notarius/internal/modules/integration/dnd_spells_runner_test.go

package integration_test

import (
	"context"
	"encoding/json"
	"os"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/core/config"
	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/pipeline"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/spells"
	spellnormalize "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/normalize/spells"
	spellshape "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/validate/spells/shape"
	spellsourcerefs "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/validate/spells/source_refs"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/seriatim/input/transcript"
)

func runPreparedPipeline(t *testing.T, registries pipeline.Registries, resolved pipeline.ResolvedPipeline, llmClient contracts.StructuredLLMClient, input pipeline.RunInput) (pipeline.RunOutput, error) {
	t.Helper()
	prepared, err := pipeline.Prepare(resolved, registries, pipeline.ModuleDependencies{LLM: llmClient})
	if err != nil {
		return pipeline.RunOutput{}, err
	}
	input.Prepared = prepared
	return pipeline.New().Run(context.Background(), input)
}

func TestRunnerProcessesSeriatimInputWithDNDSpellsExtractor(t *testing.T) {
	raw := readDNDSpellsFixture(t)
	expectedDoc := parseDNDSpellsFixture(t, raw)
	resolved := resolveDNDSpellsPipeline(t)
	llmClient := &fakeSpellsLLMClient{
		response: extractionResponse{
			SpellCasts: []spellCastResponse{
				{
					Caster:               "Aria",
					Spell:                "Cure Wounds",
					Effect:               "Heals an injured ally.",
					NarrativeDescription: "Aria restores the fighter after the fight.",
					SourceRefs:           responseSourceRefs(expectedDoc.ID, 1, 1),
				},
				{
					Caster:               "Borin",
					Spell:                "Fire Bolt",
					Effect:               "Scorches the wight.",
					NarrativeDescription: "Borin hurls fire at the wight.",
					SourceRefs:           responseSourceRefs(expectedDoc.ID, 3, 3),
				},
			},
		},
	}

	output, err := runPreparedPipeline(t, dndSpellsRunnerRegistries(t), resolved.ResolvedPipeline, llmClient, pipeline.RunInput{
		RawInput: raw,
	})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}

	if len(output.NormalizeOutputs) != 1 {
		t.Fatalf("len(NormalizeOutputs) = %d, want 1", len(output.NormalizeOutputs))
	}
	serializedOutput := output.NormalizeOutputs[0]
	if serializedOutput.LaneID != "spells" || serializedOutput.Artifact.Schema.ID != spells.ResponseSchemaID || serializedOutput.Artifact.Schema.Version != spells.SchemaVersion {
		t.Fatalf("serialized output envelope = %#v, want dnd spells schema on spells lane", serializedOutput)
	}
	response := decodeRunnerSpellResponse(t, serializedOutput.Artifact.Content)
	if len(response.SpellCasts) != 2 {
		t.Fatalf("len(spell_casts) = %d, want 2", len(response.SpellCasts))
	}
	first, second := response.SpellCasts[0], response.SpellCasts[1]
	if first.Spell != "Cure Wounds" || second.Spell != "Fire Bolt" {
		t.Fatalf("spell order = %q, %q; want source-unit order", first.Spell, second.Spell)
	}
	if first.Caster != "Aria" || second.Caster != "Borin" {
		t.Fatalf("casters = %q, %q; want spell data", first.Caster, second.Caster)
	}
	for _, spell := range response.SpellCasts {
		if len(spell.SourceRefs) != 1 {
			t.Fatalf("len(SourceRefs) = %d, want 1", len(spell.SourceRefs))
		}
		if spell.SourceRefs[0].SourceID != expectedDoc.ID {
			t.Fatalf("SourceID = %q, want fixture document ID", spell.SourceRefs[0].SourceID)
		}
	}

	if output.Manifest.InputModule != transcript.Key {
		t.Fatalf("manifest input module = %q, want %q", output.Manifest.InputModule, transcript.Key)
	}
	if output.Manifest.ValidationStatus != "approved" {
		t.Fatalf("ValidationStatus = %q, want approved", output.Manifest.ValidationStatus)
	}
	if len(output.Manifest.ArtifactLanes) != 1 {
		t.Fatalf("len(ArtifactLanes) = %d, want 1", len(output.Manifest.ArtifactLanes))
	}
	lane := output.Manifest.ArtifactLanes[0]
	if lane.ID != "spells" || lane.Extractor != spells.Key {
		t.Fatalf("manifest lane = %#v, want spells lane with dnd/spells extractor", lane)
	}
	extractorMetadata, ok := lane.Metadata["extractor"].(map[string]any)
	if !ok {
		t.Fatalf("manifest lane metadata = %#v, want extractor metadata", lane.Metadata)
	}
	if extractorMetadata["prompt_id"] != spells.PromptID ||
		extractorMetadata["response_schema_key"] != string(spells.ResponseSchemaKey) ||
		extractorMetadata["response_schema_name"] != spells.ResponseSchemaName {
		t.Fatalf("extractor metadata = %#v, want prompt/schema identifiers", extractorMetadata)
	}
	if len(output.OutputFiles) != 1 {
		t.Fatalf("len(OutputFiles) = %d, want 1", len(output.OutputFiles))
	}
	if output.OutputFiles[0].ContentType != "application/json" {
		t.Fatalf("ContentType = %q, want application/json", output.OutputFiles[0].ContentType)
	}
}

