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

Implement runner retries and raw validation

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-07 19:19:04 +00:00
parent bcedf19a08
commit cc6b050367
8 changed files with 805 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

414
internal/framework/pipeline/runner.go
View File

@@ -18,13 +18,14 @@ import (
)
type Registries struct {
Inputs *InputAdapterRegistry
Chunkers *ChunkerRegistry
Extractors *ExtractorRegistry
Mergers *MergerRegistry
Normalizers *NormalizerRegistry
Validators *ValidatorRegistry
Outputs *OutputEncoderRegistry
Inputs *InputAdapterRegistry
Chunkers *ChunkerRegistry
Extractors *ExtractorRegistry
Mergers *MergerRegistry
Normalizers *NormalizerRegistry
Validators *ValidatorRegistry
RawValidators *RawValidationRegistry
Outputs *OutputEncoderRegistry
}
type Runner struct {
@@ -103,31 +104,51 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
return failOutput(output), fmt.Errorf("build chunker %q: %w", input.Pipeline.Chunk.Module, err)
}
attachModuleManifestMetadata(&output, "chunker", chunker)
chunkResult, err := chunker.Chunk(ctx, contracts.ChunkRequest{
Source: doc,
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(input.Pipeline.ChunkReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: input.Pipeline.Chunk.LLMProfile,
Options: cloneOptions(input.Pipeline.Chunk.Options),
Metadata: input.Metadata,
var canonicalChunks []contracts.SourceChunk
var chunkWarnings []contracts.Warning
chunksAccepted, chunkRejection, err := runWithRetry(ctx, input.Pipeline.Chunk.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
chunkResult, err := chunker.Chunk(ctx, contracts.ChunkRequest{
Source: doc,
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(input.Pipeline.ChunkReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: input.Pipeline.Chunk.LLMProfile,
Options: cloneOptions(input.Pipeline.Chunk.Options),
Metadata: input.Metadata,
})
if err != nil {
return false, nil, fmt.Errorf("chunk source with chunker %q: %w", chunker.Key(), err)
}
if len(chunkResult.Chunks) == 0 {
return false, nil, fmt.Errorf("chunker %q returned no chunks", chunker.Key())
}
chunks, err := validateAndCanonicalizeChunkResult(doc, chunkResult.Chunks)
if err != nil {
return false, nil, fmt.Errorf("validate chunks from chunker %q: %w", chunker.Key(), err)
}
rejection, err := r.validateChunksRaw(ctx, doc, chunker.Key(), chunks, input.Metadata, attempt)
if err != nil || rejection != nil {
return false, rejection, err
}
canonicalChunks = chunks
chunkWarnings = cloneWarnings(chunkResult.Warnings)
return true, nil, nil
})
output.Warnings = append(output.Warnings, chunkResult.Warnings...)
if err != nil {
return failOutput(output), fmt.Errorf("chunk source with chunker %q: %w", chunker.Key(), err)
return failOutput(output), err
}
if len(chunkResult.Chunks) == 0 {
return failOutput(output), fmt.Errorf("chunker %q returned no chunks", chunker.Key())
}
canonicalChunks, err := validateAndCanonicalizeChunkResult(doc, chunkResult.Chunks)
if err != nil {
return failOutput(output), fmt.Errorf("validate chunks from chunker %q: %w", chunker.Key(), err)
if !chunksAccepted {
output.Rejected = append(output.Rejected, *chunkRejection)
} else {
output.Warnings = append(output.Warnings, chunkWarnings...)
}
for _, lane := range input.Pipeline.ArtifactLanes {
if err := r.runLane(ctx, input, doc, sourceInput, sessionID, canonicalChunks, lane, &output); err != nil {
return failOutput(output), err
if chunksAccepted {
for _, lane := range input.Pipeline.ArtifactLanes {
if err := r.runLane(ctx, input, doc, sourceInput, sessionID, canonicalChunks, lane, &output); err != nil {
return failOutput(output), err
}
}
}
@@ -183,70 +204,165 @@ func (r *Runner) runLane(ctx context.Context, input RunInput, doc *source.Source
extractOutputs := make([]contracts.ExtractOutput, 0, len(chunks))
for index := range chunks {
chunk := chunks[index]
result, err := extractor.Extract(ctx, contracts.ExtractionRequest{
var acceptedOutput contracts.ExtractOutput
var acceptedWarnings []contracts.Warning
accepted, rejection, err := runWithRetry(ctx, lane.Extract.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
result, err := extractor.Extract(ctx, contracts.ExtractionRequest{
Source: doc,
Chunk: &chunk,
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.ExtractReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Extract.LLMProfile,
Options: cloneOptions(lane.Extract.Options),
Metadata: input.Metadata,
})
if err != nil {
return false, nil, fmt.Errorf("extract lane %q chunk %q with extractor %q: %w", lane.ID, chunk.ID, extractor.Key(), err)
}
extractOutput := result.Output
extractOutput.LaneID = lane.ID
extractOutput.ExtractorKey = extractor.Key()
extractOutput.SourceID = doc.ID
extractOutput.ChunkID = chunk.ID
extractOutput.ChunkIndex = chunk.Index
extractOutput.Payload.Warnings = append(extractOutput.Payload.Warnings, result.Warnings...)
validationWarnings, rejection, err := r.validateRaw(ctx, rawValidationTarget{
stage: StageExtract,
laneID: lane.ID,
moduleKey: extractor.Key(),
source: doc,
sourceID: doc.ID,
chunkID: chunk.ID,
chunkIndex: chunk.Index,
schema: extractOutput.Schema,
payload: extractOutput.Payload,
metadata: input.Metadata,
attempt: attempt,
})
if err != nil || rejection != nil {
return false, rejection, err
}
acceptedOutput = cloneExtractOutput(extractOutput)
acceptedWarnings = append(cloneWarnings(result.Warnings), validationWarnings...)
return true, nil, nil
})
if err != nil {
return err
}
if !accepted {
output.Rejected = append(output.Rejected, *rejection)
continue
}
output.Warnings = append(output.Warnings, acceptedWarnings...)
extractOutputs = append(extractOutputs, acceptedOutput)
}
if len(extractOutputs) == 0 {
return nil
}
var acceptedMerge contracts.MergeOutput
var mergeWarnings []contracts.Warning
mergeAccepted, mergeRejection, err := runWithRetry(ctx, lane.Merge.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
mergeResult, err := merger.Merge(ctx, contracts.MergeRequest{
Source: doc,
LaneID: lane.ID,
ExtractOutputs: cloneExtractOutputs(extractOutputs),
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.MergeReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Merge.LLMProfile,
Options: cloneOptions(lane.Merge.Options),
Metadata: input.Metadata,
})
if err != nil {
return false, nil, fmt.Errorf("merge lane %q with merger %q: %w", lane.ID, merger.Key(), err)
}
mergeOutput := mergeResult.Output
mergeOutput.LaneID = lane.ID
mergeOutput.MergerKey = merger.Key()
mergeOutput.SourceID = doc.ID
mergeOutput.Payload.Warnings = append(mergeOutput.Payload.Warnings, mergeResult.Warnings...)
validationWarnings, rejection, err := r.validateRaw(ctx, rawValidationTarget{
stage: StageMerge,
laneID: lane.ID,
moduleKey: merger.Key(),
source: doc,
sourceID: doc.ID,
schema: mergeOutput.Schema,
payload: mergeOutput.Payload,
metadata: input.Metadata,
attempt: attempt,
})
if err != nil || rejection != nil {
return false, rejection, err
}
acceptedMerge = cloneMergeOutput(mergeOutput)
mergeWarnings = append(cloneWarnings(mergeResult.Warnings), validationWarnings...)
return true, nil, nil
})
if err != nil {
return err
}
if !mergeAccepted {
output.Rejected = append(output.Rejected, *mergeRejection)
return nil
}
output.Warnings = append(output.Warnings, mergeWarnings...)
var acceptedNormalize contracts.NormalizeOutput
var normalizeWarnings []contracts.Warning
normalizeAccepted, normalizeRejection, err := runWithRetry(ctx, lane.Normalize.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
normalizeResult, err := normalizer.Normalize(ctx, contracts.NormalizeRequest{
Source: doc,
Chunk: &chunk,
LaneID: lane.ID,
MergeOutput: cloneMergeOutput(acceptedMerge),
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.ExtractReferences.ReferenceSet),
References: CloneReferenceSet(lane.NormalizeReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Extract.LLMProfile,
Options: cloneOptions(lane.Extract.Options),
LLMProfile: lane.Normalize.LLMProfile,
Options: cloneOptions(lane.Normalize.Options),
Metadata: input.Metadata,
})
output.Warnings = append(output.Warnings, result.Warnings...)
if err != nil {
return fmt.Errorf("extract lane %q chunk %q with extractor %q: %w", lane.ID, chunk.ID, extractor.Key(), err)
return false, nil, fmt.Errorf("normalize lane %q with normalizer %q: %w", lane.ID, normalizer.Key(), err)
}
extractOutput := result.Output
extractOutput.LaneID = lane.ID
extractOutput.ExtractorKey = extractor.Key()
extractOutput.SourceID = doc.ID
extractOutput.ChunkID = chunk.ID
extractOutput.ChunkIndex = chunk.Index
extractOutput.Payload.Warnings = append(extractOutput.Payload.Warnings, result.Warnings...)
extractOutputs = append(extractOutputs, cloneExtractOutput(extractOutput))
}
mergeResult, err := merger.Merge(ctx, contracts.MergeRequest{
Source: doc,
LaneID: lane.ID,
ExtractOutputs: cloneExtractOutputs(extractOutputs),
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.MergeReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Merge.LLMProfile,
Options: cloneOptions(lane.Merge.Options),
Metadata: input.Metadata,
normalizeOutput := normalizeResult.Output
normalizeOutput.LaneID = lane.ID
normalizeOutput.NormalizerKey = normalizer.Key()
normalizeOutput.SourceID = doc.ID
normalizeOutput.Payload.Warnings = append(normalizeOutput.Payload.Warnings, normalizeResult.Warnings...)
validationWarnings, rejection, err := r.validateRaw(ctx, rawValidationTarget{
stage: StageNormalize,
laneID: lane.ID,
moduleKey: normalizer.Key(),
source: doc,
sourceID: doc.ID,
schema: normalizeOutput.Schema,
payload: normalizeOutput.Payload,
metadata: input.Metadata,
attempt: attempt,
})
if err != nil || rejection != nil {
return false, rejection, err
}
acceptedNormalize = cloneNormalizeOutput(normalizeOutput)
normalizeWarnings = append(cloneWarnings(normalizeResult.Warnings), validationWarnings...)
return true, nil, nil
})
output.Warnings = append(output.Warnings, mergeResult.Warnings...)
if err != nil {
return fmt.Errorf("merge lane %q with merger %q: %w", lane.ID, merger.Key(), err)
return err
}
normalizeResult, err := normalizer.Normalize(ctx, contracts.NormalizeRequest{
Source: doc,
LaneID: lane.ID,
MergeOutput: cloneMergeOutput(mergeResult.Output),
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(lane.NormalizeReferences.ReferenceSet),
LLMClient: input.LLMClient,
LLMProfile: lane.Normalize.LLMProfile,
Options: cloneOptions(lane.Normalize.Options),
Metadata: input.Metadata,
})
output.Warnings = append(output.Warnings, normalizeResult.Warnings...)
if err != nil {
return fmt.Errorf("normalize lane %q with normalizer %q: %w", lane.ID, normalizer.Key(), err)
if !normalizeAccepted {
output.Rejected = append(output.Rejected, *normalizeRejection)
return nil
}
normalizeOutput := normalizeResult.Output
normalizeOutput.LaneID = lane.ID
normalizeOutput.NormalizerKey = normalizer.Key()
normalizeOutput.SourceID = doc.ID
normalizeOutput.Payload.Warnings = append(normalizeOutput.Payload.Warnings, normalizeResult.Warnings...)
output.NormalizeOutputs = append(output.NormalizeOutputs, cloneNormalizeOutput(normalizeOutput))
output.Warnings = append(output.Warnings, normalizeWarnings...)
output.NormalizeOutputs = append(output.NormalizeOutputs, acceptedNormalize)
return nil
}
@@ -255,6 +371,146 @@ type validatorExecution struct {
binding ModuleBinding
}
type rawValidationTarget struct {
stage ModuleStage
laneID string
moduleKey string
source *source.SourceDocument
sourceID string
chunkID string
chunkIndex int
schema contracts.ResponseSchema
payload contracts.RawPayload
metadata map[string]any
attempt int
}
func runWithRetry(ctx context.Context, retries int, run func(attempt int) (bool, *contracts.RejectedOutput, error)) (bool, *contracts.RejectedOutput, error) {
attempts := 1
if retries > 0 {
attempts += retries
}
var lastRejection *contracts.RejectedOutput
for attempt := 1; attempt <= attempts; attempt++ {
if err := ctx.Err(); err != nil {
return false, nil, err
}
accepted, rejection, err := run(attempt)
if err != nil {
if ctxErr := ctx.Err(); ctxErr != nil {
return false, nil, ctxErr
}
if attempt == attempts {
return false, nil, fmt.Errorf("failed after %d attempt(s): %w", attempt, err)
}
continue
}
if accepted {
return true, nil, nil
}
if rejection != nil {
rejection.AttemptCount = attempt
lastRejection = rejection
}
if ctxErr := ctx.Err(); ctxErr != nil {
return false, nil, ctxErr
}
if attempt == attempts {
if lastRejection == nil {
lastRejection = &contracts.RejectedOutput{
ReasonCode: "raw_output_rejected",
Message: "raw output rejected",
AttemptCount: attempt,
}
}
return false, lastRejection, nil
}
}
return false, lastRejection, nil
}
func (r *Runner) validateChunksRaw(ctx context.Context, doc *source.SourceDocument, moduleKey string, chunks []contracts.SourceChunk, metadata map[string]any, attempt int) (*contracts.RejectedOutput, error) {
for _, chunk := range chunks {
_, rejection, err := r.validateRaw(ctx, rawValidationTarget{
stage: StageChunk,
moduleKey: moduleKey,
source: doc,
sourceID: doc.ID,
chunkID: chunk.ID,
chunkIndex: chunk.Index,
payload: contracts.RawPayload{
Content: append([]byte(nil), chunk.Content...),
MediaType: chunk.MediaType,
Metadata: cloneMetadata(chunk.Metadata),
},
metadata: metadata,
attempt: attempt,
})
if err != nil {
return nil, err
}
if rejection != nil {
return rejection, nil
}
}
return nil, nil
}
func (r *Runner) validateRaw(ctx context.Context, target rawValidationTarget) ([]contracts.Warning, *contracts.RejectedOutput, error) {
validators := r.registries.RawValidators.Validators(target.stage, target.moduleKey)
if len(validators) == 0 {
return nil, nil, nil
}
request := contracts.RawValidationRequest{
Stage: string(target.stage),
LaneID: target.laneID,
ModuleKey: target.moduleKey,
Source: target.source,
SourceID: target.sourceID,
ChunkID: target.chunkID,
ChunkIndex: target.chunkIndex,
Schema: target.schema,
Payload: cloneRawPayload(target.payload),
Metadata: cloneMetadata(target.metadata),
}
var warnings []contracts.Warning
for _, validator := range validators {
result, err := validator.ValidateRaw(ctx, request)
if err != nil {
return nil, nil, fmt.Errorf("validate raw %s output with validator %q: %w", target.stage, validator.Name(), err)
}
if !result.Approved {
reasonCode := strings.TrimSpace(result.ReasonCode)
if reasonCode == "" {
reasonCode = "raw_output_rejected"
}
message := strings.TrimSpace(result.Message)
if message == "" {
message = "raw output rejected"
}
return nil, &contracts.RejectedOutput{
Stage: string(target.stage),
LaneID: target.laneID,
ModuleKey: target.moduleKey,
ChunkID: target.chunkID,
ChunkIndex: target.chunkIndex,
ValidatorName: validator.Name(),
ReasonCode: reasonCode,
Message: message,
AttemptCount: target.attempt,
DiagnosticArtifactPath: result.DiagnosticArtifactPath,
}, nil
}
warnings = append(warnings, result.Warnings...)
}
return warnings, nil, nil
}
func (r *Runner) buildConfiguredValidators(lane ResolvedArtifactLane) ([]validatorExecution, error) {
validators := make([]validatorExecution, 0, len(lane.Validators))
for _, binding := range lane.Validators {