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

Generate and materialize canonical chunk plans

Browse Source
This commit is contained in:
Eric Rakestraw
2026-07-17 23:57:59 +00:00
parent 3bfac14397
commit 7844c0a93f
23 changed files with 396 additions and 540 deletions
Download Patch File Download Diff File Expand all files Collapse all files

143
internal/modules/dnd/chunk/scenes/chunker.go
View File

@@ -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...)
}

125
internal/modules/dnd/chunk/scenes/chunker_test.go
View File

@@ -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())
}
}

82
internal/modules/generic/chunk/units/chunker.go
View File

@@ -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...)
}

111
internal/modules/generic/chunk/units/chunker_test.go
View File

@@ -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
}

9
internal/modules/integration/concurrent_runner_test.go
View File

@@ -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 {

16
internal/modules/integration/dnd_spells_config_test.go
View File

@@ -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
}

4
internal/modules/seriatim/input/transcript/config_test.go
View File

@@ -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{}

16
internal/modules/seriatim/input/transcript/runner_test.go
View File

@@ -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
}