promptkit/examples/go-library/observe/main.go

package main

import (
	"context"
	"embed"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"

	"gitea.maximumdirect.net/eric/promptkit"
)

//go:embed prompts/*.yaml schemas/*.json
var assets embed.FS

type scriptedTransport struct {
	proposalCalls int
}

func (t *scriptedTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	if err := req.Context().Err(); err != nil {
		return nil, err
	}
	var request struct {
		Model string `json:"model"`
	}
	if err := json.NewDecoder(req.Body).Decode(&request); err != nil {
		return nil, err
	}
	switch request.Model {
	case "proposal-model":
		t.proposalCalls++
		if t.proposalCalls == 1 {
			return httpResponse(http.StatusServiceUnavailable, `{"error":{}}`), nil
		}
		return modelResponse(`{"title":"Queue cleanup","risk":"low"}`, 10, 8), nil
	case "assessment-model":
		return modelResponse(`{"recommendation":"accept","confidence":0.9}`, 7, 5), nil
	default:
		return nil, fmt.Errorf("unexpected model in offline transport")
	}
}

func modelResponse(content string, promptTokens, completionTokens int) *http.Response {
	body, _ := json.Marshal(map[string]any{
		"choices": []any{map[string]any{"message": map[string]string{"content": content}}},
		"usage":   map[string]int{"prompt_tokens": promptTokens, "completion_tokens": completionTokens},
	})
	return httpResponse(http.StatusOK, string(body))
}

func httpResponse(status int, body string) *http.Response {
	return &http.Response{StatusCode: status, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}
}

type eventLog struct{ events []promptkit.ExecutionEvent }

func (l *eventLog) record(event promptkit.ExecutionEvent) { l.events = append(l.events, event) }

func (l *eventLog) terminal() (promptkit.ExecutionEvent, error) {
	if len(l.events) == 0 || l.events[len(l.events)-1].Kind != promptkit.ExecutionFinished {
		return promptkit.ExecutionEvent{}, fmt.Errorf("execution did not finish")
	}
	return l.events[len(l.events)-1], nil
}

type summary struct {
	ProposalAttempts           uint64 `json:"proposal_attempts"`
	ProposalTransientRetries   uint64 `json:"proposal_transient_retries"`
	ProposalUsageComplete      bool   `json:"proposal_usage_complete"`
	AssessmentAttempts         uint64 `json:"assessment_attempts"`
	AssessmentTransientRetries uint64 `json:"assessment_transient_retries"`
	AssessmentOutcome          string `json:"assessment_outcome"`
}

func run() (summary, error) {
	transport := &scriptedTransport{}
	engine, err := promptkit.NewEngine(promptkit.Config{HTTPClient: &http.Client{Transport: transport}},
		promptkit.WithPromptFS(assets, "prompts"),
		promptkit.WithSchemaFS(assets, "schemas"),
		promptkit.WithBackend(promptkit.Backend{ID: "offline-backend", Endpoint: "https://example.invalid/v1", ConcurrencyLimit: 1}),
		promptkit.WithProfiles(
			promptkit.Profile{ID: "proposal-profile", BackendID: "offline-backend", Model: "proposal-model"},
			promptkit.Profile{ID: "assessment-profile", BackendID: "offline-backend", Model: "assessment-model"},
		))
	if err != nil {
		return summary{}, err
	}

	var proposalLog eventLog
	proposal, err := engine.Run(context.Background(), promptkit.RunRequest{
		PromptID: "example.proposal", Inputs: map[string]promptkit.ArtifactRef{"idea": promptkit.Inline("Reduce queue cleanup toil.")}, Observer: proposalLog.record,
	})
	if err != nil {
		return summary{}, err
	}
	if !proposal.Validation.IsValid {
		return summary{}, fmt.Errorf("proposal failed structural validation")
	}
	proposalTerminal, err := proposalLog.terminal()
	if err != nil {
		return summary{}, err
	}
	var assessmentLog eventLog
	assessment, err := engine.Run(context.Background(), promptkit.RunRequest{
		PromptID: "example.assessment", Inputs: map[string]promptkit.ArtifactRef{"proposal": promptkit.Inline(proposal.RawOutput)}, Observer: assessmentLog.record,
	})
	if err != nil {
		return summary{}, err
	}
	if !assessment.Validation.IsValid {
		return summary{}, fmt.Errorf("assessment failed structural validation")
	}
	assessmentTerminal, err := assessmentLog.terminal()
	if err != nil {
		return summary{}, err
	}
	var proposalValue struct {
		Risk string `json:"risk"`
	}
	var assessmentValue struct {
		Recommendation string `json:"recommendation"`
	}
	if err := json.Unmarshal([]byte(proposal.RawOutput), &proposalValue); err != nil {
		return summary{}, err
	}
	if err := json.Unmarshal([]byte(assessment.RawOutput), &assessmentValue); err != nil {
		return summary{}, err
	}
	decision := "rejected"
	if proposalValue.Risk == "low" && assessmentValue.Recommendation == "accept" {
		decision = "accepted"
	}
	return summary{
		ProposalAttempts:           proposalTerminal.Attempts,
		ProposalTransientRetries:   proposalTerminal.TransientRetries,
		ProposalUsageComplete:      proposalTerminal.Usage.PromptTokens.Complete,
		AssessmentAttempts:         assessmentTerminal.Attempts,
		AssessmentTransientRetries: assessmentTerminal.TransientRetries,
		AssessmentOutcome:          decision,
	}, nil
}

func main() {
	value, err := run()
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
	if err := json.NewEncoder(os.Stdout).Encode(value); err != nil {
		fmt.Fprintln(os.Stderr, err)
		os.Exit(1)
	}
}