Implement integer source units and chunk payloads
This commit is contained in:
@@ -9,8 +9,12 @@ import (
|
||||
)
|
||||
|
||||
func validateAndCanonicalizeChunkResult(doc *source.SourceDocument, chunks []contracts.SourceChunk) ([]contracts.SourceChunk, error) {
|
||||
sourceUnitIndexes := make(map[string]int, len(doc.Units))
|
||||
sourceUnits := make(map[string]source.SourceUnit, len(doc.Units))
|
||||
if len(chunks) == 0 {
|
||||
return nil, fmt.Errorf("chunks must not be empty")
|
||||
}
|
||||
|
||||
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
|
||||
@@ -33,25 +37,42 @@ func validateAndCanonicalizeChunkResult(doc *source.SourceDocument, chunks []con
|
||||
if chunk.Index != chunkIndex {
|
||||
return nil, fmt.Errorf("chunk %q index %d does not match returned order %d", chunk.ID, chunk.Index, chunkIndex)
|
||||
}
|
||||
startIndex, ok := sourceUnitIndexes[chunk.StartUnitID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("chunk %q start_unit_id %d was not found in source document %q", chunk.ID, chunk.StartUnitID, doc.ID)
|
||||
}
|
||||
endIndex, ok := sourceUnitIndexes[chunk.EndUnitID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("chunk %q end_unit_id %d was not found in source document %q", chunk.ID, chunk.EndUnitID, doc.ID)
|
||||
}
|
||||
if startIndex > endIndex {
|
||||
return nil, fmt.Errorf("chunk %q start_unit_id %d appears after end_unit_id %d", chunk.ID, chunk.StartUnitID, chunk.EndUnitID)
|
||||
}
|
||||
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)
|
||||
}
|
||||
|
||||
seenUnitIDs := make(map[string]struct{}, len(chunk.Units))
|
||||
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 strings.TrimSpace(unit.ID) == "" {
|
||||
return nil, fmt.Errorf("chunk %q unit[%d].id must not be empty", chunk.ID, unitIndex)
|
||||
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 %q", chunk.ID, unit.ID)
|
||||
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 %q was not found in source document %q", chunk.ID, unit.ID, doc.ID)
|
||||
return nil, fmt.Errorf("chunk %q source unit %d was not found in source document %q", chunk.ID, unit.ID, doc.ID)
|
||||
}
|
||||
if sourceIndex <= previousSourceIndex {
|
||||
return nil, fmt.Errorf("chunk %q source units must appear in source document order", chunk.ID)
|
||||
@@ -61,11 +82,15 @@ func validateAndCanonicalizeChunkResult(doc *source.SourceDocument, chunks []con
|
||||
}
|
||||
|
||||
canonicalChunks = append(canonicalChunks, contracts.SourceChunk{
|
||||
ID: chunk.ID,
|
||||
SourceID: chunk.SourceID,
|
||||
Index: chunk.Index,
|
||||
Units: canonicalUnits,
|
||||
Metadata: cloneMetadata(chunk.Metadata),
|
||||
ID: chunk.ID,
|
||||
SourceID: chunk.SourceID,
|
||||
Index: chunk.Index,
|
||||
StartUnitID: chunk.StartUnitID,
|
||||
EndUnitID: chunk.EndUnitID,
|
||||
Content: append([]byte(nil), chunk.Content...),
|
||||
MediaType: chunk.MediaType,
|
||||
Units: canonicalUnits,
|
||||
Metadata: cloneMetadata(chunk.Metadata),
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -301,7 +301,7 @@ func (adapter fakeAdapter) Parse(ctx context.Context, req contracts.ParseRequest
|
||||
Format: "text/plain",
|
||||
Digest: "sha256:abc123",
|
||||
Units: []source.SourceUnit{
|
||||
{ID: "u1", Kind: "unit", Text: "Source unit."},
|
||||
{ID: 1, Kind: "unit", Text: "Source unit."},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
||||
@@ -126,10 +126,14 @@ func (chunker integrationChunker) Chunk(ctx context.Context, req contracts.Chunk
|
||||
return contracts.ChunkResult{
|
||||
Chunks: []contracts.SourceChunk{
|
||||
{
|
||||
ID: "chunk-0",
|
||||
SourceID: req.Source.ID,
|
||||
Index: 0,
|
||||
Units: req.Source.Units,
|
||||
ID: "chunk-0",
|
||||
SourceID: req.Source.ID,
|
||||
Index: 0,
|
||||
StartUnitID: req.Source.Units[0].ID,
|
||||
EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID,
|
||||
Content: []byte(`{"units":[1]}`),
|
||||
MediaType: "application/json",
|
||||
Units: req.Source.Units,
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
@@ -267,7 +271,7 @@ func integrationSourceDocument() *source.SourceDocument {
|
||||
Format: "text/plain",
|
||||
Digest: "sha256:abc123",
|
||||
Units: []source.SourceUnit{
|
||||
{ID: "u1", Kind: "unit", Text: "Source unit."},
|
||||
{ID: 1, Kind: "unit", Text: "Source unit."},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -264,45 +264,70 @@ func TestRunRejectsInvalidChunks(t *testing.T) {
|
||||
}{
|
||||
{
|
||||
name: "empty chunk id",
|
||||
chunks: []contracts.SourceChunk{{ID: "", SourceID: "source-1", Index: 0, Units: []source.SourceUnit{unitWithID("u1")}}},
|
||||
chunks: []contracts.SourceChunk{chunkWithUnits("", "source-1", 0, unitWithID("u1"))},
|
||||
want: "id must not be empty",
|
||||
},
|
||||
{
|
||||
name: "duplicate chunk id",
|
||||
chunks: []contracts.SourceChunk{
|
||||
{ID: "chunk-0", SourceID: "source-1", Index: 0, Units: []source.SourceUnit{unitWithID("u1")}},
|
||||
{ID: "chunk-0", SourceID: "source-1", Index: 1, Units: []source.SourceUnit{unitWithID("u2")}},
|
||||
chunkWithUnits("chunk-0", "source-1", 0, unitWithID("u1")),
|
||||
chunkWithUnits("chunk-0", "source-1", 1, unitWithID("u2")),
|
||||
},
|
||||
want: "duplicated",
|
||||
},
|
||||
{
|
||||
name: "wrong source id",
|
||||
chunks: []contracts.SourceChunk{{ID: "chunk-0", SourceID: "other-source", Index: 0, Units: []source.SourceUnit{unitWithID("u1")}}},
|
||||
chunks: []contracts.SourceChunk{chunkWithUnits("chunk-0", "other-source", 0, unitWithID("u1"))},
|
||||
want: "source_id",
|
||||
},
|
||||
{
|
||||
name: "wrong index",
|
||||
chunks: []contracts.SourceChunk{{ID: "chunk-0", SourceID: "source-1", Index: 1, Units: []source.SourceUnit{unitWithID("u1")}}},
|
||||
chunks: []contracts.SourceChunk{chunkWithUnits("chunk-0", "source-1", 1, unitWithID("u1"))},
|
||||
want: "index",
|
||||
},
|
||||
{
|
||||
name: "unknown start id",
|
||||
chunks: []contracts.SourceChunk{chunkWithBounds("chunk-0", "source-1", 0, 9, 1, unitWithID("u1"))},
|
||||
want: "start_unit_id",
|
||||
},
|
||||
{
|
||||
name: "unknown end id",
|
||||
chunks: []contracts.SourceChunk{chunkWithBounds("chunk-0", "source-1", 0, 1, 9, unitWithID("u1"))},
|
||||
want: "end_unit_id",
|
||||
},
|
||||
{
|
||||
name: "reversed bounds",
|
||||
chunks: []contracts.SourceChunk{chunkWithBounds("chunk-0", "source-1", 0, 2, 1, unitWithID("u1"), unitWithID("u2"))},
|
||||
want: "appears after",
|
||||
},
|
||||
{
|
||||
name: "empty units",
|
||||
chunks: []contracts.SourceChunk{{ID: "chunk-0", SourceID: "source-1", Index: 0}},
|
||||
chunks: []contracts.SourceChunk{{ID: "chunk-0", SourceID: "source-1", Index: 0, StartUnitID: 1, EndUnitID: 1, Content: []byte(`{"units":[]}`), MediaType: "application/json"}},
|
||||
want: "units must not be empty",
|
||||
},
|
||||
{
|
||||
name: "empty content",
|
||||
chunks: []contracts.SourceChunk{{ID: "chunk-0", SourceID: "source-1", Index: 0, StartUnitID: 1, EndUnitID: 1, MediaType: "application/json", Units: []source.SourceUnit{unitWithID("u1")}}},
|
||||
want: "content must not be empty",
|
||||
},
|
||||
{
|
||||
name: "empty media type",
|
||||
chunks: []contracts.SourceChunk{{ID: "chunk-0", SourceID: "source-1", Index: 0, StartUnitID: 1, EndUnitID: 1, Content: []byte(`{"units":[1]}`), Units: []source.SourceUnit{unitWithID("u1")}}},
|
||||
want: "media_type must not be empty",
|
||||
},
|
||||
{
|
||||
name: "repeated unit inside chunk",
|
||||
chunks: []contracts.SourceChunk{{ID: "chunk-0", SourceID: "source-1", Index: 0, Units: []source.SourceUnit{unitWithID("u1"), unitWithID("u1")}}},
|
||||
chunks: []contracts.SourceChunk{chunkWithUnits("chunk-0", "source-1", 0, unitWithID("u1"), unitWithID("u1"))},
|
||||
want: "repeats source unit",
|
||||
},
|
||||
{
|
||||
name: "unknown unit",
|
||||
chunks: []contracts.SourceChunk{{ID: "chunk-0", SourceID: "source-1", Index: 0, Units: []source.SourceUnit{unitWithID("u9")}}},
|
||||
chunks: []contracts.SourceChunk{chunkWithUnits("chunk-0", "source-1", 0, unitWithID("u9"))},
|
||||
want: "was not found",
|
||||
},
|
||||
{
|
||||
name: "units out of source order",
|
||||
chunks: []contracts.SourceChunk{{ID: "chunk-0", SourceID: "source-1", Index: 0, Units: []source.SourceUnit{unitWithID("u2"), unitWithID("u1")}}},
|
||||
chunks: []contracts.SourceChunk{chunkWithBounds("chunk-0", "source-1", 0, 1, 2, unitWithID("u2"), unitWithID("u1"))},
|
||||
want: "source document order",
|
||||
},
|
||||
}
|
||||
@@ -328,8 +353,8 @@ func TestRunRejectsInvalidChunks(t *testing.T) {
|
||||
func TestRunAllowsPartialCoverageAndOverlappingChunks(t *testing.T) {
|
||||
modules := defaultRunnerModules()
|
||||
modules.chunker.chunks = []contracts.SourceChunk{
|
||||
{ID: "chunk-0", SourceID: "source-1", Index: 0, Units: []source.SourceUnit{unitWithID("u1"), unitWithID("u2")}},
|
||||
{ID: "chunk-1", SourceID: "source-1", Index: 1, Units: []source.SourceUnit{unitWithID("u2")}},
|
||||
chunkWithUnits("chunk-0", "source-1", 0, unitWithID("u1"), unitWithID("u2")),
|
||||
chunkWithUnits("chunk-1", "source-1", 1, unitWithID("u2")),
|
||||
}
|
||||
|
||||
output, err := New(newRunnerRegistries(t, modules)).Run(context.Background(), RunInput{Pipeline: resolvedPipeline()})
|
||||
@@ -346,12 +371,16 @@ func TestRunCanonicalizesChunkUnitsBeforeExtraction(t *testing.T) {
|
||||
modules.input.doc = sourceDocumentWithUnitMetadata()
|
||||
modules.chunker.chunks = []contracts.SourceChunk{
|
||||
{
|
||||
ID: "chunk-0",
|
||||
SourceID: "source-1",
|
||||
Index: 0,
|
||||
ID: "chunk-0",
|
||||
SourceID: "source-1",
|
||||
Index: 0,
|
||||
StartUnitID: 1,
|
||||
EndUnitID: 1,
|
||||
Content: []byte(`{"units":[{"id":1}]}`),
|
||||
MediaType: "application/json",
|
||||
Units: []source.SourceUnit{
|
||||
{
|
||||
ID: "u1",
|
||||
ID: 1,
|
||||
Kind: "mutated-kind",
|
||||
Text: "mutated text",
|
||||
Metadata: map[string]any{
|
||||
@@ -380,7 +409,7 @@ func TestRunCanonicalizesChunkUnitsBeforeExtraction(t *testing.T) {
|
||||
if chunk == nil {
|
||||
t.Fatal("extractor chunk = nil, want canonical chunk")
|
||||
}
|
||||
if chunk.Units[0].ID != "u1" || chunk.Units[0].Kind != "source-kind" || chunk.Units[0].Text != "source text" {
|
||||
if chunk.Units[0].ID != 1 || chunk.Units[0].Kind != "source-kind" || chunk.Units[0].Text != "source text" {
|
||||
t.Fatalf("chunk unit = %#v, want source document unit values", chunk.Units[0])
|
||||
}
|
||||
if got := chunk.Units[0].Metadata["speaker"]; got != "source-speaker" {
|
||||
@@ -402,9 +431,13 @@ func TestRunPreservesChunkMetadataDuringCanonicalization(t *testing.T) {
|
||||
modules := defaultRunnerModules()
|
||||
modules.chunker.chunks = []contracts.SourceChunk{
|
||||
{
|
||||
ID: "chunk-0",
|
||||
SourceID: "source-1",
|
||||
Index: 0,
|
||||
ID: "chunk-0",
|
||||
SourceID: "source-1",
|
||||
Index: 0,
|
||||
StartUnitID: 1,
|
||||
EndUnitID: 1,
|
||||
Content: []byte(`{"units":[{"id":1}]}`),
|
||||
MediaType: "application/json",
|
||||
Units: []source.SourceUnit{
|
||||
unitWithID("u1"),
|
||||
},
|
||||
@@ -1739,9 +1772,9 @@ func validSourceDocument() *source.SourceDocument {
|
||||
Format: "text/plain",
|
||||
Digest: "sha256:source",
|
||||
Units: []source.SourceUnit{
|
||||
{ID: "u1", Kind: "unit", Text: "Source unit."},
|
||||
{ID: "u2", Kind: "unit", Text: "Second source unit."},
|
||||
{ID: "u3", Kind: "unit", Text: "Third source unit."},
|
||||
{ID: 1, Kind: "unit", Text: "Source unit."},
|
||||
{ID: 2, Kind: "unit", Text: "Second source unit."},
|
||||
{ID: 3, Kind: "unit", Text: "Third source unit."},
|
||||
},
|
||||
}
|
||||
}
|
||||
@@ -1754,7 +1787,7 @@ func sourceDocumentWithUnitMetadata() *source.SourceDocument {
|
||||
Digest: "sha256:source",
|
||||
Units: []source.SourceUnit{
|
||||
{
|
||||
ID: "u1",
|
||||
ID: 1,
|
||||
Kind: "source-kind",
|
||||
Text: "source text",
|
||||
Metadata: map[string]any{
|
||||
@@ -1763,7 +1796,7 @@ func sourceDocumentWithUnitMetadata() *source.SourceDocument {
|
||||
},
|
||||
},
|
||||
{
|
||||
ID: "u2",
|
||||
ID: 2,
|
||||
Kind: "source-kind",
|
||||
Text: "second source text",
|
||||
Metadata: map[string]any{
|
||||
@@ -1775,26 +1808,53 @@ func sourceDocumentWithUnitMetadata() *source.SourceDocument {
|
||||
}
|
||||
|
||||
func sourceChunkWithID(id string, index int) contracts.SourceChunk {
|
||||
unit := unitWithID("u1")
|
||||
return contracts.SourceChunk{
|
||||
ID: id,
|
||||
SourceID: "source-1",
|
||||
Index: index,
|
||||
Units: []source.SourceUnit{
|
||||
{ID: "u1", Kind: "unit", Text: "Source unit."},
|
||||
},
|
||||
ID: id,
|
||||
SourceID: "source-1",
|
||||
Index: index,
|
||||
StartUnitID: unit.ID,
|
||||
EndUnitID: unit.ID,
|
||||
Content: []byte(`{"units":[{"id":1,"kind":"unit","text":"Source unit."}]}`),
|
||||
MediaType: "application/json",
|
||||
Units: []source.SourceUnit{unit},
|
||||
}
|
||||
}
|
||||
|
||||
func unitWithID(id string) source.SourceUnit {
|
||||
switch id {
|
||||
case "u1":
|
||||
return source.SourceUnit{ID: "u1", Kind: "unit", Text: "Source unit."}
|
||||
return source.SourceUnit{ID: 1, Kind: "unit", Text: "Source unit."}
|
||||
case "u2":
|
||||
return source.SourceUnit{ID: "u2", Kind: "unit", Text: "Second source unit."}
|
||||
return source.SourceUnit{ID: 2, Kind: "unit", Text: "Second source unit."}
|
||||
case "u3":
|
||||
return source.SourceUnit{ID: "u3", Kind: "unit", Text: "Third source unit."}
|
||||
return source.SourceUnit{ID: 3, Kind: "unit", Text: "Third source unit."}
|
||||
case "u9":
|
||||
return source.SourceUnit{ID: 9, Kind: "unit", Text: "Unknown source unit."}
|
||||
default:
|
||||
return source.SourceUnit{ID: id, Kind: "unit", Text: "Unknown source unit."}
|
||||
return source.SourceUnit{ID: 99, Kind: "unit", Text: "Unknown source unit."}
|
||||
}
|
||||
}
|
||||
|
||||
func chunkWithUnits(id string, sourceID string, index int, units ...source.SourceUnit) contracts.SourceChunk {
|
||||
startUnitID, endUnitID := 1, 1
|
||||
if len(units) > 0 {
|
||||
startUnitID = units[0].ID
|
||||
endUnitID = units[len(units)-1].ID
|
||||
}
|
||||
return chunkWithBounds(id, sourceID, index, startUnitID, endUnitID, units...)
|
||||
}
|
||||
|
||||
func chunkWithBounds(id string, sourceID string, index int, startUnitID int, endUnitID int, units ...source.SourceUnit) contracts.SourceChunk {
|
||||
return contracts.SourceChunk{
|
||||
ID: id,
|
||||
SourceID: sourceID,
|
||||
Index: index,
|
||||
StartUnitID: startUnitID,
|
||||
EndUnitID: endUnitID,
|
||||
Content: []byte(`{"units":[1]}`),
|
||||
MediaType: "application/json",
|
||||
Units: append([]source.SourceUnit(nil), units...),
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@@ -2,15 +2,15 @@
|
||||
"id": "fixture-source",
|
||||
"units": [
|
||||
{
|
||||
"id": "u1",
|
||||
"id": 1,
|
||||
"text": "First event."
|
||||
},
|
||||
{
|
||||
"id": "u2",
|
||||
"id": 2,
|
||||
"text": "Second event."
|
||||
},
|
||||
{
|
||||
"id": "u3",
|
||||
"id": 3,
|
||||
"text": "Third event."
|
||||
}
|
||||
]
|
||||
|
||||
@@ -25,8 +25,8 @@
|
||||
"source_refs": [
|
||||
{
|
||||
"source_id": "fixture-source",
|
||||
"start_unit_id": "u1",
|
||||
"end_unit_id": "u2"
|
||||
"start_unit_id": 1,
|
||||
"end_unit_id": 2
|
||||
}
|
||||
]
|
||||
},
|
||||
@@ -42,8 +42,8 @@
|
||||
"source_refs": [
|
||||
{
|
||||
"source_id": "fixture-source",
|
||||
"start_unit_id": "u3",
|
||||
"end_unit_id": "u3"
|
||||
"start_unit_id": 3,
|
||||
"end_unit_id": 3
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
@@ -185,7 +185,7 @@ func (input walkingSkeletonInput) Parse(ctx context.Context, req contracts.Parse
|
||||
var fixture struct {
|
||||
ID string `json:"id"`
|
||||
Units []struct {
|
||||
ID string `json:"id"`
|
||||
ID int `json:"id"`
|
||||
Text string `json:"text"`
|
||||
} `json:"units"`
|
||||
}
|
||||
@@ -227,16 +227,24 @@ func (chunker walkingSkeletonChunker) Chunk(ctx context.Context, req contracts.C
|
||||
return contracts.ChunkResult{
|
||||
Chunks: []contracts.SourceChunk{
|
||||
{
|
||||
ID: req.Source.ID + ":chunk:0",
|
||||
SourceID: req.Source.ID,
|
||||
Index: 0,
|
||||
Units: append([]source.SourceUnit(nil), req.Source.Units[:2]...),
|
||||
ID: req.Source.ID + ":chunk:0",
|
||||
SourceID: req.Source.ID,
|
||||
Index: 0,
|
||||
StartUnitID: req.Source.Units[0].ID,
|
||||
EndUnitID: req.Source.Units[1].ID,
|
||||
Content: []byte(`{"units":[1,2]}`),
|
||||
MediaType: "application/json",
|
||||
Units: append([]source.SourceUnit(nil), req.Source.Units[:2]...),
|
||||
},
|
||||
{
|
||||
ID: req.Source.ID + ":chunk:1",
|
||||
SourceID: req.Source.ID,
|
||||
Index: 1,
|
||||
Units: append([]source.SourceUnit(nil), req.Source.Units[2:]...),
|
||||
ID: req.Source.ID + ":chunk:1",
|
||||
SourceID: req.Source.ID,
|
||||
Index: 1,
|
||||
StartUnitID: req.Source.Units[2].ID,
|
||||
EndUnitID: req.Source.Units[len(req.Source.Units)-1].ID,
|
||||
Content: []byte(`{"units":[3]}`),
|
||||
MediaType: "application/json",
|
||||
Units: append([]source.SourceUnit(nil), req.Source.Units[2:]...),
|
||||
},
|
||||
},
|
||||
}, nil
|
||||
|
||||
Reference in New Issue
Block a user