notarius/internal/modules/dnd/register/prompt_cache_test.go

package register

import (
	"context"
	"reflect"
	"strings"
	"testing"
	"time"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
	combatextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/combatturns"
	enemyeventextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/enemyevents"
	itemoccurrenceextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/itemoccurrences"
	itemregistryextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/itemregistry"
	locationoccurrenceextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/locationoccurrences"
	locationextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/locationregistry"
	occurrenceextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/npcoccurrences"
	npcextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/npcregistry"
	scenedescriptionextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/scenedescriptions"
	spellextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/spells"
	"gitea.maximumdirect.net/eric/promptkit"
)

func TestExtractionPromptComposition(t *testing.T) {
	const (
		transcriptSentinel   = "shared-transcript-sentinel"
		playersSentinel      = "shared-players-sentinel"
		partySentinel        = "shared-party-sentinel"
		glossarySentinel     = "shared-glossary-sentinel"
		npcSentinel          = "npc-registry-sentinel"
		catalogSentinel      = "spell-catalog-sentinel"
		evidenceSentinel     = "Transcript units are the only evidence"
		combatPolicySentinel = "substantive active combat materially organizes the"
	)
	registry := llm.NewAssetRegistry()
	if err := registerPromptAssets(registry); err != nil {
		t.Fatalf("registerPromptAssets() error = %v", err)
	}
	engine := newPromptCacheEngine(t, registry)
	commonInputs := map[string]promptkit.ArtifactRef{
		"transcript": promptkit.InlineWithURI("file:///session.json", `{"sentinel":"`+transcriptSentinel+`"}`),
		"players":    promptkit.Inline(playersSentinel),
		"party":      promptkit.Inline(partySentinel),
		"glossary":   promptkit.Inline(glossarySentinel),
	}
	cases := []struct {
		name           string
		promptID       string
		promptVersion  string
		inputs         map[string]promptkit.ArtifactRef
		suffixGroups   [][]string
		middleSentinel string
	}{
		{name: "npcs", promptID: npcextract.PromptID, promptVersion: npcextract.SchemaVersion, inputs: commonInputs, suffixGroups: [][]string{{evidenceSentinel}}},
		{name: "locations", promptID: locationextract.PromptID, promptVersion: locationextract.SchemaVersion, inputs: commonInputs, suffixGroups: [][]string{{evidenceSentinel}}},
		{name: "item occurrences", promptID: itemoccurrenceextract.PromptID, promptVersion: itemoccurrenceextract.SchemaVersion, inputs: withPromptInputs(commonInputs, map[string]promptkit.ArtifactRef{
			"item_registry": promptkit.Inline(`{"items":[{"id":"item-registry-sentinel","name":"Torch"}]}`),
		}), suffixGroups: [][]string{{evidenceSentinel}, {"item-registry-sentinel"}}},
		{name: "item registry", promptID: itemregistryextract.PromptID, promptVersion: itemregistryextract.SchemaVersion, inputs: commonInputs, suffixGroups: [][]string{{evidenceSentinel}}},
		{name: "scene descriptions", promptID: scenedescriptionextract.PromptID, promptVersion: scenedescriptionextract.SchemaVersion, inputs: commonInputs, middleSentinel: combatPolicySentinel},
		{
			name:          "combat turns",
			promptID:      combatextract.PromptID,
			promptVersion: combatextract.SchemaVersion,
			inputs: withPromptInputs(commonInputs, map[string]promptkit.ArtifactRef{
				"npc_registry": promptkit.Inline(`{"sentinel":"` + npcSentinel + `"}`),
			}),
			suffixGroups: [][]string{{evidenceSentinel}, {npcSentinel}},
		},
		{
			name:          "enemy events",
			promptID:      enemyeventextract.PromptID,
			promptVersion: enemyeventextract.SchemaVersion,
			inputs: withPromptInputs(commonInputs, map[string]promptkit.ArtifactRef{
				"npc_registry":    promptkit.Inline(`{"sentinel":"` + npcSentinel + `"}`),
				"combat_turns":    promptkit.Inline(`{"sentinel":"combat-turns-sentinel"}`),
				"npc_occurrences": promptkit.Inline(`{"sentinel":"npc-occurrences-sentinel"}`),
			}),
			suffixGroups: [][]string{
				{evidenceSentinel},
				{npcSentinel},
				{"combat-turns-sentinel", "npc-occurrences-sentinel"},
			},
		},
		{
			name:          "npc occurrences",
			promptID:      occurrenceextract.PromptID,
			promptVersion: occurrenceextract.SchemaVersion,
			inputs: withPromptInputs(commonInputs, map[string]promptkit.ArtifactRef{
				"npc_registry": promptkit.Inline(`{"sentinel":"` + npcSentinel + `"}`),
			}),
			suffixGroups: [][]string{{evidenceSentinel}, {npcSentinel}},
		},
		{
			name:          "location occurrences",
			promptID:      locationoccurrenceextract.PromptID,
			promptVersion: locationoccurrenceextract.SchemaVersion,
			inputs: withPromptInputs(commonInputs, map[string]promptkit.ArtifactRef{
				"location_registry": promptkit.Inline(`{"sentinel":"location-registry-sentinel"}`),
			}),
			suffixGroups: [][]string{{evidenceSentinel}, {"location-registry-sentinel"}},
		},
		{
			name:          "spells",
			promptID:      spellextract.PromptID,
			promptVersion: spellextract.SchemaVersion,
			inputs: withPromptInputs(commonInputs, map[string]promptkit.ArtifactRef{
				"npc_registry":  promptkit.Inline(`{"sentinel":"` + npcSentinel + `"}`),
				"spell_catalog": promptkit.Inline(`{"sentinel":"` + catalogSentinel + `"}`),
			}),
			suffixGroups: [][]string{{evidenceSentinel}, {npcSentinel}, {catalogSentinel}},
		},
	}

	var sharedPrefix []promptkit.RenderedMessage
	for _, testCase := range cases {
		t.Run(testCase.name, func(t *testing.T) {
			prepared, err := engine.Prepare(context.Background(), promptkit.RunRequest{
				PromptID:      testCase.promptID,
				PromptVersion: testCase.promptVersion,
				ProfileID:     "prompt-cache-test-profile",
				Inputs:        testCase.inputs,
			})
			if err != nil {
				t.Fatalf("Prepare() error = %v", err)
			}
			transcriptIndex := renderedInputMessageIndex(t, prepared.Messages, transcriptSentinel)
			referenceIndex := renderedInputMessageIndex(t, prepared.Messages, playersSentinel)
			if referenceIndex >= transcriptIndex {
				t.Fatalf("campaign references rendered at message %d, want before transcript message %d", referenceIndex, transcriptIndex)
			}
			assertEphemeralCache(t, prepared.Messages, referenceIndex)
			assertEphemeralCache(t, prepared.Messages, transcriptIndex)
			prefix := prepared.Messages[:transcriptIndex+1]
			if len(prepared.Messages) <= len(prefix) {
				t.Fatalf("prepared prompt has %d messages, want lane-specific suffix after transcript", len(prepared.Messages))
			}
			if sharedPrefix == nil {
				sharedPrefix = append([]promptkit.RenderedMessage(nil), prefix...)
			} else if !reflect.DeepEqual(prefix, sharedPrefix) {
				t.Fatalf("rendered prefix = %#v, want %#v", prefix, sharedPrefix)
			}
			previousIndex := transcriptIndex
			for _, group := range testCase.suffixGroups {
				groupIndex := -1
				for _, sentinel := range group {
					inputIndex := renderedInputMessageIndex(t, prepared.Messages, sentinel)
					if groupIndex < 0 {
						groupIndex = inputIndex
					} else if inputIndex != groupIndex {
						t.Fatalf("input sentinel %q rendered at message %d, want grouped at message %d", sentinel, inputIndex, groupIndex)
					}
				}
				if groupIndex <= previousIndex {
					t.Fatalf("prompt suffix group rendered at message %d, want after message %d", groupIndex, previousIndex)
				}
				previousIndex = groupIndex
			}
			instructionIndex := len(prepared.Messages) - 1
			if testCase.middleSentinel != "" {
				policyIndex := renderedInputMessageIndex(t, prepared.Messages, testCase.middleSentinel)
				if policyIndex <= transcriptIndex || policyIndex >= instructionIndex {
					t.Fatalf("combat policy rendered at message %d, want after transcript message %d and before instructions message %d", policyIndex, transcriptIndex, instructionIndex)
				}
			}
			if instructionIndex <= previousIndex {
				t.Fatalf("instructions rendered at message %d, want after lane input message %d", instructionIndex, previousIndex)
			}
			assertEphemeralCache(t, prepared.Messages, instructionIndex)
		})
	}
}