func TestRunnerPassesPartyAndGlossaryReferencesToDNDSpellsPrompt(t *testing.T) {
	raw := readDNDSpellsFixture(t)
	expectedDoc := parseDNDSpellsFixture(t, raw)
	resolved := resolveDNDSpellsPipeline(t)
	resolved.ResolvedPipeline.Steps[0].ArtifactLanes[0].ExtractReferences.ReferenceSet = dndSpellsReferenceSet(
		"Aria: party cleric\nBorin: fighter",
		"Fire Bolt: evocation cantrip",
	)
	llmClient := &fakeSpellsLLMClient{
		response: extractionResponse{
			SpellCasts: []spellCastResponse{
				{
					Caster:               "Borin",
					Spell:                "Fire Bolt",
					Effect:               "Scorches the wight.",
					NarrativeDescription: "Borin hurls fire at the wight.",
					SourceRefs:           responseSourceRefs(expectedDoc.ID, 3, 3),
				},
			},
		},
	}

	output, err := runPreparedPipeline(t, dndSpellsRunnerRegistries(t), resolved.ResolvedPipeline, llmClient, pipeline.RunInput{
		RawInput: raw,
	})
	if err != nil {
		t.Fatalf("Run() error = %v, want nil", err)
	}

	if len(output.NormalizeOutputs) != 1 {
		t.Fatalf("len(NormalizeOutputs) = %d, want 1", len(output.NormalizeOutputs))
	}
	if len(output.Manifest.References) != 2 {
		t.Fatalf("manifest references = %#v, want party and glossary provenance", output.Manifest.References)
	}
	if len(llmClient.requests) != 1 {
		t.Fatalf("LLM calls = %d, want 1", len(llmClient.requests))
	}
	request := llmClient.requests[0]
	if request.PromptID != spells.PromptID || request.PromptVersion != spells.SchemaVersion {
		t.Fatalf("prompt = %q/%q, want %q/%q", request.PromptID, request.PromptVersion, spells.PromptID, spells.SchemaVersion)
	}
	if got := string(request.Inputs["party"].Content); got != "Aria: party cleric\nBorin: fighter" {
		t.Fatalf("party input = %q, want reference text", got)
	}
	if got := string(request.Inputs["glossary"].Content); got != "Fire Bolt: evocation cantrip" {
		t.Fatalf("glossary input = %q, want reference text", got)
	}
}

func TestProductionSpellPipelineRoutesSemanticCandidatesToDeterministicValidators(t *testing.T) {
	registries := productionNPCRegistries(t)
	configValue := config.Default()
	configValue.Pipelines["dnd-spells-shape"] = pipeline.PipelineProfile{
		Input: pipeline.Binding(transcript.Key),
		Chunk: pipeline.ModuleBinding{Module: pipeline.DefaultChunkModule, Options: map[string]any{"max_units": 100}},
		Artifacts: map[string]pipeline.ArtifactLaneProfile{
			"spells": {
				Extract:   pipeline.ModuleBinding{Module: spells.Key, Retries: 2},
				Normalize: pipeline.Binding(spellnormalize.Key),
			},
		},
	}
	effective, err := configValue.Resolve(config.ResolveInput{
		PipelineID: "dnd-spells-shape",
		Catalog:    moduleCatalog(registries),
	})
	if err != nil {
		t.Fatalf("Resolve() error = %v, want nil", err)
	}
	for _, test := range []struct {
		name          string
		response      []byte
		reasonCode    string
		validatorName string
	}{
		{
			name:          "blank string",
			response:      []byte(`{"spell_casts":[{"caster":"Aria","spell":"Cure Wounds","effect":"","narrative_description":"Aria casts the spell.","source_refs":[{"start_unit_id":1,"end_unit_id":1}]}]}`),
			reasonCode:    spellshape.ReasonCode,
			validatorName: spellshape.Key,
		},
		{
			name:          "empty evidence",
			response:      []byte(`{"spell_casts":[{"caster":"Aria","spell":"Cure Wounds","effect":"Heals.","narrative_description":"Aria casts the spell.","source_refs":[]}]}`),
			reasonCode:    spellshape.ReasonCode,
			validatorName: spellshape.Key,
		},
		{
			name:          "nonpositive unit candidate",
			response:      []byte(`{"spell_casts":[{"caster":"Aria","spell":"Cure Wounds","effect":"Heals.","narrative_description":"Aria casts the spell.","source_refs":[{"start_unit_id":0,"end_unit_id":1}]}]}`),
			reasonCode:    spellsourcerefs.ReasonCode,
			validatorName: spellsourcerefs.Key,
		},
		{
			name:          "unknown unit candidate",
			response:      []byte(`{"spell_casts":[{"caster":"Aria","spell":"Cure Wounds","effect":"Heals.","narrative_description":"Aria casts the spell.","source_refs":[{"start_unit_id":99,"end_unit_id":99}]}]}`),
			reasonCode:    spellsourcerefs.ReasonCode,
			validatorName: spellsourcerefs.Key,
		},
	} {
		t.Run(test.name, func(t *testing.T) {
			client := &fakeSpellsLLMClient{rawResponses: [][]byte{test.response, test.response, test.response}}
			output, err := runPreparedPipeline(t, registries, effective.ResolvedPipeline, client, pipeline.RunInput{RawInput: readDNDSpellsFixture(t)})
			if err != nil {
				t.Fatalf("Run() error = %v, want non-fatal rejected output", err)
			}
			if len(client.requests) != 3 || len(output.Rejected) != 1 || len(output.NormalizeOutputs) != 0 {
				t.Fatalf("LLM requests = %d rejected = %#v normalized = %#v, want exhausted rejection", len(client.requests), output.Rejected, output.NormalizeOutputs)
			}
			rejection := output.Rejected[0]
			if rejection.ReasonCode != test.reasonCode || rejection.ValidatorName != test.validatorName || rejection.AttemptCount != 3 {
				t.Fatalf("rejection = %#v, want exhausted %s rejection", rejection, test.validatorName)
			}
		})
	}
}

