Canonicalize chunk units before extraction
This commit is contained in:
@@ -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`,
|
||||
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:
|
||||
|
||||
- `max_units`: positive integer, default `50`;
|
||||
|
||||
@@ -97,6 +97,11 @@ invariants before running extractors:
|
||||
- every chunk source unit must exist in the source document;
|
||||
- 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
|
||||
overlap between different chunks. Stricter policies, such as full coverage or
|
||||
non-overlap, belong to individual chunk modules when they are part of that
|
||||
|
||||
@@ -8,53 +8,75 @@ import (
|
||||
"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))
|
||||
sourceUnits := make(map[string]source.SourceUnit, len(doc.Units))
|
||||
for index, unit := range doc.Units {
|
||||
sourceUnitIndexes[unit.ID] = index
|
||||
sourceUnits[unit.ID] = unit
|
||||
}
|
||||
|
||||
canonicalChunks := make([]contracts.SourceChunk, 0, len(chunks))
|
||||
seenChunkIDs := make(map[string]struct{}, len(chunks))
|
||||
for chunkIndex, chunk := range chunks {
|
||||
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 {
|
||||
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{}{}
|
||||
|
||||
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 {
|
||||
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 {
|
||||
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))
|
||||
previousSourceIndex := -1
|
||||
canonicalUnits := make([]source.SourceUnit, 0, len(chunk.Units))
|
||||
for unitIndex, unit := range chunk.Units {
|
||||
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 {
|
||||
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{}{}
|
||||
|
||||
sourceIndex, ok := sourceUnitIndexes[unit.ID]
|
||||
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 {
|
||||
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
|
||||
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),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -105,13 +105,14 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (RunOutput, error) {
|
||||
if len(chunkResult.Chunks) == 0 {
|
||||
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)
|
||||
}
|
||||
|
||||
nextCandidateIndex := 0
|
||||
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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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) {
|
||||
modules := defaultRunnerModules()
|
||||
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 {
|
||||
return contracts.SourceChunk{
|
||||
ID: id,
|
||||
|
||||
Reference in New Issue
Block a user