func newPromptCacheEngine(t *testing.T, registry *llm.AssetRegistry) *promptkit.Engine {
	t.Helper()
	options, err := registry.PromptKitOptions()
	if err != nil {
		t.Fatalf("PromptKitOptions() error = %v", err)
	}
	options = append(options, promptkit.WithProfiles(promptkit.OpenAICompatibleProfile(promptkit.OpenAICompatibleProfileConfig{
		ID: "prompt-cache-test-profile", Endpoint: "http://127.0.0.1:1/v1", Model: "prompt-cache-test-model",
	})))
	engine, err := promptkit.NewEngine(promptkit.Config{Timeout: time.Second}, options...)
	if err != nil {
		t.Fatalf("NewEngine() error = %v", err)
	}
	return engine
}

func withPromptInputs(inputs, extras map[string]promptkit.ArtifactRef) map[string]promptkit.ArtifactRef {
	merged := make(map[string]promptkit.ArtifactRef, len(inputs)+len(extras))
	for name, input := range inputs {
		merged[name] = input
	}
	for name, input := range extras {
		merged[name] = input
	}
	return merged
}

func assertEphemeralCache(t *testing.T, messages []promptkit.RenderedMessage, index int) {
	t.Helper()
	if cache := messages[index].CacheControl; cache == nil || cache.Type != promptkit.CacheControlEphemeral {
		t.Fatalf("message %d cache control = %#v, want ephemeral", index, cache)
	}
}

func renderedInputMessageIndex(t *testing.T, messages []promptkit.RenderedMessage, sentinel string) int {
	t.Helper()
	index := -1
	occurrences := 0
	for messageIndex, message := range messages {
		count := strings.Count(message.Content, sentinel)
		if count > 0 {
			index = messageIndex
			occurrences += count
		}
	}
	if occurrences != 1 {
		t.Fatalf("input sentinel %q rendered %d times, want exactly once", sentinel, occurrences)
	}
	return index
}