Implement integer source units and chunk payloads
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user