Canonicalize chunk units before extraction

This commit is contained in:
2026-07-05 13:28:09 +00:00
parent e700df82d8
commit 95218218e2
5 changed files with 173 additions and 13 deletions

View File

@@ -49,6 +49,10 @@ The `generic` chunker splits source units into ordered chunks. It validates the
source document, clones source units, assigns chunk IDs such as `chunk-000001`, source document, clones source units, assigns chunk IDs such as `chunk-000001`,
and records chunk metadata for start unit, end unit, and unit count. and records chunk metadata for start unit, end unit, and unit count.
The pipeline runner canonicalizes chunk units from the source document by ID
before extractors and mergers run. Chunker-owned context should stay in
`SourceChunk.Metadata`.
Options: Options:
- `max_units`: positive integer, default `50`; - `max_units`: positive integer, default `50`;

View File

@@ -97,6 +97,11 @@ invariants before running extractors:
- every chunk source unit must exist in the source document; - every chunk source unit must exist in the source document;
- source units inside each chunk must appear in source-document order. - source units inside each chunk must appear in source-document order.
After validation, the runner rebuilds each chunk from source-document units by
ID, preserving the chunk boundary order and cloning chunk metadata. Extractors
and downstream stages therefore see canonical source units, while
`SourceChunk.Metadata` remains the supported place for chunker-owned context.
The framework does not require complete source-unit coverage and does not reject The framework does not require complete source-unit coverage and does not reject
overlap between different chunks. Stricter policies, such as full coverage or overlap between different chunks. Stricter policies, such as full coverage or
non-overlap, belong to individual chunk modules when they are part of that non-overlap, belong to individual chunk modules when they are part of that

View File

@@ -8,53 +8,75 @@ import (
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
) )
func validateChunkResult(doc *source.SourceDocument, chunks []contracts.SourceChunk) error { func validateAndCanonicalizeChunkResult(doc *source.SourceDocument, chunks []contracts.SourceChunk) ([]contracts.SourceChunk, error) {
sourceUnitIndexes := make(map[string]int, len(doc.Units)) sourceUnitIndexes := make(map[string]int, len(doc.Units))
sourceUnits := make(map[string]source.SourceUnit, len(doc.Units))
for index, unit := range doc.Units { for index, unit := range doc.Units {
sourceUnitIndexes[unit.ID] = index sourceUnitIndexes[unit.ID] = index
sourceUnits[unit.ID] = unit
} }
canonicalChunks := make([]contracts.SourceChunk, 0, len(chunks))
seenChunkIDs := make(map[string]struct{}, len(chunks)) seenChunkIDs := make(map[string]struct{}, len(chunks))
for chunkIndex, chunk := range chunks { for chunkIndex, chunk := range chunks {
if strings.TrimSpace(chunk.ID) == "" { if strings.TrimSpace(chunk.ID) == "" {
return fmt.Errorf("chunk[%d].id must not be empty", chunkIndex) return nil, fmt.Errorf("chunk[%d].id must not be empty", chunkIndex)
} }
if _, ok := seenChunkIDs[chunk.ID]; ok { if _, ok := seenChunkIDs[chunk.ID]; ok {
return fmt.Errorf("chunk id %q is duplicated", chunk.ID) return nil, fmt.Errorf("chunk id %q is duplicated", chunk.ID)
} }
seenChunkIDs[chunk.ID] = struct{}{} seenChunkIDs[chunk.ID] = struct{}{}
if chunk.SourceID != doc.ID { if chunk.SourceID != doc.ID {
return fmt.Errorf("chunk %q source_id %q does not match source document id %q", chunk.ID, chunk.SourceID, doc.ID) return nil, fmt.Errorf("chunk %q source_id %q does not match source document id %q", chunk.ID, chunk.SourceID, doc.ID)
} }
if chunk.Index != chunkIndex { if chunk.Index != chunkIndex {
return fmt.Errorf("chunk %q index %d does not match returned order %d", chunk.ID, chunk.Index, chunkIndex) return nil, fmt.Errorf("chunk %q index %d does not match returned order %d", chunk.ID, chunk.Index, chunkIndex)
} }
if len(chunk.Units) == 0 { if len(chunk.Units) == 0 {
return fmt.Errorf("chunk %q units must not be empty", chunk.ID) return nil, fmt.Errorf("chunk %q units must not be empty", chunk.ID)
} }
seenUnitIDs := make(map[string]struct{}, len(chunk.Units)) seenUnitIDs := make(map[string]struct{}, len(chunk.Units))
previousSourceIndex := -1 previousSourceIndex := -1
canonicalUnits := make([]source.SourceUnit, 0, len(chunk.Units))
for unitIndex, unit := range chunk.Units { for unitIndex, unit := range chunk.Units {
if strings.TrimSpace(unit.ID) == "" { if strings.TrimSpace(unit.ID) == "" {
return fmt.Errorf("chunk %q unit[%d].id must not be empty", chunk.ID, unitIndex) return nil, fmt.Errorf("chunk %q unit[%d].id must not be empty", chunk.ID, unitIndex)
} }
if _, ok := seenUnitIDs[unit.ID]; ok { if _, ok := seenUnitIDs[unit.ID]; ok {
return fmt.Errorf("chunk %q repeats source unit %q", chunk.ID, unit.ID) return nil, fmt.Errorf("chunk %q repeats source unit %q", chunk.ID, unit.ID)
} }
seenUnitIDs[unit.ID] = struct{}{} seenUnitIDs[unit.ID] = struct{}{}
sourceIndex, ok := sourceUnitIndexes[unit.ID] sourceIndex, ok := sourceUnitIndexes[unit.ID]
if !ok { if !ok {
return fmt.Errorf("chunk %q source unit %q was not found in source document %q", chunk.ID, unit.ID, doc.ID) return nil, fmt.Errorf("chunk %q source unit %q was not found in source document %q", chunk.ID, unit.ID, doc.ID)
} }
if sourceIndex <= previousSourceIndex { if sourceIndex <= previousSourceIndex {
return fmt.Errorf("chunk %q source units must appear in source document order", chunk.ID) return nil, fmt.Errorf("chunk %q source units must appear in source document order", chunk.ID)
} }
previousSourceIndex = sourceIndex previousSourceIndex = sourceIndex
canonicalUnits = append(canonicalUnits, cloneSourceUnit(sourceUnits[unit.ID]))
} }
canonicalChunks = append(canonicalChunks, contracts.SourceChunk{
ID: chunk.ID,
SourceID: chunk.SourceID,
Index: chunk.Index,
Units: canonicalUnits,
Metadata: cloneMetadata(chunk.Metadata),
})
} }
return nil return canonicalChunks, nil
}
func cloneSourceUnit(unit source.SourceUnit) source.SourceUnit {
return source.SourceUnit{
ID: unit.ID,
Kind: unit.Kind,
Text: unit.Text,
Metadata: cloneMetadata(unit.Metadata),
}
} }

