From 4bb4582695eaec1b18f92b7731949a05ed64f9d9 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Thu, 27 Aug 2026 01:06:17 +0000 Subject: [PATCH] Integrate extraction validation retries --- docs/roadmap/implementation.md | 2 +- .../assembled_spell_pipeline_contract_test.go | 136 +++++++++++++++++- .../cli/spell_catalog_retry_contract_test.go | 5 +- .../framework/pipeline/producer_attempts.go | 2 +- .../pipeline/runner_concurrency_test.go | 1 + .../framework/pipeline/runner_concurrent.go | 120 ++++++++++------ .../runner_extract_correction_test.go | 97 +++++++++++++ .../pipeline/runner_extract_handoff_test.go | 6 +- .../runner_rejection_warnings_test.go | 3 +- .../pipeline/runner_terminal_debug_test.go | 10 ++ internal/framework/pipeline/runner_typed.go | 20 ++- .../integration/concurrent_runner_test.go | 2 +- .../integration/dnd_combat_runner_test.go | 1 + .../dnd_npc_occurrences_runner_test.go | 5 + .../dnd_npc_registry_runner_test.go | 4 +- .../integration/dnd_spells_runner_test.go | 4 +- 16 files changed, 358 insertions(+), 60 deletions(-) create mode 100644 internal/framework/pipeline/runner_extract_correction_test.go diff --git a/docs/roadmap/implementation.md b/docs/roadmap/implementation.md index 2cdc6068..1ec2926d 100644 --- a/docs/roadmap/implementation.md +++ b/docs/roadmap/implementation.md @@ -550,7 +550,7 @@ Chunk is the first complete production stage using feedback-aware retries, with cache and terminal behavior matching the roadmap. This stage is one Terra prompt. -## Stage 12 — Integrate Per-Chunk Extraction +## Stage 12 ✅ — Integrate Per-Chunk Extraction ### Goal diff --git a/internal/cli/assembled_spell_pipeline_contract_test.go b/internal/cli/assembled_spell_pipeline_contract_test.go index 2891bac4..f56589dc 100644 --- a/internal/cli/assembled_spell_pipeline_contract_test.go +++ b/internal/cli/assembled_spell_pipeline_contract_test.go @@ -21,6 +21,10 @@ import ( const assembledSpellExtractorKey = "test/dnd/spell-casts" +const assembledCorrectingSpellExtractorKey = "test/dnd/correcting-spell-casts" + +const assembledDirectSpellValidatorKey = "test/dnd/direct-spell-correction" + func TestAssembledSpellPipelineNormalizesMergedCasts(t *testing.T) { registries, resolved, extractor := assembledSpellPipeline(t, assembledSpellPipelineOptions{}) prepared, err := pipeline.Prepare(resolved, registries, pipeline.ModuleDependencies{}) @@ -101,6 +105,30 @@ func TestAssembledSpellPipelineNormalizesMergedCasts(t *testing.T) { } } +func TestAssembledSpellPipelineCorrectsRejectedDirectExtraction(t *testing.T) { + registries, resolved, extractor := assembledCorrectingSpellPipeline(t) + prepared, err := pipeline.Prepare(resolved, registries, pipeline.ModuleDependencies{}) + if err != nil { + t.Fatalf("Prepare() error = %v, want nil", err) + } + + output, err := pipeline.New().Run(context.Background(), pipeline.RunInput{ + Prepared: prepared, + RawInput: readRepositoryFile(t, "examples", "seriatim-minimal-transcript.json"), + ChunkCacheMode: pipeline.ChunkCacheBypass, + }) + if err != nil { + t.Fatalf("Run() error = %v, want nil", err) + } + if len(output.Rejected) != 0 || output.Manifest.ValidationStatus != "approved" || len(output.NormalizeOutputs) != 1 { + t.Fatalf("run output = %#v, want corrected accepted spell output", output) + } + correction := extractor.correctionSnapshot() + if correction == nil || string(correction.AssistantResponse) != `{"spell":"Mysterious Burst"}` || correction.UserGuidance != "use a known spell name" { + t.Fatalf("extract correction = %#v, want exact rejected model response and validator guidance", correction) + } +} + func TestAssembledSpellPipelineHonorsNormalizeValidatorOverride(t *testing.T) { registries, resolved, _ := assembledSpellPipeline(t, assembledSpellPipelineOptions{normalizeValidatorOverride: true}) var normalizeChain *pipeline.ResolvedValidatorChain @@ -161,8 +189,8 @@ func TestAssembledSpellPipelinePromotesTerminalUnknownSpellWarning(t *testing.T) if !reflect.DeepEqual(rejectedFile.Rejected, output.Rejected) { t.Fatalf("rejected file = %#v, run rejections = %#v, want durable rejection diagnostic", rejectedFile.Rejected, output.Rejected) } - if len(output.Warnings) != 1 || output.Warnings[0].ReasonCode != spellnormalize.ReasonCodeSpellNameUnresolved || output.Warnings[0].Scope != "spell_casts[0]" { - t.Fatalf("warnings = %#v, want terminal normalize catalog warning", output.Warnings) + if len(output.Warnings) != 2 || output.Warnings[0].ReasonCode != spellnormalize.ReasonCodeSpellNameUnresolved || output.Warnings[0].Scope != "spell_casts[0]" || output.Warnings[1].ReasonCode != "spell_not_near_source" { + t.Fatalf("warnings = %#v, want complete terminal normalize validation warnings", output.Warnings) } } @@ -206,6 +234,47 @@ type assembledSpellPipelineOptions struct { unknownSpell bool } +func assembledCorrectingSpellPipeline(t *testing.T) (pipeline.Registries, pipeline.ResolvedPipeline, *assembledCorrectingSpellExtractor) { + t.Helper() + components := productionTestComponents(t) + extractor := &assembledCorrectingSpellExtractor{} + if err := pipeline.RegisterExtractor[dnd.SpellList](components.registries.Extractors, pipeline.ModuleSpec{ + Key: assembledCorrectingSpellExtractorKey, + Stage: pipeline.StageExtract, + ExecutionClass: contracts.ExecutionClassLLMBacked, + CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1, + Requires: []string{"chunks", "source.transcript"}, + Provides: []string{"dnd.spell_casts"}, + ArtifactKind: dnd.SpellListKind, + }, func() (contracts.Extractor[dnd.SpellList], error) { + return extractor, nil + }); err != nil { + t.Fatalf("register correcting extractor: %v", err) + } + if err := pipeline.RegisterTypedValidator[dnd.SpellList](components.registries.Validators, dnd.SpellListKind, pipeline.ValidatorSpec{Key: assembledDirectSpellValidatorKey, ExecutionClass: contracts.ExecutionClassDeterministic}, func() (contracts.TypedValidator[dnd.SpellList], error) { + return assembledDirectSpellValidator{}, nil + }); err != nil { + t.Fatalf("register direct spell validator: %v", err) + } + + extract := pipeline.Binding(assembledCorrectingSpellExtractorKey) + extract.Retries = 1 + extract.Validators = pipeline.ValidatorOverride{Set: true, Validators: []pipeline.ModuleBinding{{Module: assembledDirectSpellValidatorKey}}} + resolved, err := pipeline.ResolvePipeline(pipeline.PipelineProfile{ + ID: "assembled-dnd-correcting-spells", + Input: pipeline.Binding("seriatim"), + Chunk: pipeline.ModuleBinding{Module: "generic", Options: map[string]any{"max_units": 1}}, + Artifacts: map[string]pipeline.ArtifactLaneProfile{ + "spells": {Extract: extract, Normalize: pipeline.Binding(spellnormalize.Key)}, + }, + Output: pipeline.Binding("json"), + }, pipeline.ResolveOptions{}, catalogFromRegistries(components.registries)) + if err != nil { + t.Fatalf("ResolvePipeline() error = %v, want nil", err) + } + return components.registries, resolved, extractor +} + func assembledSpellPipeline(t *testing.T, options assembledSpellPipelineOptions) (pipeline.Registries, pipeline.ResolvedPipeline, *assembledSpellExtractor) { t.Helper() components := productionTestComponents(t) @@ -251,6 +320,69 @@ type assembledSpellExtractor struct { unknownSpell bool } +type assembledCorrectingSpellExtractor struct { + mu sync.Mutex + correction *contracts.SemanticCorrection +} + +func (*assembledCorrectingSpellExtractor) Key() string { return assembledCorrectingSpellExtractorKey } + +func (*assembledCorrectingSpellExtractor) ReferenceSlots() []contracts.ReferenceSlot { return nil } + +func (e *assembledCorrectingSpellExtractor) Extract(_ context.Context, req contracts.TypedExtractionRequest) (contracts.TypedExtractionResult[dnd.SpellList], error) { + if req.Source == nil || req.Chunk == nil { + return contracts.TypedExtractionResult[dnd.SpellList]{}, fmt.Errorf("correcting assembled extractor requires source and chunk") + } + response := `{"spell":"accepted"}` + value := dnd.SpellList{SpellCasts: []dnd.SpellCast{}} + if req.Chunk.Index == 0 && req.Correction == nil { + response = `{"spell":"Mysterious Burst"}` + value.SpellCasts = []dnd.SpellCast{{Caster: "Aria", Spell: "Mysterious Burst", SourceRefs: []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 1, EndUnitID: 1}}}} + } + if req.Chunk.Index == 0 && req.Correction != nil { + correction, err := contracts.CloneSemanticCorrection(req.Correction) + if err != nil { + return contracts.TypedExtractionResult[dnd.SpellList]{}, err + } + e.mu.Lock() + e.correction = correction + e.mu.Unlock() + value.SpellCasts = []dnd.SpellCast{{Caster: "Aria", Spell: "Cure Wounds", SourceRefs: []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 1, EndUnitID: 1}}}} + } + candidate, err := contracts.NewModelCandidate([]byte(response), contracts.CorrectionProtocolSingleResponseV1) + if err != nil { + return contracts.TypedExtractionResult[dnd.SpellList]{}, err + } + return contracts.TypedExtractionResult[dnd.SpellList]{Value: value, ModelCandidate: candidate}, nil +} + +func (e *assembledCorrectingSpellExtractor) correctionSnapshot() *contracts.SemanticCorrection { + e.mu.Lock() + defer e.mu.Unlock() + correction, err := contracts.CloneSemanticCorrection(e.correction) + if err != nil { + return nil + } + return correction +} + +type assembledDirectSpellValidator struct{} + +func (assembledDirectSpellValidator) Name() string { return assembledDirectSpellValidatorKey } + +func (assembledDirectSpellValidator) ExecutionClass() contracts.ExecutionClass { + return contracts.ExecutionClassDeterministic +} + +func (assembledDirectSpellValidator) Validate(_ context.Context, req contracts.TypedValidationRequest[dnd.SpellList]) (contracts.ValidationResult, error) { + for _, cast := range req.Value.SpellCasts { + if cast.Spell == "Mysterious Burst" { + return contracts.ValidationResult{Approved: false, ReasonCode: "unknown_spell", Message: "spell is not in the catalog", CorrectionGuidance: "use a known spell name"}, nil + } + } + return contracts.ValidationResult{Approved: true}, nil +} + func (e *assembledSpellExtractor) Key() string { return assembledSpellExtractorKey } func (*assembledSpellExtractor) ReferenceSlots() []contracts.ReferenceSlot { return nil } diff --git a/internal/cli/spell_catalog_retry_contract_test.go b/internal/cli/spell_catalog_retry_contract_test.go index 2a246b48..bda25b92 100644 --- a/internal/cli/spell_catalog_retry_contract_test.go +++ b/internal/cli/spell_catalog_retry_contract_test.go @@ -65,6 +65,7 @@ func TestProductionSpellCatalogValidationRetries(t *testing.T) { t.Fatalf("materialize production references: %v", err) } materialized.Steps[0].ArtifactLanes[0].Extract.Retries = retries + materialized.Steps[0].ArtifactLanes[0].ExtractValidationPolicy.SemanticRejection = pipeline.SemanticRejectionRejectOutput llmClient := &catalogRetryLLMClient{responses: tt.responses} prepared, err := pipeline.Prepare(materialized, components.registries, pipeline.ModuleDependencies{LLM: llmClient}) @@ -91,8 +92,8 @@ func TestProductionSpellCatalogValidationRetries(t *testing.T) { if rejection.ReasonCode != "unknown_spell" || rejection.AttemptCount != retries+1 { t.Fatalf("rejection = %#v, want exhausted unknown-spell rejection", rejection) } - if len(output.Warnings) != 0 { - t.Fatalf("warnings = %#v, want no emitted warnings from rejected attempts", output.Warnings) + if len(output.Warnings) != 1 || output.Warnings[0].ReasonCode != "spell_not_near_source" { + t.Fatalf("warnings = %#v, want complete terminal validation warnings", output.Warnings) } return } diff --git a/internal/framework/pipeline/producer_attempts.go b/internal/framework/pipeline/producer_attempts.go index bb0fef81..2c484a88 100644 --- a/internal/framework/pipeline/producer_attempts.go +++ b/internal/framework/pipeline/producer_attempts.go @@ -278,7 +278,7 @@ func applyStructuralTerminalPolicy(policy ValidationPolicy, provenance []produce } func applySemanticTerminalPolicy(policy ValidationPolicy, provenance []producerAttemptProvenance, number int, output producerAttemptOutput, report validationReport, rejection validationRecord) (producerAttemptTerminal, error) { - rejected := contracts.RejectedOutput{ReasonCode: rejection.reasonCode, Message: rejection.message, AttemptCount: number} + rejected := contracts.RejectedOutput{ValidatorName: rejection.validatorName, ReasonCode: rejection.reasonCode, Message: rejection.message, AttemptCount: number, DiagnosticArtifactPath: rejection.diagnosticPath} switch policy.SemanticRejection { case SemanticRejectionRejectOutput: return producerAttemptTerminal{Action: producerTerminalRejected, Warnings: terminalWarnings(output, report), Rejection: &rejected, Validation: report, ValidationIncomplete: firstIncompleteValidation(report) != nil, Provenance: cloneProducerAttemptProvenance(provenance)}, nil diff --git a/internal/framework/pipeline/runner_concurrency_test.go b/internal/framework/pipeline/runner_concurrency_test.go index 02a74180..9ba8ccd9 100644 --- a/internal/framework/pipeline/runner_concurrency_test.go +++ b/internal/framework/pipeline/runner_concurrency_test.go @@ -558,6 +558,7 @@ func TestRunnerSelectsMergeErrorBeforeEarlierLaneNormalizeError(t *testing.T) { func TestRunnerKeepsSuccessfulChunksWhenAnotherExtractIsRejected(t *testing.T) { prepared := preparedConcurrentPipeline(t, 2) + prepared.Steps[0].lanes[0].resolved.ExtractValidationPolicy.SemanticRejection = SemanticRejectionRejectOutput for laneIndex := range prepared.Steps[0].lanes { lane := laneIndex installExtractOperation(prepared, lane, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) { diff --git a/internal/framework/pipeline/runner_concurrent.go b/internal/framework/pipeline/runner_concurrent.go index e63f34f6..6341d98b 100644 --- a/internal/framework/pipeline/runner_concurrent.go +++ b/internal/framework/pipeline/runner_concurrent.go @@ -24,6 +24,7 @@ type laneExtractState struct { serialized []CheckpointArtifact warnings []contracts.Warning rejected []contracts.RejectedOutput + incomplete []int results map[int]extractJobResult remaining int failed bool @@ -36,6 +37,7 @@ type finalizedExtractResults struct { serialized []CheckpointArtifact warnings []contracts.Warning rejected []contracts.RejectedOutput + incomplete []int decision CheckpointDecision } @@ -56,13 +58,20 @@ type extractJob struct { } type extractJobResult struct { - laneIndex int - chunkIndex int - value erasedExtractArtifact + laneIndex int + chunkIndex int + value erasedExtractArtifact + serialized CheckpointArtifact + warnings []contracts.Warning + rejected *contracts.RejectedOutput + validationIncomplete bool + err error +} + +type extractAttemptValue struct { + artifact erasedExtractArtifact serialized CheckpointArtifact - warnings []contracts.Warning - rejected *contracts.RejectedOutput - err error + terminal *attemptTerminalRecorder } type laneCompletion struct { @@ -407,23 +416,21 @@ func (r *Runner) runExtractJob(ctx context.Context, input RunInput, doc *source. result.err = fmt.Errorf("clone chunk %q for extraction: %w", job.chunk.ID, cloneErr) return result } - var accepted erasedExtractArtifact - var serialized CheckpointArtifact - var acceptedWarnings []contracts.Warning - retryResult, err := runSimpleRetry(ctx, lane.Extract.Retries, func(attempt int) (retryAttemptResult, error) { + terminalResult, err := runProducerAttempts(ctx, producerAttemptConfig{Retries: lane.Extract.Retries, Policy: lane.ExtractValidationPolicy, AllowStructuralRetry: lane.ExtractExecutionClass == contracts.ExecutionClassLLMBacked}, func(attemptCtx context.Context, request producerAttemptRequest) (producerAttemptOutput, error) { + attempt := request.Number started := time.Now().UTC() attemptPath := path.Join("extract", fileio.EncodePathComponent(lane.ID), fmt.Sprintf("chunk-%06d", chunk.Index+1), fmt.Sprintf("attempt-%02d", attempt)) - attemptCtx, llmScope := withDebugLLMScope(ctx, attemptPath) + attemptCtx, llmScope := withDebugLLMScope(attemptCtx, attemptPath) terminal := newAttemptTerminalRecorder(input.Debug, attemptPath, "extract", llmScope, debugTimedEnvelope{Stage: string(StageExtract), StepID: input.stepID, LaneID: lane.ID, ModuleKey: lane.Extract.Module, Attempt: attempt, StartedAt: started}) requestMetadata, metadataErr := cloneMetadata(input.Metadata) if metadataErr != nil { - return retryAttemptResult{}, terminal.record(nil, fmt.Errorf("clone extract request metadata: %w", metadataErr)) + return producerAttemptOutput{}, terminal.record(nil, fmt.Errorf("clone extract request metadata: %w", metadataErr)) } extractReferences := operationReferenceSet(input, lane.ExtractReferences) - extracted, callErr := typed.extract(attemptCtx, typed.extractor, contracts.TypedExtractionRequest{Source: doc, Chunk: &chunk, SourceInput: chunkInputMaterial(sourceInput, chunk), SessionID: sessionID, References: CloneReferenceSet(extractReferences), LLMProfile: lane.Extract.LLMProfile, StructuredOutputRepairAttempts: cloneStructuredOutputRepairAttempts(lane.Extract.StructuredOutputRepairAttempts), Metadata: requestMetadata}) + extracted, callErr := typed.extract(attemptCtx, typed.extractor, contracts.TypedExtractionRequest{Source: doc, Chunk: &chunk, SourceInput: chunkInputMaterial(sourceInput, chunk), SessionID: sessionID, References: CloneReferenceSet(extractReferences), LLMProfile: lane.Extract.LLMProfile, StructuredOutputRepairAttempts: cloneStructuredOutputRepairAttempts(lane.Extract.StructuredOutputRepairAttempts), Correction: request.Correction, Metadata: requestMetadata}) if callErr != nil { attemptErr := fmt.Errorf("extract lane %q chunk %q with extractor %q: %w", lane.ID, chunk.ID, lane.Extract.Module, callErr) - return retryAttemptResult{}, terminal.record(nil, attemptErr) + return producerAttemptOutput{}, terminal.record(nil, attemptErr) } artifact := erasedExtractArtifact{LaneID: lane.ID, ExtractorKey: lane.Extract.Module, SourceID: doc.ID, ChunkID: chunk.ID, ChunkIndex: chunk.Index, ChunkRef: chunk.Ref, Value: extracted.Value} attemptWarnings := cloneWarnings(extracted.Warnings) @@ -431,41 +438,65 @@ func (r *Runner) runExtractJob(ctx context.Context, input RunInput, doc *source. if encodeErr != nil { attemptErr := fmt.Errorf("serialize extract candidate for lane %q chunk %q: %w", lane.ID, chunk.ID, encodeErr) payload := map[string]any{"warnings": debugWarningEnvelopes(attemptWarnings)} - return retryAttemptResult{}, terminal.record(payload, attemptErr) + return producerAttemptOutput{}, terminal.record(payload, attemptErr) } serializedCandidate.ChunkID, serializedCandidate.ChunkIndex, serializedCandidate.ChunkRef = artifact.ChunkID, artifact.ChunkIndex, artifact.ChunkRef - warnings, rejected, validateErr := r.validateTypedArtifact(attemptCtx, typed.codec, typedValidationTarget{stage: StageExtract, stepID: input.stepID, laneID: lane.ID, moduleKey: lane.Extract.Module, source: doc, sourceID: doc.ID, sourceInput: chunkInputMaterial(sourceInput, chunk), sessionID: sessionID, references: extractReferences, metadata: input.Metadata, chunk: &chunk, ref: chunk.Ref, value: extracted.Value, candidate: &serializedCandidate}, state.prepared.extractValidators, attempt, input.Debug) - attemptWarnings = append(attemptWarnings, warnings...) - payload := map[string]any{"output": debugCheckpointArtifact(serializedCandidate), "warnings": debugWarningEnvelopes(attemptWarnings), "rejection": debugRejectedOutputPtr(rejected)} - if validateErr != nil { - return retryAttemptResult{}, terminal.record(payload, validateErr) + return producerAttemptOutput{Value: extractAttemptValue{artifact: artifact, serialized: serializedCandidate, terminal: &terminal}, Candidate: extracted.ModelCandidate, Warnings: attemptWarnings}, nil + }, func(validationCtx context.Context, output producerAttemptOutput) (validationReport, error) { + candidate, ok := output.Value.(extractAttemptValue) + if !ok { + return validationReport{}, fmt.Errorf("extract attempt has incompatible value") } - if rejected != nil { - if debugErr := terminal.record(payload, nil); debugErr != nil { - return retryAttemptResult{}, debugErr + report, validationErr := r.validateTypedReport(validationCtx, typed.codec, typedValidationTarget{stage: StageExtract, stepID: input.stepID, laneID: lane.ID, moduleKey: lane.Extract.Module, source: doc, sourceID: doc.ID, sourceInput: chunkInputMaterial(sourceInput, chunk), sessionID: sessionID, references: operationReferenceSet(input, lane.ExtractReferences), metadata: input.Metadata, chunk: &chunk, ref: chunk.Ref, value: candidate.artifact.Value, candidate: &candidate.serialized}, state.prepared.extractValidators, candidate.terminal.envelope.Attempt, input.Debug) + payload := map[string]any{ + "output": debugCheckpointArtifact(candidate.serialized), + "warnings": debugWarningEnvelopes(append(cloneWarnings(output.Warnings), report.Warnings()...)), + "rejection": debugRejectedOutputPtr(typedRejection(report, typedValidationTarget{stage: StageExtract, stepID: input.stepID, laneID: lane.ID, moduleKey: lane.Extract.Module, chunk: &chunk}, candidate.terminal.envelope.Attempt)), + } + if validationErr != nil { + return report, candidate.terminal.record(payload, validationErr) + } + if report.FirstRejection() != nil { + return report, candidate.terminal.record(payload, nil) + } + if report.FirstRejection() == nil && lane.ExtractValidationPolicy.ValidatorFailure == ValidatorFailureFailRun { + if failure := firstIncompleteValidation(report); failure != nil { + return report, candidate.terminal.record(payload, validatorFailureError(*failure)) } - return retryAttemptResult{rejection: rejected, warnings: attemptWarnings}, nil } - stored, encodeErr := checkpointArtifact(typed.codec, artifact.LaneID, artifact.ExtractorKey, artifact.SourceID, artifact.Value) - if encodeErr != nil { - attemptErr := fmt.Errorf("serialize accepted extract output for lane %q chunk %q: %w", lane.ID, chunk.ID, encodeErr) - return retryAttemptResult{}, terminal.record(payload, attemptErr) - } - stored.ChunkID, stored.ChunkIndex, stored.ChunkRef = artifact.ChunkID, artifact.ChunkIndex, artifact.ChunkRef - accepted, serialized = artifact, stored - acceptedWarnings = attemptWarnings - if debugErr := terminal.record(payload, nil); debugErr != nil { - return retryAttemptResult{}, debugErr - } - return retryAttemptResult{accepted: true}, nil + return report, candidate.terminal.record(payload, nil) }) result.err = err - if err == nil && !retryResult.accepted { - result.rejected = retryResult.rejection - result.warnings = cloneWarnings(retryResult.warnings) + if err == nil && terminalResult.Action == producerTerminalRejected { + result.rejected = terminalResult.Rejection + if result.rejected != nil { + result.rejected.Stage, result.rejected.StepID, result.rejected.LaneID, result.rejected.ModuleKey, result.rejected.ChunkID, result.rejected.ChunkIndex = string(StageExtract), input.stepID, lane.ID, lane.Extract.Module, chunk.ID, chunk.Index + } + result.warnings = cloneWarnings(terminalResult.Warnings) return result } - result.value, result.serialized, result.warnings = accepted, serialized, acceptedWarnings + if err == nil { + candidate, ok := terminalResult.Value.(extractAttemptValue) + if !ok { + result.err = fmt.Errorf("extract attempt terminal has incompatible value") + return result + } + stored, encodeErr := checkpointArtifact(typed.codec, candidate.artifact.LaneID, candidate.artifact.ExtractorKey, candidate.artifact.SourceID, candidate.artifact.Value) + payload := map[string]any{"output": debugCheckpointArtifact(candidate.serialized), "warnings": debugWarningEnvelopes(terminalResult.Warnings), "rejection": debugRejectedOutputPtr(nil)} + if encodeErr != nil { + attemptErr := fmt.Errorf("serialize accepted extract output for lane %q chunk %q: %w", lane.ID, chunk.ID, encodeErr) + result.err = candidate.terminal.record(payload, attemptErr) + return result + } + stored.ChunkID, stored.ChunkIndex, stored.ChunkRef = candidate.artifact.ChunkID, candidate.artifact.ChunkIndex, candidate.artifact.ChunkRef + if debugErr := candidate.terminal.record(payload, nil); debugErr != nil { + result.err = debugErr + return result + } + result.value, result.serialized = candidate.artifact, stored + result.warnings = cloneWarnings(terminalResult.Warnings) + result.validationIncomplete = terminalResult.ValidationIncomplete + } return result } @@ -486,11 +517,15 @@ func finalizeLaneExtract(checkpoints CheckpointRecorder, stepID string, state *l state.values = append(state.values, result.value) state.serialized = append(state.serialized, result.serialized) state.warnings = append(state.warnings, result.warnings...) + if result.validationIncomplete { + state.incomplete = append(state.incomplete, result.chunkIndex) + } } sort.SliceStable(state.values, func(i, j int) bool { return state.values[i].ChunkIndex < state.values[j].ChunkIndex }) sort.SliceStable(state.serialized, func(i, j int) bool { return state.serialized[i].ChunkIndex < state.serialized[j].ChunkIndex }) sort.SliceStable(state.rejected, func(i, j int) bool { return state.rejected[i].ChunkIndex < state.rejected[j].ChunkIndex }) - if !state.decision.Reused { + sort.Ints(state.incomplete) + if !state.decision.Reused && len(state.incomplete) == 0 { if err := recordExtract(checkpoints, stepID, lane.ID, lane.Extract.Module, state.deps, state.serialized, state.rejected, state.warnings); err != nil { return fmt.Errorf("write extract checkpoint for lane %q: %w", lane.ID, err) } @@ -506,6 +541,7 @@ func (r *Runner) continueLane(ctx context.Context, input RunInput, checkpoints C serialized: state.serialized, warnings: state.warnings, rejected: state.rejected, + incomplete: state.incomplete, decision: state.decision, } local.Warnings = append(local.Warnings, cloneWarnings(results.warnings)...) @@ -513,7 +549,7 @@ func (r *Runner) continueLane(ctx context.Context, input RunInput, checkpoints C if err := writeDebugTimed(input.Debug, path.Join("extract", fileio.EncodePathComponent(lane.ID), "input.json"), debugTimedEnvelope{Stage: string(StageExtract), StepID: input.stepID, LaneID: lane.ID, ModuleKey: lane.Extract.Module, StartedAt: time.Now().UTC(), Payload: map[string]any{"reused": results.decision.Reused, "decision": results.decision, "source": debugSourceDocumentEnvelope(doc), "chunks": debugSourceChunkEnvelopes(chunks), "options": redactSensitiveMap(lane.Extract.Options), "metadata": redactSensitiveMap(input.Metadata)}}); err != nil { return local, &laneRunError{stage: StageExtract, err: err} } - if err := writeDebugTimed(input.Debug, path.Join("extract", fileio.EncodePathComponent(lane.ID), "output.json"), debugTimedEnvelope{Stage: string(StageExtract), StepID: input.stepID, LaneID: lane.ID, ModuleKey: lane.Extract.Module, StartedAt: time.Now().UTC(), Payload: map[string]any{"reused": results.decision.Reused, "outputs": debugCheckpointArtifacts(results.serialized), "rejected": debugRejectedOutputEnvelopes(results.rejected), "warnings": debugWarningEnvelopes(results.warnings)}}); err != nil { + if err := writeDebugTimed(input.Debug, path.Join("extract", fileio.EncodePathComponent(lane.ID), "output.json"), debugTimedEnvelope{Stage: string(StageExtract), StepID: input.stepID, LaneID: lane.ID, ModuleKey: lane.Extract.Module, StartedAt: time.Now().UTC(), Payload: map[string]any{"reused": results.decision.Reused, "outputs": debugCheckpointArtifacts(results.serialized), "rejected": debugRejectedOutputEnvelopes(results.rejected), "warnings": debugWarningEnvelopes(results.warnings), "validation_incomplete_chunks": append([]int(nil), results.incomplete...)}}); err != nil { return local, &laneRunError{stage: StageExtract, err: err} } if len(results.accepted) == 0 { diff --git a/internal/framework/pipeline/runner_extract_correction_test.go b/internal/framework/pipeline/runner_extract_correction_test.go new file mode 100644 index 00000000..8419f4c0 --- /dev/null +++ b/internal/framework/pipeline/runner_extract_correction_test.go @@ -0,0 +1,97 @@ +package pipeline + +import ( + "context" + "fmt" + "strings" + "sync" + "testing" + + "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" +) + +func TestRunnerCorrectsExtractCandidatesIndependentlyPerChunk(t *testing.T) { + prepared := preparedConcurrentPipeline(t, 2) + lane := &prepared.Steps[0].lanes[0] + lane.resolved.Extract.Retries = 1 + + var mu sync.Mutex + corrections := make(map[int]*contracts.SemanticCorrection) + installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) { + chunkIndex := request.Chunk.Index + if request.Correction != nil { + correction, err := contracts.CloneSemanticCorrection(request.Correction) + if err != nil { + return erasedTypedResult{}, err + } + mu.Lock() + corrections[chunkIndex] = correction + mu.Unlock() + candidate, err := contracts.NewModelCandidate([]byte(fmt.Sprintf("corrected-response-%d", chunkIndex)), contracts.CorrectionProtocolSingleResponseV1) + if err != nil { + return erasedTypedResult{}, err + } + return erasedTypedResult{Value: codecNotes{Items: []string{fmt.Sprintf("corrected-%d", chunkIndex)}}, ModelCandidate: candidate}, nil + } + candidate, err := contracts.NewModelCandidate([]byte(fmt.Sprintf("initial-response-%d", chunkIndex)), contracts.CorrectionProtocolSingleResponseV1) + if err != nil { + return erasedTypedResult{}, err + } + return erasedTypedResult{Value: codecNotes{Items: []string{fmt.Sprintf("invalid-%d", chunkIndex)}}, ModelCandidate: candidate}, nil + }) + lane.extractValidators.validators[0].typedValidate = func(_ context.Context, _ any, target typedValidationTarget) (contracts.ValidationResult, error) { + value := target.value.(codecNotes) + if value.Items[0] == fmt.Sprintf("invalid-%d", target.chunk.Index) { + return contracts.ValidationResult{Approved: false, ReasonCode: "incorrect_extract", Message: "candidate needs correction", CorrectionGuidance: fmt.Sprintf("correct chunk %d", target.chunk.Index)}, nil + } + return contracts.ValidationResult{Approved: true}, nil + } + + output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), ExtractWorkers: 2}) + if err != nil { + t.Fatalf("Run() error = %v, want nil", err) + } + if len(output.Rejected) != 0 || len(output.NormalizeOutputs) != 2 { + t.Fatalf("run output = %#v, want corrected accepted extracts", output) + } + mu.Lock() + defer mu.Unlock() + if len(corrections) != 2 { + t.Fatalf("corrections = %#v, want one correction per chunk", corrections) + } + for index := 0; index < 2; index++ { + correction := corrections[index] + if correction == nil || string(correction.AssistantResponse) != fmt.Sprintf("initial-response-%d", index) || correction.UserGuidance != fmt.Sprintf("correct chunk %d", index) { + t.Fatalf("chunk %d correction = %#v, want its exact initial response and guidance", index, correction) + } + } +} + +func TestRunnerContinuesValidationIncompleteExtractWithoutCheckpoint(t *testing.T) { + prepared := preparedAttemptDebugPipeline(t) + lane := &prepared.Steps[0].lanes[0] + lane.resolved.ExtractValidationPolicy.ValidatorFailure = ValidatorFailureWarnContinue + installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) { + return erasedTypedResult{Value: typedValueForLane(0, request.Chunk.Index)}, nil + }) + lane.extractValidators.validators[0].typedValidate = func(context.Context, any, typedValidationTarget) (contracts.ValidationResult, error) { + return contracts.ValidationResult{}, fmt.Errorf("validator unavailable") + } + recorder := &extractCaptureRecorder{CheckpointRecorder: NoopCheckpointRecorder()} + debug := newCapturedDebugRecorder() + + output, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Checkpoints: recorder, Debug: debug}) + if err != nil { + t.Fatalf("Run() error = %v, want nil", err) + } + if len(output.NormalizeOutputs) != 1 || len(output.Rejected) != 0 { + t.Fatalf("run output = %#v, want accepted incomplete extract", output) + } + if len(recorder.checkpoint.Outputs) != 0 || len(recorder.checkpoint.Rejected) != 0 { + t.Fatalf("extract checkpoint = %#v, want no persisted incomplete output", recorder.checkpoint) + } + encoded := string(debug.json["extract/notes/output.json"]) + if !strings.Contains(encoded, `"validation_incomplete_chunks":[0]`) { + t.Fatalf("extract output debug = %s, want incomplete chunk marker", encoded) + } +} diff --git a/internal/framework/pipeline/runner_extract_handoff_test.go b/internal/framework/pipeline/runner_extract_handoff_test.go index bd927968..0fe750d0 100644 --- a/internal/framework/pipeline/runner_extract_handoff_test.go +++ b/internal/framework/pipeline/runner_extract_handoff_test.go @@ -3,6 +3,7 @@ package pipeline import ( "context" "encoding/json" + "fmt" "reflect" "strings" "testing" @@ -221,8 +222,9 @@ func TestRunnerPromotesOnlyAcceptedExtractRetryWarnings(t *testing.T) { scope = "accepted" } return erasedTypedResult{ - Value: typedValueForLane(0, request.Chunk.Index), - Warnings: []contracts.Warning{{Scope: scope, ReasonCode: "observed", Message: scope}}, + Value: typedValueForLane(0, request.Chunk.Index), + Warnings: []contracts.Warning{{Scope: scope, ReasonCode: "observed", Message: scope}}, + ModelCandidate: attemptCandidate(t, fmt.Sprintf(`{"items":["%s"]}`, scope)), }, nil }) validatorCalls := 0 diff --git a/internal/framework/pipeline/runner_rejection_warnings_test.go b/internal/framework/pipeline/runner_rejection_warnings_test.go index abf503ca..2e93ef73 100644 --- a/internal/framework/pipeline/runner_rejection_warnings_test.go +++ b/internal/framework/pipeline/runner_rejection_warnings_test.go @@ -75,9 +75,10 @@ func TestRunnerPromotesOnlyTerminalRejectionWarnings(t *testing.T) { } case StageExtract: lane.resolved.Extract.Retries = 1 + lane.resolved.ExtractValidationPolicy.SemanticRejection = SemanticRejectionRejectOutput installExtractOperation(prepared, 0, func(context.Context, contracts.TypedExtractionRequest) (erasedTypedResult, error) { attempts++ - return erasedTypedResult{Value: codecNotes{Items: []string{"extract"}}, Warnings: []contracts.Warning{{Scope: fmt.Sprintf("operation-%d", attempts), ReasonCode: "operation", Message: "operation warning"}}}, nil + return erasedTypedResult{Value: codecNotes{Items: []string{"extract"}}, Warnings: []contracts.Warning{{Scope: fmt.Sprintf("operation-%d", attempts), ReasonCode: "operation", Message: "operation warning"}}, ModelCandidate: attemptCandidate(t, fmt.Sprintf(`{"items":["%d"]}`, attempts))}, nil }) lane.extractValidators.validators = rejectionWarningTypedValidators(first, reject) case StageMerge: diff --git a/internal/framework/pipeline/runner_terminal_debug_test.go b/internal/framework/pipeline/runner_terminal_debug_test.go index e76e343d..6edf6fec 100644 --- a/internal/framework/pipeline/runner_terminal_debug_test.go +++ b/internal/framework/pipeline/runner_terminal_debug_test.go @@ -128,6 +128,11 @@ func preparedTerminalDebugPipeline(t *testing.T) (*PreparedPipeline, source.Chun SemanticRejection: SemanticRejectionRejectOutput, ValidatorFailure: ValidatorFailureFailRun, } + prepared.Steps[0].lanes[0].resolved.ExtractValidationPolicy = ValidationPolicy{ + ProducerStructuralFailure: ProducerStructuralFailureFailRun, + SemanticRejection: SemanticRejectionRejectOutput, + ValidatorFailure: ValidatorFailureFailRun, + } chunker, ok := prepared.chunker.(*typedTestChunker) if !ok { t.Fatalf("prepared chunker = %T, want *typedTestChunker", prepared.chunker) @@ -295,6 +300,11 @@ func TestRunnerRecordsExtractTerminalOutcomes(t *testing.T) { for _, tc := range tests { t.Run(tc.name, func(t *testing.T) { prepared := preparedAttemptDebugPipeline(t) + prepared.Steps[0].lanes[0].resolved.ExtractValidationPolicy = ValidationPolicy{ + ProducerStructuralFailure: ProducerStructuralFailureFailRun, + SemanticRejection: SemanticRejectionRejectOutput, + ValidatorFailure: ValidatorFailureFailRun, + } value := "extract-terminal" codec := &observedNotesCodec{} if tc.candidateFail { diff --git a/internal/framework/pipeline/runner_typed.go b/internal/framework/pipeline/runner_typed.go index f2f0f068..de0bae8d 100644 --- a/internal/framework/pipeline/runner_typed.go +++ b/internal/framework/pipeline/runner_typed.go @@ -551,16 +551,24 @@ func (r *Runner) validateTypedArtifact(ctx context.Context, codec artifactCodecE if failure := report.FirstFailure(); failure != nil { return report.Warnings(), nil, validatorFailureError(*failure) } - if rejection := report.FirstRejection(); rejection != nil { - chunkID, chunkIndex := "", 0 - if target.chunk != nil { - chunkID, chunkIndex = target.chunk.ID, target.chunk.Index - } - return report.Warnings(), &contracts.RejectedOutput{Stage: string(target.stage), StepID: target.stepID, LaneID: target.laneID, ModuleKey: target.moduleKey, ChunkID: chunkID, ChunkIndex: chunkIndex, ValidatorName: rejection.validatorName, ReasonCode: rejection.reasonCode, Message: rejection.message, AttemptCount: attempt, DiagnosticArtifactPath: rejection.diagnosticPath}, nil + if rejected := typedRejection(report, target, attempt); rejected != nil { + return report.Warnings(), rejected, nil } return report.Warnings(), nil, nil } +func typedRejection(report validationReport, target typedValidationTarget, attempt int) *contracts.RejectedOutput { + rejection := report.FirstRejection() + if rejection == nil { + return nil + } + chunkID, chunkIndex := "", 0 + if target.chunk != nil { + chunkID, chunkIndex = target.chunk.ID, target.chunk.Index + } + return &contracts.RejectedOutput{Stage: string(target.stage), StepID: target.stepID, LaneID: target.laneID, ModuleKey: target.moduleKey, ChunkID: chunkID, ChunkIndex: chunkIndex, ValidatorName: rejection.validatorName, ReasonCode: rejection.reasonCode, Message: rejection.message, AttemptCount: attempt, DiagnosticArtifactPath: rejection.diagnosticPath} +} + func validationCandidateArtifact(codec artifactCodecEntry, target typedValidationTarget) (CheckpointArtifact, error) { if target.candidate != nil { return cloneCheckpointArtifact(*target.candidate), nil diff --git a/internal/modules/integration/concurrent_runner_test.go b/internal/modules/integration/concurrent_runner_test.go index 40773f72..03f7bff8 100644 --- a/internal/modules/integration/concurrent_runner_test.go +++ b/internal/modules/integration/concurrent_runner_test.go @@ -127,7 +127,7 @@ func TestRunnerIndependentlyBoundsWorkersAndProviderCallsAcrossRegisteredModules t.Fatalf("register chunker: %v", err) } validateOptions := func(options map[string]any) error { return pipeline.RejectUnknownOptions(options) } - if err := pipeline.RegisterExtractorBuilder(catalog.Extractors, pipeline.ModuleSpec{Key: concurrentExtractorKey, Stage: pipeline.StageExtract, ExecutionClass: contracts.ExecutionClassLLMBacked, ArtifactKind: dnd.SpellListKind, Requires: []string{"chunks", "source.transcript"}, Provides: []string{"dnd.spell_casts"}}, validateOptions, func(request pipeline.BuildRequest) (contracts.Extractor[dnd.SpellList], error) { + if err := pipeline.RegisterExtractorBuilder(catalog.Extractors, pipeline.ModuleSpec{Key: concurrentExtractorKey, Stage: pipeline.StageExtract, ExecutionClass: contracts.ExecutionClassLLMBacked, CorrectionProtocol: contracts.CorrectionProtocolSingleResponseV1, ArtifactKind: dnd.SpellListKind, Requires: []string{"chunks", "source.transcript"}, Provides: []string{"dnd.spell_casts"}}, validateOptions, func(request pipeline.BuildRequest) (contracts.Extractor[dnd.SpellList], error) { return &concurrentExtractor{client: request.Dependencies.LLM, tracker: tracker}, nil }); err != nil { t.Fatalf("register extractor: %v", err) diff --git a/internal/modules/integration/dnd_combat_runner_test.go b/internal/modules/integration/dnd_combat_runner_test.go index 827e7363..2d73d333 100644 --- a/internal/modules/integration/dnd_combat_runner_test.go +++ b/internal/modules/integration/dnd_combat_runner_test.go @@ -157,6 +157,7 @@ func TestProductionCombatPipelineAttributesExhaustedInvalidEnumsToShapeValidatio if err != nil { t.Fatalf("MaterializeReferences() error = %v", err) } + materialized.Steps[0].ArtifactLanes[0].ExtractValidationPolicy.SemanticRejection = pipeline.SemanticRejectionRejectOutput prepared, err := pipeline.Prepare(materialized, registries, pipeline.ModuleDependencies{LLM: client}) if err != nil { t.Fatalf("Prepare() error = %v, want nil", err) diff --git a/internal/modules/integration/dnd_npc_occurrences_runner_test.go b/internal/modules/integration/dnd_npc_occurrences_runner_test.go index df3d6128..64a4d0e0 100644 --- a/internal/modules/integration/dnd_npc_occurrences_runner_test.go +++ b/internal/modules/integration/dnd_npc_occurrences_runner_test.go @@ -91,6 +91,11 @@ func TestSemanticNPCNormalizationCrossesOrderedRegistryHandoff(t *testing.T) { if err != nil || len(warnings) != 0 { t.Fatalf("MaterializeReferences() error = %v warnings = %#v", err, warnings) } + for stepIndex := range resolved.Steps { + for laneIndex := range resolved.Steps[stepIndex].ArtifactLanes { + resolved.Steps[stepIndex].ArtifactLanes[laneIndex].ExtractValidationPolicy.SemanticRejection = pipeline.SemanticRejectionRejectOutput + } + } client := &semanticNPCOccurrenceClient{} raw := []byte(`{"metadata":{"id":"semantic-session","title":"Semantic NPC session"},"segments":[{"id":1,"start":0,"end":1,"speaker":"DM","text":"Mira Thorn enters."},{"id":2,"start":1,"end":2,"speaker":"DM","text":"Mira Thorn, the Greencloak, waves."}]}`) output, err := runPreparedPipeline(t, registries, resolved, client, pipeline.RunInput{RawInput: raw}) diff --git a/internal/modules/integration/dnd_npc_registry_runner_test.go b/internal/modules/integration/dnd_npc_registry_runner_test.go index c84cf2d8..e2796dce 100644 --- a/internal/modules/integration/dnd_npc_registry_runner_test.go +++ b/internal/modules/integration/dnd_npc_registry_runner_test.go @@ -157,7 +157,9 @@ func TestProductionNPCPipelineRoutesSemanticCandidatesToDeterministicValidators( } { t.Run(test.name, func(t *testing.T) { client := &fakeNPCProductionLLMClient{rawResponses: [][]byte{test.response, test.response, test.response}} - output, err := runPreparedPipeline(t, registries, effective.ResolvedPipeline, client, pipeline.RunInput{RawInput: readNPCFixture(t)}) + resolved := effective.ResolvedPipeline + resolved.Steps[0].ArtifactLanes[0].ExtractValidationPolicy.SemanticRejection = pipeline.SemanticRejectionRejectOutput + output, err := runPreparedPipeline(t, registries, resolved, client, pipeline.RunInput{RawInput: readNPCFixture(t)}) if err != nil { t.Fatalf("Run() error = %v, want non-fatal rejected output", err) } diff --git a/internal/modules/integration/dnd_spells_runner_test.go b/internal/modules/integration/dnd_spells_runner_test.go index 6f75b5c5..c5540145 100644 --- a/internal/modules/integration/dnd_spells_runner_test.go +++ b/internal/modules/integration/dnd_spells_runner_test.go @@ -214,7 +214,9 @@ func TestProductionSpellPipelineRoutesSemanticCandidatesToDeterministicValidator } { t.Run(test.name, func(t *testing.T) { client := &fakeSpellsLLMClient{rawResponses: [][]byte{test.response, test.response, test.response}} - output, err := runPreparedPipeline(t, registries, effective.ResolvedPipeline, client, pipeline.RunInput{RawInput: readDNDSpellsFixture(t)}) + resolved := effective.ResolvedPipeline + resolved.Steps[0].ArtifactLanes[0].ExtractValidationPolicy.SemanticRejection = pipeline.SemanticRejectionRejectOutput + output, err := runPreparedPipeline(t, registries, resolved, client, pipeline.RunInput{RawInput: readDNDSpellsFixture(t)}) if err != nil { t.Fatalf("Run() error = %v, want non-fatal rejected output", err) }