Move chunks into the canonical source model

This commit is contained in:
2026-07-17 05:32:11 +00:00
parent 40709e4ad8
commit 075888c97f
33 changed files with 548 additions and 325 deletions

View File

@@ -141,17 +141,20 @@ func (chunker compositionChunker) Chunk(ctx context.Context, req contracts.Chunk
}
return contracts.ChunkResult{
Chunks: []contracts.SourceChunk{
Chunks: []source.Chunk{
{
ID: req.Source.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":[{"id":1,"kind":"unit","text":"First source unit."},{"id":2,"kind":"unit","text":"Second source unit."}]}`),
MediaType: "application/json",
Units: append([]source.SourceUnit(nil), req.Source.Units...),
Metadata: map[string]any{"strategy": "whole-document"},
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":"unit","text":"First source unit."},{"id":2,"kind":"unit","text":"Second source unit."}]}`),
MediaType: "application/json",
Units: append([]source.SourceUnit(nil), req.Source.Units...),
Metadata: map[string]any{"strategy": "whole-document"},
},
},
}, nil

View File

@@ -138,18 +138,6 @@ type InputAdapter interface {
Parse(ctx context.Context, req ParseRequest) (*source.SourceDocument, error)
}
type SourceChunk struct {
ID string `json:"id"`
SourceID string `json:"source_id"`
Index int `json:"index"`
StartUnitID int `json:"start_unit_id"`
EndUnitID int `json:"end_unit_id"`
Content []byte `json:"-"`
MediaType string `json:"media_type"`
Units []source.SourceUnit `json:"units"`
Metadata map[string]any `json:"metadata,omitempty"`
}
type ChunkRequest struct {
Source *source.SourceDocument `json:"-"`
SourceInput LLMInputMaterial `json:"source_input,omitempty"`
@@ -162,8 +150,8 @@ type ChunkRequest struct {
}
type ChunkResult struct {
Chunks []SourceChunk `json:"chunks"`
Warnings []Warning `json:"warnings,omitempty"`
Chunks []source.Chunk `json:"chunks"`
Warnings []Warning `json:"warnings,omitempty"`
}
type Chunker interface {
@@ -224,7 +212,7 @@ type ReferenceSet struct {
type ExtractionRequest struct {
Source *source.SourceDocument `json:"-"`
Chunk *SourceChunk `json:"chunk,omitempty"`
Chunk *source.Chunk `json:"chunk,omitempty"`
AmbientContext map[string]any `json:"ambient_context,omitempty"`
SourceInput LLMInputMaterial `json:"source_input,omitempty"`
SessionID string `json:"session_id,omitempty"`
@@ -277,8 +265,8 @@ type ValidationRequest struct {
Payload RawPayload `json:"payload"`
ChunkID string `json:"chunk_id,omitempty"`
ChunkIndex int `json:"chunk_index,omitempty"`
Chunk *SourceChunk `json:"chunk,omitempty"`
Chunks []SourceChunk `json:"chunks,omitempty"`
Chunk *source.Chunk `json:"chunk,omitempty"`
Chunks []source.Chunk `json:"chunks,omitempty"`
ExtractOutputs []ExtractOutput `json:"extract_outputs,omitempty"`
MergeOutput MergeOutput `json:"merge_output,omitempty"`
}

View File

@@ -78,22 +78,22 @@ func TestFakeChunkerReturnsSourceChunks(t *testing.T) {
chunk := result.Chunks[0]
if chunk.ID != "source-1:chunk:0" {
t.Fatalf("SourceChunk.ID = %q, want source-1:chunk:0", chunk.ID)
t.Fatalf("source.Chunk.ID = %q, want source-1:chunk:0", chunk.ID)
}
if chunk.SourceID != doc.ID {
t.Fatalf("SourceChunk.SourceID = %q, want %q", chunk.SourceID, doc.ID)
t.Fatalf("source.Chunk.SourceID = %q, want %q", chunk.SourceID, doc.ID)
}
if chunk.Index != 0 {
t.Fatalf("SourceChunk.Index = %d, want 0", chunk.Index)
t.Fatalf("source.Chunk.Index = %d, want 0", chunk.Index)
}
if chunk.StartUnitID != 1 || chunk.EndUnitID != 1 {
t.Fatalf("SourceChunk boundaries = %d-%d, want 1-1", chunk.StartUnitID, chunk.EndUnitID)
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("SourceChunk payload = %q %s, want JSON units", chunk.MediaType, chunk.Content)
t.Fatalf("source.Chunk payload = %q %s, want JSON units", chunk.MediaType, chunk.Content)
}
if len(chunk.Units) != 1 {
t.Fatalf("len(SourceChunk.Units) = %d, want 1", len(chunk.Units))
t.Fatalf("len(source.Chunk.Units) = %d, want 1", len(chunk.Units))
}
}
@@ -130,15 +130,14 @@ func TestFakeExtractorReceivesChunkAndAmbientContext(t *testing.T) {
{ID: 2, Kind: "section", Text: "Second source text."},
},
}
chunk := SourceChunk{
ID: "source-1:chunk:1",
SourceID: doc.ID,
Index: 1,
StartUnitID: 2,
EndUnitID: 2,
Content: []byte(`{"units":[{"id":2,"kind":"section","text":"Second source text."}]}`),
MediaType: "application/json",
Units: []source.SourceUnit{doc.Units[1]},
chunk := source.Chunk{
ID: "source-1:chunk:1",
SourceID: doc.ID,
Index: 1,
Ref: source.SourceRef{SourceID: doc.ID, StartUnitID: 2, EndUnitID: 2},
Content: []byte(`{"units":[{"id":2,"kind":"section","text":"Second source text."}]}`),
MediaType: "application/json",
Units: []source.SourceUnit{doc.Units[1]},
}
result, err := extractor.Extract(context.Background(), ExtractionRequest{
@@ -472,16 +471,19 @@ func (chunker fakeChunker) ReferenceSlots() []ReferenceSlot {
func (chunker fakeChunker) Chunk(ctx context.Context, req ChunkRequest) (ChunkResult, error) {
return ChunkResult{
Chunks: []SourceChunk{
Chunks: []source.Chunk{
{
ID: req.Source.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":[{"id":1,"kind":"section","text":"Source text."}]}`),
MediaType: "application/json",
Units: append([]source.SourceUnit(nil), req.Source.Units...),
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...),
},
},
}, nil