292 lines
13 KiB
Go
292 lines
13 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\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)
|
|
}
|
|
var messages [][]structuredoutput.Message
|
|
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 []structuredoutput.Message `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 := initial.Exchange.Messages
|
|
if len(messages) != 3 || !reflect.DeepEqual(messages[0], prefix) || !reflect.DeepEqual(messages[1][:len(prefix)], prefix) || !reflect.DeepEqual(messages[2][:len(prefix)], prefix) {
|
|
t.Fatalf("prefixes: %+v", messages)
|
|
}
|
|
if len(messages[1]) != len(prefix)+2 || len(messages[2]) != len(prefix)+4 || messages[2][len(prefix)+2].Content != outputs[1] {
|
|
t.Fatalf("turns: %+v", messages)
|
|
}
|
|
_, 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(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))
|
|
}
|
|
}
|