Add bounded structured-output repair

2026-05-04 21:35:31 -05:00
parent 328703df36
commit ea05945457
3 changed files with 431 additions and 11 deletions
--- a/internal/usecase/runner_test.go
+++ b/internal/usecase/runner_test.go
@@ -87,6 +87,29 @@ func (f *fakeValidator) Validate(ctx context.Context, artifact *domain.Artifact,
 	return f.result, nil
 }

+type fakeRepairer struct {
+	responses []*domain.GenerateResponse
+	err       error
+	calls     int
+	lastReq   RepairRequest
+}
+
+func (f *fakeRepairer) Repair(ctx context.Context, req RepairRequest) (*domain.GenerateResponse, error) {
+	f.calls++
+	f.lastReq = req
+	if f.err != nil {
+		return nil, f.err
+	}
+	if len(f.responses) == 0 {
+		return nil, errors.New("no repair response configured")
+	}
+	idx := f.calls - 1
+	if idx >= len(f.responses) {
+		idx = len(f.responses) - 1
+	}
+	return f.responses[idx], nil
+}
+
 func TestRunnerRunSuccessful(t *testing.T) {
 	repo := &fakeProfileRepo{
 		profile: &domain.PromptProfile{
@@ -377,6 +400,267 @@ func TestRunnerRunValidationFailureWithRealValidatorPreservesRawOutput(t *testin
 	}
 }

+func TestRunnerRunNoRepairWhenDisabled(t *testing.T) {
+	repairer := &fakeRepairer{
+		responses: []*domain.GenerateResponse{{Content: `{"ok":true}`}},
+	}
+
+	runner := NewRunnerWithRepairer(
+		&fakeProfileRepo{profile: &domain.PromptProfile{
+			ID:           "p-json",
+			Version:      "1",
+			OutputFormat: domain.FormatJSON,
+			ModelDefaults: domain.ModelTarget{
+				Endpoint: "ep",
+				Model:    "m",
+			},
+			Validation: domain.OutputContract{
+				ValidationMode: domain.ValidationJSON,
+				Format:         domain.FormatJSON,
+				RepairAttempts: 0,
+			},
+		}},
+		&fakeArtifactReader{artifactsByURI: map[string]*domain.Artifact{"a://ok": {Body: []byte("x"), Hash: hashString("x")}}},
+		&fakeRenderer{rendered: &domain.RenderedPrompt{}},
+		&fakeLLM{resp: &domain.GenerateResponse{Content: `{"broken":`}},
+		validate.NewStandardValidator(t.TempDir()),
+		repairer,
+	)
+
+	res, err := runner.Run(context.Background(), domain.RunRequest{
+		ProfileID: "p-json",
+		Inputs: map[string]domain.ArtifactRef{
+			"transcript": {Type: domain.ArtifactRefFile, URI: "a://ok"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if repairer.calls != 0 {
+		t.Fatalf("expected no repair calls, got %d", repairer.calls)
+	}
+	if res.Validation.Status != domain.ValidationFailed {
+		t.Fatalf("expected failed validation, got %q", res.Validation.Status)
+	}
+	if res.RawOutput != `{"broken":` {
+		t.Fatalf("expected original output preserved, got %q", res.RawOutput)
+	}
+	if res.Validation.RepairAttempts != 0 {
+		t.Fatalf("expected repair attempts 0, got %d", res.Validation.RepairAttempts)
+	}
+}
+
+func TestRunnerRunSuccessfulRepairAfterInvalidJSON(t *testing.T) {
+	repairer := &fakeRepairer{
+		responses: []*domain.GenerateResponse{{Content: `{"ok":true}`, Usage: domain.TokenUsage{TotalTokens: 5}}},
+	}
+
+	runner := NewRunnerWithRepairer(
+		&fakeProfileRepo{profile: &domain.PromptProfile{
+			ID:           "p-json",
+			Version:      "1",
+			OutputFormat: domain.FormatJSON,
+			ModelDefaults: domain.ModelTarget{
+				Endpoint: "ep",
+				Model:    "m",
+			},
+			Validation: domain.OutputContract{
+				ValidationMode: domain.ValidationJSON,
+				Format:         domain.FormatJSON,
+				RepairAttempts: 1,
+			},
+		}},
+		&fakeArtifactReader{artifactsByURI: map[string]*domain.Artifact{"a://ok": {Body: []byte("x"), Hash: hashString("x")}}},
+		&fakeRenderer{rendered: &domain.RenderedPrompt{}},
+		&fakeLLM{resp: &domain.GenerateResponse{Content: `{"broken":`, Usage: domain.TokenUsage{TotalTokens: 3}}},
+		validate.NewStandardValidator(t.TempDir()),
+		repairer,
+	)
+
+	res, err := runner.Run(context.Background(), domain.RunRequest{
+		ProfileID: "p-json",
+		Inputs: map[string]domain.ArtifactRef{
+			"transcript": {Type: domain.ArtifactRefFile, URI: "a://ok"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if repairer.calls != 1 {
+		t.Fatalf("expected one repair call, got %d", repairer.calls)
+	}
+	if res.Validation.Status != domain.ValidationPassed {
+		t.Fatalf("expected passed validation, got %q", res.Validation.Status)
+	}
+	if res.Validation.RepairAttempts != 1 {
+		t.Fatalf("expected repair attempts 1, got %d", res.Validation.RepairAttempts)
+	}
+	if res.RawOutput != `{"ok":true}` {
+		t.Fatalf("expected repaired output, got %q", res.RawOutput)
+	}
+}
+
+func TestRunnerRunSuccessfulRepairAfterSchemaFailure(t *testing.T) {
+	tmp := t.TempDir()
+	if err := os.WriteFile(filepath.Join(tmp, "schema.json"), []byte(`{
+  "$schema": "https://json-schema.org/draft/2020-12/schema",
+  "type": "object",
+  "required": ["name"],
+  "properties": {
+    "name": {"type": "string"}
+  }
+}`), 0644); err != nil {
+		t.Fatal(err)
+	}
+
+	repairer := &fakeRepairer{
+		responses: []*domain.GenerateResponse{{Content: `{"name":"eris"}`}},
+	}
+
+	runner := NewRunnerWithRepairer(
+		&fakeProfileRepo{profile: &domain.PromptProfile{
+			ID:           "p-json",
+			Version:      "1",
+			OutputFormat: domain.FormatJSON,
+			ModelDefaults: domain.ModelTarget{
+				Endpoint: "ep",
+				Model:    "m",
+			},
+			Validation: domain.OutputContract{
+				ValidationMode: domain.ValidationJSONSchema,
+				SchemaPath:     "schema.json",
+				Format:         domain.FormatJSON,
+				RepairAttempts: 1,
+			},
+		}},
+		&fakeArtifactReader{artifactsByURI: map[string]*domain.Artifact{"a://ok": {Body: []byte("x"), Hash: hashString("x")}}},
+		&fakeRenderer{rendered: &domain.RenderedPrompt{}},
+		&fakeLLM{resp: &domain.GenerateResponse{Content: `{"count":1}`}},
+		validate.NewStandardValidator(tmp),
+		repairer,
+	)
+
+	res, err := runner.Run(context.Background(), domain.RunRequest{
+		ProfileID: "p-json",
+		Inputs: map[string]domain.ArtifactRef{
+			"transcript": {Type: domain.ArtifactRefFile, URI: "a://ok"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if res.Validation.Status != domain.ValidationPassed {
+		t.Fatalf("expected passed validation, got %q", res.Validation.Status)
+	}
+	if res.Validation.RepairAttempts != 1 {
+		t.Fatalf("expected repair attempts 1, got %d", res.Validation.RepairAttempts)
+	}
+	if res.RawOutput != `{"name":"eris"}` {
+		t.Fatalf("expected repaired output, got %q", res.RawOutput)
+	}
+}
+
+func TestRunnerRunFailedRepairPreservesRawOutputAndErrors(t *testing.T) {
+	repairer := &fakeRepairer{
+		responses: []*domain.GenerateResponse{
+			{Content: `{"repair1":`},
+			{Content: `{"repair2":`},
+		},
+	}
+
+	runner := NewRunnerWithRepairer(
+		&fakeProfileRepo{profile: &domain.PromptProfile{
+			ID:           "p-json",
+			Version:      "1",
+			OutputFormat: domain.FormatJSON,
+			ModelDefaults: domain.ModelTarget{
+				Endpoint: "ep",
+				Model:    "m",
+			},
+			Validation: domain.OutputContract{
+				ValidationMode: domain.ValidationJSON,
+				Format:         domain.FormatJSON,
+				RepairAttempts: 2,
+			},
+		}},
+		&fakeArtifactReader{artifactsByURI: map[string]*domain.Artifact{"a://ok": {Body: []byte("x"), Hash: hashString("x")}}},
+		&fakeRenderer{rendered: &domain.RenderedPrompt{}},
+		&fakeLLM{resp: &domain.GenerateResponse{Content: `{"initial":`}},
+		validate.NewStandardValidator(t.TempDir()),
+		repairer,
+	)
+
+	res, err := runner.Run(context.Background(), domain.RunRequest{
+		ProfileID: "p-json",
+		Inputs: map[string]domain.ArtifactRef{
+			"transcript": {Type: domain.ArtifactRefFile, URI: "a://ok"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if res.Validation.Status != domain.ValidationFailed {
+		t.Fatalf("expected failed validation, got %q", res.Validation.Status)
+	}
+	if len(res.Validation.Errors) == 0 {
+		t.Fatal("expected validation errors after failed repair")
+	}
+	if res.Validation.RepairAttempts != 2 {
+		t.Fatalf("expected repair attempts 2, got %d", res.Validation.RepairAttempts)
+	}
+	if res.RawOutput != `{"repair2":` {
+		t.Fatalf("expected final repaired output preserved, got %q", res.RawOutput)
+	}
+}
+
+func TestRunnerRunRepairAttemptsBounded(t *testing.T) {
+	repairer := &fakeRepairer{
+		responses: []*domain.GenerateResponse{
+			{Content: `{"repair1":`},
+			{Content: `{"repair2":`},
+			{Content: `{"repair3":`},
+		},
+	}
+
+	runner := NewRunnerWithRepairer(
+		&fakeProfileRepo{profile: &domain.PromptProfile{
+			ID:           "p-json",
+			Version:      "1",
+			OutputFormat: domain.FormatJSON,
+			ModelDefaults: domain.ModelTarget{
+				Endpoint: "ep",
+				Model:    "m",
+			},
+			Validation: domain.OutputContract{
+				ValidationMode: domain.ValidationJSON,
+				Format:         domain.FormatJSON,
+				RepairAttempts: 1,
+			},
+		}},
+		&fakeArtifactReader{artifactsByURI: map[string]*domain.Artifact{"a://ok": {Body: []byte("x"), Hash: hashString("x")}}},
+		&fakeRenderer{rendered: &domain.RenderedPrompt{}},
+		&fakeLLM{resp: &domain.GenerateResponse{Content: `{"initial":`}},
+		validate.NewStandardValidator(t.TempDir()),
+		repairer,
+	)
+
+	res, err := runner.Run(context.Background(), domain.RunRequest{
+		ProfileID: "p-json",
+		Inputs: map[string]domain.ArtifactRef{
+			"transcript": {Type: domain.ArtifactRefFile, URI: "a://ok"},
+		},
+	})
+	if err != nil {
+		t.Fatalf("expected no error, got %v", err)
+	}
+	if repairer.calls != 1 {
+		t.Fatalf("expected repair calls bounded to 1, got %d", repairer.calls)
+	}
+	if res.Validation.RepairAttempts != 1 {
+		t.Fatalf("expected repair attempts 1, got %d", res.Validation.RepairAttempts)
+	}
+}
+
 func minimalProfile() *domain.PromptProfile {
 	return &domain.PromptProfile{
 		ID:           "p",