Move chunks into the canonical source model
This commit is contained in:
@@ -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
|
||||
|
||||
@@ -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"`
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user