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"
	itemeventextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/itemevents"
	interactionextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/npcinteractions"
	npcextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/npcs"
	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 TestExtractionPromptsShareRenderedPrefix(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"
	)
	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
		npcInput          bool
		spellCatalogInput bool
	}{
		{name: "npcs", promptID: npcextract.PromptID, promptVersion: npcextract.SchemaVersion, inputs: commonInputs},
		{name: "item events", promptID: itemeventextract.PromptID, promptVersion: itemeventextract.SchemaVersion, inputs: commonInputs},
		{name: "scene descriptions", promptID: scenedescriptionextract.PromptID, promptVersion: scenedescriptionextract.SchemaVersion, inputs: commonInputs},
		{
			name:          "combat turns",
			promptID:      combatextract.PromptID,
			promptVersion: combatextract.SchemaVersion,
			inputs: withPromptInputs(commonInputs, map[string]promptkit.ArtifactRef{
				"npcs": promptkit.Inline(`{"sentinel":"` + npcSentinel + `"}`),
			}),
			npcInput: true,
		},
		{
			name:          "npc interactions",
			promptID:      interactionextract.PromptID,
			promptVersion: interactionextract.SchemaVersion,
			inputs: withPromptInputs(commonInputs, map[string]promptkit.ArtifactRef{
				"npcs": promptkit.Inline(`{"sentinel":"` + npcSentinel + `"}`),
			}),
			npcInput: true,
		},
		{
			name:          "spells",
			promptID:      spellextract.PromptID,
			promptVersion: spellextract.SchemaVersion,
			inputs: withPromptInputs(commonInputs, map[string]promptkit.ArtifactRef{
				"npcs":          promptkit.Inline(`{"sentinel":"` + npcSentinel + `"}`),
				"spell_catalog": promptkit.Inline(`{"sentinel":"` + catalogSentinel + `"}`),
			}),
			npcInput:          true,
			spellCatalogInput: true,
		},
	}

	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)
			prefix := prepared.Messages[:transcriptIndex+1]
			if len(prefix) != 4 {
				t.Fatalf("messages through transcript = %d, want 4", len(prefix))
			}
			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)
			}
			if testCase.npcInput {
				assertRenderedInputAfter(t, prepared.Messages, npcSentinel, transcriptIndex)
			}
			if testCase.spellCatalogInput {
				assertRenderedInputAfter(t, prepared.Messages, catalogSentinel, transcriptIndex)
			}
		})
	}
}

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 assertRenderedInputAfter(t *testing.T, messages []promptkit.RenderedMessage, sentinel string, index int) {
	t.Helper()
	if inputIndex := renderedInputMessageIndex(t, messages, sentinel); inputIndex <= index {
		t.Fatalf("input sentinel %q rendered at message %d, want after transcript message %d", sentinel, inputIndex, index)
	}
}

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
}