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

Integrate extraction validation retries

Browse Source
This commit is contained in:
Eric Rakestraw
2026-08-27 01:06:17 +00:00
parent 7d3434e5f6
commit 4bb4582695
16 changed files with 358 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

120
internal/framework/pipeline/runner_concurrent.go
View File

@@ -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 {