Introduce typed D&D spell artifacts

internal/modules/dnd/extract/spells/test_helpers_test.go

@@ -1,7 +1,9 @@
 package spells
 
 import (
+	"context"
 	"encoding/json"
+	"errors"
 	"testing"
 
 	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
@@ -9,7 +11,7 @@ import (
 	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
 )
 
-func promptExtractionRequest() contracts.ExtractionRequest {
+func promptExtractionRequest() contracts.TypedExtractionRequest {
 	doc := promptSourceDocument()
 	chunk := &source.Chunk{
 		ID:       "session-alpha:chunk:0",
@@ -25,7 +27,7 @@ func promptExtractionRequest() contracts.ExtractionRequest {
 		Units:     append([]source.SourceUnit(nil), doc.Units...),
 		Metadata:  map[string]any{"ignored": "chunk metadata"},
 	}
-	return contracts.ExtractionRequest{
+	return contracts.TypedExtractionRequest{
 		Source: doc,
 		Chunk:  chunk,
 	}
@@ -70,22 +72,90 @@ func mustJSON(t *testing.T, value any) string {
 	return string(encoded)
 }
 
-func responseSourceRefs(sourceID string, startUnitID int, endUnitID int) []shared.SourceRefResponse {
-	return []shared.SourceRefResponse{
+func responseSourceRefs(startUnitID int, endUnitID int) []spellSourceRefResponse {
+	return []spellSourceRefResponse{
 		{
-			SourceID:    sourceID,
 			StartUnitID: shared.UnitRefFromInt(startUnitID),
 			EndUnitID:   shared.UnitRefFromInt(endUnitID),
 		},
 	}
 }
 
-func responseSourceRefsInt(sourceID string, startUnitID int, endUnitID int) []shared.SourceRefResponse {
-	return []shared.SourceRefResponse{
-		{
-			SourceID:    sourceID,
-			StartUnitID: shared.UnitRefFromInt(startUnitID),
-			EndUnitID:   shared.UnitRefFromInt(endUnitID),
-		},
-	}
+const spellTranscriptJSON = `{"id":"session-alpha","segments":[{"id":1,"text":"Aria raises her hand and casts Cure Wounds."}]}`
+
+func spellSourceInput() contracts.LLMInputMaterial {
+	return contracts.NewLLMInputMaterial("source", "application/json", []byte(spellTranscriptJSON), "sha256:transcript", "file:///session-alpha.json")
+}
+
+func spellChunkInput(chunk *source.Chunk) contracts.LLMInputMaterial {
+	return contracts.NewLLMInputMaterial("source", chunk.MediaType, chunk.Content, "sha256:chunk", "file:///session-alpha.json")
+}
+
+func extractionRequest() contracts.TypedExtractionRequest {
+	req := promptExtractionRequest()
+	req.SourceInput = spellChunkInput(req.Chunk)
+	req.SessionID = "session-123"
+	req.LLMProfile = "profile-spells"
+	return req
+}
+
+func emptyChunkRequest(req contracts.TypedExtractionRequest) contracts.TypedExtractionRequest {
+	req.Chunk = &source.Chunk{ID: req.Chunk.ID, SourceID: req.Chunk.SourceID, Index: req.Chunk.Index}
+	return req
+}
+
+func mismatchedSourceInputRequest(req contracts.TypedExtractionRequest) contracts.TypedExtractionRequest {
+	req.SourceInput = spellSourceInput()
+	return req
+}
+
+func newExtractor(t *testing.T, client contracts.StructuredLLMClient) *Extractor {
+	t.Helper()
+	extractor, err := New(client, Options{})
+	if err != nil {
+		t.Fatalf("New() error = %v, want nil", err)
+	}
+	return extractor
+}
+
+type fakeSpellsLLMClient struct {
+	response extractionResponse
+	content  []byte
+	err      error
+	requests []contracts.StructuredCompletionRequest
+}
+
+func (client *fakeSpellsLLMClient) CompleteStructured(_ context.Context, req contracts.StructuredCompletionRequest, out any) (contracts.StructuredCompletionResponse, error) {
+	client.requests = append(client.requests, cloneStructuredCompletionRequest(req))
+	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 := append([]byte(nil), client.content...)
+	if len(content) == 0 {
+		var err error
+		content, err = json.Marshal(client.response)
+		if err != nil {
+			return contracts.StructuredCompletionResponse{}, err
+		}
+	}
+	return contracts.StructuredCompletionResponse{Content: content}, nil
+}
+
+func cloneStructuredCompletionRequest(req contracts.StructuredCompletionRequest) contracts.StructuredCompletionRequest {
+	req.Inputs = req.Inputs.Clone()
+	if len(req.Vars) == 0 {
+		req.Vars = nil
+		return req
+	}
+	vars := make(map[string]any, len(req.Vars))
+	for key, value := range req.Vars {
+		vars[key] = value
+	}
+	req.Vars = vars
+	return req
 }