216 lines
9.3 KiB
Go
216 lines
9.3 KiB
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"
|
|
itemeventextract "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/extract/itemevents"
|
|
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"
|
|
)
|
|
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
|
|
}{
|
|
{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: itemeventextract.PromptID, promptVersion: itemeventextract.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},
|
|
{
|
|
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 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
|
|
}
|