package spells import ( "context" "encoding/json" "errors" "testing" "gitea.maximumdirect.net/eric/notarius/internal/core/source" "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" spellcatalog "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/spells/catalog" ) func promptExtractionRequest() contracts.TypedExtractionRequest { doc := promptSourceDocument() chunk := &source.Chunk{ ID: "session-alpha:chunk:0", SourceID: doc.ID, Index: 0, Ref: source.SourceRef{ SourceID: doc.ID, StartUnitID: doc.Units[0].ID, EndUnitID: doc.Units[len(doc.Units)-1].ID, }, Content: []byte(`{"units":[1,2]}`), MediaType: "application/json", Units: append([]source.SourceUnit(nil), doc.Units...), Metadata: map[string]any{"ignored": "chunk metadata"}, } return contracts.TypedExtractionRequest{ Source: doc, Chunk: chunk, } } func promptSourceDocument() *source.SourceDocument { return &source.SourceDocument{ ID: "session-alpha", Kind: "transcript", Format: "application/vnd.seriatim.minimal+json", Digest: "sha256:test", Units: []source.SourceUnit{ { ID: 1, Kind: "transcript_segment", Text: "Aria raises her hand and casts Cure Wounds.", Ref: source.SourceRef{SourceID: "session-alpha", StartUnitID: 1, EndUnitID: 1}, Metadata: map[string]any{ "speaker": "Alice", "start": json.Number("1.25"), "end": json.Number("3.5"), "ignored": "not rendered", }, }, { ID: 2, Kind: "transcript_segment", Text: "The fighter's wounds begin to close.", Ref: source.SourceRef{SourceID: "session-alpha", StartUnitID: 2, EndUnitID: 2}, Metadata: map[string]any{"ignored": "not rendered"}, }, }, } } func mustJSON(t *testing.T, value any) string { t.Helper() encoded, err := json.Marshal(value) if err != nil { t.Fatalf("Marshal() error = %v, want nil", err) } return string(encoded) } func responseSourceRefs(startUnitID int, endUnitID int) []spellSourceRefResponse { return []spellSourceRefResponse{ { StartUnitID: startUnitID, EndUnitID: 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 mismatchedSourceInputRequest(req contracts.TypedExtractionRequest) contracts.TypedExtractionRequest { req.SourceInput = spellSourceInput() return req } func newExtractor(t *testing.T, client contracts.StructuredLLMClient, references ...contracts.ReferenceSet) *Extractor { t.Helper() extractor, err := New(client, Options{}, references...) if err != nil { t.Fatalf("New() error = %v, want nil", err) } return extractor } func spellCatalogReference(content string) contracts.ReferenceSet { return contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{ spellcatalog.SpellCatalogReferenceSlot: { Items: []contracts.ReferenceItem{{ SlotName: spellcatalog.SpellCatalogReferenceSlot, MediaType: "application/json", Content: []byte(content), }}, }, }} } func overlaySpellCatalogReference() contracts.ReferenceSet { return spellCatalogReference(`{"schema_version":"notarius.dnd.spell-catalog-overlay.v1","catalogs":[{"id":"campaign.example","ruleset":"dnd-5e-2014","source":{"title":"Private campaign source","version":"1","url":"file:///private-source.json","license":"private"},"spells":[{"name":"Aegis of Emberfall","aliases":["Emberfall Aegis"]}]}]}`) } 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") } content := append([]byte(nil), client.content...) if len(content) != 0 { if err := json.Unmarshal(content, target); err != nil { return contracts.StructuredCompletionResponse{}, err } } else { *target = client.response 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 }