notarius/internal/modules/dnd/shared/prompt_inputs_test.go

package shared

import (
	"reflect"
	"strings"
	"testing"

	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

func TestPromptInputsIdentity(t *testing.T) {
	source := contracts.NewLLMInputMaterial("source", "application/json", []byte("source text"), "sha256:source", "file:///source.json")
	baseReferences := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"players":  slotWithContent("players", "Alice: Aria"),
		"party":    slotWithContent("party", "Aria: cleric"),
		"glossary": slotWithContent("glossary", "Brightmantle: temple"),
	}}
	orderedReferences := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"players": {Items: []contracts.ReferenceItem{
			referenceItem("players", "file:///b.txt", "sha256:bbb", "second"),
			referenceItem("players", "file:///a.txt", "sha256:aaa", "first"),
		}},
	}}
	reversedReferences := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"players": {Items: []contracts.ReferenceItem{
			referenceItem("players", "file:///a.txt", "sha256:aaa", "first"),
			referenceItem("players", "file:///b.txt", "sha256:bbb", "second"),
		}},
	}}
	rosterReferences := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"party":  {},
		"roster": slotWithContent("roster", "Legacy roster text"),
	}}
	explicitPartyReferences := contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{
		"party":  slotWithContent("party", "Current party text"),
		"roster": slotWithContent("roster", "Legacy roster text"),
	}}

	tests := []struct {
		name       string
		references contracts.ReferenceSet
		other      *contracts.ReferenceSet
		want       map[string]string
	}{
		{
			name:       "identical materials produce deeply equal inputs",
			references: baseReferences,
			other:      &baseReferences,
			want: map[string]string{
				"transcript": "source text",
				"players":    "Alice: Aria",
				"party":      "Aria: cleric",
				"glossary":   "Brightmantle: temple",
			},
		},
		{
			name:       "reference insertion order does not change bytes",
			references: orderedReferences,
			other:      &reversedReferences,
			want:       map[string]string{"players": "Reference 1\nOrigin-URI: file:///a.txt\nDigest: sha256:aaa\n\nfirst\n\nReference 2\nOrigin-URI: file:///b.txt\nDigest: sha256:bbb\n\nsecond"},
		},
		{
			name:       "roster becomes canonical party input",
			references: rosterReferences,
			want:       map[string]string{"party": "Legacy roster text"},
		},
		{
			name:       "explicit party wins over roster",
			references: explicitPartyReferences,
			want:       map[string]string{"party": "Current party text"},
		},
		{
			name:       "missing optional slots use placeholders",
			references: contracts.ReferenceSet{},
			want: map[string]string{
				"players":  " ",
				"party":    " ",
				"glossary": " ",
			},
		},
	}

	for _, test := range tests {
		t.Run(test.name, func(t *testing.T) {
			inputs := PromptInputs(source, test.references)
			if test.other != nil {
				other := PromptInputs(source, *test.other)
				if !reflect.DeepEqual(inputs, other) {
					t.Fatalf("PromptInputs() differs for equivalent references:\nfirst=%#v\nsecond=%#v", inputs, other)
				}
			}
			for name, want := range test.want {
				got, ok := inputs[name]
				if !ok {
					t.Fatalf("PromptInputs() missing %q: %#v", name, inputs)
				}
				if got := string(got.Content); got != want {
					t.Fatalf("%s content = %q, want %q", name, got, want)
				}
			}
			if _, ok := inputs["roster"]; ok {
				t.Fatalf("PromptInputs() included roster input: %#v", inputs)
			}
		})
	}
}

func TestTranscriptPromptMaterialClonesSource(t *testing.T) {
	source := contracts.NewLLMInputMaterial("source", "text/plain", []byte("source text"), "sha256:source", "file:///source.txt")
	got := TranscriptPromptMaterial(source)

	if got.Name != "transcript" {
		t.Fatalf("Name = %q, want transcript", got.Name)
	}
	if got.MediaType != source.MediaType || got.Digest != source.Digest || got.OriginURI != source.OriginURI || got.SizeBytes != source.SizeBytes {
		t.Fatalf("TranscriptPromptMaterial() = %#v, want cloned metadata from %#v", got, source)
	}
	source.Content[0] = 'X'
	if string(got.Content) != "source text" {
		t.Fatalf("TranscriptPromptMaterial() reused content slice: %q", got.Content)
	}
}

