eric/notarius
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enable structural output repair by default

Browse Source
This commit is contained in:
Eric Rakestraw
2026-08-25 20:04:56 +00:00
parent 63c397d86a
commit 3d3f16db4a
25 changed files with 241 additions and 52 deletions
Download Patch File Download Diff File Expand all files Collapse all files

131
internal/framework/llm/promptkit_client_test.go
View File

@@ -4,6 +4,7 @@ import (
"context"
"encoding/json"
"errors"
"fmt"
"io"
"io/fs"
"net/http"
@@ -922,12 +923,15 @@ func TestLLMProfileRecorderDistinguishesEffectiveTargets(t *testing.T) {
}
func TestPromptKitClientValidationFailureReturnsError(t *testing.T) {
client := newTestPromptKitClient(t, &fakePromptKitLLM{content: `{"bad":true}`})
attempts := 1
fake := &fakePromptKitLLM{responses: []promptkit.GenerateResponse{{Content: `{"bad":true}`}, {Content: `{"bad":true}`}}}
client := newTestPromptKitClient(t, fake)
var out map[string]any
resp, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
PromptID: "adapter.test",
SessionID: "session-123",
PromptID: "adapter.test",
SessionID: "session-123",
StructuredOutputRepairAttempts: &attempts,
Inputs: contracts.LLMInputSet{
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
},
@@ -935,6 +939,9 @@ func TestPromptKitClientValidationFailureReturnsError(t *testing.T) {
if err == nil || !errors.Is(err, contracts.ErrInvalidStructuredOutput) || !strings.Contains(err.Error(), "validation failed") {
t.Fatalf("CompleteStructured() error = %v, want validation failure", err)
}
if got := atomic.LoadInt32(&fake.calls); got != 2 {
t.Fatalf("provider calls = %d, want exhausted repair budget", got)
}
if got := string(resp.Content); got != `{"bad":true}` {
t.Fatalf("response content = %q, want raw failed output", got)
}
@@ -946,6 +953,45 @@ func TestPromptKitClientValidationFailureReturnsError(t *testing.T) {
}
}
func TestPromptKitClientRepairsStructuredOutputAndReportsCumulativeUsage(t *testing.T) {
attempts := 1
fake := &fakePromptKitLLM{responses: []promptkit.GenerateResponse{
{Content: `{"bad":true}`, Usage: promptkit.TokenUsage{PromptTokens: 3, CompletionTokens: 5, TotalTokens: 8}},
{Content: `{"ok":true}`, Usage: promptkit.TokenUsage{PromptTokens: 7, CompletionTokens: 11, TotalTokens: 18}},
}}
client := newTestPromptKitClient(t, fake)
var out struct {
OK bool `json:"ok"`
}
response, err := client.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
PromptID: "adapter.test",
StructuredOutputRepairAttempts: &attempts,
SessionID: "repair-test",
Inputs: contracts.LLMInputSet{
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
},
}, &out)
if err != nil || !out.OK {
t.Fatalf("CompleteStructured() = (%#v, %v), want repaired success", response, err)
}
if got := atomic.LoadInt32(&fake.calls); got != 2 {
t.Fatalf("provider calls = %d, want initial generation and one repair", got)
}
if response.RepairAttempts != 1 || response.PromptTokens != 10 || response.CompletionTokens != 16 || response.TotalTokens != 26 {
t.Fatalf("response repair and usage = %#v, want one repair and PromptKit cumulative usage", response)
}
if response.Debug == nil || response.Debug.Prompt == nil || response.Debug.Response == nil {
t.Fatalf("debug = %#v, want prompt and response details", response.Debug)
}
if response.Debug.Prompt.OutputContract["repair_attempts"] != float64(1) ||
response.Debug.Response.Validation["repair_attempts"] != float64(1) ||
response.Debug.Response.Content != `{"ok":true}` ||
response.Debug.Response.Usage.TotalTokens != 26 {
t.Fatalf("debug repair result = %#v, want configured contract and repaired response", response.Debug)
}
}
func TestPromptKitClientDecodeFailureReturnsRawResponse(t *testing.T) {
client := newTestPromptKitClient(t, &fakePromptKitLLM{content: `{"ok":true}`})
@@ -1221,6 +1267,66 @@ func TestScheduledPromptKitClientBoundsConcurrentCalls(t *testing.T) {
wg.Wait()
}
func TestScheduledPromptKitClientHoldsPermitAcrossStructuredOutputRepair(t *testing.T) {
attempts := 1
fake := &fakePromptKitLLM{
responses: []promptkit.GenerateResponse{
{Content: `{"bad":true}`},
{Content: `{"ok":true}`},
{Content: `{"ok":true}`},
},
block: make(chan struct{}),
}
scheduler, err := NewScheduler(1)
if err != nil {
t.Fatalf("NewScheduler() error = %v", err)
}
scheduled := NewScheduledClient(newTestPromptKitClient(t, fake), scheduler)
firstDone := make(chan error, 1)
go func() {
var out map[string]any
_, callErr := scheduled.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
PromptID: "adapter.test",
SessionID: "repair-session",
StructuredOutputRepairAttempts: &attempts,
Inputs: contracts.LLMInputSet{
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
},
}, &out)
firstDone <- callErr
}()
waitForAtomicAtLeast(t, &fake.calls, 1)
secondDone := make(chan error, 1)
go func() {
var out map[string]any
_, callErr := scheduled.CompleteStructured(context.Background(), contracts.StructuredCompletionRequest{
PromptID: "adapter.test",
SessionID: "queued-session",
Inputs: contracts.LLMInputSet{
"transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""),
},
}, &out)
secondDone <- callErr
}()
close(fake.block)
if err := <-firstDone; err != nil {
t.Fatalf("repaired completion error = %v", err)
}
if err := <-secondDone; err != nil {
t.Fatalf("queued completion error = %v", err)
}
requests := fake.requestsSnapshot()
if len(requests) != 3 || requests[1].Prompt.SessionID != "repair-session" || requests[2].Prompt.SessionID != "queued-session" {
t.Fatalf("generation order = %#v, want repair before queued completion", requests)
}
if got := atomic.LoadInt32(&fake.maxInFlight); got > 1 {
t.Fatalf("max in-flight calls = %d, want one scheduled logical completion", got)
}
}
func TestPromptKitClientValidatesRequest(t *testing.T) {
client := newTestPromptKitClient(t, &fakePromptKitLLM{content: `{"ok":true}`})
var out map[string]any
@@ -1375,11 +1481,13 @@ func (f *switchingPromptFS) resumePromptRead() {
type fakePromptKitLLM struct {
content string
responses []promptkit.GenerateResponse
allowEmpty bool
err error
block chan struct{}
mu sync.Mutex
last promptkit.GenerateRequest
requests []promptkit.GenerateRequest
calls int32
inFlight int32
maxInFlight int32
@@ -1394,8 +1502,9 @@ func (*credentialBearingProviderError) Error() string {
func (f *fakePromptKitLLM) Generate(ctx context.Context, req promptkit.GenerateRequest) (*promptkit.GenerateResponse, error) {
f.mu.Lock()
f.last = req
f.requests = append(f.requests, req)
f.mu.Unlock()
atomic.AddInt32(&f.calls, 1)
call := atomic.AddInt32(&f.calls, 1)
current := atomic.AddInt32(&f.inFlight, 1)
for {
seen := atomic.LoadInt32(&f.maxInFlight)
@@ -1414,6 +1523,14 @@ func (f *fakePromptKitLLM) Generate(ctx context.Context, req promptkit.GenerateR
if f.err != nil {
return nil, f.err
}
if len(f.responses) > 0 {
index := int(call - 1)
if index >= len(f.responses) {
return nil, fmt.Errorf("unexpected provider call %d", call)
}
response := f.responses[index]
return &response, nil
}
content := f.content
if content == "" && !f.allowEmpty {
content = `{"ok":true}`
@@ -1444,3 +1561,9 @@ func (f *fakePromptKitLLM) lastRequest() promptkit.GenerateRequest {
defer f.mu.Unlock()
return f.last
}
func (f *fakePromptKitLLM) requestsSnapshot() []promptkit.GenerateRequest {
f.mu.Lock()
defer f.mu.Unlock()
return append([]promptkit.GenerateRequest(nil), f.requests...)
}

4
internal/framework/semanticreconcile/assets_test.go
View File

@@ -45,8 +45,8 @@ func TestRegisterAssetsPreparesGenericPromptOffline(t *testing.T) {
if prepared.SelectedProfileID != "semantic-reconciliation-test" {
t.Fatalf("selected profile = %q, want explicit test profile", prepared.SelectedProfileID)
}
if contract := prepared.OutputContract; contract.SchemaPath != "semantic_reconciliation_llm.v1.json" || contract.RepairAttempts != 0 {
t.Fatalf("output contract = %#v, want generic schema without repair", contract)
if contract := prepared.OutputContract; contract.SchemaPath != "semantic_reconciliation_llm.v1.json" || contract.RepairAttempts != 1 {
t.Fatalf("output contract = %#v, want generic schema with one repair", contract)
}
if len(prepared.Messages) != 5 || prepared.Messages[0].Role != "system" {
t.Fatalf("prepared messages = %#v, want five ordered messages beginning with system", prepared.Messages)