audita/internal/framework/llm/conversation_test.go

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)
		}
	}
}