package spells import ( "context" "encoding/json" "errors" "testing" "gitea.maximumdirect.net/eric/notarius/internal/core/source" "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared" ) 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: 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 }