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

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"
	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: 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, 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")
	}
	*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
}