Add feedback-aware correction contracts

This commit is contained in:
2026-08-26 23:24:20 +00:00
parent 9abd93502f
commit 1c3da3e869
33 changed files with 703 additions and 37 deletions

View File

@@ -9,14 +9,15 @@ import (
)
type StructuredCompletionRequest struct {
StageName string `json:"stage_name"`
PromptID string `json:"prompt_id,omitempty"`
PromptVersion string `json:"prompt_version,omitempty"`
ProfileID string `json:"profile_id,omitempty"`
SessionID string `json:"session_id,omitempty"`
Inputs LLMInputSet `json:"inputs,omitempty"`
Vars map[string]any `json:"vars,omitempty"`
StructuredOutputRepairAttempts *int `json:"structured_output_repair_attempts,omitempty"`
StageName string `json:"stage_name"`
PromptID string `json:"prompt_id,omitempty"`
PromptVersion string `json:"prompt_version,omitempty"`
ProfileID string `json:"profile_id,omitempty"`
SessionID string `json:"session_id,omitempty"`
Inputs LLMInputSet `json:"inputs,omitempty"`
Vars map[string]any `json:"vars,omitempty"`
StructuredOutputRepairAttempts *int `json:"structured_output_repair_attempts,omitempty"`
Correction *SemanticCorrection `json:"-"`
}
type StructuredCompletionResponse struct {
@@ -156,12 +157,14 @@ type ChunkRequest struct {
References ReferenceSet `json:"references,omitempty"`
LLMProfile string `json:"llm_profile,omitempty"`
StructuredOutputRepairAttempts *int `json:"structured_output_repair_attempts,omitempty"`
Correction *SemanticCorrection `json:"-"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type ChunkPlanResult struct {
Plan source.ChunkPlan `json:"plan"`
Warnings []Warning `json:"warnings,omitempty"`
Plan source.ChunkPlan `json:"plan"`
Warnings []Warning `json:"warnings,omitempty"`
ModelCandidate *ModelCandidate `json:"-"`
}
type Chunker interface {
@@ -283,6 +286,7 @@ type ValidationResult struct {
Approved bool `json:"approved"`
ReasonCode string `json:"reason_code,omitempty"`
Message string `json:"message,omitempty"`
CorrectionGuidance string `json:"-"`
DiagnosticArtifactPath string `json:"diagnostic_artifact_path,omitempty"`
Warnings []Warning `json:"warnings,omitempty"`
}

View File

@@ -0,0 +1,169 @@
package contracts
import (
"errors"
"fmt"
"strings"
"unicode/utf8"
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
)
// CorrectionProtocol identifies how a producer can represent the model output
// that directly controlled a candidate.
type CorrectionProtocol string
const (
CorrectionProtocolSingleResponseV1 CorrectionProtocol = "single_response_v1"
)
const (
MaxValidationReasonCodeBytes = 128
MaxValidationCorrectionGuidanceBytes = 4 * 1024
MaxAssistantResponseBytes = 1 << 20
MaxCorrectionGuidanceBytes = 64 * 1024
MaxCorrectionContentBytes = MaxAssistantResponseBytes + MaxCorrectionGuidanceBytes
)
func (protocol CorrectionProtocol) Validate() error {
if protocol == CorrectionProtocolSingleResponseV1 {
return nil
}
return fmt.Errorf("unsupported correction protocol %q", protocol)
}
// SemanticCorrection carries the latest model response and application-owned
// guidance for one fresh corrected request. Its content is sensitive and is
// deliberately excluded from ordinary JSON serialization.
type SemanticCorrection struct {
AssistantResponse []byte `json:"-"`
UserGuidance string `json:"-"`
}
func NewSemanticCorrection(assistantResponse []byte, userGuidance string) (*SemanticCorrection, error) {
correction := &SemanticCorrection{
AssistantResponse: append([]byte(nil), assistantResponse...),
UserGuidance: userGuidance,
}
if err := correction.Validate(); err != nil {
return nil, err
}
return correction, nil
}
func CloneSemanticCorrection(correction *SemanticCorrection) (*SemanticCorrection, error) {
if correction == nil {
return nil, nil
}
return NewSemanticCorrection(correction.AssistantResponse, correction.UserGuidance)
}
// CloneStructuredCompletionRequest returns a request whose mutable values are
// owned by the caller. It is suitable for clients that retain requests after
// CompleteStructured returns.
func CloneStructuredCompletionRequest(request StructuredCompletionRequest) (StructuredCompletionRequest, error) {
correction, err := CloneSemanticCorrection(request.Correction)
if err != nil {
return StructuredCompletionRequest{}, fmt.Errorf("clone correction: %w", err)
}
vars, err := source.CloneMetadata(request.Vars)
if err != nil {
return StructuredCompletionRequest{}, fmt.Errorf("clone variables: %w", err)
}
request.Inputs = request.Inputs.Clone()
request.Vars = vars
request.Correction = correction
if request.StructuredOutputRepairAttempts != nil {
attempts := *request.StructuredOutputRepairAttempts
request.StructuredOutputRepairAttempts = &attempts
}
return request, nil
}
func (correction SemanticCorrection) Validate() error {
if err := validateAssistantResponse(correction.AssistantResponse); err != nil {
return fmt.Errorf("semantic correction assistant response: %w", err)
}
if err := validateBoundedText(correction.UserGuidance, MaxCorrectionGuidanceBytes, "semantic correction user guidance", false); err != nil {
return err
}
if len(correction.AssistantResponse)+len(correction.UserGuidance) > MaxCorrectionContentBytes {
return errors.New("semantic correction content exceeds maximum length")
}
return nil
}
// ModelCandidate preserves the exact single model response that directly
// controlled a producer result. It is attempt-local and never durable data.
type ModelCandidate struct {
Response []byte `json:"-"`
Protocol CorrectionProtocol `json:"-"`
}
func NewModelCandidate(response []byte, protocol CorrectionProtocol) (*ModelCandidate, error) {
candidate := &ModelCandidate{Response: append([]byte(nil), response...), Protocol: protocol}
if err := candidate.Validate(); err != nil {
return nil, err
}
return candidate, nil
}
func CloneModelCandidate(candidate *ModelCandidate) (*ModelCandidate, error) {
if candidate == nil {
return nil, nil
}
return NewModelCandidate(candidate.Response, candidate.Protocol)
}
func (candidate ModelCandidate) Validate() error {
if err := candidate.Protocol.Validate(); err != nil {
return err
}
if err := validateAssistantResponse(candidate.Response); err != nil {
return fmt.Errorf("model candidate response: %w", err)
}
return nil
}
func ValidateValidationResult(result ValidationResult) error {
if result.ReasonCode != "" {
if err := validateBoundedText(result.ReasonCode, MaxValidationReasonCodeBytes, "validation reason code", false); err != nil {
return err
}
}
if result.CorrectionGuidance != "" {
if err := validateBoundedText(result.CorrectionGuidance, MaxValidationCorrectionGuidanceBytes, "validation correction guidance", false); err != nil {
return err
}
}
return nil
}
func validateAssistantResponse(response []byte) error {
if len(response) > MaxAssistantResponseBytes {
return errors.New("exceeds maximum length")
}
if !utf8.Valid(response) {
return errors.New("must be valid UTF-8")
}
if len(strings.TrimSpace(string(response))) == 0 {
return errors.New("must not be blank")
}
return nil
}
func validateBoundedText(value string, maximum int, name string, optional bool) error {
if value == "" && optional {
return nil
}
if !utf8.ValidString(value) {
return fmt.Errorf("%s must be valid UTF-8", name)
}
if strings.TrimSpace(value) == "" {
return fmt.Errorf("%s must not be blank", name)
}
if len(value) > maximum {
return fmt.Errorf("%s exceeds maximum length", name)
}
return nil
}

View File

@@ -0,0 +1,141 @@
package contracts
import (
"bytes"
"encoding/json"
"strings"
"testing"
)
func TestSemanticCorrectionOwnsValidatedContent(t *testing.T) {
assistant := []byte(`{"items":["original"]}`)
correction, err := NewSemanticCorrection(assistant, "Return one corrected replacement.")
if err != nil {
t.Fatalf("NewSemanticCorrection() error = %v", err)
}
assistant[0] = '['
if got := string(correction.AssistantResponse); got != `{"items":["original"]}` {
t.Fatalf("assistant response = %q, want owned original content", got)
}
clone, err := CloneSemanticCorrection(correction)
if err != nil {
t.Fatalf("CloneSemanticCorrection() error = %v", err)
}
clone.AssistantResponse[0] = '['
if got := string(correction.AssistantResponse); got != `{"items":["original"]}` {
t.Fatalf("source correction changed through clone = %q", got)
}
if nilClone, err := CloneSemanticCorrection(nil); err != nil || nilClone != nil {
t.Fatalf("CloneSemanticCorrection(nil) = %#v, %v; want nil, nil", nilClone, err)
}
encoded, err := json.Marshal(correction)
if err != nil {
t.Fatalf("marshal correction: %v", err)
}
if string(encoded) != "{}" {
t.Fatalf("correction JSON = %s, want no sensitive content", encoded)
}
}
func TestCorrectionContractsRejectInvalidContent(t *testing.T) {
tooLongAssistant := bytes.Repeat([]byte("a"), MaxAssistantResponseBytes+1)
tooLongGuidance := strings.Repeat("a", MaxCorrectionGuidanceBytes+1)
tooLongReason := strings.Repeat("a", MaxValidationReasonCodeBytes+1)
tooLongValidationGuidance := strings.Repeat("a", MaxValidationCorrectionGuidanceBytes+1)
for _, test := range []struct {
name string
call func() error
}{
{"blank assistant", func() error { _, err := NewSemanticCorrection([]byte(" \n"), "guidance"); return err }},
{"invalid assistant utf8", func() error { _, err := NewSemanticCorrection([]byte{0xff}, "guidance"); return err }},
{"oversized assistant", func() error { _, err := NewSemanticCorrection(tooLongAssistant, "guidance"); return err }},
{"blank guidance", func() error { _, err := NewSemanticCorrection([]byte("response"), " \t"); return err }},
{"invalid guidance utf8", func() error { _, err := NewSemanticCorrection([]byte("response"), string([]byte{0xff})); return err }},
{"oversized guidance", func() error { _, err := NewSemanticCorrection([]byte("response"), tooLongGuidance); return err }},
{"unsupported protocol", func() error { _, err := NewModelCandidate([]byte("response"), "multiple_responses"); return err }},
{"missing candidate protocol", func() error { _, err := NewModelCandidate([]byte("response"), ""); return err }},
{"blank candidate response", func() error { _, err := NewModelCandidate([]byte(" "), CorrectionProtocolSingleResponseV1); return err }},
{"oversized reason code", func() error { return ValidateValidationResult(ValidationResult{ReasonCode: tooLongReason}) }},
{"blank reason code", func() error { return ValidateValidationResult(ValidationResult{ReasonCode: " \t"}) }},
{"invalid correction guidance utf8", func() error {
return ValidateValidationResult(ValidationResult{CorrectionGuidance: string([]byte{0xff})})
}},
{"oversized correction guidance", func() error {
return ValidateValidationResult(ValidationResult{CorrectionGuidance: tooLongValidationGuidance})
}},
} {
t.Run(test.name, func(t *testing.T) {
if err := test.call(); err == nil {
t.Fatal("validation error = nil, want error")
}
})
}
}
func TestModelCandidateOwnsValidatedResponse(t *testing.T) {
response := []byte(`{"items":["original"]}`)
candidate, err := NewModelCandidate(response, CorrectionProtocolSingleResponseV1)
if err != nil {
t.Fatalf("NewModelCandidate() error = %v", err)
}
response[0] = '['
if got := string(candidate.Response); got != `{"items":["original"]}` {
t.Fatalf("candidate response = %q, want owned original content", got)
}
clone, err := CloneModelCandidate(candidate)
if err != nil {
t.Fatalf("CloneModelCandidate() error = %v", err)
}
clone.Response[0] = '['
if got := string(candidate.Response); got != `{"items":["original"]}` {
t.Fatalf("source candidate changed through clone = %q", got)
}
if nilClone, err := CloneModelCandidate(nil); err != nil || nilClone != nil {
t.Fatalf("CloneModelCandidate(nil) = %#v, %v; want nil, nil", nilClone, err)
}
}
func TestValidationResultAllowsAbsentOptionalCorrectionFields(t *testing.T) {
if err := ValidateValidationResult(ValidationResult{Approved: true}); err != nil {
t.Fatalf("ValidateValidationResult() error = %v, want nil", err)
}
if err := ValidateValidationResult(ValidationResult{ReasonCode: "invalid-evidence", CorrectionGuidance: "Provide source-backed evidence."}); err != nil {
t.Fatalf("ValidateValidationResult() error = %v, want nil", err)
}
if err := CorrectionProtocol("").Validate(); err == nil {
t.Fatal("empty correction protocol validation error = nil, want error")
}
}
func TestCloneStructuredCompletionRequestOwnsCorrection(t *testing.T) {
correction, err := NewSemanticCorrection([]byte(`{"value":"original"}`), "Correct the value.")
if err != nil {
t.Fatalf("NewSemanticCorrection() error = %v", err)
}
attempts := 2
request := StructuredCompletionRequest{
Inputs: LLMInputSet{"source": NewLLMInputMaterial("source", "application/json", []byte(`{"source":true}`), "", "")},
Vars: map[string]any{"labels": []string{"original"}},
StructuredOutputRepairAttempts: &attempts,
Correction: correction,
}
clone, err := CloneStructuredCompletionRequest(request)
if err != nil {
t.Fatalf("CloneStructuredCompletionRequest() error = %v", err)
}
correction.AssistantResponse[0] = '['
request.Inputs["source"] = NewLLMInputMaterial("source", "application/json", []byte(`{"source":false}`), "", "")
*request.StructuredOutputRepairAttempts = 7
if got := string(clone.Correction.AssistantResponse); got != `{"value":"original"}` {
t.Fatalf("cloned correction response = %q, want owned original content", got)
}
if got := string(clone.Inputs["source"].Content); got != `{"source":true}` {
t.Fatalf("cloned input = %q, want owned original content", got)
}
if clone.StructuredOutputRepairAttempts == nil || *clone.StructuredOutputRepairAttempts != 2 {
t.Fatalf("cloned repair attempts = %v, want 2", clone.StructuredOutputRepairAttempts)
}
}

View File

@@ -42,12 +42,14 @@ type TypedExtractionRequest struct {
References ReferenceSet
LLMProfile string
StructuredOutputRepairAttempts *int
Correction *SemanticCorrection
Metadata map[string]any
}
type TypedExtractionResult[T any] struct {
Value T
Warnings []Warning
Value T
Warnings []Warning
ModelCandidate *ModelCandidate
}
type Extractor[T any] interface {
@@ -65,12 +67,14 @@ type TypedMergeRequest[T any] struct {
References ReferenceSet
LLMProfile string
StructuredOutputRepairAttempts *int
Correction *SemanticCorrection
Metadata map[string]any
}
type TypedMergeResult[T any] struct {
Value T
Warnings []Warning
Value T
Warnings []Warning
ModelCandidate *ModelCandidate
}
type Merger[T any] interface {
@@ -87,13 +91,15 @@ type TypedNormalizeRequest[T any] struct {
References ReferenceSet
LLMProfile string
StructuredOutputRepairAttempts *int
Correction *SemanticCorrection
Metadata map[string]any
}
type TypedNormalizeResult[T any] struct {
Value T
Warnings []Warning
Retry *NormalizeRetry
Value T
Warnings []Warning
Retry *NormalizeRetry
ModelCandidate *ModelCandidate
}
// Normalize retry diagnostic limits bound module-provided values before the

View File

@@ -691,6 +691,7 @@ func debugResponseModel(response contracts.StructuredCompletionResponse) string
func debugValidationResultEnvelope(result contracts.ValidationResult) contracts.ValidationResult {
result.Message = string(redactSecretBytes([]byte(result.Message)))
result.CorrectionGuidance = ""
result.DiagnosticArtifactPath = string(redactSecretBytes([]byte(result.DiagnosticArtifactPath)))
for i := range result.Warnings {
result.Warnings[i].Message = string(redactSecretBytes([]byte(result.Warnings[i].Message)))

View File

@@ -58,11 +58,20 @@ func RegisterExtractorBuilder[T any](registry *ExtractorRegistry, spec ModuleSpe
if !ok {
return erasedTypedResult{}, fmt.Errorf("extractor %q has incompatible implementation %T", normalized.Key, implementation)
}
correction, err := contracts.CloneSemanticCorrection(request.Correction)
if err != nil {
return erasedTypedResult{}, fmt.Errorf("clone extraction correction: %w", err)
}
request.Correction = correction
result, err := extractor.Extract(ctx, request)
if err != nil {
return erasedTypedResult{}, err
}
return erasedTypedResult{Value: result.Value, Warnings: result.Warnings}, nil
candidate, err := contracts.CloneModelCandidate(result.ModelCandidate)
if err != nil {
return erasedTypedResult{}, fmt.Errorf("clone extraction model candidate: %w", err)
}
return erasedTypedResult{Value: result.Value, Warnings: cloneWarnings(result.Warnings), ModelCandidate: candidate}, nil
}}
if registry.typedEntries == nil {
registry.typedEntries = map[string]typedExtractorEntry{}

View File

@@ -85,11 +85,19 @@ func RegisterMergerBuilder[T any](registry *MergerRegistry, spec ModuleSpec, val
}
outputs[i] = contracts.ExtractArtifact[T]{LaneID: output.LaneID, ExtractorKey: output.ExtractorKey, SourceID: output.SourceID, ChunkID: output.ChunkID, ChunkIndex: output.ChunkIndex, ChunkRef: output.ChunkRef, Value: value}
}
result, err := merger.Merge(ctx, contracts.TypedMergeRequest[T]{Source: request.Source, LaneID: request.LaneID, ExtractOutputs: outputs, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, StructuredOutputRepairAttempts: request.StructuredOutputRepairAttempts, Metadata: request.Metadata})
correction, err := contracts.CloneSemanticCorrection(request.Correction)
if err != nil {
return erasedTypedResult{}, fmt.Errorf("clone merge correction: %w", err)
}
result, err := merger.Merge(ctx, contracts.TypedMergeRequest[T]{Source: request.Source, LaneID: request.LaneID, ExtractOutputs: outputs, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, StructuredOutputRepairAttempts: request.StructuredOutputRepairAttempts, Correction: correction, Metadata: request.Metadata})
if err != nil {
return erasedTypedResult{}, err
}
return erasedTypedResult{Value: result.Value, Warnings: result.Warnings}, nil
candidate, err := contracts.CloneModelCandidate(result.ModelCandidate)
if err != nil {
return erasedTypedResult{}, fmt.Errorf("clone merge model candidate: %w", err)
}
return erasedTypedResult{Value: result.Value, Warnings: cloneWarnings(result.Warnings), ModelCandidate: candidate}, nil
},
}
return nil

View File

@@ -76,11 +76,19 @@ func RegisterNormalizerBuilder[T any](registry *NormalizerRegistry, spec ModuleS
if err != nil {
return erasedTypedResult{}, err
}
result, err := normalizer.Normalize(ctx, contracts.TypedNormalizeRequest[T]{Source: request.Source, LaneID: request.LaneID, MergeOutput: contracts.MergeArtifact[T]{LaneID: request.MergeOutput.LaneID, MergerKey: request.MergeOutput.MergerKey, SourceID: request.MergeOutput.SourceID, Value: value}, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, StructuredOutputRepairAttempts: request.StructuredOutputRepairAttempts, Metadata: request.Metadata})
correction, err := contracts.CloneSemanticCorrection(request.Correction)
if err != nil {
return erasedTypedResult{}, fmt.Errorf("clone normalize correction: %w", err)
}
result, err := normalizer.Normalize(ctx, contracts.TypedNormalizeRequest[T]{Source: request.Source, LaneID: request.LaneID, MergeOutput: contracts.MergeArtifact[T]{LaneID: request.MergeOutput.LaneID, MergerKey: request.MergeOutput.MergerKey, SourceID: request.MergeOutput.SourceID, Value: value}, SourceInput: request.SourceInput, SessionID: request.SessionID, References: request.References, LLMProfile: request.LLMProfile, StructuredOutputRepairAttempts: request.StructuredOutputRepairAttempts, Correction: correction, Metadata: request.Metadata})
if err != nil {
return erasedTypedResult{}, err
}
return erasedTypedResult{Value: result.Value, Warnings: cloneWarnings(result.Warnings), Retry: cloneNormalizeRetry(result.Retry)}, nil
candidate, err := contracts.CloneModelCandidate(result.ModelCandidate)
if err != nil {
return erasedTypedResult{}, fmt.Errorf("clone normalize model candidate: %w", err)
}
return erasedTypedResult{Value: result.Value, Warnings: cloneWarnings(result.Warnings), Retry: cloneNormalizeRetry(result.Retry), ModelCandidate: candidate}, nil
},
}
return nil

View File

@@ -499,6 +499,9 @@ func (r *Runner) validateChunks(ctx context.Context, doc *source.SourceDocument,
default:
return nil, nil, fmt.Errorf("validator %q is incompatible with chunk validation", binding.Module)
}
if err == nil {
err = contracts.ValidateValidationResult(result)
}
debugContent := debugContentEnvelope(content, "application/json", nil, nil)
debugContent.ContentDigest = debugContentDigest(content)
debugCall := debugValidationCall{ValidatorName: binding.Module, Request: map[string]any{"stage": string(StageChunk), "module_key": moduleKey, "source_id": doc.ID, "schema": schema, "schema_digest": contracts.DigestArtifactSchema(schema), "content": debugContent, "metadata": redactSensitiveMap(metadata)}, Result: debugValidationResultEnvelope(result)}

View File

@@ -523,6 +523,9 @@ func (r *Runner) validateTypedArtifact(ctx context.Context, codec artifactCodecE
default:
return nil, nil, fmt.Errorf("validator %q is incompatible with typed artifact validation", binding.Module)
}
if err == nil {
err = contracts.ValidateValidationResult(result)
}
artifact, _ := validationCandidateArtifact(codec, target)
debugCall := debugValidationCall{ValidatorName: binding.Module, Request: map[string]any{"stage": string(target.stage), "lane_id": target.laneID, "module_key": target.moduleKey, "source_id": target.sourceID, "artifact": debugCheckpointArtifact(artifact), "metadata": redactSensitiveMap(target.metadata)}, Result: debugValidationResultEnvelope(result)}
if err != nil {

View File

@@ -22,9 +22,10 @@ type erasedMergeArtifact struct {
}
type erasedTypedResult struct {
Value any
Warnings []contracts.Warning
Retry *contracts.NormalizeRetry
Value any
Warnings []contracts.Warning
Retry *contracts.NormalizeRetry
ModelCandidate *contracts.ModelCandidate
}
type typedValidationTarget struct {

View File

@@ -1,12 +1,70 @@
package pipeline
import (
"bytes"
"context"
"testing"
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)
type correctionObservingNotesExtractor struct {
correction *contracts.SemanticCorrection
candidate *contracts.ModelCandidate
}
func (*correctionObservingNotesExtractor) Key() string { return "test/correction-observing-extract" }
func (*correctionObservingNotesExtractor) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (e *correctionObservingNotesExtractor) Extract(_ context.Context, request contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[codecNotes], error) {
e.correction = request.Correction
candidate, err := contracts.NewModelCandidate([]byte(`{"items":["extract"]}`), contracts.CorrectionProtocolSingleResponseV1)
if err != nil {
return contracts.TypedExtractionResult[codecNotes]{}, err
}
e.candidate = candidate
return contracts.TypedExtractionResult[codecNotes]{Value: codecNotes{Items: []string{"extract"}}, ModelCandidate: candidate}, nil
}
type correctionObservingNotesMerger struct {
correction *contracts.SemanticCorrection
candidate *contracts.ModelCandidate
}
func (*correctionObservingNotesMerger) Key() string { return "test/correction-observing-merge" }
func (m *correctionObservingNotesMerger) Merge(_ context.Context, request contracts.TypedMergeRequest[codecNotes]) (contracts.TypedMergeResult[codecNotes], error) {
m.correction = request.Correction
candidate, err := contracts.NewModelCandidate([]byte(`{"items":["merge"]}`), contracts.CorrectionProtocolSingleResponseV1)
if err != nil {
return contracts.TypedMergeResult[codecNotes]{}, err
}
m.candidate = candidate
return contracts.TypedMergeResult[codecNotes]{Value: codecNotes{Items: []string{"merge"}}, ModelCandidate: candidate}, nil
}
type correctionObservingNotesNormalizer struct {
correction *contracts.SemanticCorrection
candidate *contracts.ModelCandidate
}
func (*correctionObservingNotesNormalizer) Key() string { return "test/correction-observing-normalize" }
func (*correctionObservingNotesNormalizer) ReferenceSlots() []contracts.ReferenceSlot {
return nil
}
func (n *correctionObservingNotesNormalizer) Normalize(_ context.Context, request contracts.TypedNormalizeRequest[codecNotes]) (contracts.TypedNormalizeResult[codecNotes], error) {
n.correction = request.Correction
candidate, err := contracts.NewModelCandidate([]byte(`{"items":["normalize"]}`), contracts.CorrectionProtocolSingleResponseV1)
if err != nil {
return contracts.TypedNormalizeResult[codecNotes]{}, err
}
n.candidate = candidate
return contracts.TypedNormalizeResult[codecNotes]{Value: codecNotes{Items: []string{"normalize"}}, ModelCandidate: candidate}, nil
}
type repairObservingNotesMerger struct {
attempts *int
}
@@ -93,3 +151,103 @@ func TestTypedRegistryErasurePreservesStructuredOutputRepairAttempts(t *testing.
}
})
}
func TestTypedRegistryErasurePreservesCorrectionAndCandidateOwnership(t *testing.T) {
const artifactKind contracts.ArtifactKind = "test/notes"
t.Run("extract", func(t *testing.T) {
correction := newTestCorrection(t)
implementation := &correctionObservingNotesExtractor{}
registry := NewExtractorRegistry()
if err := RegisterExtractor(registry, ModuleSpec{Key: implementation.Key(), Stage: StageExtract, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Extractor[codecNotes], error) {
return implementation, nil
}); err != nil {
t.Fatalf("RegisterExtractor() error = %v", err)
}
entry, ok := registry.typedEntry(implementation.Key())
if !ok {
t.Fatal("typed extractor entry missing")
}
built, err := entry.builder(BuildRequest{})
if err != nil {
t.Fatalf("builder() error = %v", err)
}
result, err := entry.extract(context.Background(), built, contracts.TypedExtractionRequest{Correction: correction})
if err != nil {
t.Fatalf("extract() error = %v", err)
}
assertCorrectionAndCandidateOwnership(t, correction, implementation.correction, implementation.candidate, result.ModelCandidate)
})
t.Run("merge", func(t *testing.T) {
correction := newTestCorrection(t)
implementation := &correctionObservingNotesMerger{}
registry := NewMergerRegistry()
if err := RegisterMerger(registry, ModuleSpec{Key: implementation.Key(), Stage: StageMerge, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Merger[codecNotes], error) {
return implementation, nil
}); err != nil {
t.Fatalf("RegisterMerger() error = %v", err)
}
entry, ok := registry.typedEntry(implementation.Key(), artifactKind)
if !ok {
t.Fatal("typed merger entry missing")
}
built, err := entry.builder(BuildRequest{})
if err != nil {
t.Fatalf("builder() error = %v", err)
}
result, err := entry.merge(context.Background(), built, contracts.TypedMergeRequest[any]{Correction: correction, ExtractOutputs: []contracts.ExtractArtifact[any]{{Value: codecNotes{Items: []string{"extract"}}}}})
if err != nil {
t.Fatalf("merge() error = %v", err)
}
assertCorrectionAndCandidateOwnership(t, correction, implementation.correction, implementation.candidate, result.ModelCandidate)
})
t.Run("normalize", func(t *testing.T) {
correction := newTestCorrection(t)
implementation := &correctionObservingNotesNormalizer{}
registry := NewNormalizerRegistry()
if err := RegisterNormalizer(registry, ModuleSpec{Key: implementation.Key(), Stage: StageNormalize, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: artifactKind}, func() (contracts.Normalizer[codecNotes], error) {
return implementation, nil
}); err != nil {
t.Fatalf("RegisterNormalizer() error = %v", err)
}
entry, ok := registry.typedEntry(implementation.Key(), artifactKind)
if !ok {
t.Fatal("typed normalizer entry missing")
}
built, err := entry.builder(BuildRequest{})
if err != nil {
t.Fatalf("builder() error = %v", err)
}
result, err := entry.normalize(context.Background(), built, contracts.TypedNormalizeRequest[any]{Correction: correction, MergeOutput: contracts.MergeArtifact[any]{Value: codecNotes{Items: []string{"merge"}}}})
if err != nil {
t.Fatalf("normalize() error = %v", err)
}
assertCorrectionAndCandidateOwnership(t, correction, implementation.correction, implementation.candidate, result.ModelCandidate)
})
}
func newTestCorrection(t *testing.T) *contracts.SemanticCorrection {
t.Helper()
correction, err := contracts.NewSemanticCorrection([]byte(`{"items":["original"]}`), "Correct the response.")
if err != nil {
t.Fatalf("NewSemanticCorrection() error = %v", err)
}
return correction
}
func assertCorrectionAndCandidateOwnership(t *testing.T, callerCorrection, observedCorrection *contracts.SemanticCorrection, producerCandidate, returnedCandidate *contracts.ModelCandidate) {
t.Helper()
if observedCorrection == nil || producerCandidate == nil || returnedCandidate == nil {
t.Fatal("correction and candidates must be present")
}
callerCorrection.AssistantResponse[0] = '['
if got := string(observedCorrection.AssistantResponse); got != `{"items":["original"]}` {
t.Fatalf("observed correction = %q, want owned original content", got)
}
producerCandidate.Response[0] = '['
if bytes.Equal(producerCandidate.Response, returnedCandidate.Response) {
t.Fatalf("returned candidate aliases producer candidate: %q", returnedCandidate.Response)
}
}

View File

@@ -229,8 +229,11 @@ type recordingReconciliationClient struct {
}
func (client *recordingReconciliationClient) CompleteStructured(_ context.Context, request contracts.StructuredCompletionRequest, output any) (contracts.StructuredCompletionResponse, error) {
request.Inputs = request.Inputs.Clone()
client.requests = append(client.requests, request)
snapshot, err := contracts.CloneStructuredCompletionRequest(request)
if err != nil {
return contracts.StructuredCompletionResponse{}, fmt.Errorf("clone recording request: %w", err)
}
client.requests = append(client.requests, snapshot)
index := len(client.requests) - 1
if index < len(client.errors) && client.errors[index] != nil {
return contracts.StructuredCompletionResponse{}, client.errors[index]