Files
audita/internal/framework/llm/conversation_test.go

302 lines
14 KiB
Go

package llm
import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"net/http"
"os"
"path/filepath"
"reflect"
"strings"
"testing"
"time"
"gitea.maximumdirect.net/eric/audita/v3/internal/core/config"
"gitea.maximumdirect.net/eric/audita/v3/internal/framework/responseschema"
"gitea.maximumdirect.net/eric/audita/v3/internal/framework/structuredoutput"
"gitea.maximumdirect.net/eric/audita/v3/internal/prompts"
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)
}
}
}
func TestRoleClientCapturesAndReplaysRepairConversation(t *testing.T) {
dir := t.TempDir()
for name, content := range map[string]string{
"config.yml": "version: 3\n",
"prompts.yml": "profiles:\n - id: proposal\n endpoint: https://example.invalid/v1\n model: proposal\n - id: repair\n endpoint: https://example.invalid/v1\n model: repair\ncollections:\n audita:\n prompts:\n - id: modules.grammar.repair\n messages:\n - role: user\n content: '{{ .RepairRequestJSON }}'\n cache_control: {type: ephemeral, ttl: 1h}\n output:\n format: json\n validation_mode: json_schema\n schema_path: correction_repair.v1.json\n",
} {
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)
}
type wireMessage struct {
Role string
Content json.RawMessage
}
var messages [][]wireMessage
var models []string
outputs := []string{`{"corrections":[]}`, `{"proposal_ref":1,"action":"revise","correction":{"id":1,"original_text":"a","corrected_text":"b","confidence":0.8}}`, `{"proposal_ref":1,"action":"withdraw","correction":null}`, `{"proposal_ref":1,"action":"withdraw","correction":null}`}
transport := transportFunc(func(r *http.Request) (*http.Response, error) {
var req struct {
Model string `json:"model"`
Messages []wireMessage `json:"messages"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Fatal(err)
}
models = append(models, req.Model)
messages = append(messages, req.Messages)
content, _ := json.Marshal(outputs[len(messages)-1])
body := fmt.Sprintf(`{"model":%q,"choices":[{"message":{"content":%s}}]}`, req.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)
}
settings := config.LLM{OperationTimeout: 10 * time.Second}
proposal, err := runtime.NewRoleClient("proposal", "proposal", settings, CallHooks{})
if err != nil {
t.Fatal(err)
}
repair, err := runtime.NewRoleClient("repair", "repair", settings, CallHooks{})
if err != nil {
t.Fatal(err)
}
vars := map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}
initial, err := proposal.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarProposal, Schema: responseschema.CorrectionSetKey, Variables: vars, CaptureExchange: true}, &map[string]any{})
if err != nil || initial.Exchange == nil {
t.Fatalf("capture: %+v %v", initial, err)
}
vars["SectionJSON"] = "changed"
replay := &structuredoutput.Replay{Exchange: initial.Exchange}
for _, finding := range []string{"first", "second"} {
result, err := repair.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarRepair, Schema: responseschema.CorrectionRepairKey, Variables: map[string]string{"RepairRequestJSON": finding}, Replay: replay}, &map[string]any{})
if err != nil {
t.Fatal(err)
}
if len(result.RepairInstruction) != 1 {
t.Fatalf("instruction: %+v", result)
}
replay.Turns = append(replay.Turns, result.RepairInstruction...)
replay.Turns = append(replay.Turns, structuredoutput.Message{Role: "assistant", Content: result.FinalContent})
}
if !reflect.DeepEqual(models, []string{"proposal", "repair", "repair"}) {
t.Fatalf("models = %v", models)
}
prefix := messages[0]
if len(messages) != 3 || !reflect.DeepEqual(messages[1][:len(prefix)], prefix) || !reflect.DeepEqual(messages[2][:len(prefix)], prefix) {
t.Fatalf("prefixes: %+v", messages)
}
encodedRevision, _ := json.Marshal(outputs[1])
if len(messages[1]) != len(prefix)+2 || len(messages[2]) != len(prefix)+4 || string(messages[2][len(prefix)+2].Content) != string(encodedRevision) {
t.Fatalf("turns: %+v", messages)
}
first := messages[1][len(prefix)+1]
if first.Role != "user" || !strings.Contains(string(first.Content), `"ttl":"1h"`) ||
!strings.Contains(string(first.Content), `"text":"first"`) || !reflect.DeepEqual(first, messages[2][len(prefix)+1]) {
t.Fatalf("cached instruction changed during replay: first=%s replay=%s", first.Content, messages[2][len(prefix)+1].Content)
}
_, err = repair.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarRepair, Schema: responseschema.CorrectionRepairKey, Variables: map[string]string{"RepairRequestJSON": "independent"}, Replay: &structuredoutput.Replay{Exchange: initial.Exchange}}, &map[string]any{})
if err != nil || len(messages) != 4 || len(messages[3]) != len(prefix)+2 || !strings.Contains(string(messages[3][len(prefix)+1].Content), "independent") {
t.Fatalf("independent branch: %v %+v", err, messages)
}
bad := *initial.Exchange
bad.Messages = append([]structuredoutput.Message(nil), bad.Messages...)
bad.Messages[0].Content = "changed"
_, err = repair.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarRepair, Schema: responseschema.CorrectionRepairKey, Variables: map[string]string{"RepairRequestJSON": "failure"}, Replay: &structuredoutput.Replay{Exchange: &bad}}, &map[string]any{})
var failure *structuredoutput.Failure
if !errors.As(err, &failure) || failure.Class != structuredoutput.FailureFatal || len(messages) != 4 {
t.Fatalf("prefix failure: %v, calls %d", err, len(messages))
}
}
func TestRepairReplayUsesReplacementAndCountsStructuralRepair(t *testing.T) {
dir := t.TempDir()
files := map[string]string{
"config.yml": "version: 3\n",
"prompts.yml": "profiles:\n - id: fixture\n endpoint: https://example.invalid/v1\n model: fixture\ncollections:\n audita:\n prompts:\n - id: modules.grammar.repair\n messages:\n - role: user\n content_file: repair.txt\n output:\n format: json\n validation_mode: json_schema\n schema_path: correction_repair.v1.json\n",
"repair.txt": "Replace {{.RepairRequestJSON}}",
}
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)
}
outputs := []string{`{"corrections":[]}`, `{"bad":true}`, `{"proposal_ref":1,"action":"withdraw","correction":null}`}
var requests [][]structuredoutput.Message
transport := transportFunc(func(r *http.Request) (*http.Response, error) {
var req struct {
Messages []structuredoutput.Message `json:"messages"`
}
if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
t.Fatal(err)
}
requests = append(requests, req.Messages)
content, _ := json.Marshal(outputs[len(requests)-1])
body := fmt.Sprintf(`{"model":"fixture","choices":[{"message":{"content":%s}}]}`, 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)
}
client, err := runtime.NewRoleClient("repair", "fixture", config.LLM{OperationTimeout: 10 * time.Second, RepairAttempts: 1}, CallHooks{})
if err != nil {
t.Fatal(err)
}
base, err := client.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarProposal, Schema: responseschema.CorrectionSetKey, Variables: map[string]string{"TranscriptDescriptionBlock": "", "GlossaryJSON": "[]", "SectionJSON": "[]"}, CaptureExchange: true}, &map[string]any{})
if err != nil {
t.Fatal(err)
}
result, err := client.CompleteStructured(context.Background(), structuredoutput.Request{PromptID: prompts.PromptIDModuleGrammarRepair, Schema: responseschema.CorrectionRepairKey, Variables: map[string]string{"RepairRequestJSON": "fixture"}, Replay: &structuredoutput.Replay{Exchange: base.Exchange}}, &map[string]any{})
if err != nil {
t.Fatal(err)
}
if len(requests) != 3 || result.Execution.ProviderAttempts != 2 || result.Execution.StructuralRepairs != 1 || result.FinalContent != outputs[2] || len(result.RepairInstruction) != 1 || result.RepairInstruction[0].Content != "Replace fixture" {
t.Fatalf("result=%+v calls=%d", result, len(requests))
}
}