Use PromptKit prepared execution for structured model calls

internal/framework/llm/adapter.go

@@ -0,0 +1,138 @@
+package llm
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"strings"
+	"time"
+
+	"gitea.maximumdirect.net/eric/audita/internal/framework/responseschema"
+	"gitea.maximumdirect.net/eric/audita/internal/framework/structuredoutput"
+	"gitea.maximumdirect.net/eric/audita/internal/prompts"
+	pk "gitea.maximumdirect.net/eric/promptkit"
+)
+
+// CallHooks receive one logical start and finish record and the execution
+// events for this call. Call records can contain transcript content.
+type CallHooks struct {
+	OnCall  func(structuredoutput.CallRecord)
+	OnEvent pk.ExecutionObserver
+}
+
+type RoleClient struct {
+	runtime *Runtime
+	role    string
+	hooks   CallHooks
+}
+
+var _ structuredoutput.Client = (*RoleClient)(nil)
+
+func (r *Runtime) NewRoleClient(role string, hooks CallHooks) (*RoleClient, error) {
+	if _, ok := r.Profiles[role]; !ok {
+		return nil, fmt.Errorf("unknown LLM role %q", role)
+	}
+	return &RoleClient{runtime: r, role: role, hooks: hooks}, nil
+}
+
+func (c *RoleClient) CompleteStructured(parent context.Context, req structuredoutput.Request, out any) (response structuredoutput.Response, err error) {
+	ctx, cancel := context.WithTimeout(parent, c.runtime.Config.LLM.OperationTimeout)
+	defer cancel()
+	record := structuredoutput.CallRecord{Context: req.Context, Started: time.Now().UTC()}
+	record.Context.Role = c.role
+	if c.hooks.OnCall != nil {
+		c.hooks.OnCall(record)
+	}
+	defer func() {
+		record.Finished = time.Now().UTC()
+		if err != nil {
+			if structuredoutput.IsMalformedError(err) {
+				record.Failure = structuredoutput.FailureMalformedOutput
+			} else {
+				record.Failure = structuredoutput.FailureRuntime
+			}
+		}
+		if c.hooks.OnCall != nil {
+			c.hooks.OnCall(record)
+		}
+	}()
+	if err = validateOutputTarget(out); err != nil {
+		return response, &structuredoutput.Failure{Kind: structuredoutput.FailureRuntime, Message: "invalid structured output target", Cause: err}
+	}
+	schemaPath, ok := responseschema.AssetPath(req.Schema)
+	if !ok {
+		return response, &structuredoutput.Failure{Kind: structuredoutput.FailureRuntime, Message: "unknown structured response schema", Cause: fmt.Errorf("schema %q", req.Schema)}
+	}
+	profile := c.runtime.Profiles[c.role]
+	backend := c.runtime.Config.Backends[profile.EffectiveModelParams.BackendID]
+	vars := make(map[string]string, len(req.Variables)+1)
+	for key, value := range req.Variables {
+		vars[key] = value
+	}
+	vars["HardeningPolicy"] = prompts.HardeningText()
+	apiKey := ""
+	if backend.Credential != nil {
+		apiKey = backend.Credential.Value()
+	}
+	handle, prepErr := c.runtime.Engine.PrepareExecution(ctx, pk.RunRequest{
+		PromptID: req.PromptID, PromptVersion: req.PromptVersion,
+		ProfileID: profile.ProfileID, APIKey: apiKey, Vars: vars,
+		Validation: &pk.OutputContract{Format: pk.FormatJSON, ValidationMode: pk.ValidationJSONSchema, SchemaPath: schemaPath, RepairAttempts: c.runtime.Config.LLM.RepairAttempts},
+		Observer:   c.hooks.OnEvent,
+	})
+	if prepErr != nil {
+		return response, &structuredoutput.Failure{Kind: structuredoutput.FailureRuntime, Message: "could not prepare structured model call", Cause: prepErr}
+	}
+	defer handle.Discard()
+	details := handle.Details()
+	response.Prompt = promptMetadata(details)
+	response.Model = details.EffectiveModelParams.Model
+	response.Prepared = preparedSnapshot(details)
+	record.Prepared = response.Prepared
+	selected, exists := c.runtime.Config.Backends[details.SelectedBackendID]
+	endpoint := func(s string) string { return strings.TrimRight(strings.TrimSpace(s), "/") }
+	if !exists || details.SelectedProfileID != profile.ProfileID || details.SelectedBackendID != profile.EffectiveModelParams.BackendID || endpoint(details.EffectiveModelParams.Endpoint) != endpoint(selected.Endpoint) || details.EffectiveModelParams.APIKeyEnv != "" {
+		return response, &structuredoutput.Failure{Kind: structuredoutput.FailureRuntime, Message: "structured model target changed after configuration inspection", Cause: fmt.Errorf("prepared target does not match registered role backend")}
+	}
+	result, runErr := c.runtime.Engine.RunPrepared(ctx, handle)
+	if result != nil {
+		response.FinalContent = result.RawOutput
+		record.FinalContent = result.RawOutput
+		response.Model = result.ModelName
+		response.PromptTokens = result.Usage.PromptTokens
+		response.CompletionTokens = result.Usage.CompletionTokens
+		response.TotalTokens = result.Usage.TotalTokens
+	}
+	if runErr != nil {
+		return response, &structuredoutput.Failure{Kind: structuredoutput.FailureRuntime, Message: "structured model call failed", Cause: runErr}
+	}
+	if result == nil {
+		return response, &structuredoutput.Failure{Kind: structuredoutput.FailureRuntime, Message: "structured model call returned no result"}
+	}
+	if !result.Validation.IsValid {
+		return response, &structuredoutput.Failure{Kind: structuredoutput.FailureMalformedOutput, Message: "malformed structured output"}
+	}
+	if decodeErr := json.Unmarshal([]byte(result.RawOutput), out); decodeErr != nil {
+		return response, &structuredoutput.Failure{Kind: structuredoutput.FailureRuntime, Message: "validated structured output could not be decoded", Cause: decodeErr}
+	}
+	response.Content = json.RawMessage(result.RawOutput)
+	response.Provider = "openai-compatible"
+	return response, nil
+}
+
+func promptMetadata(details pk.PreparedRun) structuredoutput.PromptMetadata {
+	return structuredoutput.PromptMetadata{
+		PromptID: details.PromptID, PromptVersion: details.PromptVersion,
+		DefinitionHash: details.PromptHash, RenderedHash: details.RenderedPromptHash,
+		ProfileID: details.SelectedProfileID, BackendID: details.SelectedBackendID,
+		Model: details.EffectiveModelParams.Model,
+	}
+}
+
+func preparedSnapshot(details pk.PreparedRun) *structuredoutput.PreparedSnapshot {
+	messages := make([]structuredoutput.Message, len(details.Messages))
+	for i, message := range details.Messages {
+		messages[i] = structuredoutput.Message{Role: string(message.Role), Content: message.Content}
+	}
+	return &structuredoutput.PreparedSnapshot{Messages: messages, Prompt: promptMetadata(details)}
+}