Implement integer source units and chunk payloads

This commit is contained in:
2026-07-07 18:34:23 +00:00
parent 4f057b99ac
commit 9e3f8809b3
45 changed files with 618 additions and 451 deletions

View File

@@ -68,11 +68,19 @@ func (c *Chunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contra
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, contracts.SourceChunk{
ID: fmt.Sprintf("chunk-%06d", len(chunks)+1),
SourceID: req.Source.ID,
Index: len(chunks),
Units: units,
ID: fmt.Sprintf("chunk-%06d", len(chunks)+1),
SourceID: req.Source.ID,
Index: len(chunks),
StartUnitID: units[0].ID,
EndUnitID: units[len(units)-1].ID,
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,
@@ -87,6 +95,18 @@ func (c *Chunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contra
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
}
func ModuleSpec() pipeline.ModuleSpec {
return pipeline.ModuleSpec{
Key: Key,

View File

@@ -59,10 +59,16 @@ func TestChunkUsesDefaultsForSingleChunk(t *testing.T) {
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, []string{"u001", "u002", "u003"}) {
if got := unitIDs(chunk.Units); !reflect.DeepEqual(got, []int{1, 2, 3}) {
t.Fatalf("unit IDs = %#v, want all units", got)
}
if chunk.Metadata["start_unit_id"] != "u001" || chunk.Metadata["end_unit_id"] != "u003" || chunk.Metadata["unit_count"] != 3 {
if chunk.StartUnitID != 1 || chunk.EndUnitID != 3 {
t.Fatalf("chunk boundaries = %d-%d, want 1-3", chunk.StartUnitID, chunk.EndUnitID)
}
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)
}
}
@@ -79,8 +85,8 @@ func TestChunkExactBoundaries(t *testing.T) {
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 := [][]string{unitIDs(result.Chunks[0].Units), unitIDs(result.Chunks[1].Units), unitIDs(result.Chunks[2].Units)}
wantUnits := [][]string{{"u001", "u002"}, {"u003", "u004"}, {"u005", "u006"}}
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)
}
@@ -95,11 +101,11 @@ func TestChunkOverlap(t *testing.T) {
t.Fatalf("Chunk() error = %v, want nil", err)
}
gotUnits := make([][]string, 0, len(result.Chunks))
gotUnits := make([][]int, 0, len(result.Chunks))
for _, chunk := range result.Chunks {
gotUnits = append(gotUnits, unitIDs(chunk.Units))
}
wantUnits := [][]string{{"u001", "u002", "u003"}, {"u003", "u004", "u005"}, {"u005", "u006", "u007"}}
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)
}
@@ -162,10 +168,10 @@ func TestChunkDefensivelyCopiesUnits(t *testing.T) {
t.Fatalf("len(Chunks) = %d, want 2", len(result.Chunks))
}
doc.Units[0].ID = "changed"
doc.Units[0].ID = 99
doc.Units[0].Metadata["speaker"] = "changed"
if result.Chunks[0].Units[0].ID != "u001" {
if result.Chunks[0].Units[0].ID != 1 {
t.Fatalf("chunk unit ID changed after source mutation: %#v", result.Chunks[0].Units[0])
}
if result.Chunks[0].Units[0].Metadata["speaker"] != "speaker-001" {
@@ -176,11 +182,10 @@ func TestChunkDefensivelyCopiesUnits(t *testing.T) {
func testSource(count int) *source.SourceDocument {
units := make([]source.SourceUnit, 0, count)
for i := 1; i <= count; i++ {
id := "u" + zeroPad3(i)
units = append(units, source.SourceUnit{
ID: id,
ID: i,
Kind: "unit",
Text: "Text for " + id,
Text: "Text for " + zeroPad3(i),
Metadata: map[string]any{
"speaker": "speaker-" + zeroPad3(i),
},
@@ -207,8 +212,8 @@ func chunkIDs(chunks []contracts.SourceChunk) []string {
return ids
}
func unitIDs(units []source.SourceUnit) []string {
ids := make([]string, 0, len(units))
func unitIDs(units []source.SourceUnit) []int {
ids := make([]int, 0, len(units))
for _, unit := range units {
ids = append(ids, unit.ID)
}