Implement D&D scene chunker
This commit is contained in:
303
internal/modules/chunk/dnd/scenes/chunker.go
Normal file
303
internal/modules/chunk/dnd/scenes/chunker.go
Normal file
@@ -0,0 +1,303 @@
|
||||
package scenes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
||||
)
|
||||
|
||||
const Key = "dnd/scenes"
|
||||
|
||||
var requiredCapabilities = []string{
|
||||
"source.transcript",
|
||||
}
|
||||
|
||||
var providedCapabilities = []string{
|
||||
"chunks",
|
||||
"chunks.scenes",
|
||||
}
|
||||
|
||||
var _ contracts.Chunker = (*Chunker)(nil)
|
||||
var _ contracts.ManifestMetadataProvider = (*Chunker)(nil)
|
||||
|
||||
type Chunker struct{}
|
||||
|
||||
func New() *Chunker {
|
||||
return &Chunker{}
|
||||
}
|
||||
|
||||
func (c *Chunker) Key() string {
|
||||
return Key
|
||||
}
|
||||
|
||||
func (c *Chunker) ManifestMetadata() map[string]any {
|
||||
promptMetadata := scenesPromptBundle.Metadata()
|
||||
metadata := map[string]any{
|
||||
"prompt_id": PromptID,
|
||||
"prompt_version": promptMetadata.PromptVersion,
|
||||
"prompt_sha256": promptMetadata.SHA256,
|
||||
"response_schema_key": string(ResponseSchemaKey),
|
||||
"response_schema_id": ResponseSchemaID,
|
||||
"response_schema_name": ResponseSchemaName,
|
||||
}
|
||||
if schema, err := loadResponseSchema(); err == nil {
|
||||
metadata["response_schema_version"] = schema.Version
|
||||
metadata["response_schema_sha256"] = schema.SHA256
|
||||
}
|
||||
return metadata
|
||||
}
|
||||
|
||||
func (c *Chunker) Chunk(ctx context.Context, req contracts.ChunkRequest) (contracts.ChunkResult, error) {
|
||||
if c == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("chunker must not be nil")
|
||||
}
|
||||
if ctx == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("context must not be nil")
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("context error before chunking: %w", err)
|
||||
}
|
||||
if req.Source == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("source must not be nil")
|
||||
}
|
||||
if len(req.Source.Units) == 0 {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("source units must not be empty")
|
||||
}
|
||||
if err := source.ValidateDocument(req.Source); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("validate source document: %w", err)
|
||||
}
|
||||
if req.LLMClient == nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("LLM client must not be nil")
|
||||
}
|
||||
if len(req.Options) > 0 {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("options are not supported")
|
||||
}
|
||||
|
||||
system, user, _, err := renderPrompt(req)
|
||||
if err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("render prompt: %w", err)
|
||||
}
|
||||
schema, err := loadResponseSchema()
|
||||
if err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("load response schema %q: %w", ResponseSchemaKey, err)
|
||||
}
|
||||
|
||||
var response chunkResponse
|
||||
if _, err := req.LLMClient.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
|
||||
StageName: Key,
|
||||
Messages: []contracts.LLMMessage{
|
||||
{Role: "system", Content: system},
|
||||
{Role: "user", Content: user},
|
||||
},
|
||||
ResponseSchemaName: schema.Name,
|
||||
ResponseSchema: schema.JSONSchema,
|
||||
}, &response); err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("complete structured output: %w", err)
|
||||
}
|
||||
|
||||
chunks, err := chunksFromResponse(req.Source, response)
|
||||
if err != nil {
|
||||
return contracts.ChunkResult{}, chunkerErrorf("malformed structured output: %w", err)
|
||||
}
|
||||
return contracts.ChunkResult{
|
||||
Chunks: chunks,
|
||||
Warnings: warningsFromCaveats(response.BoundaryCaveats),
|
||||
}, nil
|
||||
}
|
||||
|
||||
func ModuleSpec() pipeline.ModuleSpec {
|
||||
return pipeline.ModuleSpec{
|
||||
Key: Key,
|
||||
Stage: pipeline.StageChunk,
|
||||
Requires: append([]string(nil), requiredCapabilities...),
|
||||
Provides: append([]string(nil), providedCapabilities...),
|
||||
}
|
||||
}
|
||||
|
||||
func Register(registry *pipeline.ChunkerRegistry) error {
|
||||
return registry.RegisterWithSpec(ModuleSpec(), func() (contracts.Chunker, error) {
|
||||
return New(), nil
|
||||
})
|
||||
}
|
||||
|
||||
func chunksFromResponse(doc *source.SourceDocument, response chunkResponse) ([]contracts.SourceChunk, error) {
|
||||
if response.Scenes == nil {
|
||||
return nil, fmt.Errorf("scenes must be present")
|
||||
}
|
||||
if len(response.Scenes) == 0 {
|
||||
return nil, fmt.Errorf("scenes must not be empty")
|
||||
}
|
||||
|
||||
unitIndexes := make(map[string]int, len(doc.Units))
|
||||
for i, unit := range doc.Units {
|
||||
unitIndexes[unit.ID] = i
|
||||
}
|
||||
|
||||
chunks := make([]contracts.SourceChunk, 0, len(response.Scenes))
|
||||
previousEnd := -1
|
||||
for i, scene := range response.Scenes {
|
||||
normalized, err := normalizeScene(i, scene)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
startIndex, ok := unitIndexes[normalized.StartUnitID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("scene[%d] start_unit_id %q was not found", i, normalized.StartUnitID)
|
||||
}
|
||||
endIndex, ok := unitIndexes[normalized.EndUnitID]
|
||||
if !ok {
|
||||
return nil, fmt.Errorf("scene[%d] end_unit_id %q was not found", i, normalized.EndUnitID)
|
||||
}
|
||||
if startIndex > endIndex {
|
||||
return nil, fmt.Errorf("scene[%d] start_unit_id %q appears after end_unit_id %q", i, normalized.StartUnitID, normalized.EndUnitID)
|
||||
}
|
||||
|
||||
if i == 0 && startIndex != 0 {
|
||||
return nil, fmt.Errorf("first scene must start at first source unit %q", doc.Units[0].ID)
|
||||
}
|
||||
if i > 0 {
|
||||
if startIndex <= previousEnd {
|
||||
return nil, 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)
|
||||
}
|
||||
}
|
||||
previousEnd = endIndex
|
||||
|
||||
units := cloneUnits(doc.Units[startIndex : endIndex+1])
|
||||
chunks = append(chunks, contracts.SourceChunk{
|
||||
ID: fmt.Sprintf("scene-%06d", i+1),
|
||||
SourceID: doc.ID,
|
||||
Index: i,
|
||||
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),
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
if previousEnd != len(doc.Units)-1 {
|
||||
return nil, fmt.Errorf("final scene must end at final source unit %q", doc.Units[len(doc.Units)-1].ID)
|
||||
}
|
||||
return chunks, nil
|
||||
}
|
||||
|
||||
func normalizeScene(index int, scene sceneResponse) (sceneResponse, error) {
|
||||
out := sceneResponse{
|
||||
StartUnitID: strings.TrimSpace(scene.StartUnitID),
|
||||
EndUnitID: strings.TrimSpace(scene.EndUnitID),
|
||||
ShortTitle: strings.TrimSpace(scene.ShortTitle),
|
||||
PrimaryMode: strings.TrimSpace(scene.PrimaryMode),
|
||||
Summary: strings.TrimSpace(scene.Summary),
|
||||
BoundaryNote: strings.TrimSpace(scene.BoundaryNote),
|
||||
BoundaryConfidence: strings.TrimSpace(scene.BoundaryConfidence),
|
||||
}
|
||||
|
||||
required := map[string]string{
|
||||
"start_unit_id": out.StartUnitID,
|
||||
"end_unit_id": out.EndUnitID,
|
||||
"short_title": out.ShortTitle,
|
||||
"primary_mode": out.PrimaryMode,
|
||||
"summary": out.Summary,
|
||||
"boundary_note": out.BoundaryNote,
|
||||
"boundary_confidence": out.BoundaryConfidence,
|
||||
}
|
||||
for field, value := range required {
|
||||
if value == "" {
|
||||
return sceneResponse{}, fmt.Errorf("scene[%d] %s must not be empty", index, field)
|
||||
}
|
||||
}
|
||||
if !validPrimaryMode(out.PrimaryMode) {
|
||||
return sceneResponse{}, fmt.Errorf("scene[%d] primary_mode %q is not supported", index, out.PrimaryMode)
|
||||
}
|
||||
if !validBoundaryConfidence(out.BoundaryConfidence) {
|
||||
return sceneResponse{}, fmt.Errorf("scene[%d] boundary_confidence %q is not supported", index, out.BoundaryConfidence)
|
||||
}
|
||||
if len(scene.MainParticipants) == 0 {
|
||||
return sceneResponse{}, fmt.Errorf("scene[%d] main_participants must not be empty", index)
|
||||
}
|
||||
out.MainParticipants = make([]string, 0, len(scene.MainParticipants))
|
||||
for participantIndex, participant := range scene.MainParticipants {
|
||||
trimmed := strings.TrimSpace(participant)
|
||||
if trimmed == "" {
|
||||
return sceneResponse{}, fmt.Errorf("scene[%d] main_participants[%d] must not be empty", index, participantIndex)
|
||||
}
|
||||
out.MainParticipants = append(out.MainParticipants, trimmed)
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func validPrimaryMode(value string) bool {
|
||||
switch value {
|
||||
case "Recap", "Discussion", "Combat", "Narrative":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func validBoundaryConfidence(value string) bool {
|
||||
switch value {
|
||||
case "High", "Medium", "Low":
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
func warningsFromCaveats(caveats []string) []contracts.Warning {
|
||||
if len(caveats) == 0 {
|
||||
return nil
|
||||
}
|
||||
warnings := make([]contracts.Warning, 0, len(caveats))
|
||||
for _, caveat := range caveats {
|
||||
warnings = append(warnings, contracts.Warning{
|
||||
Scope: Key,
|
||||
ReasonCode: "scene_boundary_caveat",
|
||||
Message: caveat,
|
||||
})
|
||||
}
|
||||
return warnings
|
||||
}
|
||||
|
||||
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,
|
||||
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...)
|
||||
}
|
||||
465
internal/modules/chunk/dnd/scenes/chunker_test.go
Normal file
465
internal/modules/chunk/dnd/scenes/chunker_test.go
Normal file
@@ -0,0 +1,465 @@
|
||||
package scenes
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"reflect"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
|
||||
"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
|
||||
)
|
||||
|
||||
func TestNewModuleSpecAndRegister(t *testing.T) {
|
||||
chunker := New()
|
||||
if chunker == nil {
|
||||
t.Fatal("New() = nil, want chunker")
|
||||
}
|
||||
if chunker.Key() != Key {
|
||||
t.Fatalf("Key() = %q, want %q", chunker.Key(), Key)
|
||||
}
|
||||
|
||||
want := pipeline.ModuleSpec{
|
||||
Key: Key,
|
||||
Stage: pipeline.StageChunk,
|
||||
Requires: []string{"source.transcript"},
|
||||
Provides: []string{"chunks", "chunks.scenes"},
|
||||
}
|
||||
if got := ModuleSpec(); !reflect.DeepEqual(got, want) {
|
||||
t.Fatalf("ModuleSpec() = %#v, want %#v", got, want)
|
||||
}
|
||||
got := ModuleSpec()
|
||||
got.Requires[0] = "changed"
|
||||
got.Provides[0] = "changed"
|
||||
if again := ModuleSpec(); !reflect.DeepEqual(again, want) {
|
||||
t.Fatalf("ModuleSpec() after caller mutation = %#v, want %#v", again, want)
|
||||
}
|
||||
|
||||
registry := pipeline.NewChunkerRegistry()
|
||||
if err := Register(registry); err != nil {
|
||||
t.Fatalf("Register() error = %v, want nil", err)
|
||||
}
|
||||
registered, ok := registry.Spec(Key)
|
||||
if !ok {
|
||||
t.Fatalf("Spec(%q) ok = false, want true", Key)
|
||||
}
|
||||
if !reflect.DeepEqual(registered, want) {
|
||||
t.Fatalf("registered spec = %#v, want %#v", registered, want)
|
||||
}
|
||||
built, err := registry.Build(Key)
|
||||
if err != nil {
|
||||
t.Fatalf("Build(%q) error = %v, want nil", Key, err)
|
||||
}
|
||||
if built.Key() != Key {
|
||||
t.Fatalf("built Key() = %q, want %q", built.Key(), Key)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRegisterNilRegistryReturnsError(t *testing.T) {
|
||||
err := Register(nil)
|
||||
if err == nil {
|
||||
t.Fatal("Register(nil) error = nil, want error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "chunker registry") {
|
||||
t.Fatalf("Register(nil) error = %q, want registry context", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkReturnsSceneChunksFromStructuredOutput(t *testing.T) {
|
||||
client := &fakeScenesLLMClient{
|
||||
response: chunkResponse{
|
||||
Scenes: []sceneResponse{
|
||||
{
|
||||
StartUnitID: "seg-001",
|
||||
EndUnitID: "seg-002",
|
||||
ShortTitle: " Goblin parley ",
|
||||
PrimaryMode: "Discussion",
|
||||
MainParticipants: []string{" Aria ", "Goblin scout"},
|
||||
Summary: " The party negotiates with a scout. ",
|
||||
BoundaryNote: " The scene covers the discussion before fighting starts. ",
|
||||
BoundaryConfidence: "High",
|
||||
},
|
||||
{
|
||||
StartUnitID: "seg-003",
|
||||
EndUnitID: "seg-004",
|
||||
ShortTitle: "Ambush at the gate",
|
||||
PrimaryMode: "Combat",
|
||||
MainParticipants: []string{"Aria", "Goblin ambushers"},
|
||||
Summary: "The goblins attack at the gate.",
|
||||
BoundaryNote: "Combat begins and resolves the immediate threat.",
|
||||
BoundaryConfidence: "Medium",
|
||||
},
|
||||
},
|
||||
BoundaryCaveats: []string{"The transition into combat is gradual."},
|
||||
},
|
||||
}
|
||||
|
||||
result, err := New().Chunk(context.Background(), chunkRequestWithClient(client))
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
}
|
||||
|
||||
if len(client.requests) != 1 {
|
||||
t.Fatalf("LLM calls = %d, want 1", len(client.requests))
|
||||
}
|
||||
req := client.requests[0]
|
||||
if req.StageName != Key {
|
||||
t.Fatalf("StageName = %q, want %q", req.StageName, Key)
|
||||
}
|
||||
schema, err := loadResponseSchema()
|
||||
if err != nil {
|
||||
t.Fatalf("loadResponseSchema() error = %v, want nil", err)
|
||||
}
|
||||
if req.ResponseSchemaName != schema.Name {
|
||||
t.Fatalf("ResponseSchemaName = %q, want %q", req.ResponseSchemaName, schema.Name)
|
||||
}
|
||||
if !bytes.Equal(req.ResponseSchema, schema.JSONSchema) {
|
||||
t.Fatal("ResponseSchema does not match D&D scenes schema")
|
||||
}
|
||||
if len(req.Messages) != 2 || req.Messages[0].Role != "system" || req.Messages[1].Role != "user" {
|
||||
t.Fatalf("Messages = %#v, want system then user", req.Messages)
|
||||
}
|
||||
for _, want := range []string{"session-alpha", "seg-001", "seg-004", "start_unit_id", "boundary_confidence"} {
|
||||
if !strings.Contains(req.Messages[1].Content, want) {
|
||||
t.Fatalf("user message = %q, want substring %q", req.Messages[1].Content, want)
|
||||
}
|
||||
}
|
||||
|
||||
if got := chunkIDs(result.Chunks); !reflect.DeepEqual(got, []string{"scene-000001", "scene-000002"}) {
|
||||
t.Fatalf("chunk IDs = %#v, want deterministic scene IDs", got)
|
||||
}
|
||||
gotUnits := [][]string{unitIDs(result.Chunks[0].Units), unitIDs(result.Chunks[1].Units)}
|
||||
wantUnits := [][]string{{"seg-001", "seg-002"}, {"seg-003", "seg-004"}}
|
||||
if !reflect.DeepEqual(gotUnits, wantUnits) {
|
||||
t.Fatalf("chunk units = %#v, want %#v", gotUnits, wantUnits)
|
||||
}
|
||||
first := result.Chunks[0]
|
||||
if first.SourceID != "session-alpha" || first.Index != 0 {
|
||||
t.Fatalf("first chunk = %#v, want source and index fields", first)
|
||||
}
|
||||
if first.Metadata["scene_title"] != "Goblin parley" ||
|
||||
first.Metadata["primary_mode"] != "Discussion" ||
|
||||
first.Metadata["summary"] != "The party negotiates with a scout." ||
|
||||
first.Metadata["boundary_note"] != "The scene covers the discussion before fighting starts." ||
|
||||
first.Metadata["boundary_confidence"] != "High" ||
|
||||
first.Metadata["start_unit_id"] != "seg-001" ||
|
||||
first.Metadata["end_unit_id"] != "seg-002" ||
|
||||
first.Metadata["unit_count"] != 2 {
|
||||
t.Fatalf("first metadata = %#v, want scene metadata", first.Metadata)
|
||||
}
|
||||
if got, ok := first.Metadata["main_participants"].([]string); !ok || !reflect.DeepEqual(got, []string{"Aria", "Goblin scout"}) {
|
||||
t.Fatalf("main_participants = %#v, want trimmed participant slice", first.Metadata["main_participants"])
|
||||
}
|
||||
if got := result.Warnings; len(got) != 1 ||
|
||||
got[0].Scope != Key ||
|
||||
got[0].ReasonCode != "scene_boundary_caveat" ||
|
||||
got[0].Message != "The transition into combat is gradual." {
|
||||
t.Fatalf("Warnings = %#v, want boundary caveat warning", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkDefensivelyCopiesSourceUnitsAndMetadata(t *testing.T) {
|
||||
doc := sceneSourceDocument()
|
||||
client := &fakeScenesLLMClient{response: validSceneResponse()}
|
||||
|
||||
result, err := New().Chunk(context.Background(), contracts.ChunkRequest{
|
||||
Source: doc,
|
||||
LLMClient: client,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("Chunk() error = %v, want nil", err)
|
||||
}
|
||||
|
||||
doc.Units[0].ID = "mutated"
|
||||
doc.Units[0].Metadata["speaker"] = "mutated"
|
||||
client.response.Scenes[0].MainParticipants[0] = "mutated"
|
||||
|
||||
if result.Chunks[0].Units[0].ID != "seg-001" {
|
||||
t.Fatalf("chunk unit ID changed after source mutation: %#v", result.Chunks[0].Units[0])
|
||||
}
|
||||
if result.Chunks[0].Units[0].Metadata["speaker"] != "Alice" {
|
||||
t.Fatalf("chunk unit metadata changed after source mutation: %#v", result.Chunks[0].Units[0].Metadata)
|
||||
}
|
||||
participants, ok := result.Chunks[0].Metadata["main_participants"].([]string)
|
||||
if !ok || participants[0] != "Aria" {
|
||||
t.Fatalf("participants = %#v, want defensive copy", result.Chunks[0].Metadata["main_participants"])
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkerManifestMetadataIncludesPromptAndSchemaProvenance(t *testing.T) {
|
||||
metadata := New().ManifestMetadata()
|
||||
|
||||
tests := map[string]string{
|
||||
"prompt_id": PromptID,
|
||||
"prompt_version": ResponseSchemaVersion,
|
||||
"response_schema_key": string(ResponseSchemaKey),
|
||||
"response_schema_id": ResponseSchemaID,
|
||||
"response_schema_name": ResponseSchemaName,
|
||||
"response_schema_version": ResponseSchemaVersion,
|
||||
}
|
||||
for key, want := range tests {
|
||||
if metadata[key] != want {
|
||||
t.Fatalf("metadata[%q] = %#v, want %q", key, metadata[key], want)
|
||||
}
|
||||
}
|
||||
for _, key := range []string{"prompt_sha256", "response_schema_sha256"} {
|
||||
value, ok := metadata[key].(string)
|
||||
if !ok || !strings.HasPrefix(value, "sha256:") {
|
||||
t.Fatalf("metadata[%q] = %#v, want sha256 value", key, metadata[key])
|
||||
}
|
||||
}
|
||||
for _, forbidden := range []string{"prompt", "schema", "source", "text"} {
|
||||
if _, ok := metadata[forbidden]; ok {
|
||||
t.Fatalf("metadata includes raw %q field: %#v", forbidden, metadata)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRejectsInvalidRequests(t *testing.T) {
|
||||
validClient := &fakeScenesLLMClient{response: validSceneResponse()}
|
||||
validReq := chunkRequestWithClient(validClient)
|
||||
canceledCtx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
invalidDoc := sceneSourceDocument()
|
||||
invalidDoc.Units[0].ID = ""
|
||||
emptyDoc := sceneSourceDocument()
|
||||
emptyDoc.Units = nil
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
chunker *Chunker
|
||||
ctx context.Context
|
||||
req contracts.ChunkRequest
|
||||
want string
|
||||
}{
|
||||
{name: "nil chunker", chunker: nil, ctx: context.Background(), req: validReq, want: "chunker"},
|
||||
{name: "nil context", chunker: New(), ctx: nil, req: validReq, want: "context"},
|
||||
{name: "canceled context", chunker: New(), ctx: canceledCtx, req: validReq, want: "context"},
|
||||
{name: "nil source", chunker: New(), ctx: context.Background(), req: contracts.ChunkRequest{LLMClient: validClient}, want: "source"},
|
||||
{name: "empty source units", chunker: New(), ctx: context.Background(), req: contracts.ChunkRequest{Source: emptyDoc, LLMClient: validClient}, want: "units"},
|
||||
{name: "invalid source", chunker: New(), ctx: context.Background(), req: contracts.ChunkRequest{Source: invalidDoc, LLMClient: validClient}, want: "validate source document"},
|
||||
{name: "nil LLM client", chunker: New(), ctx: context.Background(), req: contracts.ChunkRequest{Source: sceneSourceDocument()}, want: "LLM client"},
|
||||
{name: "unsupported options", chunker: New(), ctx: context.Background(), req: requestWithOptions(validReq), want: "options"},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
_, err := tt.chunker.Chunk(tt.ctx, tt.req)
|
||||
if err == nil {
|
||||
t.Fatal("Chunk() error = nil, want error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "dnd scenes") || !strings.Contains(err.Error(), tt.want) {
|
||||
t.Fatalf("Chunk() error = %q, want module context and %q", err.Error(), tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkRejectsMalformedStructuredOutput(t *testing.T) {
|
||||
tests := []struct {
|
||||
name string
|
||||
response chunkResponse
|
||||
want string
|
||||
}{
|
||||
{name: "missing scenes", response: chunkResponse{}, want: "scenes"},
|
||||
{name: "empty scenes", response: chunkResponse{Scenes: []sceneResponse{}}, want: "scenes"},
|
||||
{
|
||||
name: "unknown boundary id",
|
||||
response: replaceScenes(validSceneResponse(), []sceneResponse{
|
||||
scene("seg-001", "seg-999"),
|
||||
}),
|
||||
want: "was not found",
|
||||
},
|
||||
{
|
||||
name: "out of order boundaries",
|
||||
response: replaceScenes(validSceneResponse(), []sceneResponse{
|
||||
scene("seg-003", "seg-002"),
|
||||
}),
|
||||
want: "appears after",
|
||||
},
|
||||
{
|
||||
name: "gap",
|
||||
response: replaceScenes(validSceneResponse(), []sceneResponse{
|
||||
scene("seg-001", "seg-001"),
|
||||
scene("seg-003", "seg-004"),
|
||||
}),
|
||||
want: "gap",
|
||||
},
|
||||
{
|
||||
name: "overlap",
|
||||
response: replaceScenes(validSceneResponse(), []sceneResponse{
|
||||
scene("seg-001", "seg-002"),
|
||||
scene("seg-002", "seg-004"),
|
||||
}),
|
||||
want: "overlap",
|
||||
},
|
||||
{
|
||||
name: "incomplete coverage",
|
||||
response: replaceScenes(validSceneResponse(), []sceneResponse{
|
||||
scene("seg-001", "seg-003"),
|
||||
}),
|
||||
want: "final scene",
|
||||
},
|
||||
{
|
||||
name: "empty metadata field",
|
||||
response: replaceScenes(validSceneResponse(), []sceneResponse{
|
||||
{
|
||||
StartUnitID: "seg-001",
|
||||
EndUnitID: "seg-004",
|
||||
ShortTitle: " ",
|
||||
PrimaryMode: "Narrative",
|
||||
MainParticipants: []string{"Aria"},
|
||||
Summary: "Summary.",
|
||||
BoundaryNote: "Note.",
|
||||
BoundaryConfidence: "High",
|
||||
},
|
||||
}),
|
||||
want: "short_title",
|
||||
},
|
||||
{
|
||||
name: "empty participant",
|
||||
response: replaceScenes(validSceneResponse(), []sceneResponse{
|
||||
{
|
||||
StartUnitID: "seg-001",
|
||||
EndUnitID: "seg-004",
|
||||
ShortTitle: "Title",
|
||||
PrimaryMode: "Narrative",
|
||||
MainParticipants: []string{"Aria", " "},
|
||||
Summary: "Summary.",
|
||||
BoundaryNote: "Note.",
|
||||
BoundaryConfidence: "High",
|
||||
},
|
||||
}),
|
||||
want: "main_participants",
|
||||
},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
client := &fakeScenesLLMClient{response: tt.response}
|
||||
_, err := New().Chunk(context.Background(), chunkRequestWithClient(client))
|
||||
if err == nil {
|
||||
t.Fatal("Chunk() error = nil, want error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "dnd scenes") || !strings.Contains(err.Error(), tt.want) {
|
||||
t.Fatalf("Chunk() error = %q, want module context and %q", err.Error(), tt.want)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestChunkWrapsLLMClientError(t *testing.T) {
|
||||
client := &fakeScenesLLMClient{err: errors.New("provider unavailable")}
|
||||
|
||||
_, err := New().Chunk(context.Background(), chunkRequestWithClient(client))
|
||||
if err == nil {
|
||||
t.Fatal("Chunk() error = nil, want LLM error")
|
||||
}
|
||||
if !strings.Contains(err.Error(), "dnd scenes") || !strings.Contains(err.Error(), "provider unavailable") {
|
||||
t.Fatalf("Chunk() error = %q, want wrapped LLM context", err.Error())
|
||||
}
|
||||
}
|
||||
|
||||
func chunkRequestWithClient(client contracts.StructuredLLMClient) contracts.ChunkRequest {
|
||||
return contracts.ChunkRequest{
|
||||
Source: sceneSourceDocument(),
|
||||
LLMClient: client,
|
||||
}
|
||||
}
|
||||
|
||||
func requestWithOptions(req contracts.ChunkRequest) contracts.ChunkRequest {
|
||||
req.Options = map[string]any{"max_units": 2}
|
||||
return req
|
||||
}
|
||||
|
||||
func sceneSourceDocument() *source.SourceDocument {
|
||||
return &source.SourceDocument{
|
||||
ID: "session-alpha",
|
||||
Kind: "transcript",
|
||||
Format: "application/vnd.seriatim.minimal+json",
|
||||
Digest: "sha256:source",
|
||||
Units: []source.SourceUnit{
|
||||
{ID: "seg-001", Kind: "transcript_segment", Text: "Aria asks whether the goblin will parley.", Metadata: map[string]any{"speaker": "Alice"}},
|
||||
{ID: "seg-002", Kind: "transcript_segment", Text: "The goblin scout describes the gate guards."},
|
||||
{ID: "seg-003", Kind: "transcript_segment", Text: "The guards rush out with blades drawn."},
|
||||
{ID: "seg-004", Kind: "transcript_segment", Text: "The party defeats the ambushers."},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func validSceneResponse() chunkResponse {
|
||||
return chunkResponse{
|
||||
Scenes: []sceneResponse{
|
||||
scene("seg-001", "seg-004"),
|
||||
},
|
||||
BoundaryCaveats: []string{},
|
||||
}
|
||||
}
|
||||
|
||||
func replaceScenes(response chunkResponse, scenes []sceneResponse) chunkResponse {
|
||||
response.Scenes = scenes
|
||||
return response
|
||||
}
|
||||
|
||||
func scene(startUnitID string, endUnitID string) sceneResponse {
|
||||
return sceneResponse{
|
||||
StartUnitID: startUnitID,
|
||||
EndUnitID: endUnitID,
|
||||
ShortTitle: "Scene title",
|
||||
PrimaryMode: "Narrative",
|
||||
MainParticipants: []string{"Aria"},
|
||||
Summary: "A compact summary.",
|
||||
BoundaryNote: "The source units form one coherent scene.",
|
||||
BoundaryConfidence: "High",
|
||||
}
|
||||
}
|
||||
|
||||
func chunkIDs(chunks []contracts.SourceChunk) []string {
|
||||
ids := make([]string, 0, len(chunks))
|
||||
for _, chunk := range chunks {
|
||||
ids = append(ids, chunk.ID)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
func unitIDs(units []source.SourceUnit) []string {
|
||||
ids := make([]string, 0, len(units))
|
||||
for _, unit := range units {
|
||||
ids = append(ids, unit.ID)
|
||||
}
|
||||
return ids
|
||||
}
|
||||
|
||||
type fakeScenesLLMClient struct {
|
||||
response chunkResponse
|
||||
err error
|
||||
requests []contracts.StructuredCompletionRequest
|
||||
}
|
||||
|
||||
func (client *fakeScenesLLMClient) CompleteStructured(ctx context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
|
||||
client.requests = append(client.requests, contracts.StructuredCompletionRequest{
|
||||
StageName: req.StageName,
|
||||
Messages: append([]contracts.LLMMessage(nil), req.Messages...),
|
||||
Model: req.Model,
|
||||
ResponseSchemaName: req.ResponseSchemaName,
|
||||
ResponseSchema: append(json.RawMessage(nil), req.ResponseSchema...),
|
||||
})
|
||||
if client.err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, client.err
|
||||
}
|
||||
|
||||
target, ok := out.(*chunkResponse)
|
||||
if !ok {
|
||||
return contracts.StructuredCompletionResponse{}, errors.New("unexpected output target")
|
||||
}
|
||||
*target = client.response
|
||||
content, err := json.Marshal(client.response)
|
||||
if err != nil {
|
||||
return contracts.StructuredCompletionResponse{}, err
|
||||
}
|
||||
return contracts.StructuredCompletionResponse{Content: content}, nil
|
||||
}
|
||||
Reference in New Issue
Block a user