View File

@@ -105,13 +105,14 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
if len(chunkResult.Chunks) == 0 { if len(chunkResult.Chunks) == 0 {
return failOutput(output), fmt.Errorf("chunker %q returned no chunks", chunker.Key()) return failOutput(output), fmt.Errorf("chunker %q returned no chunks", chunker.Key())
} }
if err := validateChunkResult(doc, chunkResult.Chunks); err != nil { canonicalChunks, err := validateAndCanonicalizeChunkResult(doc, chunkResult.Chunks)
if err != nil {
return failOutput(output), fmt.Errorf("validate chunks from chunker %q: %w", chunker.Key(), err) return failOutput(output), fmt.Errorf("validate chunks from chunker %q: %w", chunker.Key(), err)
} }
nextCandidateIndex := 0 nextCandidateIndex := 0
for _, lane := range input.Pipeline.ArtifactLanes { for _, lane := range input.Pipeline.ArtifactLanes {
if err := r.runLane(ctx, input, doc, chunkResult.Chunks, lane, &output, &nextCandidateIndex); err != nil { if err := r.runLane(ctx, input, doc, canonicalChunks, lane, &output, &nextCandidateIndex); err != nil {
return failOutput(output), err return failOutput(output), err
} }
} }

View File

@@ -341,6 +341,106 @@ func TestRunAllowsPartialCoverageAndOverlappingChunks(t *testing.T) {
} }
} }
func TestRunCanonicalizesChunkUnitsBeforeExtraction(t *testing.T) {
modules := defaultRunnerModules()
modules.input.doc = sourceDocumentWithUnitMetadata()
modules.chunker.chunks = []contracts.SourceChunk{
{
ID: "chunk-0",
SourceID: "source-1",
Index: 0,
Units: []source.SourceUnit{
{
ID: "u1",
Kind: "mutated-kind",
Text: "mutated text",
Metadata: map[string]any{
"speaker": "chunker-speaker",
"note": "chunker note",
},
},
},
},
}
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: resolvedPipeline()})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
if len(output.Approved) != 1 {
t.Fatalf("len(Approved) = %d, want 1", len(output.Approved))
}
extractor := modules.extractors["extract-alpha"]
if len(extractor.requests) != 1 {
t.Fatalf("len(extractor requests) = %d, want 1", len(extractor.requests))
}
chunk := extractor.requests[0].Chunk
if chunk == nil {
t.Fatal("extractor chunk = nil, want canonical chunk")
}
if chunk.Units[0].ID != "u1" || chunk.Units[0].Kind != "source-kind" || chunk.Units[0].Text != "source text" {
t.Fatalf("chunk unit = %#v, want source document unit values", chunk.Units[0])
}
if got := chunk.Units[0].Metadata["speaker"]; got != "source-speaker" {
t.Fatalf("chunk unit metadata = %#v, want source document metadata", chunk.Units[0].Metadata)
}
if got := chunk.Units[0].Metadata["topic"]; got != "source-topic" {
t.Fatalf("chunk unit metadata = %#v, want cloned source document metadata", chunk.Units[0].Metadata)
}
modules.input.doc.Units[0].Kind = "changed-kind"
modules.input.doc.Units[0].Text = "changed text"
modules.input.doc.Units[0].Metadata["speaker"] = "changed-speaker"
if chunk.Units[0].Kind != "source-kind" || chunk.Units[0].Text != "source text" || chunk.Units[0].Metadata["speaker"] != "source-speaker" {
t.Fatalf("chunk unit changed after source mutation: %#v", chunk.Units[0])
}
}
func TestRunPreservesChunkMetadataDuringCanonicalization(t *testing.T) {
modules := defaultRunnerModules()
modules.chunker.chunks = []contracts.SourceChunk{
{
ID: "chunk-0",
SourceID: "source-1",
Index: 0,
Units: []source.SourceUnit{
unitWithID("u1"),
},
Metadata: map[string]any{
"scene_title": "Original scene",
"boundary_note": "Chunker note",
},
},
}
_, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: resolvedPipeline()})
if err != nil {
t.Fatalf("Run() error = %v, want nil", err)
}
extractor := modules.extractors["extract-alpha"]
if len(extractor.requests) != 1 || extractor.requests[0].Chunk == nil {
t.Fatalf("extractor requests = %#v, want one canonical chunk", extractor.requests)
}
if got := extractor.requests[0].Chunk.Metadata["scene_title"]; got != "Original scene" {
t.Fatalf("chunk metadata = %#v, want chunker metadata", extractor.requests[0].Chunk.Metadata)
}
if got := extractor.requests[0].Chunk.Metadata["boundary_note"]; got != "Chunker note" {
t.Fatalf("chunk metadata = %#v, want chunker metadata", extractor.requests[0].Chunk.Metadata)
}
modules.chunker.chunks[0].Metadata["scene_title"] = "changed"
modules.chunker.chunks[0].Metadata["boundary_note"] = "changed"
if got := extractor.requests[0].Chunk.Metadata["scene_title"]; got != "Original scene" {
t.Fatalf("chunk metadata aliased to chunker map: %#v", extractor.requests[0].Chunk.Metadata)
}
if got := extractor.requests[0].Chunk.Metadata["boundary_note"]; got != "Chunker note" {
t.Fatalf("chunk metadata aliased to chunker map: %#v", extractor.requests[0].Chunk.Metadata)
}
}
func TestRunExecutesChunksAndPassesChunkAndLLMClient(t *testing.T) { func TestRunExecutesChunksAndPassesChunkAndLLMClient(t *testing.T) {
modules := defaultRunnerModules() modules := defaultRunnerModules()
llmClient := fakeLLMClient{} llmClient := fakeLLMClient{}
@@ -1334,6 +1434,34 @@ func validSourceDocument() *source.SourceDocument {
} }
} }
func sourceDocumentWithUnitMetadata() *source.SourceDocument {
return &source.SourceDocument{
ID: "source-1",
Kind: "document",
Format: "text/plain",
Digest: "sha256:source",
Units: []source.SourceUnit{
{
ID: "u1",
Kind: "source-kind",
Text: "source text",
Metadata: map[string]any{
"speaker": "source-speaker",
"topic": "source-topic",
},
},
{
ID: "u2",
Kind: "source-kind",
Text: "second source text",
Metadata: map[string]any{
"speaker": "source-speaker-2",
},
},
},
}
}
func sourceChunkWithID(id string, index int) contracts.SourceChunk { func sourceChunkWithID(id string, index int) contracts.SourceChunk {
return contracts.SourceChunk{ return contracts.SourceChunk{
ID: id, ID: id,