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

Integrate chunk plan caching into the runner

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-18 00:20:56 +00:00
parent ebd449d847
commit 51a36efb6b
15 changed files with 590 additions and 441 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -38,19 +38,21 @@ func New() *Runner {
}
type RunInput struct {
Prepared *PreparedPipeline
SourceID string
Path string
RawInput []byte
SessionID string
RunID string
StartedAt time.Time
LLMProfiles []artifacts.LLMProfileManifest
Metadata map[string]any
Warnings []contracts.Warning
Checkpoints CheckpointRecorder
Checkpoint CheckpointLoader
Debug DebugRecorder
Prepared *PreparedPipeline
SourceID string
Path string
RawInput []byte
SessionID string
RunID string
StartedAt time.Time
LLMProfiles []artifacts.LLMProfileManifest
Metadata map[string]any
Warnings []contracts.Warning
ChunkCacheMode ChunkCacheMode
ChunkPlans ChunkPlanStore
Checkpoints CheckpointRecorder
Checkpoint CheckpointLoader
Debug DebugRecorder
// ExtractWorkers bounds run-wide extract jobs. Values less than one use a
// single worker so direct framework callers retain deterministic behavior.
ExtractWorkers int
@@ -174,19 +176,14 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
chunker := input.Prepared.chunker
attachModuleManifestMetadata(&output, "chunker", chunker)
var canonicalChunks []source.Chunk
var generatedPlan *source.ChunkPlan
var chunkWarnings []contracts.Warning
chunkCheckpoint, chunkDecision := checkpointLoader.Chunk(chunker.Key(), doc.Digest)
recordCheckpointEvent(&output, checkpointLoader, string(StageChunk), "", chunker.Key(), chunkDecision)
chunkStarted := time.Now().UTC()
chunkMode := effectiveChunkCacheMode(input.ChunkCacheMode)
if err := writeDebugTimed(debugRecorder, "chunk/input.json", debugTimedEnvelope{
Stage: string(StageChunk),
ModuleKey: chunker.Key(),
StartedAt: chunkStarted,
Payload: map[string]any{
"reused": chunkDecision.Reused,
"decision": chunkDecision,
"cache_mode": chunkMode,
"source": debugSourceDocumentEnvelope(doc),
"source_input": debugContentEnvelope(sourceInput.Content, sourceInput.MediaType, nil, nil),
"options": redactSensitiveMap(input.pipeline.Chunk.Options),
@@ -195,85 +192,28 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
}); err != nil {
return failOutput(output), fmt.Errorf("write chunk debug artifact: %w", err)
}
chunksAccepted := chunkDecision.Reused
var chunkRejection *contracts.RejectedOutput
if chunkDecision.Reused {
canonicalChunks = cloneSourceChunks(chunkCheckpoint.Chunks)
chunkWarnings = cloneWarnings(chunkCheckpoint.Warnings)
output.Warnings = append(output.Warnings, chunkWarnings...)
} else {
if err := checkpoints.ChunkRunning(chunker.Key(), doc.Digest); err != nil {
return failOutput(output), fmt.Errorf("write chunk checkpoint: %w", err)
}
chunksAccepted, chunkRejection, err = runWithRetry(ctx, input.pipeline.Chunk.Retries, func(attempt int) (bool, *contracts.RejectedOutput, error) {
attemptStarted := time.Now().UTC()
attemptPath := path.Join("chunk", fmt.Sprintf("attempt-%02d", attempt))
attemptCtx, llmScope := withDebugLLMScope(ctx, attemptPath)
terminal := newAttemptTerminalRecorder(debugRecorder, attemptPath, "chunk", llmScope, debugTimedEnvelope{Stage: string(StageChunk), ModuleKey: chunker.Key(), Attempt: attempt, StartedAt: attemptStarted})
chunkResult, err := chunker.Plan(attemptCtx, contracts.ChunkRequest{
Source: doc,
SourceInput: sourceInput.Clone(),
SessionID: sessionID,
References: CloneReferenceSet(input.pipeline.ChunkReferences.ReferenceSet),
LLMProfile: input.pipeline.Chunk.LLMProfile,
Metadata: input.Metadata,
})
if err != nil {
attemptErr := fmt.Errorf("chunk source with chunker %q: %w", chunker.Key(), err)
return false, nil, terminal.record(nil, attemptErr)
}
plan, chunks, err := validateAndMaterializeChunkPlan(doc, chunkResult.Plan)
if err != nil {
attemptErr := fmt.Errorf("validate chunk plan from chunker %q: %w", chunker.Key(), err)
payload := map[string]any{"plan": debugChunkPlanEnvelope(chunkResult.Plan), "warnings": debugWarningEnvelopes(chunkResult.Warnings)}
return false, nil, terminal.record(payload, attemptErr)
}
validationWarnings, rejection, err := r.validateChunks(attemptCtx, doc, chunker.Key(), chunks, sourceInput, sessionID, input.pipeline.ChunkReferences.ReferenceSet, input.Metadata, input.Prepared.chunkValidators, attempt, input.Debug)
attemptWarnings := append(cloneWarnings(chunkResult.Warnings), validationWarnings...)
payload := map[string]any{
"plan": debugChunkPlanEnvelope(plan),
"materialized_chunks": debugSourceChunkEnvelopes(chunks),
"warnings": debugWarningEnvelopes(attemptWarnings),
"rejection": debugRejectedOutputPtr(rejection),
}
if err != nil || rejection != nil {
return false, rejection, terminal.record(payload, err)
}
canonicalChunks = chunks
generatedPlan = &plan
chunkWarnings = attemptWarnings
if err := terminal.record(payload, nil); err != nil {
return false, nil, err
}
return true, nil, nil
})
if err != nil {
_ = checkpoints.ChunkFailed(chunker.Key(), doc.Digest, err)
return failOutput(output), err
}
if !chunksAccepted {
output.Rejected = append(output.Rejected, *chunkRejection)
if err := checkpoints.ChunkRejected(chunker.Key(), doc.Digest, *chunkRejection); err != nil {
return failOutput(output), fmt.Errorf("write chunk checkpoint: %w", err)
}
} else {
output.Warnings = append(output.Warnings, chunkWarnings...)
if err := checkpoints.ChunkSucceeded(chunker.Key(), doc.Digest, canonicalChunks, chunkWarnings); err != nil {
return failOutput(output), fmt.Errorf("write chunk checkpoint: %w", err)
}
}
chunkResult, err := r.runChunkPlan(ctx, input, doc, sourceInput, sessionID)
if err != nil {
return failOutput(output), err
}
if chunkResult.rejection != nil {
output.Rejected = append(output.Rejected, *chunkResult.rejection)
}
if chunkResult.accepted || chunkResult.lookup.Status == ChunkPlanHit {
output.Warnings = append(output.Warnings, chunkResult.warnings...)
}
chunkDebugPayload := map[string]any{
"reused": chunkDecision.Reused,
"accepted": chunksAccepted,
"materialized_chunks": debugSourceChunkEnvelopes(canonicalChunks),
"warnings": chunkWarnings,
"cache_mode": chunkMode,
"lookup": chunkResult.lookup,
"accepted": chunkResult.accepted,
"materialized_chunks": debugSourceChunkEnvelopes(chunkResult.chunks),
"warnings": chunkResult.warnings,
}
if generatedPlan != nil {
chunkDebugPayload["plan"] = debugChunkPlanEnvelope(*generatedPlan)
if chunkResult.plan != nil {
chunkDebugPayload["plan"] = debugChunkPlanEnvelope(*chunkResult.plan)
}
if chunkRejection != nil {
chunkDebugPayload["rejection"] = debugRejectedOutputEnvelope(*chunkRejection)
if chunkResult.rejection != nil {
chunkDebugPayload["rejection"] = debugRejectedOutputEnvelope(*chunkResult.rejection)
}
if err := writeDebugTimed(debugRecorder, "chunk/output.json", debugTimedEnvelope{
Stage: string(StageChunk),
@@ -284,8 +224,8 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
return failOutput(output), fmt.Errorf("write chunk debug artifact: %w", err)
}
if chunksAccepted {
laneOutput, laneErr := r.runLanes(ctx, input, checkpoints, checkpointLoader, doc, sourceInput, sessionID, canonicalChunks)
if chunkResult.accepted {
laneOutput, laneErr := r.runLanes(ctx, input, checkpoints, checkpointLoader, doc, sourceInput, sessionID, chunkResult.chunks)
mergeLaneOutput(&output, laneOutput)
if laneErr != nil {
return failOutput(output), laneErr
@@ -473,6 +413,13 @@ func validateRunInput(input RunInput) error {
if input.Prepared == nil {
return fmt.Errorf("prepared pipeline must not be nil")
}
mode := effectiveChunkCacheMode(input.ChunkCacheMode)
if err := mode.Validate(); err != nil {
return err
}
if mode != ChunkCacheBypass && input.ChunkPlans == nil {
return fmt.Errorf("chunk plan store is required for %q mode", mode)
}
return validateResolvedPipeline(input.Prepared.resolved)
}
@@ -546,10 +493,6 @@ func manifestFromPipeline(input RunInput) artifacts.RunManifest {
References: ReferenceProvenance(pipeline),
LLMProfiles: cloneLLMProfiles(input.LLMProfiles),
}
// The runner does not currently maintain a cache or idempotency key. Reference
// digests are recorded in manifest provenance and intentionally kept separate
// from source_digests.
for _, lane := range pipeline.ArtifactLanes {
laneManifest := artifacts.ArtifactLaneManifest{
ID: lane.ID,