func dndSpellsReferenceSet(party string, glossary string) contracts.ReferenceSet {
	slots := make(map[string]contracts.ResolvedReferenceSlot)
	if strings.TrimSpace(party) != "" {
		slots["party"] = contracts.ResolvedReferenceSlot{
			Slot: contracts.ReferenceSlot{Name: "party"},
			Items: []contracts.ReferenceItem{
				{
					SlotName:      "party",
					MediaType:     "text/plain; charset=utf-8",
					Content:       []byte(party),
					Digest:        "sha256:party",
					Origin:        contracts.ReferenceOrigin{Type: "file", URI: "file:///tmp/party.txt"},
					SizeBytes:     int64(len(party)),
					BindingSource: contracts.ReferenceBindingSourceConfig,
				},
			},
		}
	}
	if strings.TrimSpace(glossary) != "" {
		slots["glossary"] = contracts.ResolvedReferenceSlot{
			Slot: contracts.ReferenceSlot{Name: "glossary"},
			Items: []contracts.ReferenceItem{
				{
					SlotName:      "glossary",
					MediaType:     "text/plain; charset=utf-8",
					Content:       []byte(glossary),
					Digest:        "sha256:glossary",
					Origin:        contracts.ReferenceOrigin{Type: "file", URI: "file:///tmp/glossary.txt"},
					SizeBytes:     int64(len(glossary)),
					BindingSource: contracts.ReferenceBindingSourceConfig,
				},
			},
		}
	}
	return contracts.ReferenceSet{Slots: slots}
}

func resolveDNDSpellsPipeline(t *testing.T) config.EffectiveConfig {
	t.Helper()

	resolved, err := loadDNDSpellsPipelineConfig(t).Resolve(config.ResolveInput{
		PipelineID: "dnd-spells-fixture",
		Catalog:    dndSpellsTestCatalog(t, dndSpellsCatalogSpecs{}),
	})
	if err != nil {
		t.Fatalf("Resolve() error = %v, want nil", err)
	}
	return resolved
}

func dndSpellsRunnerRegistries(t *testing.T) pipeline.Registries {
	t.Helper()

	catalog := dndSpellsTestCatalog(t, dndSpellsCatalogSpecs{})
	return pipeline.Registries{
		Inputs:         catalog.Inputs,
		Chunkers:       catalog.Chunkers,
		ArtifactCodecs: catalog.ArtifactCodecs,
		Extractors:     catalog.Extractors,
		Mergers:        catalog.Mergers,
		Normalizers:    catalog.Normalizers,
		Outputs:        catalog.Outputs,
	}
}

func readDNDSpellsFixture(t *testing.T) []byte {
	t.Helper()

	raw, err := os.ReadFile("testdata/seriatim_spell_session.json")
	if err != nil {
		t.Fatalf("ReadFile(seriatim_spell_session.json) error = %v, want nil", err)
	}
	return raw
}

func parseDNDSpellsFixture(t *testing.T, raw []byte) *source.SourceDocument {
	t.Helper()

	doc, err := transcript.New().Parse(context.Background(), contracts.ParseRequest{Raw: raw})
	if err != nil {
		t.Fatalf("Parse() error = %v, want nil", err)
	}
	return doc
}

func decodeRunnerSpellResponse(t *testing.T, raw []byte) extractionResponse {
	t.Helper()

	var response extractionResponse
	if err := json.Unmarshal(raw, &response); err != nil {
		t.Fatalf("Unmarshal(raw output) error = %v, want nil", err)
	}
	return response
}