Define correction repair response and conversation replay

internal/framework/llm/conversation_test.go

@@ -0,0 +1,150 @@
+package llm
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"path/filepath"
+	"reflect"
+	"strings"
+	"testing"
+
+	"gitea.maximumdirect.net/eric/audita/v3/internal/core/config"
+	"gitea.maximumdirect.net/eric/audita/v3/internal/framework/responseschema"
+	ck "go.maximumdirect.net/configkit"
+	pk "go.maximumdirect.net/promptkit"
+)
+
+func TestNativeConversationReplayWithRepairContract(t *testing.T) {
+	dir := t.TempDir()
+	files := map[string]string{
+		"config.yml": "version: 3\n",
+		"prompts.yml": `profiles:
+  - id: proposal-profile
+    endpoint: https://example.invalid/v1
+    model: proposal-model
+  - id: repair-profile
+    endpoint: https://example.invalid/v1
+    model: repair-model
+prompts:
+  - id: proposal-fixture
+    messages:
+      - role: system
+        content: 'System {{.Subject}}'
+      - role: user
+        content: 'Proposal {{.Subject}}'
+    output:
+      format: json
+      validation_mode: json_schema
+      schema_path: correction_set.v1.json
+  - id: repair-fixture
+    messages:
+      - role: user
+        content: 'Repair {{.Finding}}'
+    output:
+      format: text
+      validation_mode: none
+`,
+	}
+	for name, content := range files {
+		if err := os.WriteFile(filepath.Join(dir, name), []byte(content), 0644); err != nil {
+			t.Fatal(err)
+		}
+	}
+	loaded, err := config.Load(ck.Input{Source: ck.File(filepath.Join(dir, "config.yml"))}, "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	var requests []struct {
+		Model          string                           `json:"model"`
+		Messages       []struct{ Role, Content string } `json:"messages"`
+		ResponseFormat map[string]any                   `json:"response_format"`
+	}
+	outputs := []string{`{"corrections":[]}`, `{"proposal_ref":1,"action":"revise","correction":{"id":1,"original_text":" ","corrected_text":"","confidence":0.8}}`, `{"proposal_ref":1,"action":"withdraw","correction":null}`}
+	transport := transportFunc(func(r *http.Request) (*http.Response, error) {
+		var request struct {
+			Model          string                           `json:"model"`
+			Messages       []struct{ Role, Content string } `json:"messages"`
+			ResponseFormat map[string]any                   `json:"response_format"`
+		}
+		if err := json.NewDecoder(r.Body).Decode(&request); err != nil {
+			t.Fatal(err)
+		}
+		requests = append(requests, request)
+		content, _ := json.Marshal(outputs[len(requests)-1])
+		body := fmt.Sprintf(`{"model":%q,"choices":[{"message":{"content":%s}}]}`, request.Model, content)
+		return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body)), Request: r}, nil
+	})
+	runtime, err := NewRuntime(context.Background(), loaded, &http.Client{Transport: transport})
+	if err != nil {
+		t.Fatal(err)
+	}
+	ctx := context.Background()
+	base := pk.RunRequest{PromptID: "proposal-fixture", ProfileID: "proposal-profile", Vars: map[string]string{"Subject": "source"}}
+	original, err := runtime.Engine.PrepareExecution(ctx, base)
+	if err != nil {
+		t.Fatal(err)
+	}
+	prefix := original.Details().Messages
+	initial, err := runtime.Engine.RunPrepared(ctx, original)
+	if err != nil {
+		t.Fatal(err)
+	}
+	suffix := []pk.RenderedMessage{{Role: pk.RoleAssistant, Content: initial.RawOutput}}
+	contract, err := OutputContract(responseschema.CorrectionRepairKey, 0)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for _, finding := range []string{"first", "second"} {
+		instruction, err := runtime.Engine.Prepare(ctx, pk.RunRequest{PromptID: "repair-fixture", ProfileID: "repair-profile", Vars: map[string]string{"Finding": finding}})
+		if err != nil {
+			t.Fatal(err)
+		}
+		suffix = append(suffix, instruction.Messages...)
+		req := base
+		req.ProfileID = "repair-profile"
+		req.AppendedMessages = suffix
+		req.Validation = contract
+		handle, err := runtime.Engine.PrepareExecution(ctx, req)
+		if err != nil {
+			t.Fatal(err)
+		}
+		details := handle.Details()
+		if !reflect.DeepEqual(details.Messages[:len(prefix)], prefix) || !reflect.DeepEqual(details.Messages[len(prefix):], suffix) {
+			t.Fatalf("prepared conversation = %#v", details.Messages)
+		}
+		result, err := runtime.Engine.RunPrepared(ctx, handle)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if !result.Validation.IsValid {
+			t.Fatalf("invalid result: %+v", result.Validation)
+		}
+		suffix = append(suffix, pk.RenderedMessage{Role: pk.RoleAssistant, Content: result.RawOutput})
+	}
+	if len(requests) != 3 || requests[0].Model != "proposal-model" || requests[1].Model != "repair-model" || requests[2].Model != "repair-model" {
+		t.Fatalf("models: %+v", requests)
+	}
+	for i, wantLen := range []int{2, 4, 6} {
+		req := requests[i]
+		if len(req.Messages) != wantLen {
+			t.Fatalf("request %d messages: %+v", i, req.Messages)
+		}
+		for j, message := range req.Messages {
+			if message.Role != string(append(prefix, suffix...)[j].Role) || message.Content != append(prefix, suffix...)[j].Content {
+				t.Fatalf("request %d message %d: %+v", i, j, message)
+			}
+		}
+		encoded, _ := json.Marshal(req.ResponseFormat)
+		expected := "corrections"
+		if i > 0 {
+			expected = "proposal_ref"
+		}
+		if !strings.Contains(string(encoded), expected) {
+			t.Fatalf("request %d contract: %s", i, encoded)
+		}
+	}
+}