Replace the Scriptorium adapter with PromptKit

2026-07-28 16:31:36 +00:00
parent 7cfab8ada0
commit 53a330587b
20 changed files with 374 additions and 277 deletions
--- a/internal/framework/llm/promptkit_client.go
+++ b/internal/framework/llm/promptkit_client.go
@@ -0,0 +1,379 @@
+package llm
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"regexp"
+	"sort"
+	"strings"
+	"sync"
+	"time"
+
+	"gitea.maximumdirect.net/eric/notarius/internal/core/artifacts"
+	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
+	"gitea.maximumdirect.net/eric/promptkit"
+)
+
+const promptKitProviderName = "promptkit"
+
+type PromptKitClientConfig struct {
+	ProfileDir    string
+	ProfileFile   string
+	Assets        *AssetRegistry
+	Timeout       time.Duration
+	HTTPClient    *http.Client
+	EngineOptions []promptkit.Option
+	Recorder      *LLMProfileRecorder
+}
+
+type PromptKitClient struct {
+	engine   *promptkit.Engine
+	recorder *LLMProfileRecorder
+}
+
+type LLMProfileRecorder struct {
+	mu       sync.Mutex
+	profiles map[string]artifacts.LLMProfileManifest
+}
+
+var _ contracts.StructuredLLMClient = (*PromptKitClient)(nil)
+var _ contracts.LLMProfileManifestProvider = (*PromptKitClient)(nil)
+
+func NewPromptKitClient(cfg PromptKitClientConfig) (*PromptKitClient, error) {
+	if cfg.Assets == nil {
+		return nil, fmt.Errorf("PromptKit client assets must not be nil")
+	}
+	if strings.TrimSpace(cfg.ProfileDir) != "" && strings.TrimSpace(cfg.ProfileFile) != "" {
+		return nil, fmt.Errorf("PromptKit profile_dir and profile_file are mutually exclusive")
+	}
+	options, err := cfg.Assets.PromptKitOptions()
+	if err != nil {
+		return nil, err
+	}
+	if profileFile := strings.TrimSpace(cfg.ProfileFile); profileFile != "" {
+		options = append(options, promptkit.WithProfileFile(profileFile))
+	}
+	options = append(options, cfg.EngineOptions...)
+
+	engine, err := promptkit.NewEngine(promptkit.Config{
+		ProfileDir: strings.TrimSpace(cfg.ProfileDir),
+		Timeout:    cfg.Timeout,
+		HTTPClient: cfg.HTTPClient,
+	}, options...)
+	if err != nil {
+		return nil, fmt.Errorf("create PromptKit engine: %w", err)
+	}
+	recorder := cfg.Recorder
+	if recorder == nil {
+		recorder = NewLLMProfileRecorder()
+	}
+	return &PromptKitClient{
+		engine:   engine,
+		recorder: recorder,
+	}, nil
+}
+
+func (c *PromptKitClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
+	if c == nil {
+		return contracts.StructuredCompletionResponse{}, fmt.Errorf("PromptKit client must not be nil")
+	}
+	if c.engine == nil {
+		return contracts.StructuredCompletionResponse{}, fmt.Errorf("PromptKit client engine must not be nil")
+	}
+	if err := validateOutputTarget(out); err != nil {
+		return contracts.StructuredCompletionResponse{}, err
+	}
+	promptID := strings.TrimSpace(req.PromptID)
+	if promptID == "" {
+		return contracts.StructuredCompletionResponse{}, fmt.Errorf("structured completion prompt_id must not be empty")
+	}
+
+	runReq := promptkit.RunRequest{
+		PromptID:      promptID,
+		PromptVersion: strings.TrimSpace(req.PromptVersion),
+		ProfileID:     strings.TrimSpace(req.ProfileID),
+		Inputs:        promptKitInputs(req.Inputs),
+		Vars:          promptKitVars(req),
+		Metadata:      promptKitMetadata(req),
+	}
+	prepared, err := c.engine.Prepare(ctx, runReq)
+	if err != nil {
+		if ctxErr := ctx.Err(); ctxErr != nil {
+			return contracts.StructuredCompletionResponse{}, ctxErr
+		}
+		return contracts.StructuredCompletionResponse{}, fmt.Errorf("prepare PromptKit prompt %q: %w", promptID, redactPromptKitError(err))
+	}
+	result, err := c.engine.Run(ctx, runReq)
+	if err != nil {
+		if ctxErr := ctx.Err(); ctxErr != nil {
+			return contracts.StructuredCompletionResponse{}, ctxErr
+		}
+		return contracts.StructuredCompletionResponse{}, fmt.Errorf("run PromptKit prompt %q: %w", promptID, redactPromptKitError(err))
+	}
+	if result == nil {
+		return contracts.StructuredCompletionResponse{}, fmt.Errorf("run PromptKit prompt %q: %w: empty result", promptID, contracts.ErrInvalidStructuredOutput)
+	}
+	response := c.responseFromResult(result, prepared)
+	if result.Validation.Status == promptkit.ValidationFailed || !result.Validation.IsValid {
+		return response, fmt.Errorf("run PromptKit prompt %q: %w: validation failed: %s", promptID, contracts.ErrInvalidStructuredOutput, strings.Join(result.Validation.Errors, "; "))
+	}
+	if len(strings.TrimSpace(string(response.Content))) == 0 {
+		return response, fmt.Errorf("run PromptKit prompt %q: %w: empty structured output", promptID, contracts.ErrInvalidStructuredOutput)
+	}
+	if err := json.Unmarshal(response.Content, out); err != nil {
+		return response, fmt.Errorf("decode PromptKit structured output for prompt %q: %w: %w", promptID, contracts.ErrInvalidStructuredOutput, err)
+	}
+	return response, nil
+}
+
+func (c *PromptKitClient) responseFromResult(result *promptkit.RunResult, prepared *promptkit.PreparedRun) contracts.StructuredCompletionResponse {
+	content := result.Artifact.Body
+	if len(content) == 0 {
+		content = []byte(result.RawOutput)
+	}
+	profile := artifacts.LLMProfileManifest{
+		ID:       strings.TrimSpace(result.SelectedProfileID),
+		Provider: promptKitProviderName,
+		Model:    firstNonEmpty(result.ModelName, result.EffectiveModelParams.Model),
+	}
+	if c.recorder != nil {
+		c.recorder.Record(profile)
+	}
+	return contracts.StructuredCompletionResponse{
+		Content:          append(json.RawMessage(nil), content...),
+		Provider:         profile.Provider,
+		Model:            profile.Model,
+		ProfileID:        profile.ID,
+		PromptTokens:     result.Usage.PromptTokens,
+		CompletionTokens: result.Usage.CompletionTokens,
+		TotalTokens:      result.Usage.TotalTokens,
+		Debug:            promptKitDebugMaterial(prepared, result),
+	}
+}
+
+func promptKitDebugMaterial(prepared *promptkit.PreparedRun, result *promptkit.RunResult) *contracts.LLMDebugMaterial {
+	material := &contracts.LLMDebugMaterial{}
+	if prepared != nil {
+		material.Prompt = promptKitDebugPrompt(prepared)
+	}
+	if result != nil {
+		material.Response = promptKitDebugResponse(result)
+	}
+	if material.Prompt == nil && material.Response == nil {
+		return nil
+	}
+	return material
+}
+
+func promptKitDebugPrompt(prepared *promptkit.PreparedRun) *contracts.LLMDebugPrompt {
+	if prepared == nil {
+		return nil
+	}
+	messages := make([]contracts.LLMDebugMessage, 0, len(prepared.Messages))
+	for _, message := range prepared.Messages {
+		messages = append(messages, contracts.LLMDebugMessage{
+			Role:         message.Role,
+			Content:      message.Content,
+			CacheControl: jsonObject(message.CacheControl),
+		})
+	}
+	return &contracts.LLMDebugPrompt{
+		PromptID:             prepared.PromptID,
+		PromptVersion:        prepared.PromptVersion,
+		PromptHash:           prepared.PromptHash,
+		SelectedProfileID:    prepared.SelectedProfileID,
+		SessionID:            prepared.SessionID,
+		RenderedPromptHash:   prepared.RenderedPromptHash,
+		Messages:             messages,
+		EffectiveModelParams: jsonObject(prepared.EffectiveModelParams),
+		OutputContract:       jsonObject(prepared.OutputContract),
+		StructuredOutput:     jsonObject(prepared.StructuredOutput),
+		InputHashes:          cloneStringMap(prepared.InputHashes),
+	}
+}
+
+func promptKitDebugResponse(result *promptkit.RunResult) *contracts.LLMDebugResponse {
+	if result == nil {
+		return nil
+	}
+	content := result.RawOutput
+	if content == "" {
+		content = string(result.Artifact.Body)
+	}
+	return &contracts.LLMDebugResponse{
+		Content:              content,
+		RunID:                result.RunID,
+		PromptID:             result.PromptID,
+		PromptVersion:        result.PromptVersion,
+		PromptHash:           result.PromptHash,
+		RenderedPromptHash:   result.RenderedPromptHash,
+		SelectedProfileID:    result.SelectedProfileID,
+		ModelName:            result.ModelName,
+		Endpoint:             result.Endpoint,
+		EffectiveModelParams: jsonObject(result.EffectiveModelParams),
+		InputHashes:          cloneStringMap(result.InputHashes),
+		Validation:           jsonObject(result.Validation),
+		Usage: contracts.LLMDebugUsage{
+			PromptTokens:     result.Usage.PromptTokens,
+			CompletionTokens: result.Usage.CompletionTokens,
+			TotalTokens:      result.Usage.TotalTokens,
+			CachedTokens:     result.Usage.CachedTokens,
+			CacheWriteTokens: result.Usage.CacheWriteTokens,
+		},
+	}
+}
+
+func jsonObject(value any) map[string]any {
+	if value == nil {
+		return nil
+	}
+	data, err := json.Marshal(value)
+	if err != nil || string(data) == "null" {
+		return nil
+	}
+	var out map[string]any
+	if err := json.Unmarshal(data, &out); err != nil {
+		return nil
+	}
+	if len(out) == 0 {
+		return nil
+	}
+	return out
+}
+
+func cloneStringMap(values map[string]string) map[string]string {
+	if len(values) == 0 {
+		return nil
+	}
+	out := make(map[string]string, len(values))
+	for key, value := range values {
+		out[key] = value
+	}
+	return out
+}
+
+func (c *PromptKitClient) LLMProfileManifests() []artifacts.LLMProfileManifest {
+	if c == nil || c.recorder == nil {
+		return nil
+	}
+	return c.recorder.Manifests()
+}
+
+func NewLLMProfileRecorder() *LLMProfileRecorder {
+	return &LLMProfileRecorder{profiles: map[string]artifacts.LLMProfileManifest{}}
+}
+
+func (r *LLMProfileRecorder) Record(profile artifacts.LLMProfileManifest) {
+	if r == nil {
+		return
+	}
+	profile.ID = strings.TrimSpace(profile.ID)
+	profile.Provider = strings.TrimSpace(profile.Provider)
+	profile.Model = strings.TrimSpace(profile.Model)
+	key := profile.ID + "\x00" + profile.Provider + "\x00" + profile.Model
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if r.profiles == nil {
+		r.profiles = map[string]artifacts.LLMProfileManifest{}
+	}
+	r.profiles[key] = profile
+}
+
+func (r *LLMProfileRecorder) Manifests() []artifacts.LLMProfileManifest {
+	if r == nil {
+		return nil
+	}
+	r.mu.Lock()
+	defer r.mu.Unlock()
+	if len(r.profiles) == 0 {
+		return nil
+	}
+	keys := make([]string, 0, len(r.profiles))
+	for key := range r.profiles {
+		keys = append(keys, key)
+	}
+	sort.Strings(keys)
+	out := make([]artifacts.LLMProfileManifest, 0, len(keys))
+	for _, key := range keys {
+		out = append(out, r.profiles[key])
+	}
+	return out
+}
+
+func promptKitInputs(inputs contracts.LLMInputSet) map[string]promptkit.ArtifactRef {
+	if len(inputs) == 0 {
+		return nil
+	}
+	out := make(map[string]promptkit.ArtifactRef, len(inputs))
+	for key, material := range inputs {
+		name := strings.TrimSpace(key)
+		if name == "" {
+			name = strings.TrimSpace(material.Name)
+		}
+		if name == "" {
+			continue
+		}
+		body := string(material.Content)
+		if body == "" {
+			body = " "
+		}
+		if origin := strings.TrimSpace(material.OriginURI); origin != "" {
+			out[name] = promptkit.InlineWithURI(origin, body)
+		} else {
+			out[name] = promptkit.Inline(body)
+		}
+	}
+	return out
+}
+
+func promptKitVars(req contracts.StructuredCompletionRequest) map[string]string {
+	vars := make(map[string]string, len(req.Vars)+1)
+	for key, value := range req.Vars {
+		name := strings.TrimSpace(key)
+		if name == "" || value == nil {
+			continue
+		}
+		vars[name] = fmt.Sprint(value)
+	}
+	if sessionID := strings.TrimSpace(req.SessionID); sessionID != "" {
+		vars["session_id"] = sessionID
+	}
+	if len(vars) == 0 {
+		return nil
+	}
+	return vars
+}
+
+func promptKitMetadata(req contracts.StructuredCompletionRequest) map[string]string {
+	metadata := map[string]string{}
+	if stageName := strings.TrimSpace(req.StageName); stageName != "" {
+		metadata["stage_name"] = stageName
+	}
+	if len(metadata) == 0 {
+		return nil
+	}
+	return metadata
+}
+
+var bearerTokenPattern = regexp.MustCompile(`(?i)Bearer\s+[A-Za-z0-9._~+/=-]+`)
+
+func redactPromptKitError(err error) error {
+	if err == nil {
+		return nil
+	}
+	return redactedProviderError{err: err}
+}
+
+type redactedProviderError struct {
+	err error
+}
+
+func (e redactedProviderError) Error() string {
+	return bearerTokenPattern.ReplaceAllString(e.err.Error(), "Bearer "+secretReplacement)
+}
+
+func (e redactedProviderError) Unwrap() error {
+	return e.err
+}