Consolidate D&D prompt ordering tests
This commit is contained in:
166
internal/modules/dnd/register/prompt_cache_test.go
Normal file
166
internal/modules/dnd/register/prompt_cache_test.go
Normal file
@@ -0,0 +1,166 @@
|
||||
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/scriptorium"
|
||||
)
|
||||
|
||||
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]scriptorium.ArtifactRef{
|
||||
"transcript": scriptorium.InlineWithURI("file:///session.json", `{"sentinel":"`+transcriptSentinel+`"}`),
|
||||
"players": scriptorium.Inline(playersSentinel),
|
||||
"party": scriptorium.Inline(partySentinel),
|
||||
"glossary": scriptorium.Inline(glossarySentinel),
|
||||
}
|
||||
cases := []struct {
|
||||
name string
|
||||
promptID string
|
||||
promptVersion string
|
||||
inputs map[string]scriptorium.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]scriptorium.ArtifactRef{
|
||||
"npcs": scriptorium.Inline(`{"sentinel":"` + npcSentinel + `"}`),
|
||||
}),
|
||||
npcInput: true,
|
||||
},
|
||||
{
|
||||
name: "npc interactions",
|
||||
promptID: interactionextract.PromptID,
|
||||
promptVersion: interactionextract.SchemaVersion,
|
||||
inputs: withPromptInputs(commonInputs, map[string]scriptorium.ArtifactRef{
|
||||
"npcs": scriptorium.Inline(`{"sentinel":"` + npcSentinel + `"}`),
|
||||
}),
|
||||
npcInput: true,
|
||||
},
|
||||
{
|
||||
name: "spells",
|
||||
promptID: spellextract.PromptID,
|
||||
promptVersion: spellextract.SchemaVersion,
|
||||
inputs: withPromptInputs(commonInputs, map[string]scriptorium.ArtifactRef{
|
||||
"npcs": scriptorium.Inline(`{"sentinel":"` + npcSentinel + `"}`),
|
||||
"spell_catalog": scriptorium.Inline(`{"sentinel":"` + catalogSentinel + `"}`),
|
||||
}),
|
||||
npcInput: true,
|
||||
spellCatalogInput: true,
|
||||
},
|
||||
}
|
||||
|
||||
var sharedPrefix []scriptorium.RenderedMessage
|
||||
for _, testCase := range cases {
|
||||
t.Run(testCase.name, func(t *testing.T) {
|
||||
prepared, err := engine.Prepare(context.Background(), scriptorium.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([]scriptorium.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) *scriptorium.Engine {
|
||||
t.Helper()
|
||||
options, err := registry.ScriptoriumOptions()
|
||||
if err != nil {
|
||||
t.Fatalf("ScriptoriumOptions() error = %v", err)
|
||||
}
|
||||
options = append(options, scriptorium.WithProfiles(scriptorium.OpenAICompatibleProfile(scriptorium.OpenAICompatibleProfileConfig{
|
||||
ID: "prompt-cache-test-profile", Endpoint: "http://127.0.0.1:1/v1", Model: "prompt-cache-test-model",
|
||||
})))
|
||||
engine, err := scriptorium.NewEngine(scriptorium.Config{Timeout: time.Second}, options...)
|
||||
if err != nil {
|
||||
t.Fatalf("NewEngine() error = %v", err)
|
||||
}
|
||||
return engine
|
||||
}
|
||||
|
||||
func withPromptInputs(inputs, extras map[string]scriptorium.ArtifactRef) map[string]scriptorium.ArtifactRef {
|
||||
merged := make(map[string]scriptorium.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 []scriptorium.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 []scriptorium.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
|
||||
}
|
||||
Reference in New Issue
Block a user