diff --git a/internal/framework/llm/promptkit_client.go b/internal/framework/llm/promptkit_client.go index e035d60f..d8a337f5 100644 --- a/internal/framework/llm/promptkit_client.go +++ b/internal/framework/llm/promptkit_client.go @@ -144,6 +144,9 @@ func (c *PromptKitClient) CompleteStructured(ctx context.Context, req contracts. if req.StructuredOutputRepairAttempts != nil { inspection, err := c.engine.InspectPrompt(ctx, promptID, strings.TrimSpace(req.PromptVersion)) if err != nil { + if ctxErr := ctx.Err(); ctxErr != nil { + return contracts.StructuredCompletionResponse{}, ctxErr + } return contracts.StructuredCompletionResponse{}, fmt.Errorf("inspect PromptKit prompt %q: %v", promptID, redactPromptKitError(err)) } contract := inspection.OutputContract diff --git a/internal/framework/llm/promptkit_client_test.go b/internal/framework/llm/promptkit_client_test.go index 99092e19..ce44f557 100644 --- a/internal/framework/llm/promptkit_client_test.go +++ b/internal/framework/llm/promptkit_client_test.go @@ -1147,19 +1147,31 @@ func TestPromptKitClientTranslatesBackendCapacityExhaustion(t *testing.T) { } func TestPromptKitClientContextCancellationIsRespected(t *testing.T) { - ctx, cancel := context.WithCancel(context.Background()) - cancel() - client := newTestPromptKitClient(t, &fakePromptKitLLM{content: `{"ok":true}`}) + attempts := 1 + for _, test := range []struct { + name string + attempts *int + }{ + {name: "during preparation"}, + {name: "during override inspection", attempts: &attempts}, + } { + t.Run(test.name, func(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + cancel() + client := newTestPromptKitClient(t, &fakePromptKitLLM{content: `{"ok":true}`}) - var out map[string]any - _, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{ - PromptID: "adapter.test", - Inputs: contracts.LLMInputSet{ - "transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""), - }, - }, &out) - if !errors.Is(err, context.Canceled) || errors.Is(err, contracts.ErrInvalidStructuredOutput) { - t.Fatalf("CompleteStructured() error = %v, want context canceled", err) + var out map[string]any + _, err := client.CompleteStructured(ctx, contracts.StructuredCompletionRequest{ + PromptID: "adapter.test", + StructuredOutputRepairAttempts: test.attempts, + Inputs: contracts.LLMInputSet{ + "transcript": contracts.NewLLMInputMaterial("transcript", "application/json", []byte(`{"source":true}`), "", ""), + }, + }, &out) + if !errors.Is(err, context.Canceled) || errors.Is(err, contracts.ErrInvalidStructuredOutput) { + t.Fatalf("CompleteStructured() error = %v, want context canceled", err) + } + }) } } diff --git a/internal/framework/pipeline/merger_registry.go b/internal/framework/pipeline/merger_registry.go index f1f4bf1f..c6a5107d 100644 --- a/internal/framework/pipeline/merger_registry.go +++ b/internal/framework/pipeline/merger_registry.go @@ -85,7 +85,7 @@ 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, Metadata: request.Metadata}) + 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}) if err != nil { return erasedTypedResult{}, err } diff --git a/internal/framework/pipeline/normalizer_registry.go b/internal/framework/pipeline/normalizer_registry.go index cd403b29..dddf0430 100644 --- a/internal/framework/pipeline/normalizer_registry.go +++ b/internal/framework/pipeline/normalizer_registry.go @@ -76,7 +76,7 @@ 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, Metadata: request.Metadata}) + 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}) if err != nil { return erasedTypedResult{}, err } diff --git a/internal/framework/pipeline/runner.go b/internal/framework/pipeline/runner.go index 203cfccc..aebf5cb0 100644 --- a/internal/framework/pipeline/runner.go +++ b/internal/framework/pipeline/runner.go @@ -493,9 +493,9 @@ func (r *Runner) validateChunks(ctx context.Context, doc *source.SourceDocument, } switch item.resolved.Target { case ValidatorTargetChunk: - result, err = item.chunk.Validate(validatorCtx, contracts.ChunkValidationRequest{ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, Metadata: requestMetadata, Chunks: requestChunks}) + result, err = item.chunk.Validate(validatorCtx, contracts.ChunkValidationRequest{ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, StructuredOutputRepairAttempts: cloneStructuredOutputRepairAttempts(binding.StructuredOutputRepairAttempts), Metadata: requestMetadata, Chunks: requestChunks}) case ValidatorTargetSerialized: - result, err = item.serialized.Validate(validatorCtx, contracts.SerializedValidationRequest{Stage: string(StageChunk), ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, Metadata: requestMetadata, Chunks: requestChunks, Schema: contracts.CloneArtifactSchema(schema), MediaType: "application/json", Content: append([]byte(nil), content...)}) + result, err = item.serialized.Validate(validatorCtx, contracts.SerializedValidationRequest{Stage: string(StageChunk), ModuleKey: moduleKey, Source: doc, SourceID: doc.ID, SourceInput: sourceInput.Clone(), SessionID: sessionID, References: CloneReferenceSet(references), LLMProfile: binding.LLMProfile, StructuredOutputRepairAttempts: cloneStructuredOutputRepairAttempts(binding.StructuredOutputRepairAttempts), Metadata: requestMetadata, Chunks: requestChunks, Schema: contracts.CloneArtifactSchema(schema), MediaType: "application/json", Content: append([]byte(nil), content...)}) default: return nil, nil, fmt.Errorf("validator %q is incompatible with chunk validation", binding.Module) } diff --git a/internal/framework/pipeline/runner_terminal_debug_test.go b/internal/framework/pipeline/runner_terminal_debug_test.go index 1759d6f6..ed3b47fe 100644 --- a/internal/framework/pipeline/runner_terminal_debug_test.go +++ b/internal/framework/pipeline/runner_terminal_debug_test.go @@ -41,6 +41,10 @@ type observingChunkValidator struct { request contracts.ChunkValidationRequest } +type observingSerializedValidator struct { + request contracts.SerializedValidationRequest +} + func (*observingChunkValidator) Name() string { return "observing/chunk-validator" } func (*observingChunkValidator) ExecutionClass() contracts.ExecutionClass { return contracts.ExecutionClassDeterministic @@ -50,6 +54,15 @@ func (v *observingChunkValidator) Validate(_ context.Context, request contracts. return contracts.ValidationResult{Approved: true}, nil } +func (*observingSerializedValidator) Name() string { return "observing/serialized-validator" } +func (*observingSerializedValidator) ExecutionClass() contracts.ExecutionClass { + return contracts.ExecutionClassDeterministic +} +func (v *observingSerializedValidator) Validate(_ context.Context, request contracts.SerializedValidationRequest) (contracts.ValidationResult, error) { + v.request = request + return contracts.ValidationResult{Approved: true}, nil +} + func (terminalChunkValidator) Name() string { return "terminal/chunk-validator" } func (terminalChunkValidator) ExecutionClass() contracts.ExecutionClass { @@ -168,6 +181,33 @@ func TestRunnerMaterializesAnnotatedPlanBeforeChunkValidation(t *testing.T) { } } +func TestRunnerForwardsDetachedRepairAttemptsToChunkValidators(t *testing.T) { + prepared, plan := preparedTerminalDebugPipeline(t) + prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan} + chunkValidator := &observingChunkValidator{} + serializedValidator := &observingSerializedValidator{} + chunkAttempts := 2 + serializedAttempts := 3 + chunkBinding := Binding(chunkValidator.Name()) + chunkBinding.StructuredOutputRepairAttempts = &chunkAttempts + serializedBinding := Binding(serializedValidator.Name()) + serializedBinding.StructuredOutputRepairAttempts = &serializedAttempts + prepared.chunkValidators.validators = []preparedValidator{ + {resolved: ResolvedValidator{Binding: chunkBinding, Target: ValidatorTargetChunk}, chunk: chunkValidator}, + {resolved: ResolvedValidator{Binding: serializedBinding, Target: ValidatorTargetSerialized}, serialized: serializedValidator}, + } + + if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil { + t.Fatalf("Run() error = %v, want nil", err) + } + if got := chunkValidator.request.StructuredOutputRepairAttempts; got == nil || *got != chunkAttempts || got == chunkBinding.StructuredOutputRepairAttempts { + t.Fatalf("chunk validator repair attempts = %v, want detached value %d", got, chunkAttempts) + } + if got := serializedValidator.request.StructuredOutputRepairAttempts; got == nil || *got != serializedAttempts || got == serializedBinding.StructuredOutputRepairAttempts { + t.Fatalf("serialized validator repair attempts = %v, want detached value %d", got, serializedAttempts) + } +} + func TestRunnerRetriesMalformedPlanWithDebugEnabled(t *testing.T) { prepared, plan := preparedTerminalDebugPipeline(t) plan.Ranges[0].Annotations = source.ChunkAnnotations{"broken": []byte(`{"value":`)} diff --git a/internal/framework/pipeline/typed_registry_request_test.go b/internal/framework/pipeline/typed_registry_request_test.go new file mode 100644 index 00000000..f8806f50 --- /dev/null +++ b/internal/framework/pipeline/typed_registry_request_test.go @@ -0,0 +1,95 @@ +package pipeline + +import ( + "context" + "testing" + + "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" +) + +type repairObservingNotesMerger struct { + attempts *int +} + +func (*repairObservingNotesMerger) Key() string { return "test/repair-observing-merge" } + +func (m *repairObservingNotesMerger) Merge(_ context.Context, request contracts.TypedMergeRequest[codecNotes]) (contracts.TypedMergeResult[codecNotes], error) { + m.attempts = request.StructuredOutputRepairAttempts + return contracts.TypedMergeResult[codecNotes]{Value: codecNotes{Items: []string{"merged"}}}, nil +} + +type repairObservingNotesNormalizer struct { + attempts *int +} + +func (*repairObservingNotesNormalizer) Key() string { return "test/repair-observing-normalize" } +func (*repairObservingNotesNormalizer) ReferenceSlots() []contracts.ReferenceSlot { + return nil +} + +func (n *repairObservingNotesNormalizer) Normalize(_ context.Context, request contracts.TypedNormalizeRequest[codecNotes]) (contracts.TypedNormalizeResult[codecNotes], error) { + n.attempts = request.StructuredOutputRepairAttempts + return contracts.TypedNormalizeResult[codecNotes]{Value: request.MergeOutput.Value}, nil +} + +func TestTypedRegistryErasurePreservesStructuredOutputRepairAttempts(t *testing.T) { + const artifactKind contracts.ArtifactKind = "test/notes" + + t.Run("merge", func(t *testing.T) { + implementation := &repairObservingNotesMerger{} + 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) + } + attempts := 2 + _, err = entry.merge(context.Background(), built, contracts.TypedMergeRequest[any]{ + ExtractOutputs: []contracts.ExtractArtifact[any]{{Value: codecNotes{Items: []string{"extracted"}}}}, + StructuredOutputRepairAttempts: &attempts, + }) + if err != nil { + t.Fatalf("merge() error = %v", err) + } + if implementation.attempts == nil || *implementation.attempts != attempts { + t.Fatalf("repair attempts = %v, want %d", implementation.attempts, attempts) + } + }) + + t.Run("normalize", func(t *testing.T) { + implementation := &repairObservingNotesNormalizer{} + 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) + } + attempts := 3 + _, err = entry.normalize(context.Background(), built, contracts.TypedNormalizeRequest[any]{ + MergeOutput: contracts.MergeArtifact[any]{Value: codecNotes{Items: []string{"merged"}}}, + StructuredOutputRepairAttempts: &attempts, + }) + if err != nil { + t.Fatalf("normalize() error = %v", err) + } + if implementation.attempts == nil || *implementation.attempts != attempts { + t.Fatalf("repair attempts = %v, want %d", implementation.attempts, attempts) + } + }) +}