350 lines
12 KiB
Go
350 lines
12 KiB
Go
package usecase
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
"unicode/utf8"
|
|
|
|
"gitea.maximumdirect.net/eric/promptkit/internal/domain"
|
|
)
|
|
|
|
type recordingRepairClient struct {
|
|
requests []domain.GenerateRequest
|
|
response *domain.GenerateResponse
|
|
err error
|
|
}
|
|
|
|
func (c *recordingRepairClient) Generate(_ context.Context, req domain.GenerateRequest) (*domain.GenerateResponse, error) {
|
|
c.requests = append(c.requests, req)
|
|
return c.response, c.err
|
|
}
|
|
|
|
func TestDefaultOutputRepairerBuildsFullContextRequest(t *testing.T) {
|
|
client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
|
|
repairer := NewDefaultOutputRepairer(client)
|
|
original := []domain.RenderedMessage{
|
|
{Role: "system", Content: "Follow the task.", CacheControl: &domain.CacheControl{Type: domain.CacheControlEphemeral, TTL: "1h"}},
|
|
{Role: "user", Content: "Summarize the report."},
|
|
}
|
|
before := append([]domain.RenderedMessage(nil), original...)
|
|
previous := strings.Repeat("candidate ", 12_000)
|
|
target := domain.ExecutionTarget{BackendID: "backend", Endpoint: "https://provider.example/v1", Model: "model"}
|
|
presence := domain.ExecutionTargetPresence{Temperature: true, TopP: true}
|
|
structured := &domain.StructuredOutputSpec{}
|
|
|
|
response, err := repairer.Repair(context.Background(), RepairRequest{
|
|
OriginalMessages: original,
|
|
PreviousOutput: previous,
|
|
ValidationErrors: []string{"invalid JSON"},
|
|
SessionID: "session",
|
|
Target: target,
|
|
TargetPresence: presence,
|
|
StructuredOutput: structured,
|
|
Attempt: 1,
|
|
MaxAttempts: 3,
|
|
Mode: domain.ValidationJSONSchema,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("repair: %v", err)
|
|
}
|
|
if response == nil || response.Content != "corrected" {
|
|
t.Fatalf("response = %+v", response)
|
|
}
|
|
if !reflect.DeepEqual(original, before) {
|
|
t.Fatalf("original messages changed: got %#v, want %#v", original, before)
|
|
}
|
|
if len(client.requests) != 1 {
|
|
t.Fatalf("generation requests = %d, want 1", len(client.requests))
|
|
}
|
|
|
|
request := client.requests[0]
|
|
if request.Prompt.SessionID != "session" || !reflect.DeepEqual(request.Target, target) || request.TargetPresence != presence || request.StructuredOutput != structured {
|
|
t.Fatalf("generation request fields = %+v", request)
|
|
}
|
|
if len(request.Prompt.Messages) != len(original)+2 {
|
|
t.Fatalf("message count = %d, want %d", len(request.Prompt.Messages), len(original)+2)
|
|
}
|
|
if !reflect.DeepEqual(request.Prompt.Messages[:len(original)], original) {
|
|
t.Fatalf("original messages = %#v, want %#v", request.Prompt.Messages[:len(original)], original)
|
|
}
|
|
assistant := request.Prompt.Messages[len(original)]
|
|
if assistant.Role != "assistant" || assistant.Content != previous {
|
|
t.Fatalf("assistant candidate = %+v", assistant)
|
|
}
|
|
correction := request.Prompt.Messages[len(original)+1]
|
|
if correction.Role != "user" || !strings.Contains(correction.Content, "Repair attempt 1 of 3") ||
|
|
!strings.Contains(correction.Content, "Preserve valid values") ||
|
|
!strings.Contains(correction.Content, "Return only corrected JSON") {
|
|
t.Fatalf("correction message = %q", correction.Content)
|
|
}
|
|
if diagnostics := repairDiagnosticsFromMessage(t, correction.Content); !reflect.DeepEqual(diagnostics, []string{"invalid JSON"}) {
|
|
t.Fatalf("diagnostics = %#v", diagnostics)
|
|
}
|
|
}
|
|
|
|
func TestDefaultOutputRepairerDoesNotAccumulateCandidates(t *testing.T) {
|
|
client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
|
|
repairer := NewDefaultOutputRepairer(client)
|
|
backing := make([]domain.RenderedMessage, 1, 4)
|
|
backing[0] = domain.RenderedMessage{Role: "user", Content: "Original task"}
|
|
before := append([]domain.RenderedMessage(nil), backing...)
|
|
|
|
for _, candidate := range []string{"first invalid", "second invalid"} {
|
|
_, err := repairer.Repair(context.Background(), RepairRequest{
|
|
OriginalMessages: backing,
|
|
PreviousOutput: candidate,
|
|
Attempt: 1,
|
|
MaxAttempts: 3,
|
|
Mode: domain.ValidationJSON,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("repair %q: %v", candidate, err)
|
|
}
|
|
}
|
|
if !reflect.DeepEqual(backing, before) {
|
|
t.Fatalf("caller messages changed: got %#v, want %#v", backing, before)
|
|
}
|
|
if len(client.requests) != 2 {
|
|
t.Fatalf("generation requests = %d, want 2", len(client.requests))
|
|
}
|
|
for index, request := range client.requests {
|
|
messages := request.Prompt.Messages
|
|
if len(messages) != 3 || messages[0] != backing[0] || messages[1].Role != "assistant" || messages[1].Content != []string{"first invalid", "second invalid"}[index] {
|
|
t.Fatalf("request %d messages = %#v", index, messages)
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestDefaultOutputRepairerOmitsEmptyCandidateMessage(t *testing.T) {
|
|
for _, tc := range []struct {
|
|
name string
|
|
candidate string
|
|
}{
|
|
{name: "empty", candidate: ""},
|
|
{name: "whitespace", candidate: " \n\t "},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
|
|
repairer := NewDefaultOutputRepairer(client)
|
|
_, err := repairer.Repair(context.Background(), RepairRequest{
|
|
OriginalMessages: []domain.RenderedMessage{{Role: "user", Content: "Original task"}},
|
|
PreviousOutput: tc.candidate,
|
|
Attempt: 1,
|
|
MaxAttempts: 1,
|
|
Mode: domain.ValidationBasic,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("repair: %v", err)
|
|
}
|
|
messages := client.requests[0].Prompt.Messages
|
|
if len(messages) != 2 || messages[1].Role != "user" || !strings.Contains(messages[1].Content, "previous response was empty") {
|
|
t.Fatalf("messages = %#v", messages)
|
|
}
|
|
if strings.Contains(messages[1].Content, tc.candidate) && tc.candidate != "" {
|
|
t.Fatalf("correction message repeated whitespace candidate: %q", messages[1].Content)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestDefaultOutputRepairerUsesModeSpecificGuidance(t *testing.T) {
|
|
tests := []struct {
|
|
mode domain.ValidationMode
|
|
want string
|
|
}{
|
|
{mode: domain.ValidationBasic, want: "Return a nonempty response"},
|
|
{mode: domain.ValidationJSON, want: "Return only corrected JSON"},
|
|
{mode: domain.ValidationJSONSchema, want: "Return only corrected JSON"},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(string(tc.mode), func(t *testing.T) {
|
|
client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "corrected"}}
|
|
_, err := NewDefaultOutputRepairer(client).Repair(context.Background(), RepairRequest{Mode: tc.mode, Attempt: 1, MaxAttempts: 1})
|
|
if err != nil {
|
|
t.Fatalf("repair: %v", err)
|
|
}
|
|
if !strings.Contains(client.requests[0].Prompt.Messages[0].Content, tc.want) {
|
|
t.Fatalf("correction message = %q", client.requests[0].Prompt.Messages[0].Content)
|
|
}
|
|
})
|
|
}
|
|
|
|
client := &recordingRepairClient{response: &domain.GenerateResponse{Content: "unexpected"}}
|
|
_, err := NewDefaultOutputRepairer(client).Repair(context.Background(), RepairRequest{Mode: domain.ValidationNone})
|
|
if err == nil || !strings.Contains(err.Error(), "unsupported validation mode") {
|
|
t.Fatalf("unsupported mode error = %v", err)
|
|
}
|
|
if len(client.requests) != 0 {
|
|
t.Fatalf("generation requests = %d, want 0", len(client.requests))
|
|
}
|
|
}
|
|
|
|
func TestFormatRepairDiagnosticsBoundsAndPreservesData(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
errors []string
|
|
wantOmission bool
|
|
check func(*testing.T, []string)
|
|
}{
|
|
{
|
|
name: "below limit",
|
|
errors: []string{"first", "second"},
|
|
check: func(t *testing.T, got []string) {
|
|
t.Helper()
|
|
if !reflect.DeepEqual(got, []string{"first", "second"}) {
|
|
t.Fatalf("diagnostics = %#v", got)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "at limit",
|
|
errors: []string{strings.Repeat("x", maxRepairDiagnosticBytes-4)},
|
|
check: func(t *testing.T, got []string) {
|
|
t.Helper()
|
|
if len(got) != 1 || len(got[0]) != maxRepairDiagnosticBytes-4 {
|
|
t.Fatalf("diagnostics lengths = %#v", got)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "multibyte truncation preserves prior entries",
|
|
errors: []string{"first", strings.Repeat("界", maxRepairDiagnosticBytes)},
|
|
wantOmission: true,
|
|
check: func(t *testing.T, got []string) {
|
|
t.Helper()
|
|
if len(got) != 3 || got[0] != "first" || !strings.HasPrefix(strings.Repeat("界", maxRepairDiagnosticBytes), got[1]) {
|
|
t.Fatalf("diagnostics = %#v", got)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "many diagnostics",
|
|
errors: manyRepairDiagnostics(),
|
|
wantOmission: true,
|
|
check: func(t *testing.T, got []string) {
|
|
t.Helper()
|
|
if len(got) < 2 || got[0] != manyRepairDiagnostics()[0] {
|
|
t.Fatalf("diagnostics = %#v", got)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "no room for partial diagnostic",
|
|
errors: func() []string {
|
|
omission, _ := json.Marshal(omittedRepairDiagnostics)
|
|
return []string{
|
|
strings.Repeat("x", maxRepairDiagnosticBytes-len(omission)-5),
|
|
strings.Repeat("y", 128),
|
|
}
|
|
}(),
|
|
wantOmission: true,
|
|
check: func(t *testing.T, got []string) {
|
|
t.Helper()
|
|
if len(got) != 2 || got[1] != omittedRepairDiagnostics {
|
|
t.Fatalf("diagnostics = %#v", got)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "invalid UTF-8",
|
|
errors: []string{"broken\xffinput"},
|
|
check: func(t *testing.T, got []string) {
|
|
t.Helper()
|
|
if !reflect.DeepEqual(got, []string{"broken\uFFFDinput"}) {
|
|
t.Fatalf("diagnostics = %#v", got)
|
|
}
|
|
},
|
|
},
|
|
{
|
|
name: "one huge diagnostic",
|
|
errors: []string{strings.Repeat("x", maxRepairDiagnosticBytes*2)},
|
|
wantOmission: true,
|
|
check: func(t *testing.T, got []string) {
|
|
t.Helper()
|
|
if len(got) != 2 || !strings.HasPrefix(strings.Repeat("x", maxRepairDiagnosticBytes*2), got[0]) {
|
|
t.Fatalf("diagnostics = %#v", got)
|
|
}
|
|
},
|
|
},
|
|
}
|
|
|
|
for _, tc := range tests {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
before := append([]string(nil), tc.errors...)
|
|
encoded := formatRepairDiagnostics(tc.errors)
|
|
if len(encoded) > maxRepairDiagnosticBytes || !utf8.ValidString(encoded) {
|
|
t.Fatalf("encoded diagnostic length/UTF-8 = (%d, %v)", len(encoded), utf8.ValidString(encoded))
|
|
}
|
|
var got []string
|
|
if err := json.Unmarshal([]byte(encoded), &got); err != nil {
|
|
t.Fatalf("decode diagnostics: %v; encoded=%q", err, encoded)
|
|
}
|
|
if !reflect.DeepEqual(tc.errors, before) {
|
|
t.Fatalf("input diagnostics changed: got %#v, want %#v", tc.errors, before)
|
|
}
|
|
if hasOmission := len(got) > 0 && got[len(got)-1] == omittedRepairDiagnostics; hasOmission != tc.wantOmission {
|
|
t.Fatalf("omission = %v, want %v; diagnostics=%#v", hasOmission, tc.wantOmission, got)
|
|
}
|
|
tc.check(t, got)
|
|
})
|
|
}
|
|
}
|
|
|
|
func manyRepairDiagnostics() []string {
|
|
diagnostics := make([]string, 1_000)
|
|
for index := range diagnostics {
|
|
diagnostics[index] = fmt.Sprintf("diagnostic %04d %s", index, strings.Repeat("x", 128))
|
|
}
|
|
return diagnostics
|
|
}
|
|
|
|
func TestDefaultOutputRepairerPropagatesGenerationFailures(t *testing.T) {
|
|
expected := errors.New("generation failed")
|
|
for _, tc := range []struct {
|
|
name string
|
|
client *recordingRepairClient
|
|
want error
|
|
}{
|
|
{name: "nil client", want: nil},
|
|
{name: "generation error", client: &recordingRepairClient{err: expected}, want: expected},
|
|
{name: "nil response", client: &recordingRepairClient{}, want: nil},
|
|
} {
|
|
t.Run(tc.name, func(t *testing.T) {
|
|
var repairer OutputRepairer
|
|
if tc.client == nil {
|
|
repairer = NewDefaultOutputRepairer(nil)
|
|
} else {
|
|
repairer = NewDefaultOutputRepairer(tc.client)
|
|
}
|
|
response, err := repairer.Repair(context.Background(), RepairRequest{Mode: domain.ValidationJSON})
|
|
if response != nil || err == nil {
|
|
t.Fatalf("response/error = (%+v, %v)", response, err)
|
|
}
|
|
if tc.want != nil && !errors.Is(err, tc.want) {
|
|
t.Fatalf("error = %v, want %v", err, tc.want)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func repairDiagnosticsFromMessage(t *testing.T, message string) []string {
|
|
t.Helper()
|
|
const marker = "Validation diagnostics (data):\n"
|
|
index := strings.Index(message, marker)
|
|
if index < 0 {
|
|
t.Fatalf("missing diagnostics marker in %q", message)
|
|
}
|
|
encoded := message[index+len(marker):]
|
|
var diagnostics []string
|
|
if err := json.Unmarshal([]byte(encoded), &diagnostics); err != nil {
|
|
t.Fatalf("decode diagnostics: %v", err)
|
|
}
|
|
return diagnostics
|
|
}
|