Generate and materialize canonical chunk plans
This commit is contained in:
@@ -351,7 +351,7 @@ tested; current modules and the runner still behave as before.
|
||||
|
||||
## Stage 2: Replace the chunk operation with plan generation
|
||||
|
||||
**Status:** Not started
|
||||
**Status:** Complete
|
||||
|
||||
### Objective
|
||||
|
||||
|
||||
@@ -341,7 +341,7 @@ func TestProductionLLMCallersShareScheduledClient(t *testing.T) {
|
||||
started.Done()
|
||||
chunker, err := scenes.New(client, scenes.Options{})
|
||||
if err == nil {
|
||||
_, err = chunker.Chunk(context.Background(), contracts.ChunkRequest{Source: doc})
|
||||
_, err = chunker.Plan(context.Background(), contracts.ChunkRequest{Source: doc})
|
||||
}
|
||||
errs <- err
|
||||
}()
|
||||
|
||||
@@ -3742,11 +3742,9 @@ func (fakeRunChunker) ReferenceSlots() []contracts.ReferenceSlot {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (fakeRunChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
return contracts.ChunkResult{
|
||||
Chunks: []source.Chunk{
|
||||
{ID: "chunk-1", SourceID: req.Source.ID, Index: 0, Ref: source.SourceRef{SourceID: req.Source.ID, StartUnitID: 1, EndUnitID: 1}, Content: []byte(`{"units":[1]}`), MediaType: "application/json", Units: req.Source.Units},
|
||||
},
|
||||
func (fakeRunChunker) Plan(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
return contracts.ChunkPlanResult{
|
||||
Plan: source.ChunkPlan{SourceDigest: req.Source.Digest, Ranges: []source.ChunkRange{{StartUnitID: 1, EndUnitID: 1}}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -135,7 +135,7 @@ func TestResolveCanBindSceneChunkerFromCatalog(t *testing.T) {
|
||||
Key: "dnd/scenes",
|
||||
Stage: pipeline.StageChunk,
|
||||
Requires: []string{"source.transcript"},
|
||||
Provides: []string{"chunks", "chunks.scenes"},
|
||||
Provides: []string{"chunks"},
|
||||
})
|
||||
|
||||
effective, err := cfg.Resolve(ResolveInput{PipelineID: "example", Catalog: catalog})
|
||||
|
||||
@@ -146,15 +146,15 @@ type ChunkRequest struct {
|
||||
Metadata map[string]any `json:"metadata,omitempty"`
|
||||
}
|
||||
|
||||
type ChunkResult struct {
|
||||
Chunks []source.Chunk `json:"chunks"`
|
||||
Warnings []Warning `json:"warnings,omitempty"`
|
||||
type ChunkPlanResult struct {
|
||||
Plan source.ChunkPlan `json:"plan"`
|
||||
Warnings []Warning `json:"warnings,omitempty"`
|
||||
}
|
||||
|
||||
type Chunker interface {
|
||||
Key() string
|
||||
ReferenceSlots() []ReferenceSlot
|
||||
Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error)
|
||||
Plan(ctx context.Context, req ChunkRequest) (ChunkPlanResult, error)
|
||||
}
|
||||
|
||||
const (
|
||||
|
||||
@@ -45,7 +45,7 @@ func TestFakeExtractorReturnsTypedOutput(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestFakeChunkerReturnsSourceChunks(t *testing.T) {
|
||||
func TestFakeChunkerReturnsSourcePlan(t *testing.T) {
|
||||
doc := &source.SourceDocument{
|
||||
ID: "source-1",
|
||||
Kind: "document",
|
||||
@@ -57,36 +57,19 @@ func TestFakeChunkerReturnsSourceChunks(t *testing.T) {
|
||||
}
|
||||
chunker := fakeChunker{key: "generic-chunker"}
|
||||
|
||||
result, err := chunker.Chunk(context.Background(), ChunkRequest{Source: doc})
|
||||
result, err := chunker.Plan(context.Background(), ChunkRequest{Source: doc})
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
t.Fatalf("Plan() error = %v, want nil", err)
|
||||
}
|
||||
|
||||
if chunker.Key() != "generic-chunker" {
|
||||
t.Fatalf("Key() = %q, want generic-chunker", chunker.Key())
|
||||
}
|
||||
if len(result.Chunks) != 1 {
|
||||
t.Fatalf("len(Chunks) = %d, want 1", len(result.Chunks))
|
||||
if len(result.Plan.Ranges) != 1 {
|
||||
t.Fatalf("len(Ranges) = %d, want 1", len(result.Plan.Ranges))
|
||||
}
|
||||
|
||||
chunk := result.Chunks[0]
|
||||
if chunk.ID != "source-1:chunk:0" {
|
||||
t.Fatalf("source.Chunk.ID = %q, want source-1:chunk:0", chunk.ID)
|
||||
}
|
||||
if chunk.SourceID != doc.ID {
|
||||
t.Fatalf("source.Chunk.SourceID = %q, want %q", chunk.SourceID, doc.ID)
|
||||
}
|
||||
if chunk.Index != 0 {
|
||||
t.Fatalf("source.Chunk.Index = %d, want 0", chunk.Index)
|
||||
}
|
||||
if chunk.Ref.StartUnitID != 1 || chunk.Ref.EndUnitID != 1 {
|
||||
t.Fatalf("source.Chunk.Ref = %#v, want source-1:1-1", chunk.Ref)
|
||||
}
|
||||
if chunk.MediaType != "application/json" || string(chunk.Content) != `{"units":[{"id":1,"kind":"section","text":"Source text."}]}` {
|
||||
t.Fatalf("source.Chunk payload = %q %s, want JSON units", chunk.MediaType, chunk.Content)
|
||||
}
|
||||
if len(chunk.Units) != 1 {
|
||||
t.Fatalf("len(source.Chunk.Units) = %d, want 1", len(chunk.Units))
|
||||
if result.Plan.SourceDigest != doc.Digest || result.Plan.Ranges[0].StartUnitID != 1 || result.Plan.Ranges[0].EndUnitID != 1 {
|
||||
t.Fatalf("Plan = %#v, want source digest and unit range", result.Plan)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -102,8 +85,8 @@ func TestFakeChunkerReceivesPerRunContext(t *testing.T) {
|
||||
}
|
||||
chunker := &recordingChunker{key: "llm-chunker"}
|
||||
|
||||
if _, err := chunker.Chunk(context.Background(), ChunkRequest{Source: doc, SessionID: "session", LLMProfile: "profile"}); err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
if _, err := chunker.Plan(context.Background(), ChunkRequest{Source: doc, SessionID: "session", LLMProfile: "profile"}); err != nil {
|
||||
t.Fatalf("Plan() error = %v, want nil", err)
|
||||
}
|
||||
if chunker.request.SessionID != "session" || chunker.request.LLMProfile != "profile" {
|
||||
t.Fatalf("ChunkRequest = %#v, want per-run session and profile", chunker.request)
|
||||
@@ -446,22 +429,14 @@ func (chunker fakeChunker) ReferenceSlots() []ReferenceSlot {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (chunker fakeChunker) Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error) {
|
||||
return ChunkResult{
|
||||
Chunks: []source.Chunk{
|
||||
{
|
||||
ID: req.Source.ID + ":chunk:0",
|
||||
SourceID: req.Source.ID,
|
||||
Index: 0,
|
||||
Ref: source.SourceRef{
|
||||
SourceID: req.Source.ID,
|
||||
StartUnitID: req.Source.Units[0].ID,
|
||||
EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID,
|
||||
},
|
||||
Content: []byte(`{"units":[{"id":1,"kind":"section","text":"Source text."}]}`),
|
||||
MediaType: "application/json",
|
||||
Units: append([]source.SourceUnit(nil), req.Source.Units...),
|
||||
},
|
||||
func (chunker fakeChunker) Plan(ctx context.Context, req ChunkRequest) (ChunkPlanResult, error) {
|
||||
return ChunkPlanResult{
|
||||
Plan: source.ChunkPlan{
|
||||
SourceDigest: req.Source.Digest,
|
||||
Ranges: []source.ChunkRange{{
|
||||
StartUnitID: req.Source.Units[0].ID,
|
||||
EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID,
|
||||
}},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
@@ -479,9 +454,9 @@ func (chunker *recordingChunker) ReferenceSlots() []ReferenceSlot {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (chunker *recordingChunker) Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error) {
|
||||
func (chunker *recordingChunker) Plan(ctx context.Context, req ChunkRequest) (ChunkPlanResult, error) {
|
||||
chunker.request = req
|
||||
return fakeChunker{key: chunker.key}.Chunk(ctx, req)
|
||||
return fakeChunker{key: chunker.key}.Plan(ctx, req)
|
||||
}
|
||||
|
||||
type fakeExtractor struct {
|
||||
|
||||
@@ -10,27 +10,29 @@ import (
|
||||
|
||||
func TestChunkCanonicalizationAndClonePreserveAnnotationScopes(t *testing.T) {
|
||||
doc := validSourceDocument()
|
||||
chunk := source.Chunk{
|
||||
ID: "chunk-1", SourceID: doc.ID, Index: 0,
|
||||
Ref: doc.Units[0].Ref, Content: []byte(`{"units":[1]}`), MediaType: "application/json", Units: doc.Units[:1],
|
||||
Annotations: source.ChunkAnnotations{"shared": json.RawMessage(` {"range": true} `)},
|
||||
PlanAnnotations: source.ChunkAnnotations{"shared": json.RawMessage(` {"plan": true} `)},
|
||||
plan := source.ChunkPlan{
|
||||
SourceDigest: doc.Digest,
|
||||
Ranges: []source.ChunkRange{{
|
||||
StartUnitID: doc.Units[0].ID, EndUnitID: doc.Units[0].ID,
|
||||
Annotations: source.ChunkAnnotations{"shared": json.RawMessage(` {"range": true} `)},
|
||||
}},
|
||||
Annotations: source.ChunkAnnotations{"shared": json.RawMessage(` {"plan": true} `)},
|
||||
}
|
||||
canonical, err := validateAndCanonicalizeChunkResult(doc, []source.Chunk{chunk})
|
||||
canonical, chunks, err := validateAndMaterializeChunkPlan(doc, plan)
|
||||
if err != nil {
|
||||
t.Fatalf("validateAndCanonicalizeChunkResult() error = %v", err)
|
||||
t.Fatalf("validateAndMaterializeChunkPlan() error = %v", err)
|
||||
}
|
||||
if got := string(canonical[0].Annotations["shared"]); got != `{"range":true}` {
|
||||
if got := string(canonical.Ranges[0].Annotations["shared"]); got != `{"range":true}` {
|
||||
t.Fatalf("range annotation = %q", got)
|
||||
}
|
||||
if got := string(canonical[0].PlanAnnotations["shared"]); got != `{"plan":true}` {
|
||||
if got := string(canonical.Annotations["shared"]); got != `{"plan":true}` {
|
||||
t.Fatalf("plan annotation = %q", got)
|
||||
}
|
||||
|
||||
cloned := cloneSourceChunk(canonical[0])
|
||||
cloned := cloneSourceChunk(chunks[0])
|
||||
cloned.Annotations["shared"][0] = '['
|
||||
cloned.PlanAnnotations["shared"][0] = '['
|
||||
if string(canonical[0].Annotations["shared"]) != `{"range":true}` || string(canonical[0].PlanAnnotations["shared"]) != `{"plan":true}` {
|
||||
if string(chunks[0].Annotations["shared"]) != `{"range":true}` || string(chunks[0].PlanAnnotations["shared"]) != `{"plan":true}` {
|
||||
t.Fatal("cloneSourceChunk() shares annotation bytes")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -2,110 +2,23 @@ package pipeline
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
)
|
||||
|
||||
func validateAndCanonicalizeChunkResult(doc *source.SourceDocument, chunks []source.Chunk) ([]source.Chunk, error) {
|
||||
if len(chunks) == 0 {
|
||||
return nil, fmt.Errorf("chunks must not be empty")
|
||||
func validateAndMaterializeChunkPlan(doc *source.SourceDocument, plan source.ChunkPlan) (source.ChunkPlan, []source.Chunk, error) {
|
||||
canonical, err := source.CanonicalizeChunkPlan(plan)
|
||||
if err != nil {
|
||||
return source.ChunkPlan{}, nil, fmt.Errorf("canonicalize chunk plan: %w", err)
|
||||
}
|
||||
|
||||
sourceUnitIndexes := make(map[int]int, len(doc.Units))
|
||||
sourceUnits := make(map[int]source.SourceUnit, len(doc.Units))
|
||||
for index, unit := range doc.Units {
|
||||
sourceUnitIndexes[unit.ID] = index
|
||||
sourceUnits[unit.ID] = unit
|
||||
if err := source.ValidateChunkPlan(doc, canonical); err != nil {
|
||||
return source.ChunkPlan{}, nil, err
|
||||
}
|
||||
|
||||
canonicalChunks := make([]source.Chunk, 0, len(chunks))
|
||||
seenChunkIDs := make(map[string]struct{}, len(chunks))
|
||||
for chunkIndex, chunk := range chunks {
|
||||
if strings.TrimSpace(chunk.ID) == "" {
|
||||
return nil, fmt.Errorf("chunk[%d].id must not be empty", chunkIndex)
|
||||
}
|
||||
if _, ok := seenChunkIDs[chunk.ID]; ok {
|
||||
return nil, fmt.Errorf("chunk id %q is duplicated", chunk.ID)
|
||||
}
|
||||
seenChunkIDs[chunk.ID] = struct{}{}
|
||||
|
||||
if 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 nil, fmt.Errorf("chunk %q index %d does not match returned order %d", chunk.ID, chunk.Index, chunkIndex)
|
||||
}
|
||||
if len(chunk.Units) == 0 {
|
||||
return nil, fmt.Errorf("chunk %q units must not be empty", chunk.ID)
|
||||
}
|
||||
if len(chunk.Content) == 0 {
|
||||
return nil, fmt.Errorf("chunk %q content must not be empty", chunk.ID)
|
||||
}
|
||||
if strings.TrimSpace(chunk.MediaType) == "" {
|
||||
return nil, fmt.Errorf("chunk %q media_type must not be empty", chunk.ID)
|
||||
}
|
||||
if err := source.ValidateRef(doc, chunk.Ref); err != nil {
|
||||
return nil, fmt.Errorf("chunk %q ref: %w", chunk.ID, err)
|
||||
}
|
||||
annotations, err := source.CanonicalizeChunkAnnotations(chunk.Annotations)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("chunk %q annotations: %w", chunk.ID, err)
|
||||
}
|
||||
planAnnotations, err := source.CanonicalizeChunkAnnotations(chunk.PlanAnnotations)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("chunk %q plan_annotations: %w", chunk.ID, err)
|
||||
}
|
||||
|
||||
seenUnitIDs := make(map[int]struct{}, len(chunk.Units))
|
||||
previousSourceIndex := -1
|
||||
canonicalUnits := make([]source.SourceUnit, 0, len(chunk.Units))
|
||||
for unitIndex, unit := range chunk.Units {
|
||||
if unit.ID <= 0 {
|
||||
return nil, fmt.Errorf("chunk %q unit[%d].id must be positive", chunk.ID, unitIndex)
|
||||
}
|
||||
if _, ok := seenUnitIDs[unit.ID]; ok {
|
||||
return nil, fmt.Errorf("chunk %q repeats source unit %d", chunk.ID, unit.ID)
|
||||
}
|
||||
seenUnitIDs[unit.ID] = struct{}{}
|
||||
|
||||
sourceIndex, ok := sourceUnitIndexes[unit.ID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("chunk %q source unit %d was not found in source document %q", chunk.ID, unit.ID, doc.ID)
|
||||
}
|
||||
if previousSourceIndex >= 0 && sourceIndex != previousSourceIndex+1 {
|
||||
return nil, fmt.Errorf("chunk %q source units must form a contiguous range in source document order", chunk.ID)
|
||||
}
|
||||
if unit.Ref != sourceUnits[unit.ID].Ref {
|
||||
return nil, fmt.Errorf("chunk %q source unit %d ref does not match source document", chunk.ID, unit.ID)
|
||||
}
|
||||
previousSourceIndex = sourceIndex
|
||||
canonicalUnits = append(canonicalUnits, cloneSourceUnit(sourceUnits[unit.ID]))
|
||||
}
|
||||
expectedRef := source.SourceRef{
|
||||
SourceID: doc.ID,
|
||||
StartUnitID: canonicalUnits[0].Ref.StartUnitID,
|
||||
EndUnitID: canonicalUnits[len(canonicalUnits)-1].Ref.EndUnitID,
|
||||
}
|
||||
if chunk.Ref != expectedRef {
|
||||
return nil, fmt.Errorf("chunk %q ref %#v does not match unit span %#v", chunk.ID, chunk.Ref, expectedRef)
|
||||
}
|
||||
|
||||
canonicalChunks = append(canonicalChunks, source.Chunk{
|
||||
ID: chunk.ID,
|
||||
SourceID: chunk.SourceID,
|
||||
Index: chunk.Index,
|
||||
Ref: expectedRef,
|
||||
Content: append([]byte(nil), chunk.Content...),
|
||||
MediaType: chunk.MediaType,
|
||||
Units: canonicalUnits,
|
||||
Metadata: cloneMetadata(chunk.Metadata),
|
||||
Annotations: annotations,
|
||||
PlanAnnotations: planAnnotations,
|
||||
})
|
||||
chunks, err := source.MaterializeChunkPlan(doc, canonical)
|
||||
if err != nil {
|
||||
return source.ChunkPlan{}, nil, err
|
||||
}
|
||||
|
||||
return canonicalChunks, nil
|
||||
return canonical, chunks, nil
|
||||
}
|
||||
|
||||
func cloneSourceUnit(unit source.SourceUnit) source.SourceUnit {
|
||||
|
||||
@@ -334,8 +334,8 @@ func (chunker registryChunker) ReferenceSlots() []contracts.ReferenceSlot {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (chunker registryChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
return contracts.ChunkResult{}, nil
|
||||
func (chunker registryChunker) Plan(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
return contracts.ChunkPlanResult{}, nil
|
||||
}
|
||||
|
||||
type registryOutputEncoder struct {
|
||||
|
||||
@@ -111,6 +111,18 @@ type debugSourceChunk struct {
|
||||
PlanAnnotations source.ChunkAnnotations `json:"plan_annotations,omitempty"`
|
||||
}
|
||||
|
||||
type debugChunkPlan struct {
|
||||
SourceDigest string `json:"source_digest,omitempty"`
|
||||
Ranges []debugChunkRange `json:"ranges,omitempty"`
|
||||
Annotations map[string]any `json:"annotations,omitempty"`
|
||||
}
|
||||
|
||||
type debugChunkRange struct {
|
||||
StartUnitID int `json:"start_unit_id"`
|
||||
EndUnitID int `json:"end_unit_id"`
|
||||
Annotations map[string]any `json:"annotations,omitempty"`
|
||||
}
|
||||
|
||||
type debugSerializedOutput struct {
|
||||
LaneID string `json:"lane_id"`
|
||||
NormalizerKey string `json:"normalizer_key"`
|
||||
@@ -489,6 +501,37 @@ func debugSourceChunkEnvelopes(chunks []source.Chunk) []debugSourceChunk {
|
||||
return out
|
||||
}
|
||||
|
||||
func debugChunkPlanEnvelope(plan source.ChunkPlan) debugChunkPlan {
|
||||
envelope := debugChunkPlan{
|
||||
SourceDigest: plan.SourceDigest,
|
||||
Ranges: make([]debugChunkRange, len(plan.Ranges)),
|
||||
Annotations: debugChunkAnnotations(plan.Annotations),
|
||||
}
|
||||
for i, chunkRange := range plan.Ranges {
|
||||
envelope.Ranges[i] = debugChunkRange{
|
||||
StartUnitID: chunkRange.StartUnitID,
|
||||
EndUnitID: chunkRange.EndUnitID,
|
||||
Annotations: debugChunkAnnotations(chunkRange.Annotations),
|
||||
}
|
||||
}
|
||||
return envelope
|
||||
}
|
||||
|
||||
func debugChunkAnnotations(annotations source.ChunkAnnotations) map[string]any {
|
||||
if len(annotations) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]any, len(annotations))
|
||||
for namespace, raw := range annotations {
|
||||
if json.Valid(raw) {
|
||||
out[namespace] = append(json.RawMessage(nil), raw...)
|
||||
} else {
|
||||
out[namespace] = string(raw)
|
||||
}
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func debugSerializedOutputEnvelope(output contracts.SerializedOutput) debugSerializedOutput {
|
||||
schema := contracts.CloneArtifactSchema(output.Artifact.Schema)
|
||||
digest := contracts.DigestArtifactSchema(schema)
|
||||
|
||||
@@ -175,6 +175,7 @@ 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)
|
||||
@@ -209,7 +210,7 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
|
||||
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.Chunk(attemptCtx, contracts.ChunkRequest{
|
||||
chunkResult, err := chunker.Plan(attemptCtx, contracts.ChunkRequest{
|
||||
Source: doc,
|
||||
SourceInput: sourceInput.Clone(),
|
||||
SessionID: sessionID,
|
||||
@@ -221,27 +222,25 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
|
||||
attemptErr := fmt.Errorf("chunk source with chunker %q: %w", chunker.Key(), err)
|
||||
return false, nil, terminal.record(nil, attemptErr)
|
||||
}
|
||||
if len(chunkResult.Chunks) == 0 {
|
||||
attemptErr := fmt.Errorf("chunker %q returned no chunks", chunker.Key())
|
||||
return false, nil, terminal.record(nil, attemptErr)
|
||||
}
|
||||
chunks, err := validateAndCanonicalizeChunkResult(doc, chunkResult.Chunks)
|
||||
plan, chunks, err := validateAndMaterializeChunkPlan(doc, chunkResult.Plan)
|
||||
if err != nil {
|
||||
attemptErr := fmt.Errorf("validate chunks from chunker %q: %w", chunker.Key(), err)
|
||||
payload := map[string]any{"warnings": debugWarningEnvelopes(chunkResult.Warnings)}
|
||||
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{
|
||||
"chunks": debugSourceChunkEnvelopes(chunks),
|
||||
"warnings": debugWarningEnvelopes(attemptWarnings),
|
||||
"rejection": debugRejectedOutputPtr(rejection),
|
||||
"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
|
||||
@@ -265,10 +264,13 @@ func (r *Runner) Run(ctx context.Context, input RunInput) (output RunOutput, err
|
||||
}
|
||||
}
|
||||
chunkDebugPayload := map[string]any{
|
||||
"reused": chunkDecision.Reused,
|
||||
"accepted": chunksAccepted,
|
||||
"chunks": debugSourceChunkEnvelopes(canonicalChunks),
|
||||
"warnings": chunkWarnings,
|
||||
"reused": chunkDecision.Reused,
|
||||
"accepted": chunksAccepted,
|
||||
"materialized_chunks": debugSourceChunkEnvelopes(canonicalChunks),
|
||||
"warnings": chunkWarnings,
|
||||
}
|
||||
if generatedPlan != nil {
|
||||
chunkDebugPayload["plan"] = debugChunkPlanEnvelope(*generatedPlan)
|
||||
}
|
||||
if chunkRejection != nil {
|
||||
chunkDebugPayload["rejection"] = debugRejectedOutputEnvelope(*chunkRejection)
|
||||
|
||||
@@ -11,7 +11,6 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
frameworkllm "gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
||||
)
|
||||
@@ -51,12 +50,13 @@ func preparedConcurrentPipeline(t *testing.T, chunkCount int) *PreparedPipeline
|
||||
if err != nil {
|
||||
t.Fatalf("Prepare() error = %v", err)
|
||||
}
|
||||
doc := typedTestDocument()
|
||||
chunks := make([]source.Chunk, chunkCount)
|
||||
for i := range chunks {
|
||||
chunks[i] = source.Chunk{ID: fmt.Sprintf("chunk-%d", i+1), SourceID: doc.ID, Index: i, Ref: doc.Units[0].Ref, Content: []byte(fmt.Sprintf(`{"chunk":%d}`, i+1)), MediaType: "application/json", Units: []source.SourceUnit{doc.Units[0]}}
|
||||
doc := typedTestDocumentWithUnits(chunkCount)
|
||||
if adapter, ok := prepared.input.(*typedTestInput); ok {
|
||||
adapter.doc = doc
|
||||
} else {
|
||||
t.Fatalf("prepared input = %T, want *typedTestInput", prepared.input)
|
||||
}
|
||||
prepared.chunker = &typedTestChunker{key: "typed/chunk", chunks: chunks}
|
||||
prepared.chunker = &typedTestChunker{key: "typed/chunk", plan: typedTestPlan(doc)}
|
||||
return prepared
|
||||
}
|
||||
|
||||
|
||||
@@ -55,8 +55,8 @@ func cloneCheckpointArtifacts(values []CheckpointArtifact) []CheckpointArtifact
|
||||
func TestRunnerPassesIndependentAnnotationScopesToExtractors(t *testing.T) {
|
||||
prepared := preparedAttemptDebugPipeline(t)
|
||||
chunker := prepared.chunker.(*typedTestChunker)
|
||||
chunker.chunks[0].Annotations = source.ChunkAnnotations{"same": json.RawMessage(`{"range":1}`)}
|
||||
chunker.chunks[0].PlanAnnotations = source.ChunkAnnotations{"same": json.RawMessage(`{"plan":2}`)}
|
||||
chunker.plan.Ranges[0].Annotations = source.ChunkAnnotations{"same": json.RawMessage(`{"range":1}`)}
|
||||
chunker.plan.Annotations = source.ChunkAnnotations{"same": json.RawMessage(`{"plan":2}`)}
|
||||
|
||||
var gotRange, gotPlan string
|
||||
installExtractOperation(prepared, 0, func(_ context.Context, request contracts.TypedExtractionRequest) (erasedTypedResult, error) {
|
||||
@@ -72,7 +72,7 @@ func TestRunnerPassesIndependentAnnotationScopesToExtractors(t *testing.T) {
|
||||
if gotRange != `{"range":1}` || gotPlan != `{"plan":2}` {
|
||||
t.Fatalf("extract annotations = %q / %q", gotRange, gotPlan)
|
||||
}
|
||||
if string(chunker.chunks[0].Annotations["same"]) != `{"range":1}` || string(chunker.chunks[0].PlanAnnotations["same"]) != `{"plan":2}` {
|
||||
if string(chunker.plan.Ranges[0].Annotations["same"]) != `{"range":1}` || string(chunker.plan.Annotations["same"]) != `{"plan":2}` {
|
||||
t.Fatal("extract request shares annotation bytes with chunker output")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -15,7 +15,7 @@ import (
|
||||
|
||||
type terminalChunker struct {
|
||||
key string
|
||||
chunks []source.Chunk
|
||||
plan source.ChunkPlan
|
||||
warnings []contracts.Warning
|
||||
err error
|
||||
calls *int
|
||||
@@ -25,11 +25,11 @@ func (c terminalChunker) Key() string { return c.key }
|
||||
|
||||
func (terminalChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil }
|
||||
|
||||
func (c terminalChunker) Chunk(context.Context, contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
func (c terminalChunker) Plan(context.Context, contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
if c.calls != nil {
|
||||
(*c.calls)++
|
||||
}
|
||||
return contracts.ChunkResult{Chunks: cloneSourceChunks(c.chunks), Warnings: cloneWarnings(c.warnings)}, c.err
|
||||
return contracts.ChunkPlanResult{Plan: source.CloneChunkPlan(c.plan), Warnings: cloneWarnings(c.warnings)}, c.err
|
||||
}
|
||||
|
||||
type terminalChunkValidator struct {
|
||||
@@ -37,6 +37,19 @@ type terminalChunkValidator struct {
|
||||
err error
|
||||
}
|
||||
|
||||
type observingChunkValidator struct {
|
||||
request contracts.ChunkValidationRequest
|
||||
}
|
||||
|
||||
func (*observingChunkValidator) Name() string { return "observing/chunk-validator" }
|
||||
func (*observingChunkValidator) ExecutionClass() contracts.ExecutionClass {
|
||||
return contracts.ExecutionClassDeterministic
|
||||
}
|
||||
func (v *observingChunkValidator) Validate(_ context.Context, request contracts.ChunkValidationRequest) (contracts.ValidationResult, error) {
|
||||
v.request = request
|
||||
return contracts.ValidationResult{Approved: true}, nil
|
||||
}
|
||||
|
||||
func (terminalChunkValidator) Name() string { return "terminal/chunk-validator" }
|
||||
|
||||
func (terminalChunkValidator) ExecutionClass() contracts.ExecutionClass {
|
||||
@@ -71,14 +84,14 @@ func assertAttemptEnvelopeSequence(t *testing.T, debug *capturedDebugRecorder, p
|
||||
}
|
||||
}
|
||||
|
||||
func preparedTerminalDebugPipeline(t *testing.T) (*PreparedPipeline, []source.Chunk) {
|
||||
func preparedTerminalDebugPipeline(t *testing.T) (*PreparedPipeline, source.ChunkPlan) {
|
||||
t.Helper()
|
||||
prepared := preparedAttemptDebugPipeline(t)
|
||||
chunker, ok := prepared.chunker.(*typedTestChunker)
|
||||
if !ok {
|
||||
t.Fatalf("prepared chunker = %T, want *typedTestChunker", prepared.chunker)
|
||||
}
|
||||
return prepared, cloneSourceChunks(chunker.chunks)
|
||||
return prepared, source.CloneChunkPlan(chunker.plan)
|
||||
}
|
||||
|
||||
func TestRunnerRecordsChunkTerminalOutcomes(t *testing.T) {
|
||||
@@ -96,8 +109,8 @@ func TestRunnerRecordsChunkTerminalOutcomes(t *testing.T) {
|
||||
}
|
||||
for _, tc := range tests {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
prepared, chunks := preparedTerminalDebugPipeline(t)
|
||||
prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, chunks: chunks, warnings: []contracts.Warning{{Scope: "chunk", ReasonCode: "observed", Message: "chunk warning"}}, err: tc.moduleError}
|
||||
prepared, plan := preparedTerminalDebugPipeline(t)
|
||||
prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan, warnings: []contracts.Warning{{Scope: "chunk", ReasonCode: "observed", Message: "chunk warning"}}, err: tc.moduleError}
|
||||
prepared.chunkValidators.validators = []preparedValidator{{resolved: ResolvedValidator{Binding: Binding("terminal/chunk-validator"), Target: ValidatorTargetChunk}, chunk: tc.validator}}
|
||||
debug := newCapturedDebugRecorder()
|
||||
|
||||
@@ -116,10 +129,63 @@ func TestRunnerRecordsChunkTerminalOutcomes(t *testing.T) {
|
||||
t.Fatalf("chunk rejection = envelope %#v, outputs %#v", envelope, output.Rejected)
|
||||
}
|
||||
}
|
||||
if tc.name == "accepted" {
|
||||
encoded := string(debug.json["chunk/attempt-01.json"])
|
||||
if !strings.Contains(encoded, `"plan":`) || !strings.Contains(encoded, `"materialized_chunks":`) || strings.Contains(encoded, `"chunks":`) {
|
||||
t.Fatalf("chunk attempt debug = %s, want separate plan and materialized_chunks", encoded)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerMaterializesAnnotatedPlanBeforeChunkValidation(t *testing.T) {
|
||||
prepared, plan := preparedTerminalDebugPipeline(t)
|
||||
plan.Annotations = source.ChunkAnnotations{"same": []byte(`{"plan":1}`)}
|
||||
plan.Ranges[0].Annotations = source.ChunkAnnotations{"same": []byte(`{"range":2}`)}
|
||||
prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan}
|
||||
validator := &observingChunkValidator{}
|
||||
prepared.chunkValidators.validators = []preparedValidator{{
|
||||
resolved: ResolvedValidator{Binding: Binding(validator.Name()), Target: ValidatorTargetChunk},
|
||||
chunk: validator,
|
||||
}}
|
||||
|
||||
if _, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input")}); err != nil {
|
||||
t.Fatalf("Run() error = %v, want nil", err)
|
||||
}
|
||||
if len(validator.request.Chunks) != 1 {
|
||||
t.Fatalf("validator chunks = %#v, want one", validator.request.Chunks)
|
||||
}
|
||||
chunk := validator.request.Chunks[0]
|
||||
if chunk.ID != "chunk-000001" || chunk.Index != 0 || chunk.MediaType != "application/json" {
|
||||
t.Fatalf("materialized chunk identity = %#v", chunk)
|
||||
}
|
||||
if string(chunk.Annotations["same"]) != `{"range":2}` || string(chunk.PlanAnnotations["same"]) != `{"plan":1}` {
|
||||
t.Fatalf("materialized annotations = %s / %s", chunk.Annotations["same"], chunk.PlanAnnotations["same"])
|
||||
}
|
||||
if chunk.Metadata["start_unit_id"] != 1 || chunk.Metadata["end_unit_id"] != 1 || chunk.Metadata["unit_count"] != 1 {
|
||||
t.Fatalf("materialized metadata = %#v", chunk.Metadata)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRunnerRetriesMalformedPlanWithDebugEnabled(t *testing.T) {
|
||||
prepared, plan := preparedTerminalDebugPipeline(t)
|
||||
plan.Ranges[0].Annotations = source.ChunkAnnotations{"broken": []byte(`{"value":`)}
|
||||
prepared.resolved.Chunk.Retries = 1
|
||||
calls := 0
|
||||
prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan, calls: &calls}
|
||||
debug := newCapturedDebugRecorder()
|
||||
|
||||
_, err := New().Run(context.Background(), RunInput{Prepared: prepared, RawInput: []byte("input"), Debug: debug})
|
||||
if err == nil || !strings.Contains(err.Error(), "valid JSON") {
|
||||
t.Fatalf("Run() error = %v, want malformed annotation error", err)
|
||||
}
|
||||
if calls != 2 {
|
||||
t.Fatalf("plan calls = %d, want two attempts", calls)
|
||||
}
|
||||
assertAttemptEnvelopeSequence(t, debug, "chunk", 1, 2)
|
||||
}
|
||||
|
||||
func TestRunnerRecordsExtractTerminalOutcomes(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
@@ -187,9 +253,9 @@ func TestRunnerRecordsExtractTerminalOutcomes(t *testing.T) {
|
||||
|
||||
func TestRunnerJoinsPrimaryAndAttemptWriteErrors(t *testing.T) {
|
||||
t.Run("chunk", func(t *testing.T) {
|
||||
prepared, chunks := preparedTerminalDebugPipeline(t)
|
||||
prepared, plan := preparedTerminalDebugPipeline(t)
|
||||
primary := errors.New("chunk operation failed")
|
||||
prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, chunks: chunks, err: primary}
|
||||
prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan, err: primary}
|
||||
debug := newCapturedDebugRecorder()
|
||||
debug.failPath = "chunk/attempt-01.json"
|
||||
|
||||
@@ -216,10 +282,10 @@ func TestRunnerJoinsPrimaryAndAttemptWriteErrors(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestRunnerDoesNotRetryAfterTerminalAttemptWriteFailure(t *testing.T) {
|
||||
prepared, chunks := preparedTerminalDebugPipeline(t)
|
||||
prepared, plan := preparedTerminalDebugPipeline(t)
|
||||
prepared.resolved.Chunk.Retries = 1
|
||||
calls := 0
|
||||
prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, chunks: chunks, calls: &calls}
|
||||
prepared.chunker = terminalChunker{key: prepared.resolved.Chunk.Module, plan: plan, calls: &calls}
|
||||
prepared.chunkValidators.validators = []preparedValidator{{
|
||||
resolved: ResolvedValidator{Binding: Binding("terminal/chunk-validator"), Target: ValidatorTargetChunk},
|
||||
chunk: terminalChunkValidator{result: contracts.ValidationResult{Approved: true}},
|
||||
|
||||
@@ -67,14 +67,14 @@ func (v *typedTestInput) Parse(context.Context, contracts.ParseRequest) (*source
|
||||
}
|
||||
|
||||
type typedTestChunker struct {
|
||||
key string
|
||||
chunks []source.Chunk
|
||||
key string
|
||||
plan source.ChunkPlan
|
||||
}
|
||||
|
||||
func (v *typedTestChunker) Key() string { return v.key }
|
||||
func (v *typedTestChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil }
|
||||
func (v *typedTestChunker) Chunk(context.Context, contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
return contracts.ChunkResult{Chunks: v.chunks}, nil
|
||||
func (v *typedTestChunker) Plan(context.Context, contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
return contracts.ChunkPlanResult{Plan: source.CloneChunkPlan(v.plan)}, nil
|
||||
}
|
||||
|
||||
type typedTestOutput struct{ key string }
|
||||
@@ -84,12 +84,27 @@ func (v *typedTestOutput) Encode(context.Context, contracts.OutputRequest) (cont
|
||||
return contracts.OutputResult{}, nil
|
||||
}
|
||||
func typedTestDocument() *source.SourceDocument {
|
||||
ref := source.SourceRef{SourceID: "source", StartUnitID: 1, EndUnitID: 1}
|
||||
doc := &source.SourceDocument{ID: "source", Kind: "document", Format: "text/plain", Units: []source.SourceUnit{{ID: 1, Kind: "line", Text: "text", Ref: ref}}}
|
||||
return typedTestDocumentWithUnits(1)
|
||||
}
|
||||
|
||||
func typedTestDocumentWithUnits(count int) *source.SourceDocument {
|
||||
doc := &source.SourceDocument{ID: "source", Kind: "document", Format: "text/plain"}
|
||||
for id := 1; id <= count; id++ {
|
||||
ref := source.SourceRef{SourceID: doc.ID, StartUnitID: id, EndUnitID: id}
|
||||
doc.Units = append(doc.Units, source.SourceUnit{ID: id, Kind: "line", Text: fmt.Sprintf("text-%d", id), Ref: ref})
|
||||
}
|
||||
doc.Digest, _ = source.DigestDocument(doc)
|
||||
return doc
|
||||
}
|
||||
|
||||
func typedTestPlan(doc *source.SourceDocument) source.ChunkPlan {
|
||||
plan := source.ChunkPlan{SourceDigest: doc.Digest, Ranges: make([]source.ChunkRange, len(doc.Units))}
|
||||
for i, unit := range doc.Units {
|
||||
plan.Ranges[i] = source.ChunkRange{StartUnitID: unit.ID, EndUnitID: unit.ID}
|
||||
}
|
||||
return plan
|
||||
}
|
||||
|
||||
func (v typedTestSerializedValidator) Name() string { return v.key }
|
||||
func (typedTestSerializedValidator) ExecutionClass() contracts.ExecutionClass {
|
||||
return contracts.ExecutionClassDeterministic
|
||||
@@ -466,7 +481,7 @@ func mustRegisterTypedTestBase(t *testing.T, catalog ModuleCatalog) {
|
||||
}
|
||||
if err := catalog.Chunkers.RegisterWithSpec(ModuleSpec{Key: "typed/chunk", Stage: StageChunk}, func() (contracts.Chunker, error) {
|
||||
doc := typedTestDocument()
|
||||
return &typedTestChunker{key: "typed/chunk", chunks: []source.Chunk{{ID: "chunk-1", SourceID: doc.ID, Index: 0, Ref: doc.Units[0].Ref, Content: []byte(`{"chunk":1}`), MediaType: "application/json", Units: []source.SourceUnit{doc.Units[0]}}}}, nil
|
||||
return &typedTestChunker{key: "typed/chunk", plan: typedTestPlan(doc)}, nil
|
||||
}); err != nil {
|
||||
t.Fatalf("register chunker: %v", err)
|
||||
}
|
||||
|
||||
@@ -20,9 +20,10 @@ var requiredCapabilities = []string{
|
||||
|
||||
var providedCapabilities = []string{
|
||||
"chunks",
|
||||
"chunks.scenes",
|
||||
}
|
||||
|
||||
const annotationNamespace = "dnd/scenes"
|
||||
|
||||
var referenceSlotDescriptions = shared.ReferenceSlotDescriptions{
|
||||
Glossary: "Optional campaign glossary reference material used only for scene disambiguation.",
|
||||
Party: "Optional party roster reference material used only for scene disambiguation.",
|
||||
@@ -74,27 +75,27 @@ func (c *Chunker) ManifestMetadata() map[string]any {
|
||||
return metadata
|
||||
}
|
||||
|
||||
func (c *Chunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
func (c *Chunker) Plan(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
if c == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("chunker must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("chunker must not be nil")
|
||||
}
|
||||
if c.llm == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("LLM client must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("LLM client must not be nil")
|
||||
}
|
||||
if ctx == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("context must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("context must not be nil")
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("context error before chunking: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("context error before chunking: %w", err)
|
||||
}
|
||||
if req.Source == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("source must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("source must not be nil")
|
||||
}
|
||||
if len(req.Source.Units) == 0 {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("source units must not be empty")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("source units must not be empty")
|
||||
}
|
||||
if err := source.ValidateDocument(req.Source); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("validate source document: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("validate source document: %w", err)
|
||||
}
|
||||
var response chunkResponse
|
||||
if _, err := c.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
||||
@@ -105,19 +106,19 @@ func (c *Chunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contra
|
||||
SessionID: req.SessionID,
|
||||
Inputs: shared.PromptInputs(req.SourceInput, req.References),
|
||||
}, &response); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("complete structured output: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("complete structured output: %w", err)
|
||||
}
|
||||
|
||||
warnings, err := warningsFromCaveats(response.BoundaryCaveats)
|
||||
if err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("malformed structured output: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("malformed structured output: %w", err)
|
||||
}
|
||||
chunks, err := chunksFromResponse(req.Source, response)
|
||||
plan, err := planFromResponse(req.Source, response, warnings)
|
||||
if err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("malformed structured output: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("malformed structured output: %w", err)
|
||||
}
|
||||
return contracts.ChunkResult{
|
||||
Chunks: chunks,
|
||||
return contracts.ChunkPlanResult{
|
||||
Plan: plan,
|
||||
Warnings: warnings,
|
||||
}, nil
|
||||
}
|
||||
@@ -154,12 +155,12 @@ func DecodeOptions(options map[string]any) (Options, error) {
|
||||
return Options{}, nil
|
||||
}
|
||||
|
||||
func chunksFromResponse(doc *source.SourceDocument, response chunkResponse) ([]source.Chunk, error) {
|
||||
func planFromResponse(doc *source.SourceDocument, response chunkResponse, warnings []contracts.Warning) (source.ChunkPlan, error) {
|
||||
if response.Scenes == nil {
|
||||
return nil, fmt.Errorf("scenes must be present")
|
||||
return source.ChunkPlan{}, fmt.Errorf("scenes must be present")
|
||||
}
|
||||
if len(response.Scenes) == 0 {
|
||||
return nil, fmt.Errorf("scenes must not be empty")
|
||||
return source.ChunkPlan{}, fmt.Errorf("scenes must not be empty")
|
||||
}
|
||||
|
||||
unitIndexes := make(map[int]int, len(doc.Units))
|
||||
@@ -167,86 +168,75 @@ func chunksFromResponse(doc *source.SourceDocument, response chunkResponse) ([]s
|
||||
unitIndexes[unit.ID] = i
|
||||
}
|
||||
|
||||
chunks := make([]source.Chunk, 0, len(response.Scenes))
|
||||
ranges := make([]source.ChunkRange, 0, len(response.Scenes))
|
||||
previousEnd := -1
|
||||
for i, scene := range response.Scenes {
|
||||
normalized, err := normalizeScene(doc, i, scene)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return source.ChunkPlan{}, err
|
||||
}
|
||||
|
||||
startIndex, ok := unitIndexes[normalized.StartUnitID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("scene[%d] start_unit_id %d was not found", i, normalized.StartUnitID)
|
||||
return source.ChunkPlan{}, fmt.Errorf("scene[%d] start_unit_id %d was not found", i, normalized.StartUnitID)
|
||||
}
|
||||
endIndex, ok := unitIndexes[normalized.EndUnitID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("scene[%d] end_unit_id %d was not found", i, normalized.EndUnitID)
|
||||
return source.ChunkPlan{}, fmt.Errorf("scene[%d] end_unit_id %d was not found", i, normalized.EndUnitID)
|
||||
}
|
||||
if startIndex > endIndex {
|
||||
return nil, fmt.Errorf("scene[%d] start_unit_id %d appears after end_unit_id %d", i, normalized.StartUnitID, normalized.EndUnitID)
|
||||
return source.ChunkPlan{}, fmt.Errorf("scene[%d] start_unit_id %d appears after end_unit_id %d", i, normalized.StartUnitID, normalized.EndUnitID)
|
||||
}
|
||||
|
||||
if i == 0 && startIndex != 0 {
|
||||
return nil, fmt.Errorf("first scene must start at first source unit %d", doc.Units[0].ID)
|
||||
return source.ChunkPlan{}, fmt.Errorf("first scene must start at first source unit %d", doc.Units[0].ID)
|
||||
}
|
||||
if i > 0 {
|
||||
if startIndex <= previousEnd {
|
||||
return nil, fmt.Errorf("scene[%d] overlaps previous scene", i)
|
||||
return source.ChunkPlan{}, fmt.Errorf("scene[%d] overlaps previous scene", i)
|
||||
}
|
||||
if startIndex > previousEnd+1 {
|
||||
return nil, fmt.Errorf("scene[%d] leaves a gap after previous scene", i)
|
||||
return source.ChunkPlan{}, fmt.Errorf("scene[%d] leaves a gap after previous scene", i)
|
||||
}
|
||||
}
|
||||
previousEnd = endIndex
|
||||
|
||||
units := cloneUnits(doc.Units[startIndex : endIndex+1])
|
||||
content, err := chunkContent(units)
|
||||
annotation, err := json.Marshal(struct {
|
||||
ShortTitle string `json:"short_title"`
|
||||
PrimaryMode string `json:"primary_mode"`
|
||||
MainParticipants []string `json:"main_participants"`
|
||||
Summary string `json:"summary"`
|
||||
BoundaryNote string `json:"boundary_note"`
|
||||
BoundaryConfidence string `json:"boundary_confidence"`
|
||||
}{normalized.ShortTitle, normalized.PrimaryMode, normalized.MainParticipants, normalized.Summary, normalized.BoundaryNote, normalized.BoundaryConfidence})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return source.ChunkPlan{}, fmt.Errorf("encode scene[%d] annotation: %w", i, err)
|
||||
}
|
||||
chunks = append(chunks, source.Chunk{
|
||||
ID: fmt.Sprintf("scene-%06d", i+1),
|
||||
SourceID: doc.ID,
|
||||
Index: i,
|
||||
Ref: source.SourceRef{
|
||||
SourceID: doc.ID,
|
||||
StartUnitID: units[0].Ref.StartUnitID,
|
||||
EndUnitID: units[len(units)-1].Ref.EndUnitID,
|
||||
},
|
||||
Content: content,
|
||||
MediaType: "application/json",
|
||||
Units: units,
|
||||
Metadata: map[string]any{
|
||||
"scene_title": normalized.ShortTitle,
|
||||
"primary_mode": normalized.PrimaryMode,
|
||||
"main_participants": append([]string(nil), normalized.MainParticipants...),
|
||||
"summary": normalized.Summary,
|
||||
"boundary_note": normalized.BoundaryNote,
|
||||
"boundary_confidence": normalized.BoundaryConfidence,
|
||||
"start_unit_id": normalized.StartUnitID,
|
||||
"end_unit_id": normalized.EndUnitID,
|
||||
"unit_count": len(units),
|
||||
},
|
||||
ranges = append(ranges, source.ChunkRange{
|
||||
StartUnitID: normalized.StartUnitID,
|
||||
EndUnitID: normalized.EndUnitID,
|
||||
Annotations: source.ChunkAnnotations{annotationNamespace: annotation},
|
||||
})
|
||||
}
|
||||
|
||||
if previousEnd != len(doc.Units)-1 {
|
||||
return nil, fmt.Errorf("final scene must end at final source unit %d", doc.Units[len(doc.Units)-1].ID)
|
||||
return source.ChunkPlan{}, fmt.Errorf("final scene must end at final source unit %d", doc.Units[len(doc.Units)-1].ID)
|
||||
}
|
||||
return chunks, nil
|
||||
}
|
||||
|
||||
func chunkContent(units []source.SourceUnit) ([]byte, error) {
|
||||
content, err := json.Marshal(struct {
|
||||
Units []source.SourceUnit `json:"units"`
|
||||
}{
|
||||
Units: units,
|
||||
})
|
||||
caveats := make([]string, 0, len(warnings))
|
||||
for _, warning := range warnings {
|
||||
caveats = append(caveats, warning.Message)
|
||||
}
|
||||
annotation, err := json.Marshal(struct {
|
||||
BoundaryCaveats []string `json:"boundary_caveats"`
|
||||
}{BoundaryCaveats: caveats})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encode chunk content: %w", err)
|
||||
return source.ChunkPlan{}, fmt.Errorf("encode plan annotation: %w", err)
|
||||
}
|
||||
return content, nil
|
||||
return source.ChunkPlan{
|
||||
SourceDigest: doc.Digest,
|
||||
Ranges: ranges,
|
||||
Annotations: source.ChunkAnnotations{annotationNamespace: annotation},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func normalizeScene(doc *source.SourceDocument, index int, scene sceneResponse) (normalizedScene, error) {
|
||||
@@ -347,31 +337,6 @@ func warningsFromCaveats(caveats []string) ([]contracts.Warning, error) {
|
||||
return warnings, nil
|
||||
}
|
||||
|
||||
func cloneUnits(units []source.SourceUnit) []source.SourceUnit {
|
||||
out := make([]source.SourceUnit, 0, len(units))
|
||||
for _, unit := range units {
|
||||
out = append(out, source.SourceUnit{
|
||||
ID: unit.ID,
|
||||
Kind: unit.Kind,
|
||||
Text: unit.Text,
|
||||
Ref: unit.Ref,
|
||||
Metadata: cloneMetadata(unit.Metadata),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneMetadata(metadata map[string]any) map[string]any {
|
||||
if len(metadata) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]any, len(metadata))
|
||||
for key, value := range metadata {
|
||||
out[key] = value
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func chunkerErrorf(format string, args ...any) error {
|
||||
return fmt.Errorf("dnd scenes chunker: "+format, args...)
|
||||
}
|
||||
|
||||
@@ -25,7 +25,7 @@ func TestNewModuleSpecAndRegister(t *testing.T) {
|
||||
Key: Key,
|
||||
Stage: pipeline.StageChunk,
|
||||
Requires: []string{"source.transcript"},
|
||||
Provides: []string{"chunks", "chunks.scenes"},
|
||||
Provides: []string{"chunks"},
|
||||
ReferenceSlots: wantReferenceSlots(),
|
||||
}
|
||||
if got := ModuleSpec(); !reflect.DeepEqual(got, want) {
|
||||
@@ -110,7 +110,7 @@ func wantReferenceSlots() []contracts.ReferenceSlot {
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkReturnsSceneChunksFromStructuredOutput(t *testing.T) {
|
||||
func TestPlanReturnsSceneRangesAndAnnotationsFromStructuredOutput(t *testing.T) {
|
||||
client := &fakeScenesLLMClient{
|
||||
response: chunkResponse{
|
||||
Scenes: []sceneResponse{
|
||||
@@ -139,9 +139,9 @@ func TestChunkReturnsSceneChunksFromStructuredOutput(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
result, err := newChunker(t, client).Chunk(context.Background(), chunkRequest())
|
||||
result, err := newChunker(t, client).Plan(context.Background(), chunkRequest())
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
t.Fatalf("Plan() error = %v, want nil", err)
|
||||
}
|
||||
|
||||
if len(client.requests) != 1 {
|
||||
@@ -177,36 +177,41 @@ func TestChunkReturnsSceneChunksFromStructuredOutput(t *testing.T) {
|
||||
t.Fatalf("glossary input = %q, want empty reference placeholder", got)
|
||||
}
|
||||
|
||||
if got := chunkIDs(result.Chunks); !reflect.DeepEqual(got, []string{"scene-000001", "scene-000002"}) {
|
||||
t.Fatalf("chunk IDs = %#v, want deterministic scene IDs", got)
|
||||
if result.Plan.SourceDigest != "sha256:source" {
|
||||
t.Fatalf("SourceDigest = %q, want source digest", result.Plan.SourceDigest)
|
||||
}
|
||||
gotUnits := [][]int{unitIDs(result.Chunks[0].Units), unitIDs(result.Chunks[1].Units)}
|
||||
wantUnits := [][]int{{1, 2}, {3, 4}}
|
||||
if !reflect.DeepEqual(gotUnits, wantUnits) {
|
||||
t.Fatalf("chunk units = %#v, want %#v", gotUnits, wantUnits)
|
||||
wantRanges := []source.ChunkRange{{StartUnitID: 1, EndUnitID: 2}, {StartUnitID: 3, EndUnitID: 4}}
|
||||
if len(result.Plan.Ranges) != len(wantRanges) {
|
||||
t.Fatalf("ranges = %#v, want two", result.Plan.Ranges)
|
||||
}
|
||||
first := result.Chunks[0]
|
||||
if first.SourceID != "session-alpha" || first.Index != 0 {
|
||||
t.Fatalf("first chunk = %#v, want source and index fields", first)
|
||||
for i, want := range wantRanges {
|
||||
if result.Plan.Ranges[i].StartUnitID != want.StartUnitID || result.Plan.Ranges[i].EndUnitID != want.EndUnitID {
|
||||
t.Fatalf("range[%d] = %#v, want %#v", i, result.Plan.Ranges[i], want)
|
||||
}
|
||||
}
|
||||
if first.Ref != (source.SourceRef{SourceID: "session-alpha", StartUnitID: 1, EndUnitID: 2}) {
|
||||
t.Fatalf("first ref = %#v, want session-alpha:1-2", first.Ref)
|
||||
var firstAnnotation map[string]any
|
||||
if err := json.Unmarshal(result.Plan.Ranges[0].Annotations[annotationNamespace], &firstAnnotation); err != nil {
|
||||
t.Fatalf("decode first scene annotation: %v", err)
|
||||
}
|
||||
if first.MediaType != "application/json" || len(first.Content) == 0 {
|
||||
t.Fatalf("first payload = media type %q length %d, want JSON content", first.MediaType, len(first.Content))
|
||||
wantFirst := map[string]any{
|
||||
"short_title": "Goblin parley", "primary_mode": "Discussion",
|
||||
"main_participants": []any{"Aria", "Goblin scout"},
|
||||
"summary": "The party negotiates with a scout.",
|
||||
"boundary_note": "The scene covers the discussion before fighting starts.",
|
||||
"boundary_confidence": "High",
|
||||
}
|
||||
if first.Metadata["scene_title"] != "Goblin parley" ||
|
||||
first.Metadata["primary_mode"] != "Discussion" ||
|
||||
first.Metadata["summary"] != "The party negotiates with a scout." ||
|
||||
first.Metadata["boundary_note"] != "The scene covers the discussion before fighting starts." ||
|
||||
first.Metadata["boundary_confidence"] != "High" ||
|
||||
first.Metadata["start_unit_id"] != 1 ||
|
||||
first.Metadata["end_unit_id"] != 2 ||
|
||||
first.Metadata["unit_count"] != 2 {
|
||||
t.Fatalf("first metadata = %#v, want scene metadata", first.Metadata)
|
||||
if !reflect.DeepEqual(firstAnnotation, wantFirst) {
|
||||
t.Fatalf("first scene annotation = %#v, want %#v", firstAnnotation, wantFirst)
|
||||
}
|
||||
if got, ok := first.Metadata["main_participants"].([]string); !ok || !reflect.DeepEqual(got, []string{"Aria", "Goblin scout"}) {
|
||||
t.Fatalf("main_participants = %#v, want trimmed participant slice", first.Metadata["main_participants"])
|
||||
if len(firstAnnotation) != 6 {
|
||||
t.Fatalf("first scene annotation keys = %#v, want exact six fields", firstAnnotation)
|
||||
}
|
||||
var planAnnotation map[string]any
|
||||
if err := json.Unmarshal(result.Plan.Annotations[annotationNamespace], &planAnnotation); err != nil {
|
||||
t.Fatalf("decode plan annotation: %v", err)
|
||||
}
|
||||
if !reflect.DeepEqual(planAnnotation, map[string]any{"boundary_caveats": []any{"The transition into combat is gradual."}}) || len(planAnnotation) != 1 {
|
||||
t.Fatalf("plan annotation = %#v, want normalized caveats only", planAnnotation)
|
||||
}
|
||||
if got := result.Warnings; len(got) != 1 ||
|
||||
got[0].Scope != Key ||
|
||||
@@ -216,7 +221,7 @@ func TestChunkReturnsSceneChunksFromStructuredOutput(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkPassesReferencesAsPromptInputs(t *testing.T) {
|
||||
func TestPlanPassesReferencesAsPromptInputs(t *testing.T) {
|
||||
client := &fakeScenesLLMClient{response: chunkResponse{
|
||||
Scenes: []sceneResponse{
|
||||
{
|
||||
@@ -255,8 +260,8 @@ func TestChunkPassesReferencesAsPromptInputs(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
if _, err := newChunker(t, client).Chunk(context.Background(), req); err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
if _, err := newChunker(t, client).Plan(context.Background(), req); err != nil {
|
||||
t.Fatalf("Plan() error = %v, want nil", err)
|
||||
}
|
||||
request := client.requests[0]
|
||||
if got := string(request.Inputs["players"].Content); got != "Alice: Aria" {
|
||||
@@ -293,7 +298,7 @@ func TestPromptInputsMapLegacyRosterReferenceToParty(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRejectsWhitespaceOnlyBoundaryCaveats(t *testing.T) {
|
||||
func TestPlanRejectsWhitespaceOnlyBoundaryCaveats(t *testing.T) {
|
||||
client := &fakeScenesLLMClient{
|
||||
response: chunkResponse{
|
||||
Scenes: validSceneResponse().Scenes,
|
||||
@@ -303,46 +308,38 @@ func TestChunkRejectsWhitespaceOnlyBoundaryCaveats(t *testing.T) {
|
||||
},
|
||||
}
|
||||
|
||||
_, err := newChunker(t, client).Chunk(context.Background(), chunkRequest())
|
||||
_, err := newChunker(t, client).Plan(context.Background(), chunkRequest())
|
||||
if err == nil {
|
||||
t.Fatal("Chunk() error = nil, want malformed structured output error")
|
||||
t.Fatal("Plan() error = nil, want malformed structured output error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "dnd scenes chunker") || !strings.Contains(err.Error(), "malformed structured output") || !strings.Contains(err.Error(), "boundary_caveats[0]") {
|
||||
t.Fatalf("Chunk() error = %q, want malformed boundary caveat context", err.Error())
|
||||
t.Fatalf("Plan() error = %q, want malformed boundary caveat context", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkDefensivelyCopiesSourceUnitsAndMetadata(t *testing.T) {
|
||||
func TestPlanDefensivelyCopiesAnnotationValues(t *testing.T) {
|
||||
doc := sceneSourceDocument()
|
||||
client := &fakeScenesLLMClient{response: validSceneResponse()}
|
||||
|
||||
result, err := newChunker(t, client).Chunk(context.Background(), contracts.ChunkRequest{
|
||||
result, err := newChunker(t, client).Plan(context.Background(), contracts.ChunkRequest{
|
||||
Source: doc,
|
||||
SourceInput: sceneSourceInput(),
|
||||
SessionID: "session-123",
|
||||
LLMProfile: "profile-scenes",
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
t.Fatalf("Plan() error = %v, want nil", err)
|
||||
}
|
||||
|
||||
doc.Units[0].ID = 99
|
||||
doc.Units[0].Ref.SourceID = "mutated"
|
||||
doc.Units[0].Metadata["speaker"] = "mutated"
|
||||
client.response.Scenes[0].MainParticipants[0] = "mutated"
|
||||
|
||||
if result.Chunks[0].Units[0].ID != 1 {
|
||||
t.Fatalf("chunk unit ID changed after source mutation: %#v", result.Chunks[0].Units[0])
|
||||
var annotation struct {
|
||||
MainParticipants []string `json:"main_participants"`
|
||||
}
|
||||
if got := result.Chunks[0].Units[0].Ref.SourceID; got != "session-alpha" {
|
||||
t.Fatalf("chunk unit ref changed after source mutation: %#v", result.Chunks[0].Units[0].Ref)
|
||||
if err := json.Unmarshal(result.Plan.Ranges[0].Annotations[annotationNamespace], &annotation); err != nil {
|
||||
t.Fatalf("decode annotation: %v", err)
|
||||
}
|
||||
if result.Chunks[0].Units[0].Metadata["speaker"] != "Alice" {
|
||||
t.Fatalf("chunk unit metadata changed after source mutation: %#v", result.Chunks[0].Units[0].Metadata)
|
||||
}
|
||||
participants, ok := result.Chunks[0].Metadata["main_participants"].([]string)
|
||||
if !ok || participants[0] != "Aria" {
|
||||
t.Fatalf("participants = %#v, want defensive copy", result.Chunks[0].Metadata["main_participants"])
|
||||
if !reflect.DeepEqual(annotation.MainParticipants, []string{"Aria"}) {
|
||||
t.Fatalf("participants = %#v, want defensive copy", annotation.MainParticipants)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -375,7 +372,7 @@ func TestChunkerManifestMetadataIncludesPromptAndSchemaProvenance(t *testing.T)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRejectsInvalidRequests(t *testing.T) {
|
||||
func TestPlanRejectsInvalidRequests(t *testing.T) {
|
||||
validClient := &fakeScenesLLMClient{response: validSceneResponse()}
|
||||
validReq := chunkRequest()
|
||||
canceledCtx, cancel := context.WithCancel(context.Background())
|
||||
@@ -402,18 +399,18 @@ func TestChunkRejectsInvalidRequests(t *testing.T) {
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
_, err := tt.chunker.Chunk(tt.ctx, tt.req)
|
||||
_, err := tt.chunker.Plan(tt.ctx, tt.req)
|
||||
if err == nil {
|
||||
t.Fatal("Chunk() error = nil, want error")
|
||||
t.Fatal("Plan() error = nil, want error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "dnd scenes") || !strings.Contains(err.Error(), tt.want) {
|
||||
t.Fatalf("Chunk() error = %q, want module context and %q", err.Error(), tt.want)
|
||||
t.Fatalf("Plan() error = %q, want module context and %q", err.Error(), tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRejectsMalformedStructuredOutput(t *testing.T) {
|
||||
func TestPlanRejectsMalformedStructuredOutput(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
response chunkResponse
|
||||
@@ -495,26 +492,26 @@ func TestChunkRejectsMalformedStructuredOutput(t *testing.T) {
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
client := &fakeScenesLLMClient{response: tt.response}
|
||||
_, err := newChunker(t, client).Chunk(context.Background(), chunkRequest())
|
||||
_, err := newChunker(t, client).Plan(context.Background(), chunkRequest())
|
||||
if err == nil {
|
||||
t.Fatal("Chunk() error = nil, want error")
|
||||
t.Fatal("Plan() error = nil, want error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "dnd scenes") || !strings.Contains(err.Error(), tt.want) {
|
||||
t.Fatalf("Chunk() error = %q, want module context and %q", err.Error(), tt.want)
|
||||
t.Fatalf("Plan() error = %q, want module context and %q", err.Error(), tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkWrapsLLMClientError(t *testing.T) {
|
||||
func TestPlanWrapsLLMClientError(t *testing.T) {
|
||||
client := &fakeScenesLLMClient{err: errors.New("provider unavailable")}
|
||||
|
||||
_, err := newChunker(t, client).Chunk(context.Background(), chunkRequest())
|
||||
_, err := newChunker(t, client).Plan(context.Background(), chunkRequest())
|
||||
if err == nil {
|
||||
t.Fatal("Chunk() error = nil, want LLM error")
|
||||
t.Fatal("Plan() error = nil, want LLM error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "dnd scenes") || !strings.Contains(err.Error(), "provider unavailable") {
|
||||
t.Fatalf("Chunk() error = %q, want wrapped LLM context", err.Error())
|
||||
t.Fatalf("Plan() error = %q, want wrapped LLM context", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -42,77 +42,46 @@ func (c *Chunker) ReferenceSlots() []contracts.ReferenceSlot {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c *Chunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
func (c *Chunker) Plan(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
if c == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("chunker must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("chunker must not be nil")
|
||||
}
|
||||
if c.options.MaxUnits <= 0 || c.options.OverlapUnits < 0 || c.options.OverlapUnits >= c.options.MaxUnits {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("chunker options must be initialized by construction")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("chunker options must be initialized by construction")
|
||||
}
|
||||
if ctx == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("context must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("context must not be nil")
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("context error before chunking: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("context error before chunking: %w", err)
|
||||
}
|
||||
if req.Source == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("source must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("source must not be nil")
|
||||
}
|
||||
if len(req.Source.Units) == 0 {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("source units must not be empty")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("source units must not be empty")
|
||||
}
|
||||
if err := source.ValidateDocument(req.Source); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("validate source document: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("validate source document: %w", err)
|
||||
}
|
||||
|
||||
step := c.options.MaxUnits - c.options.OverlapUnits
|
||||
chunks := make([]source.Chunk, 0, (len(req.Source.Units)+step-1)/step)
|
||||
ranges := make([]source.ChunkRange, 0, (len(req.Source.Units)+step-1)/step)
|
||||
for start := 0; start < len(req.Source.Units); start += step {
|
||||
end := start + c.options.MaxUnits
|
||||
if end > len(req.Source.Units) {
|
||||
end = len(req.Source.Units)
|
||||
}
|
||||
units := cloneUnits(req.Source.Units[start:end])
|
||||
content, err := chunkContent(units)
|
||||
if err != nil {
|
||||
return contracts.ChunkResult{}, err
|
||||
}
|
||||
chunks = append(chunks, source.Chunk{
|
||||
ID: fmt.Sprintf("chunk-%06d", len(chunks)+1),
|
||||
SourceID: req.Source.ID,
|
||||
Index: len(chunks),
|
||||
Ref: source.SourceRef{
|
||||
SourceID: req.Source.ID,
|
||||
StartUnitID: units[0].Ref.StartUnitID,
|
||||
EndUnitID: units[len(units)-1].Ref.EndUnitID,
|
||||
},
|
||||
Content: content,
|
||||
MediaType: "application/json",
|
||||
Units: units,
|
||||
Metadata: map[string]any{
|
||||
"start_unit_id": units[0].ID,
|
||||
"end_unit_id": units[len(units)-1].ID,
|
||||
"unit_count": len(units),
|
||||
},
|
||||
ranges = append(ranges, source.ChunkRange{
|
||||
StartUnitID: req.Source.Units[start].ID,
|
||||
EndUnitID: req.Source.Units[end-1].ID,
|
||||
})
|
||||
if end == len(req.Source.Units) {
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
return contracts.ChunkResult{Chunks: chunks}, nil
|
||||
}
|
||||
|
||||
func chunkContent(units []source.SourceUnit) ([]byte, error) {
|
||||
content, err := json.Marshal(struct {
|
||||
Units []source.SourceUnit `json:"units"`
|
||||
}{
|
||||
Units: units,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, chunkerErrorf("encode chunk content: %w", err)
|
||||
}
|
||||
return content, nil
|
||||
return contracts.ChunkPlanResult{Plan: source.ChunkPlan{SourceDigest: req.Source.Digest, Ranges: ranges}}, nil
|
||||
}
|
||||
|
||||
func ModuleSpec() pipeline.ModuleSpec {
|
||||
@@ -251,31 +220,6 @@ func minInt() int64 {
|
||||
return -maxInt() - 1
|
||||
}
|
||||
|
||||
func cloneUnits(units []source.SourceUnit) []source.SourceUnit {
|
||||
out := make([]source.SourceUnit, 0, len(units))
|
||||
for _, unit := range units {
|
||||
out = append(out, source.SourceUnit{
|
||||
ID: unit.ID,
|
||||
Kind: unit.Kind,
|
||||
Text: unit.Text,
|
||||
Ref: unit.Ref,
|
||||
Metadata: cloneMetadata(unit.Metadata),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneMetadata(metadata map[string]any) map[string]any {
|
||||
if len(metadata) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]any, len(metadata))
|
||||
for key, value := range metadata {
|
||||
out[key] = value
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func chunkerErrorf(format string, args ...any) error {
|
||||
return fmt.Errorf("generic chunker: "+format, args...)
|
||||
}
|
||||
|
||||
@@ -48,72 +48,51 @@ func TestModuleSpecAndRegister(t *testing.T) {
|
||||
if err != nil {
|
||||
t.Fatalf("BuildWithRequest(%q) error = %v, want nil", Key, err)
|
||||
}
|
||||
result, err := configured.Chunk(context.Background(), contracts.ChunkRequest{Source: testSource(2)})
|
||||
if err != nil || len(result.Chunks) != 2 {
|
||||
t.Fatalf("constructed chunker result = %#v, %v; want two chunks", result, err)
|
||||
result, err := configured.Plan(context.Background(), contracts.ChunkRequest{Source: testSource(2)})
|
||||
if err != nil || len(result.Plan.Ranges) != 2 {
|
||||
t.Fatalf("constructed chunker result = %#v, %v; want two ranges", result, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkUsesDefaultsForSingleChunk(t *testing.T) {
|
||||
result, err := newChunker(t, nil).Chunk(context.Background(), contracts.ChunkRequest{Source: testSource(3)})
|
||||
func TestPlanUsesDefaultsForSingleRange(t *testing.T) {
|
||||
result, err := newChunker(t, nil).Plan(context.Background(), contracts.ChunkRequest{Source: testSource(3)})
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
t.Fatalf("Plan() error = %v, want nil", err)
|
||||
}
|
||||
|
||||
if got := chunkIDs(result.Chunks); !reflect.DeepEqual(got, []string{"chunk-000001"}) {
|
||||
t.Fatalf("chunk IDs = %#v, want one stable ID", got)
|
||||
if result.Plan.SourceDigest != "sha256:source" || !reflect.DeepEqual(result.Plan.Ranges, []source.ChunkRange{{StartUnitID: 1, EndUnitID: 3}}) {
|
||||
t.Fatalf("Plan = %#v, want source digest and one complete range", result.Plan)
|
||||
}
|
||||
chunk := result.Chunks[0]
|
||||
if chunk.Index != 0 || chunk.SourceID != "source-1" {
|
||||
t.Fatalf("chunk = %#v, want source and index fields", chunk)
|
||||
}
|
||||
if got := unitIDs(chunk.Units); !reflect.DeepEqual(got, []int{1, 2, 3}) {
|
||||
t.Fatalf("unit IDs = %#v, want all units", got)
|
||||
}
|
||||
if chunk.Ref != (source.SourceRef{SourceID: "source-1", StartUnitID: 1, EndUnitID: 3}) {
|
||||
t.Fatalf("chunk ref = %#v, want source-1:1-3", chunk.Ref)
|
||||
}
|
||||
if chunk.MediaType != "application/json" || len(chunk.Content) == 0 {
|
||||
t.Fatalf("chunk payload = media type %q length %d, want JSON content", chunk.MediaType, len(chunk.Content))
|
||||
}
|
||||
if chunk.Metadata["start_unit_id"] != 1 || chunk.Metadata["end_unit_id"] != 3 || chunk.Metadata["unit_count"] != 3 {
|
||||
t.Fatalf("metadata = %#v, want chunk bounds", chunk.Metadata)
|
||||
if len(result.Plan.Annotations) != 0 || len(result.Plan.Ranges[0].Annotations) != 0 {
|
||||
t.Fatalf("Plan annotations = %#v / %#v, want none", result.Plan.Annotations, result.Plan.Ranges[0].Annotations)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkExactBoundaries(t *testing.T) {
|
||||
result, err := newChunker(t, map[string]any{"max_units": 2}).Chunk(context.Background(), contracts.ChunkRequest{Source: testSource(6)})
|
||||
func TestPlanExactBoundaries(t *testing.T) {
|
||||
result, err := newChunker(t, map[string]any{"max_units": 2}).Plan(context.Background(), contracts.ChunkRequest{Source: testSource(6)})
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
t.Fatalf("Plan() error = %v, want nil", err)
|
||||
}
|
||||
|
||||
if got := chunkIDs(result.Chunks); !reflect.DeepEqual(got, []string{"chunk-000001", "chunk-000002", "chunk-000003"}) {
|
||||
t.Fatalf("chunk IDs = %#v, want stable IDs", got)
|
||||
}
|
||||
gotUnits := [][]int{unitIDs(result.Chunks[0].Units), unitIDs(result.Chunks[1].Units), unitIDs(result.Chunks[2].Units)}
|
||||
wantUnits := [][]int{{1, 2}, {3, 4}, {5, 6}}
|
||||
if !reflect.DeepEqual(gotUnits, wantUnits) {
|
||||
t.Fatalf("chunk units = %#v, want %#v", gotUnits, wantUnits)
|
||||
want := []source.ChunkRange{{StartUnitID: 1, EndUnitID: 2}, {StartUnitID: 3, EndUnitID: 4}, {StartUnitID: 5, EndUnitID: 6}}
|
||||
if !reflect.DeepEqual(result.Plan.Ranges, want) {
|
||||
t.Fatalf("ranges = %#v, want %#v", result.Plan.Ranges, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkOverlap(t *testing.T) {
|
||||
result, err := newChunker(t, map[string]any{"max_units": 3, "overlap_units": 1}).Chunk(context.Background(), contracts.ChunkRequest{Source: testSource(7)})
|
||||
func TestPlanOverlap(t *testing.T) {
|
||||
result, err := newChunker(t, map[string]any{"max_units": 3, "overlap_units": 1}).Plan(context.Background(), contracts.ChunkRequest{Source: testSource(7)})
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
t.Fatalf("Plan() error = %v, want nil", err)
|
||||
}
|
||||
|
||||
gotUnits := make([][]int, 0, len(result.Chunks))
|
||||
for _, chunk := range result.Chunks {
|
||||
gotUnits = append(gotUnits, unitIDs(chunk.Units))
|
||||
}
|
||||
wantUnits := [][]int{{1, 2, 3}, {3, 4, 5}, {5, 6, 7}}
|
||||
if !reflect.DeepEqual(gotUnits, wantUnits) {
|
||||
t.Fatalf("chunk units = %#v, want %#v", gotUnits, wantUnits)
|
||||
want := []source.ChunkRange{{StartUnitID: 1, EndUnitID: 3}, {StartUnitID: 3, EndUnitID: 5}, {StartUnitID: 5, EndUnitID: 7}}
|
||||
if !reflect.DeepEqual(result.Plan.Ranges, want) {
|
||||
t.Fatalf("ranges = %#v, want %#v", result.Plan.Ranges, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRejectsInvalidOptions(t *testing.T) {
|
||||
func TestPlanRejectsInvalidOptions(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
options map[string]any
|
||||
@@ -141,42 +120,36 @@ func TestChunkRejectsInvalidOptions(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRejectsEmptySource(t *testing.T) {
|
||||
func TestPlanRejectsEmptySource(t *testing.T) {
|
||||
doc := testSource(1)
|
||||
doc.Units = nil
|
||||
|
||||
_, err := newChunker(t, nil).Chunk(context.Background(), contracts.ChunkRequest{Source: doc})
|
||||
_, err := newChunker(t, nil).Plan(context.Background(), contracts.ChunkRequest{Source: doc})
|
||||
if err == nil {
|
||||
t.Fatal("Chunk() error = nil, want empty source error")
|
||||
t.Fatal("Plan() error = nil, want empty source error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "generic chunker") || !strings.Contains(err.Error(), "units") {
|
||||
t.Fatalf("Chunk() error = %q, want empty source context", err.Error())
|
||||
t.Fatalf("Plan() error = %q, want empty source context", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkDefensivelyCopiesUnits(t *testing.T) {
|
||||
func TestPlanDoesNotRetainSourceUnits(t *testing.T) {
|
||||
doc := testSource(2)
|
||||
|
||||
result, err := newChunker(t, map[string]any{"max_units": 1}).Chunk(context.Background(), contracts.ChunkRequest{Source: doc})
|
||||
result, err := newChunker(t, map[string]any{"max_units": 1}).Plan(context.Background(), contracts.ChunkRequest{Source: doc})
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
t.Fatalf("Plan() error = %v, want nil", err)
|
||||
}
|
||||
if len(result.Chunks) != 2 {
|
||||
t.Fatalf("len(Chunks) = %d, want 2", len(result.Chunks))
|
||||
if len(result.Plan.Ranges) != 2 {
|
||||
t.Fatalf("len(Ranges) = %d, want 2", len(result.Plan.Ranges))
|
||||
}
|
||||
|
||||
doc.Units[0].ID = 99
|
||||
doc.Units[0].Ref.SourceID = "changed"
|
||||
doc.Units[0].Metadata["speaker"] = "changed"
|
||||
|
||||
if result.Chunks[0].Units[0].ID != 1 {
|
||||
t.Fatalf("chunk unit ID changed after source mutation: %#v", result.Chunks[0].Units[0])
|
||||
}
|
||||
if got := result.Chunks[0].Units[0].Ref.SourceID; got != "source-1" {
|
||||
t.Fatalf("chunk unit ref changed after source mutation: %#v", result.Chunks[0].Units[0].Ref)
|
||||
}
|
||||
if result.Chunks[0].Units[0].Metadata["speaker"] != "speaker-001" {
|
||||
t.Fatalf("chunk unit metadata changed after source mutation: %#v", result.Chunks[0].Units[0].Metadata)
|
||||
if result.Plan.Ranges[0].StartUnitID != 1 || result.Plan.Ranges[0].EndUnitID != 1 {
|
||||
t.Fatalf("first range changed after source mutation: %#v", result.Plan.Ranges[0])
|
||||
}
|
||||
}
|
||||
|
||||
@@ -214,19 +187,3 @@ func testSource(count int) *source.SourceDocument {
|
||||
func zeroPad3(value int) string {
|
||||
return fmt.Sprintf("%03d", value)
|
||||
}
|
||||
|
||||
func chunkIDs(chunks []source.Chunk) []string {
|
||||
ids := make([]string, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
ids = append(ids, chunk.ID)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func unitIDs(units []source.SourceUnit) []int {
|
||||
ids := make([]int, 0, len(units))
|
||||
for _, unit := range units {
|
||||
ids = append(ids, unit.ID)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
@@ -4,7 +4,6 @@ import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
@@ -68,12 +67,12 @@ type concurrentChunker struct{}
|
||||
|
||||
func (concurrentChunker) Key() string { return concurrentChunkerKey }
|
||||
func (concurrentChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil }
|
||||
func (concurrentChunker) Chunk(_ context.Context, request contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
chunks := make([]source.Chunk, len(request.Source.Units))
|
||||
func (concurrentChunker) Plan(_ context.Context, request contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
ranges := make([]source.ChunkRange, len(request.Source.Units))
|
||||
for i, unit := range request.Source.Units {
|
||||
chunks[i] = source.Chunk{ID: fmt.Sprintf("%s:chunk:%d", request.Source.ID, i), SourceID: request.Source.ID, Index: i, Ref: unit.Ref, Content: []byte(fmt.Sprintf(`{"unit":%d}`, unit.ID)), MediaType: "application/json", Units: []source.SourceUnit{unit}}
|
||||
ranges[i] = source.ChunkRange{StartUnitID: unit.ID, EndUnitID: unit.ID}
|
||||
}
|
||||
return contracts.ChunkResult{Chunks: chunks}, nil
|
||||
return contracts.ChunkPlanResult{Plan: source.ChunkPlan{SourceDigest: request.Source.Digest, Ranges: ranges}}, nil
|
||||
}
|
||||
|
||||
type concurrentExtractor struct {
|
||||
|
||||
@@ -271,19 +271,9 @@ func (dndSpellsChunker) ReferenceSlots() []contracts.ReferenceSlot {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (dndSpellsChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
return contracts.ChunkResult{
|
||||
Chunks: []source.Chunk{
|
||||
{
|
||||
ID: req.Source.ID + ":chunk:0",
|
||||
SourceID: req.Source.ID,
|
||||
Index: 0,
|
||||
Ref: source.SourceRef{SourceID: req.Source.ID, StartUnitID: req.Source.Units[0].ID, EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID},
|
||||
Content: []byte(`{"units":[1,2,3]}`),
|
||||
MediaType: "application/json",
|
||||
Units: append([]source.SourceUnit(nil), req.Source.Units...),
|
||||
},
|
||||
},
|
||||
func (dndSpellsChunker) Plan(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
return contracts.ChunkPlanResult{
|
||||
Plan: source.ChunkPlan{SourceDigest: req.Source.Digest, Ranges: []source.ChunkRange{{StartUnitID: req.Source.Units[0].ID, EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID}}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
@@ -220,8 +220,8 @@ func (fakeChunker) Key() string { return "fake/chunk" }
|
||||
|
||||
func (fakeChunker) ReferenceSlots() []contracts.ReferenceSlot { return nil }
|
||||
|
||||
func (fakeChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
return contracts.ChunkResult{}, nil
|
||||
func (fakeChunker) Plan(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
return contracts.ChunkPlanResult{}, nil
|
||||
}
|
||||
|
||||
type fakeExtractor struct{}
|
||||
|
||||
@@ -177,19 +177,9 @@ func (runnerSeriatimChunker) ReferenceSlots() []contracts.ReferenceSlot {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (runnerSeriatimChunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
return contracts.ChunkResult{
|
||||
Chunks: []source.Chunk{
|
||||
{
|
||||
ID: req.Source.ID + ":chunk:0",
|
||||
SourceID: req.Source.ID,
|
||||
Index: 0,
|
||||
Ref: source.SourceRef{SourceID: req.Source.ID, StartUnitID: req.Source.Units[0].ID, EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID},
|
||||
Content: []byte(`{"units":[1,2]}`),
|
||||
MediaType: "application/json",
|
||||
Units: append([]source.SourceUnit(nil), req.Source.Units...),
|
||||
},
|
||||
},
|
||||
func (runnerSeriatimChunker) Plan(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
return contracts.ChunkPlanResult{
|
||||
Plan: source.ChunkPlan{SourceDigest: req.Source.Digest, Ranges: []source.ChunkRange{{StartUnitID: req.Source.Units[0].ID, EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID}}},
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user