func TestReferencePromptMaterialUsesTextPlainAndSingleReferenceMetadata(t *testing.T) {
	slot := contracts.ResolvedReferenceSlot{Items: []contracts.ReferenceItem{{
		Content: []byte("single reference"),
		Digest:  "sha256:reference",
		Origin:  contracts.ReferenceOrigin{Type: "file", URI: "file:///reference.md"},
	}}}

	got := ReferencePromptMaterial("party", slot)
	if got.Name != "party" || got.MediaType != "text/plain" {
		t.Fatalf("ReferencePromptMaterial() name/media = %q/%q, want party/text/plain", got.Name, got.MediaType)
	}
	if got.Digest != "sha256:reference" || got.OriginURI != "file:///reference.md" {
		t.Fatalf("ReferencePromptMaterial() digest/origin = %q/%q, want single reference metadata", got.Digest, got.OriginURI)
	}
	if string(got.Content) != "single reference" {
		t.Fatalf("ReferencePromptMaterial() content = %q, want raw single reference", got.Content)
	}
}

func TestReferencePromptMaterialOmitsAggregateMetadata(t *testing.T) {
	got := ReferencePromptMaterial("party", contracts.ResolvedReferenceSlot{Items: []contracts.ReferenceItem{
		{Content: []byte("one"), Digest: "sha256:one", Origin: contracts.ReferenceOrigin{URI: "file:///one.md"}},
		{Content: []byte("two"), Digest: "sha256:two", Origin: contracts.ReferenceOrigin{URI: "file:///two.md"}},
	}})

	if got.Digest != "" || got.OriginURI != "" {
		t.Fatalf("ReferencePromptMaterial() digest/origin = %q/%q, want empty aggregate metadata", got.Digest, got.OriginURI)
	}
}

func TestReferencePromptInputRendering(t *testing.T) {
	if got := string(ReferencePromptInput(contracts.ResolvedReferenceSlot{})); got != " " {
		t.Fatalf("empty rendering = %q, want single space", got)
	}
	if got := string(ReferencePromptInput(slotWithContent("party", "single reference"))); got != "single reference" {
		t.Fatalf("single rendering = %q, want raw content", got)
	}

	slot := contracts.ResolvedReferenceSlot{Items: []contracts.ReferenceItem{
		{
			SlotName:  "party",
			MediaType: "text/plain",
			Content:   []byte("second"),
			Digest:    "sha256:bbb",
			Origin:    contracts.ReferenceOrigin{Type: "file", URI: "file:///b.txt"},
			SizeBytes: 6,
		},
		{
			SlotName:  "party",
			MediaType: "text/plain",
			Content:   []byte("first"),
			Digest:    "sha256:aaa",
			Origin:    contracts.ReferenceOrigin{Type: "file", URI: "file:///a.txt"},
			SizeBytes: 5,
		},
	}}
	first := string(ReferencePromptInput(slot))
	second := string(ReferencePromptInput(slot))
	if first != second {
		t.Fatalf("ReferencePromptInput() was not deterministic:\nfirst=%q\nsecond=%q", first, second)
	}
	if !strings.Contains(first, "Reference 1\nOrigin-Type: file\nOrigin-URI: file:///a.txt\nDigest: sha256:aaa\nMedia-Type: text/plain\nSize-Bytes: 5\n\nfirst") {
		t.Fatalf("first reference block = %q, want sorted first reference metadata", first)
	}
	if strings.Index(first, "first") > strings.Index(first, "second") {
		t.Fatalf("references were not sorted deterministically: %q", first)
	}
}

func slotWithContent(name string, content string) contracts.ResolvedReferenceSlot {
	return contracts.ResolvedReferenceSlot{
		Slot: contracts.ReferenceSlot{Name: name},
		Items: []contracts.ReferenceItem{{
			SlotName: name,
			Content:  []byte(content),
		}},
	}
}

func referenceItem(slotName, uri, digest, content string) contracts.ReferenceItem {
	return contracts.ReferenceItem{
		SlotName: slotName,
		Content:  []byte(content),
		Digest:   digest,
		Origin:   contracts.ReferenceOrigin{URI: uri},
	}
}