Generate and materialize canonical chunk plans
This commit is contained in:
@@ -20,9 +20,10 @@ var requiredCapabilities = []string{
|
||||
|
||||
var providedCapabilities = []string{
|
||||
"chunks",
|
||||
"chunks.scenes",
|
||||
}
|
||||
|
||||
const annotationNamespace = "dnd/scenes"
|
||||
|
||||
var referenceSlotDescriptions = shared.ReferenceSlotDescriptions{
|
||||
Glossary: "Optional campaign glossary reference material used only for scene disambiguation.",
|
||||
Party: "Optional party roster reference material used only for scene disambiguation.",
|
||||
@@ -74,27 +75,27 @@ func (c *Chunker) ManifestMetadata() map[string]any {
|
||||
return metadata
|
||||
}
|
||||
|
||||
func (c *Chunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
func (c *Chunker) Plan(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkPlanResult, error) {
|
||||
if c == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("chunker must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("chunker must not be nil")
|
||||
}
|
||||
if c.llm == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("LLM client must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("LLM client must not be nil")
|
||||
}
|
||||
if ctx == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("context must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("context must not be nil")
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("context error before chunking: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("context error before chunking: %w", err)
|
||||
}
|
||||
if req.Source == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("source must not be nil")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("source must not be nil")
|
||||
}
|
||||
if len(req.Source.Units) == 0 {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("source units must not be empty")
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("source units must not be empty")
|
||||
}
|
||||
if err := source.ValidateDocument(req.Source); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("validate source document: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("validate source document: %w", err)
|
||||
}
|
||||
var response chunkResponse
|
||||
if _, err := c.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
||||
@@ -105,19 +106,19 @@ func (c *Chunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contra
|
||||
SessionID: req.SessionID,
|
||||
Inputs: shared.PromptInputs(req.SourceInput, req.References),
|
||||
}, &response); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("complete structured output: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("complete structured output: %w", err)
|
||||
}
|
||||
|
||||
warnings, err := warningsFromCaveats(response.BoundaryCaveats)
|
||||
if err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("malformed structured output: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("malformed structured output: %w", err)
|
||||
}
|
||||
chunks, err := chunksFromResponse(req.Source, response)
|
||||
plan, err := planFromResponse(req.Source, response, warnings)
|
||||
if err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("malformed structured output: %w", err)
|
||||
return contracts.ChunkPlanResult{}, chunkerErrorf("malformed structured output: %w", err)
|
||||
}
|
||||
return contracts.ChunkResult{
|
||||
Chunks: chunks,
|
||||
return contracts.ChunkPlanResult{
|
||||
Plan: plan,
|
||||
Warnings: warnings,
|
||||
}, nil
|
||||
}
|
||||
@@ -154,12 +155,12 @@ func DecodeOptions(options map[string]any) (Options, error) {
|
||||
return Options{}, nil
|
||||
}
|
||||
|
||||
func chunksFromResponse(doc *source.SourceDocument, response chunkResponse) ([]source.Chunk, error) {
|
||||
func planFromResponse(doc *source.SourceDocument, response chunkResponse, warnings []contracts.Warning) (source.ChunkPlan, error) {
|
||||
if response.Scenes == nil {
|
||||
return nil, fmt.Errorf("scenes must be present")
|
||||
return source.ChunkPlan{}, fmt.Errorf("scenes must be present")
|
||||
}
|
||||
if len(response.Scenes) == 0 {
|
||||
return nil, fmt.Errorf("scenes must not be empty")
|
||||
return source.ChunkPlan{}, fmt.Errorf("scenes must not be empty")
|
||||
}
|
||||
|
||||
unitIndexes := make(map[int]int, len(doc.Units))
|
||||
@@ -167,86 +168,75 @@ func chunksFromResponse(doc *source.SourceDocument, response chunkResponse) ([]s
|
||||
unitIndexes[unit.ID] = i
|
||||
}
|
||||
|
||||
chunks := make([]source.Chunk, 0, len(response.Scenes))
|
||||
ranges := make([]source.ChunkRange, 0, len(response.Scenes))
|
||||
previousEnd := -1
|
||||
for i, scene := range response.Scenes {
|
||||
normalized, err := normalizeScene(doc, i, scene)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return source.ChunkPlan{}, err
|
||||
}
|
||||
|
||||
startIndex, ok := unitIndexes[normalized.StartUnitID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("scene[%d] start_unit_id %d was not found", i, normalized.StartUnitID)
|
||||
return source.ChunkPlan{}, fmt.Errorf("scene[%d] start_unit_id %d was not found", i, normalized.StartUnitID)
|
||||
}
|
||||
endIndex, ok := unitIndexes[normalized.EndUnitID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("scene[%d] end_unit_id %d was not found", i, normalized.EndUnitID)
|
||||
return source.ChunkPlan{}, fmt.Errorf("scene[%d] end_unit_id %d was not found", i, normalized.EndUnitID)
|
||||
}
|
||||
if startIndex > endIndex {
|
||||
return nil, fmt.Errorf("scene[%d] start_unit_id %d appears after end_unit_id %d", i, normalized.StartUnitID, normalized.EndUnitID)
|
||||
return source.ChunkPlan{}, fmt.Errorf("scene[%d] start_unit_id %d appears after end_unit_id %d", i, normalized.StartUnitID, normalized.EndUnitID)
|
||||
}
|
||||
|
||||
if i == 0 && startIndex != 0 {
|
||||
return nil, fmt.Errorf("first scene must start at first source unit %d", doc.Units[0].ID)
|
||||
return source.ChunkPlan{}, fmt.Errorf("first scene must start at first source unit %d", doc.Units[0].ID)
|
||||
}
|
||||
if i > 0 {
|
||||
if startIndex <= previousEnd {
|
||||
return nil, fmt.Errorf("scene[%d] overlaps previous scene", i)
|
||||
return source.ChunkPlan{}, fmt.Errorf("scene[%d] overlaps previous scene", i)
|
||||
}
|
||||
if startIndex > previousEnd+1 {
|
||||
return nil, fmt.Errorf("scene[%d] leaves a gap after previous scene", i)
|
||||
return source.ChunkPlan{}, fmt.Errorf("scene[%d] leaves a gap after previous scene", i)
|
||||
}
|
||||
}
|
||||
previousEnd = endIndex
|
||||
|
||||
units := cloneUnits(doc.Units[startIndex : endIndex+1])
|
||||
content, err := chunkContent(units)
|
||||
annotation, err := json.Marshal(struct {
|
||||
ShortTitle string `json:"short_title"`
|
||||
PrimaryMode string `json:"primary_mode"`
|
||||
MainParticipants []string `json:"main_participants"`
|
||||
Summary string `json:"summary"`
|
||||
BoundaryNote string `json:"boundary_note"`
|
||||
BoundaryConfidence string `json:"boundary_confidence"`
|
||||
}{normalized.ShortTitle, normalized.PrimaryMode, normalized.MainParticipants, normalized.Summary, normalized.BoundaryNote, normalized.BoundaryConfidence})
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return source.ChunkPlan{}, fmt.Errorf("encode scene[%d] annotation: %w", i, err)
|
||||
}
|
||||
chunks = append(chunks, source.Chunk{
|
||||
ID: fmt.Sprintf("scene-%06d", i+1),
|
||||
SourceID: doc.ID,
|
||||
Index: i,
|
||||
Ref: source.SourceRef{
|
||||
SourceID: doc.ID,
|
||||
StartUnitID: units[0].Ref.StartUnitID,
|
||||
EndUnitID: units[len(units)-1].Ref.EndUnitID,
|
||||
},
|
||||
Content: content,
|
||||
MediaType: "application/json",
|
||||
Units: units,
|
||||
Metadata: map[string]any{
|
||||
"scene_title": normalized.ShortTitle,
|
||||
"primary_mode": normalized.PrimaryMode,
|
||||
"main_participants": append([]string(nil), normalized.MainParticipants...),
|
||||
"summary": normalized.Summary,
|
||||
"boundary_note": normalized.BoundaryNote,
|
||||
"boundary_confidence": normalized.BoundaryConfidence,
|
||||
"start_unit_id": normalized.StartUnitID,
|
||||
"end_unit_id": normalized.EndUnitID,
|
||||
"unit_count": len(units),
|
||||
},
|
||||
ranges = append(ranges, source.ChunkRange{
|
||||
StartUnitID: normalized.StartUnitID,
|
||||
EndUnitID: normalized.EndUnitID,
|
||||
Annotations: source.ChunkAnnotations{annotationNamespace: annotation},
|
||||
})
|
||||
}
|
||||
|
||||
if previousEnd != len(doc.Units)-1 {
|
||||
return nil, fmt.Errorf("final scene must end at final source unit %d", doc.Units[len(doc.Units)-1].ID)
|
||||
return source.ChunkPlan{}, fmt.Errorf("final scene must end at final source unit %d", doc.Units[len(doc.Units)-1].ID)
|
||||
}
|
||||
return chunks, nil
|
||||
}
|
||||
|
||||
func chunkContent(units []source.SourceUnit) ([]byte, error) {
|
||||
content, err := json.Marshal(struct {
|
||||
Units []source.SourceUnit `json:"units"`
|
||||
}{
|
||||
Units: units,
|
||||
})
|
||||
caveats := make([]string, 0, len(warnings))
|
||||
for _, warning := range warnings {
|
||||
caveats = append(caveats, warning.Message)
|
||||
}
|
||||
annotation, err := json.Marshal(struct {
|
||||
BoundaryCaveats []string `json:"boundary_caveats"`
|
||||
}{BoundaryCaveats: caveats})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("encode chunk content: %w", err)
|
||||
return source.ChunkPlan{}, fmt.Errorf("encode plan annotation: %w", err)
|
||||
}
|
||||
return content, nil
|
||||
return source.ChunkPlan{
|
||||
SourceDigest: doc.Digest,
|
||||
Ranges: ranges,
|
||||
Annotations: source.ChunkAnnotations{annotationNamespace: annotation},
|
||||
}, nil
|
||||
}
|
||||
|
||||
func normalizeScene(doc *source.SourceDocument, index int, scene sceneResponse) (normalizedScene, error) {
|
||||
@@ -347,31 +337,6 @@ func warningsFromCaveats(caveats []string) ([]contracts.Warning, error) {
|
||||
return warnings, nil
|
||||
}
|
||||
|
||||
func cloneUnits(units []source.SourceUnit) []source.SourceUnit {
|
||||
out := make([]source.SourceUnit, 0, len(units))
|
||||
for _, unit := range units {
|
||||
out = append(out, source.SourceUnit{
|
||||
ID: unit.ID,
|
||||
Kind: unit.Kind,
|
||||
Text: unit.Text,
|
||||
Ref: unit.Ref,
|
||||
Metadata: cloneMetadata(unit.Metadata),
|
||||
})
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func cloneMetadata(metadata map[string]any) map[string]any {
|
||||
if len(metadata) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make(map[string]any, len(metadata))
|
||||
for key, value := range metadata {
|
||||
out[key] = value
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func chunkerErrorf(format string, args ...any) error {
|
||||
return fmt.Errorf("dnd scenes chunker: "+format, args...)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user