From b2012b9b2c36983d90d400113f6285766fa755c3 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Fri, 3 Jul 2026 23:38:41 +0000 Subject: [PATCH] Implement D&D spells extraction mapping --- .../modules/extract/dnd/spells/extractor.go | 81 +++++- .../extract/dnd/spells/extractor_test.go | 275 ++++++++++++++++++ .../extract/dnd/spells/registry_test.go | 12 - 3 files changed, 355 insertions(+), 13 deletions(-) create mode 100644 internal/modules/extract/dnd/spells/extractor_test.go diff --git a/internal/modules/extract/dnd/spells/extractor.go b/internal/modules/extract/dnd/spells/extractor.go index cfc1094..f49eeb9 100644 --- a/internal/modules/extract/dnd/spells/extractor.go +++ b/internal/modules/extract/dnd/spells/extractor.go @@ -2,9 +2,14 @@ package spells import ( "context" + "encoding/json" "fmt" + "strings" + "gitea.maximumdirect.net/eric/notarius/internal/core/artifacts" + "gitea.maximumdirect.net/eric/notarius/internal/core/source" "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" + "gitea.maximumdirect.net/eric/notarius/internal/framework/llm" "gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline" ) @@ -46,7 +51,68 @@ func (e *Extractor) Validators() []contracts.Validator { } func (e *Extractor) Extract(ctx context.Context, req contracts.ExtractionRequest) (contracts.ExtractionResult, error) { - return contracts.ExtractionResult{}, fmt.Errorf("dnd spells extractor: extraction is not implemented") + if e == nil { + return contracts.ExtractionResult{}, extractorErrorf("extractor must not be nil") + } + if ctx == nil { + return contracts.ExtractionResult{}, extractorErrorf("context must not be nil") + } + if err := ctx.Err(); err != nil { + return contracts.ExtractionResult{}, extractorErrorf("context error before extraction: %w", err) + } + if req.Source == nil { + return contracts.ExtractionResult{}, extractorErrorf("source must not be nil") + } + if req.Chunk == nil { + return contracts.ExtractionResult{}, extractorErrorf("chunk must not be nil") + } + if len(req.Chunk.Units) == 0 { + return contracts.ExtractionResult{}, extractorErrorf("chunk %q units must not be empty", req.Chunk.ID) + } + if req.LLMClient == nil { + return contracts.ExtractionResult{}, extractorErrorf("LLM client must not be nil") + } + + system, user, _, err := renderPrompt(req) + if err != nil { + return contracts.ExtractionResult{}, extractorErrorf("render prompt: %w", err) + } + schema, ok := llm.LookupResponseSchema(llm.DNDSpellsSchemaKey) + if !ok { + return contracts.ExtractionResult{}, extractorErrorf("lookup response schema %q", llm.DNDSpellsSchemaKey) + } + + var response extractionResponse + 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.ExtractionResult{}, extractorErrorf("complete structured output: %w", err) + } + if response.SpellCasts == nil { + return contracts.ExtractionResult{}, extractorErrorf("malformed structured output: spell_casts must be present") + } + if len(response.SpellCasts) == 0 { + return contracts.ExtractionResult{}, nil + } + + candidates := make([]artifacts.ArtifactCandidate, 0, len(response.SpellCasts)) + for i, spellCast := range response.SpellCasts { + payload, err := spellCastPayload(spellCast) + if err != nil { + return contracts.ExtractionResult{}, extractorErrorf("marshal spell cast[%d]: %w", i, err) + } + candidates = append(candidates, artifacts.ArtifactCandidate{ + Payload: payload, + SourceRefs: append([]source.SourceRef(nil), spellCast.SourceRefs...), + }) + } + return contracts.ExtractionResult{Candidates: candidates}, nil } func ModuleSpec() pipeline.ModuleSpec { @@ -63,3 +129,16 @@ func Register(registry *pipeline.ExtractorRegistry) error { return New(), nil }) } + +func spellCastPayload(spellCast spellCastResponse) (json.RawMessage, error) { + return json.Marshal(SpellCast{ + Caster: strings.TrimSpace(spellCast.Caster), + Spell: strings.TrimSpace(spellCast.Spell), + Effect: strings.TrimSpace(spellCast.Effect), + NarrativeDescription: strings.TrimSpace(spellCast.NarrativeDescription), + }) +} + +func extractorErrorf(format string, args ...any) error { + return fmt.Errorf("dnd spells extractor: "+format, args...) +} diff --git a/internal/modules/extract/dnd/spells/extractor_test.go b/internal/modules/extract/dnd/spells/extractor_test.go new file mode 100644 index 0000000..08c4e6e --- /dev/null +++ b/internal/modules/extract/dnd/spells/extractor_test.go @@ -0,0 +1,275 @@ +package spells + +import ( + "bytes" + "context" + "encoding/json" + "errors" + "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/llm" +) + +func TestExtractReturnsSpellCandidateFromStructuredOutput(t *testing.T) { + client := &fakeSpellsLLMClient{ + response: extractionResponse{ + SpellCasts: []spellCastResponse{ + { + Caster: " Aria ", + Spell: " Cure Wounds ", + Effect: " Heals an injured ally. ", + NarrativeDescription: " Aria restores the fighter after the fight. ", + SourceRefs: []source.SourceRef{ + {SourceID: "session-alpha", StartUnitID: "seg-001", EndUnitID: "seg-002"}, + }, + }, + }, + }, + } + + result, err := New().Extract(context.Background(), extractionRequestWithClient(client)) + if err != nil { + t.Fatalf("Extract() 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 := llm.MustLookupResponseSchema(llm.DNDSpellsSchemaKey) + 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 registered D&D spells schema") + } + if len(req.Messages) != 2 { + t.Fatalf("len(Messages) = %d, want 2", len(req.Messages)) + } + if req.Messages[0].Role != "system" || req.Messages[1].Role != "user" { + t.Fatalf("Messages roles = %#v, want system then user", req.Messages) + } + if !strings.Contains(req.Messages[0].Content, "D&D spell-cast") { + t.Fatalf("system message = %q, want D&D spell context", req.Messages[0].Content) + } + for _, want := range []string{"session-alpha", "session-alpha:chunk:0", "seg-001", "Cure Wounds"} { + if !strings.Contains(req.Messages[1].Content, want) { + t.Fatalf("user message = %q, want substring %q", req.Messages[1].Content, want) + } + } + + if len(result.Candidates) != 1 { + t.Fatalf("len(Candidates) = %d, want 1", len(result.Candidates)) + } + candidate := result.Candidates[0] + if candidate.Index != 0 || candidate.ExtractorKey != "" || candidate.ArtifactType != "" || candidate.SchemaVersion != "" { + t.Fatalf("candidate envelope fields = %#v, want runner-normalized zero values", candidate) + } + var payload SpellCast + if err := json.Unmarshal(candidate.Payload, &payload); err != nil { + t.Fatalf("Unmarshal(Payload) error = %v, want nil", err) + } + wantPayload := SpellCast{ + Caster: "Aria", + Spell: "Cure Wounds", + Effect: "Heals an injured ally.", + NarrativeDescription: "Aria restores the fighter after the fight.", + } + if payload != wantPayload { + t.Fatalf("payload = %#v, want %#v", payload, wantPayload) + } + wantRef := source.SourceRef{SourceID: "session-alpha", StartUnitID: "seg-001", EndUnitID: "seg-002"} + if len(candidate.SourceRefs) != 1 || candidate.SourceRefs[0] != wantRef { + t.Fatalf("SourceRefs = %#v, want %#v", candidate.SourceRefs, []source.SourceRef{wantRef}) + } +} + +func TestExtractReturnsNoCandidatesForEmptyResponse(t *testing.T) { + client := &fakeSpellsLLMClient{response: extractionResponse{SpellCasts: []spellCastResponse{}}} + + result, err := New().Extract(context.Background(), extractionRequestWithClient(client)) + if err != nil { + t.Fatalf("Extract() error = %v, want nil", err) + } + if len(result.Candidates) != 0 { + t.Fatalf("Candidates = %#v, want none", result.Candidates) + } +} + +func TestExtractRejectsMissingSpellCasts(t *testing.T) { + client := &fakeSpellsLLMClient{response: extractionResponse{}} + + _, err := New().Extract(context.Background(), extractionRequestWithClient(client)) + if err == nil { + t.Fatal("Extract() error = nil, want malformed output error") + } + if !strings.Contains(err.Error(), "dnd spells") || !strings.Contains(err.Error(), "spell_casts") { + t.Fatalf("Extract() error = %q, want spell_casts context", err.Error()) + } +} + +func TestExtractWrapsLLMClientError(t *testing.T) { + client := &fakeSpellsLLMClient{err: errors.New("provider unavailable")} + + _, err := New().Extract(context.Background(), extractionRequestWithClient(client)) + if err == nil { + t.Fatal("Extract() error = nil, want LLM error") + } + if !strings.Contains(err.Error(), "dnd spells") || !strings.Contains(err.Error(), "provider unavailable") { + t.Fatalf("Extract() error = %q, want wrapped LLM context", err.Error()) + } +} + +func TestExtractRejectsInvalidRequests(t *testing.T) { + validClient := &fakeSpellsLLMClient{response: extractionResponse{SpellCasts: []spellCastResponse{}}} + validReq := extractionRequestWithClient(validClient) + canceledCtx, cancel := context.WithCancel(context.Background()) + cancel() + + tests := []struct { + name string + extractor *Extractor + ctx context.Context + req contracts.ExtractionRequest + want string + }{ + {name: "nil extractor", extractor: nil, ctx: context.Background(), req: validReq, want: "extractor"}, + {name: "nil context", extractor: New(), ctx: nil, req: validReq, want: "context"}, + {name: "canceled context", extractor: New(), ctx: canceledCtx, req: validReq, want: "context"}, + {name: "nil source", extractor: New(), ctx: context.Background(), req: contracts.ExtractionRequest{Chunk: validReq.Chunk, LLMClient: validReq.LLMClient}, want: "source"}, + {name: "nil chunk", extractor: New(), ctx: context.Background(), req: contracts.ExtractionRequest{Source: validReq.Source, LLMClient: validReq.LLMClient}, want: "chunk"}, + {name: "empty chunk units", extractor: New(), ctx: context.Background(), req: emptyChunkRequest(validReq), want: "units"}, + {name: "nil LLM client", extractor: New(), ctx: context.Background(), req: contracts.ExtractionRequest{Source: validReq.Source, Chunk: validReq.Chunk}, want: "LLM client"}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, err := tt.extractor.Extract(tt.ctx, tt.req) + if err == nil { + t.Fatal("Extract() error = nil, want error") + } + if !strings.Contains(err.Error(), "dnd spells") || !strings.Contains(err.Error(), tt.want) { + t.Fatalf("Extract() error = %q, want %q context", err.Error(), tt.want) + } + }) + } +} + +func TestExtractPreservesResponseOrder(t *testing.T) { + client := &fakeSpellsLLMClient{ + response: extractionResponse{ + SpellCasts: []spellCastResponse{ + { + Caster: "Aria", + Spell: "Cure Wounds", + Effect: "Heals.", + NarrativeDescription: "First spell.", + SourceRefs: []source.SourceRef{{SourceID: "session-alpha", StartUnitID: "seg-001", EndUnitID: "seg-001"}}, + }, + { + Caster: "Bandit Shaman", + Spell: "Fire Bolt", + Effect: "Burns.", + NarrativeDescription: "Second spell.", + SourceRefs: []source.SourceRef{{SourceID: "session-alpha", StartUnitID: "seg-002", EndUnitID: "seg-002"}}, + }, + }, + }, + } + + result, err := New().Extract(context.Background(), extractionRequestWithClient(client)) + if err != nil { + t.Fatalf("Extract() error = %v, want nil", err) + } + if len(result.Candidates) != 2 { + t.Fatalf("len(Candidates) = %d, want 2", len(result.Candidates)) + } + + var first, second SpellCast + if err := json.Unmarshal(result.Candidates[0].Payload, &first); err != nil { + t.Fatalf("Unmarshal(first) error = %v, want nil", err) + } + if err := json.Unmarshal(result.Candidates[1].Payload, &second); err != nil { + t.Fatalf("Unmarshal(second) error = %v, want nil", err) + } + if first.Spell != "Cure Wounds" || second.Spell != "Fire Bolt" { + t.Fatalf("candidate order = %q, %q; want response order", first.Spell, second.Spell) + } +} + +func TestExtractCopiesCandidateSourceRefs(t *testing.T) { + client := &fakeSpellsLLMClient{ + response: extractionResponse{ + SpellCasts: []spellCastResponse{ + { + Caster: "Aria", + Spell: "Cure Wounds", + Effect: "Heals.", + NarrativeDescription: "Aria heals.", + SourceRefs: []source.SourceRef{{SourceID: "session-alpha", StartUnitID: "seg-001", EndUnitID: "seg-002"}}, + }, + }, + }, + } + + result, err := New().Extract(context.Background(), extractionRequestWithClient(client)) + if err != nil { + t.Fatalf("Extract() error = %v, want nil", err) + } + client.response.SpellCasts[0].SourceRefs[0].StartUnitID = "mutated" + + if got := result.Candidates[0].SourceRefs[0].StartUnitID; got != "seg-001" { + t.Fatalf("candidate source ref start = %q, want copied seg-001", got) + } +} + +func extractionRequestWithClient(client contracts.StructuredLLMClient) contracts.ExtractionRequest { + req := promptExtractionRequest() + req.LLMClient = client + return req +} + +func emptyChunkRequest(req contracts.ExtractionRequest) contracts.ExtractionRequest { + req.Chunk = &contracts.SourceChunk{ + ID: req.Chunk.ID, + SourceID: req.Chunk.SourceID, + Index: req.Chunk.Index, + } + return req +} + +type fakeSpellsLLMClient struct { + response extractionResponse + err error + requests []contracts.StructuredCompletionRequest +} + +func (client *fakeSpellsLLMClient) 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.(*extractionResponse) + 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 +} diff --git a/internal/modules/extract/dnd/spells/registry_test.go b/internal/modules/extract/dnd/spells/registry_test.go index a54a66c..464f5fe 100644 --- a/internal/modules/extract/dnd/spells/registry_test.go +++ b/internal/modules/extract/dnd/spells/registry_test.go @@ -1,12 +1,10 @@ package spells import ( - "context" "reflect" "strings" "testing" - "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" "gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline" ) @@ -93,13 +91,3 @@ func TestRegisterNilRegistryReturnsError(t *testing.T) { t.Fatalf("Register(nil) error = %q, want registry context", err.Error()) } } - -func TestExtractReportsNotImplemented(t *testing.T) { - _, err := New().Extract(context.Background(), contracts.ExtractionRequest{}) - if err == nil { - t.Fatal("Extract() error = nil, want error") - } - if !strings.Contains(err.Error(), "dnd spells") || !strings.Contains(err.Error(), "not implemented") { - t.Fatalf("Extract() error = %q, want not implemented D&D spells context", err.Error()) - } -}