329 lines
15 KiB
Go
329 lines
15 KiB
Go
package entityreconcile
|
|
|
|
import (
|
|
"bytes"
|
|
"encoding/json"
|
|
"io/fs"
|
|
"reflect"
|
|
"strings"
|
|
"testing"
|
|
|
|
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
|
|
"gitea.maximumdirect.net/eric/notarius/internal/framework/llm"
|
|
"github.com/santhosh-tekuri/jsonschema/v6"
|
|
)
|
|
|
|
func TestBuildContextUsesContextualSelectorsSourceOrderAndOwnedData(t *testing.T) {
|
|
doc := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{
|
|
{ID: 40, Kind: "narration", Text: "zero"},
|
|
{ID: 10, Kind: "speech", Text: "one", Metadata: map[string]any{"speaker": map[string]any{"name": "Mira"}}},
|
|
{ID: 70, Kind: "speech", Text: "two"},
|
|
{ID: 20, Kind: "narration", Text: "three"},
|
|
{ID: 90, Kind: "speech", Text: "four"},
|
|
}}
|
|
candidates := []Candidate{
|
|
{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 20}}},
|
|
{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 90, EndUnitID: 90}}},
|
|
{Name: "Broken", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 20, EndUnitID: 10}}},
|
|
}
|
|
before := cloneCandidates(candidates)
|
|
|
|
materials, ready, err := BuildContext(doc, candidates, 1)
|
|
if err != nil || !ready {
|
|
t.Fatalf("BuildContext() = %#v, %t, %v; want ready materials", materials, ready, err)
|
|
}
|
|
if !reflect.DeepEqual(candidates, before) {
|
|
t.Fatalf("BuildContext() mutated candidates: %#v", candidates)
|
|
}
|
|
if got, want := materials.CandidateKeys(), []string{"candidate-000001", "candidate-000002", "candidate-000003"}; !reflect.DeepEqual(got, want) {
|
|
t.Fatalf("CandidateKeys() = %#v, want %#v", got, want)
|
|
}
|
|
if got, want := materials.EligibleCandidateKeys(), []string{"candidate-000001", "candidate-000002"}; !reflect.DeepEqual(got, want) {
|
|
t.Fatalf("EligibleCandidateKeys() = %#v, want %#v", got, want)
|
|
}
|
|
if !json.Valid(materials.Candidates.Content) || !json.Valid(materials.Transcript.Content) {
|
|
t.Fatalf("prompt materials are not JSON: %#v", materials)
|
|
}
|
|
if strings.Contains(string(materials.Candidates.Content), doc.ID) {
|
|
t.Fatalf("candidate material leaked source identity: %s", materials.Candidates.Content)
|
|
}
|
|
|
|
var candidatePayload candidateInput
|
|
if err := json.Unmarshal(materials.Candidates.Content, &candidatePayload); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(candidatePayload.Candidates) != 2 || candidatePayload.Candidates[0].Name != candidatePayload.Candidates[1].Name || strings.Contains(string(materials.Candidates.Content), "candidate-") {
|
|
t.Fatalf("candidate payload = %#v, want contextual descriptors without keys", candidatePayload)
|
|
}
|
|
if got := candidatePayload.Candidates[0].SourceRefs[0]; got != (SourceRange{StartUnitID: 10, EndUnitID: 20}) {
|
|
t.Fatalf("candidate reference = %#v", got)
|
|
}
|
|
|
|
var transcript transcriptInput
|
|
if err := json.Unmarshal(materials.Transcript.Content, &transcript); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(transcript.Windows) != 1 || len(transcript.Windows[0].Units) != 5 {
|
|
t.Fatalf("windows = %#v, want one bounded coalesced window", transcript.Windows)
|
|
}
|
|
units := transcript.Windows[0].Units
|
|
for index, wantID := range []int{40, 10, 70, 20, 90} {
|
|
if units[index].ID != wantID {
|
|
t.Fatalf("window unit %d = %d, want source-order %d", index, units[index].ID, wantID)
|
|
}
|
|
}
|
|
if units[0].Cited || !units[1].Cited || !units[2].Cited || !units[3].Cited || !units[4].Cited {
|
|
t.Fatalf("citation flags = %#v", units)
|
|
}
|
|
|
|
windows, err := contextWindows(doc.Units, []sourceInterval{{start: 1, end: 1}}, make([]bool, len(doc.Units)))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
windows[0].Units[0].Metadata["speaker"].(map[string]any)["name"] = "changed"
|
|
if doc.Units[1].Metadata["speaker"].(map[string]any)["name"] != "Mira" {
|
|
t.Fatal("context metadata aliases source document")
|
|
}
|
|
}
|
|
|
|
func TestBuildContextExcludesCollidingDescriptors(t *testing.T) {
|
|
doc := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{{ID: 1}, {ID: 2}, {ID: 3}}}
|
|
candidates := []Candidate{
|
|
{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 1, EndUnitID: 1}}},
|
|
{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 1, EndUnitID: 1}}},
|
|
{Name: "The Market", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 3, EndUnitID: 3}}},
|
|
}
|
|
materials, ready, err := BuildContext(doc, candidates, 0)
|
|
if err != nil || ready || len(materials.EligibleCandidateKeys()) != 1 || strings.Contains(string(materials.Candidates.Content), "The Tavern") {
|
|
t.Fatalf("BuildContext() = %#v, %t, %v", materials, ready, err)
|
|
}
|
|
assessment := materials.Assess(ProposalResponse{DuplicateGroups: []DuplicateGroup{{
|
|
Members: []Selector{{Name: "The Tavern", SourceRefs: []SourceRange{{StartUnitID: 1, EndUnitID: 1}}}, {Name: "The Market", SourceRefs: []SourceRange{{StartUnitID: 3, EndUnitID: 3}}}},
|
|
Canonical: Selector{Name: "The Market", SourceRefs: []SourceRange{{StartUnitID: 3, EndUnitID: 3}}},
|
|
}}})
|
|
if !hasIssue(assessment.Issues(), "member_ineligible") {
|
|
t.Fatalf("Assess() issues = %#v, want collided descriptor rejection", assessment.Issues())
|
|
}
|
|
}
|
|
|
|
func TestAssessmentRejectsPartialAndReorderedDescriptors(t *testing.T) {
|
|
doc := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{{ID: 10}, {ID: 20}, {ID: 30}}}
|
|
materials, ready, err := BuildContext(doc, []Candidate{
|
|
{Name: "The Tavern", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 20, EndUnitID: 20}, {SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}}},
|
|
{Name: "The Market", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 30, EndUnitID: 30}}},
|
|
}, 0)
|
|
if err != nil || !ready {
|
|
t.Fatalf("BuildContext() = %#v, %t, %v", materials, ready, err)
|
|
}
|
|
selectors := materialSelectors(t, materials)
|
|
for _, refs := range [][]SourceRange{
|
|
{{StartUnitID: 10, EndUnitID: 10}},
|
|
{{StartUnitID: 20, EndUnitID: 20}, {StartUnitID: 10, EndUnitID: 10}},
|
|
} {
|
|
assessment := materials.Assess(ProposalResponse{DuplicateGroups: []DuplicateGroup{{
|
|
Members: []Selector{{Name: "The Tavern", SourceRefs: refs}, selectors[1]},
|
|
Canonical: selectors[1],
|
|
}}})
|
|
if !hasIssue(assessment.Issues(), "member_unknown") {
|
|
t.Fatalf("Assess(%#v) issues = %#v, want descriptor mismatch rejection", refs, assessment.Issues())
|
|
}
|
|
}
|
|
}
|
|
|
|
func TestBuildContextExcludesUnsafeReferencesAndCoalescesAdjacentWindows(t *testing.T) {
|
|
doc := &source.SourceDocument{ID: "session", Units: []source.SourceUnit{{ID: 9}, {ID: 3}, {ID: 8}, {ID: 1}, {ID: 7}}}
|
|
candidates := []Candidate{
|
|
{Name: "One", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 3, EndUnitID: 3}}},
|
|
{Name: "Two", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 8, EndUnitID: 8}}},
|
|
{Name: "Missing", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: 99, EndUnitID: 99}}},
|
|
{Name: "Foreign", SourceRefs: []source.SourceRef{{SourceID: "other", StartUnitID: 1, EndUnitID: 1}}},
|
|
{Name: "Blank"},
|
|
}
|
|
materials, ready, err := BuildContext(doc, candidates, 0)
|
|
if err != nil || !ready {
|
|
t.Fatalf("BuildContext() error = %v, ready = %t", err, ready)
|
|
}
|
|
if got, want := materials.EligibleCandidateKeys(), []string{"candidate-000001", "candidate-000002"}; !reflect.DeepEqual(got, want) {
|
|
t.Fatalf("EligibleCandidateKeys() = %#v, want %#v", got, want)
|
|
}
|
|
var transcript transcriptInput
|
|
if err := json.Unmarshal(materials.Transcript.Content, &transcript); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if len(transcript.Windows) != 1 || len(transcript.Windows[0].Units) != 2 || transcript.Windows[0].Units[0].ID != 3 || transcript.Windows[0].Units[1].ID != 8 {
|
|
t.Fatalf("windows = %#v, want adjacent document-order units coalesced", transcript.Windows)
|
|
}
|
|
if _, ready, err := BuildContext(doc, candidates, -1); err == nil || ready {
|
|
t.Fatalf("BuildContext(radius=-1) = ready %t, err %v", ready, err)
|
|
}
|
|
}
|
|
|
|
func TestAssessmentRejectsEveryUnsafeProposalCategory(t *testing.T) {
|
|
materials := preparedMaterials(t, 4, true)
|
|
selectors := materialSelectors(t, materials)
|
|
unsafe := []struct {
|
|
name string
|
|
response ProposalResponse
|
|
category string
|
|
}{
|
|
{"blank member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{{}, selectors[1]}, Canonical: selectors[1]}}}, "member_blank"},
|
|
{"unknown member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{{Name: "unknown", SourceRefs: []SourceRange{{StartUnitID: 99, EndUnitID: 99}}}, selectors[1]}, Canonical: selectors[1]}}}, "member_unknown"},
|
|
{"repeated member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{selectors[0], selectors[0]}, Canonical: selectors[0]}}}, "repeated_member"},
|
|
{"too small", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{selectors[0]}, Canonical: selectors[0]}}}, "fewer_than_two_members"},
|
|
{"canonical blank", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{selectors[0], selectors[1]}, Canonical: Selector{}}}}, "canonical_blank"},
|
|
{"canonical not member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []Selector{selectors[0], selectors[1]}, Canonical: selectors[2]}}}, "canonical_not_member"},
|
|
{"overlapping", ProposalResponse{DuplicateGroups: []DuplicateGroup{
|
|
{Members: []Selector{selectors[0], selectors[1]}, Canonical: selectors[0]},
|
|
{Members: []Selector{selectors[1], selectors[2]}, Canonical: selectors[2]},
|
|
}}, "overlapping_member"},
|
|
}
|
|
for _, test := range unsafe {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
assessment := materials.Assess(test.response)
|
|
if len(assessment.SafeGroups()) != 0 || assessment.DiscardedGroups() != len(test.response.DuplicateGroups) || !hasIssue(assessment.Issues(), test.category) {
|
|
t.Fatalf("Assess() = groups %#v discarded %d issues %#v; want %q rejection", assessment.SafeGroups(), assessment.DiscardedGroups(), assessment.Issues(), test.category)
|
|
}
|
|
})
|
|
}
|
|
}
|
|
|
|
func TestAssessmentReturnsNonOverlappingSafeGroupsAndDefensiveCopies(t *testing.T) {
|
|
materials := preparedMaterials(t, 4, false)
|
|
keys := materials.CandidateKeys()
|
|
selectors := materialSelectors(t, materials)
|
|
assessment := materials.Assess(ProposalResponse{DuplicateGroups: []DuplicateGroup{
|
|
{Members: []Selector{selectors[1], selectors[0]}, Canonical: selectors[1]},
|
|
{Members: []Selector{selectors[3], selectors[2]}, Canonical: selectors[2]},
|
|
}})
|
|
groups := assessment.SafeGroups()
|
|
if assessment.DiscardedGroups() != 0 || len(assessment.Issues()) != 0 || len(groups) != 2 {
|
|
t.Fatalf("assessment = %#v, %d, %#v", groups, assessment.DiscardedGroups(), assessment.Issues())
|
|
}
|
|
if got, want := groups[0].Members(), []string{keys[0], keys[1]}; !reflect.DeepEqual(got, want) || groups[0].Canonical() != keys[1] {
|
|
t.Fatalf("first safe group = %#v / %q", got, groups[0].Canonical())
|
|
}
|
|
keys[0] = "changed"
|
|
if materials.CandidateKeys()[0] == "changed" {
|
|
t.Fatal("CandidateKeys() exposed retained keys")
|
|
}
|
|
members := groups[0].Members()
|
|
members[0] = "changed"
|
|
if groups[0].Members()[0] == "changed" || assessment.SafeGroups()[0].Members()[0] == "changed" {
|
|
t.Fatal("SafeGroups() exposed retained members")
|
|
}
|
|
}
|
|
|
|
func TestSharedResponseSchemaIsPrivateStrictAndRegisterableOnce(t *testing.T) {
|
|
schema, err := LoadResponseSchema()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if schema.Key != ResponseSchemaKey || schema.ID != ResponseSchemaID || schema.Name != ResponseSchemaName || schema.Version != SchemaVersion || !strings.HasPrefix(schema.SHA256, "sha256:") || !json.Valid(schema.JSONSchema) {
|
|
t.Fatalf("schema = %#v", schema)
|
|
}
|
|
for _, test := range []struct {
|
|
name string
|
|
value any
|
|
valid bool
|
|
}{
|
|
{"empty groups", map[string]any{"duplicate_groups": []any{}}, true},
|
|
{"semantic proposal problem", map[string]any{"duplicate_groups": []any{map[string]any{"members": []any{map[string]any{"name": "Mira", "source_refs": []any{map[string]any{"start_unit_id": 1, "end_unit_id": 1}}}}, "canonical": map[string]any{"name": "Mira", "source_refs": []any{map[string]any{"start_unit_id": 1, "end_unit_id": 1}}}}}}, true},
|
|
{"missing groups", map[string]any{}, false},
|
|
{"unknown top level", map[string]any{"duplicate_groups": []any{}, "extra": true}, false},
|
|
{"replacement name", map[string]any{"duplicate_groups": []any{map[string]any{"members": []any{}, "canonical": map[string]any{"name": "Mira", "source_refs": []any{}}, "name": "replacement"}}}, false},
|
|
{"missing selector evidence", map[string]any{"duplicate_groups": []any{map[string]any{"members": []any{}, "canonical": map[string]any{"name": "Mira"}}}}, false},
|
|
{"invalid range", map[string]any{"duplicate_groups": []any{map[string]any{"members": []any{map[string]any{"name": "Mira", "source_refs": []any{map[string]any{"start_unit_id": 0, "end_unit_id": 1}}}}, "canonical": map[string]any{"name": "Mira", "source_refs": []any{}}}}}, false},
|
|
} {
|
|
t.Run(test.name, func(t *testing.T) {
|
|
content, err := json.Marshal(test.value)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
err = validateSchema(content, schema.JSONSchema)
|
|
if (err == nil) != test.valid {
|
|
t.Fatalf("validateSchema() error = %v, want valid=%t", err, test.valid)
|
|
}
|
|
})
|
|
}
|
|
first := schema.JSONSchema
|
|
first[0] = '['
|
|
second, err := LoadResponseSchema()
|
|
if err != nil || !json.Valid(second.JSONSchema) || bytes.Equal(first, second.JSONSchema) {
|
|
t.Fatalf("LoadResponseSchema() returned shared content: %s, %v", second.JSONSchema, err)
|
|
}
|
|
registry := llm.NewAssetRegistry()
|
|
if err := RegisterSchemaAssets(registry); err != nil {
|
|
t.Fatalf("RegisterSchemaAssets() error = %v", err)
|
|
}
|
|
if schemaFS, err := registry.SchemaFS(); err != nil {
|
|
t.Fatalf("SchemaFS() error = %v", err)
|
|
} else if content, err := fs.ReadFile(schemaFS, "dnd_entity_reconcile_llm.v1.json"); err != nil || !json.Valid(content) {
|
|
t.Fatalf("shared schema asset = %s, %v", content, err)
|
|
}
|
|
}
|
|
|
|
func preparedMaterials(t *testing.T, count int, includeIneligible bool) Materials {
|
|
t.Helper()
|
|
doc := &source.SourceDocument{ID: "session", Units: make([]source.SourceUnit, count)}
|
|
candidates := make([]Candidate, count)
|
|
for index := range candidates {
|
|
doc.Units[index] = source.SourceUnit{ID: index + 1, Text: "unit"}
|
|
candidates[index] = Candidate{Name: "same display name", SourceRefs: []source.SourceRef{{SourceID: doc.ID, StartUnitID: index + 1, EndUnitID: index + 1}}}
|
|
}
|
|
if includeIneligible && count > 3 {
|
|
candidates[3].SourceRefs = []source.SourceRef{{SourceID: "other", StartUnitID: 1, EndUnitID: 1}}
|
|
}
|
|
materials, ready, err := BuildContext(doc, candidates, 0)
|
|
if err != nil || !ready {
|
|
t.Fatalf("BuildContext() = %#v, %t, %v", materials, ready, err)
|
|
}
|
|
return materials
|
|
}
|
|
|
|
func materialSelectors(t *testing.T, materials Materials) []Selector {
|
|
t.Helper()
|
|
var input candidateInput
|
|
if err := json.Unmarshal(materials.Candidates.Content, &input); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
return input.Candidates
|
|
}
|
|
|
|
func cloneCandidates(input []Candidate) []Candidate {
|
|
output := make([]Candidate, len(input))
|
|
copy(output, input)
|
|
for index := range output {
|
|
output[index].SourceRefs = append([]source.SourceRef(nil), input[index].SourceRefs...)
|
|
}
|
|
return output
|
|
}
|
|
|
|
func hasIssue(issues []Issue, want string) bool {
|
|
for _, issue := range issues {
|
|
if issue.Category == want {
|
|
return true
|
|
}
|
|
}
|
|
return false
|
|
}
|
|
|
|
func validateSchema(instanceContent, schemaContent []byte) error {
|
|
instance, err := jsonschema.UnmarshalJSON(bytes.NewReader(instanceContent))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
document, err := jsonschema.UnmarshalJSON(bytes.NewReader(schemaContent))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
compiler := jsonschema.NewCompiler()
|
|
if err := compiler.AddResource("schema.json", document); err != nil {
|
|
return err
|
|
}
|
|
compiled, err := compiler.Compile("schema.json")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
return compiled.Validate(instance)
|
|
}
|