Implement D&D spells extraction mapping

2026-07-03 23:38:41 +00:00
parent 404bfc8da5
commit b2012b9b2c
3 changed files with 355 additions and 13 deletions
--- 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...)
+}
--- a/internal/modules/extract/dnd/spells/extractor_test.go
+++ 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
+}
--- 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())
-	}
-}