eric/notarius
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use contextual descriptors for entity reconciliation

Browse Source
This commit is contained in:
Eric Rakestraw
2026-08-08 15:05:35 +00:00
parent 51d62de1f3
commit fc449863f2
18 changed files with 435 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

116
internal/modules/dnd/shared/entityreconcile/context.go
View File

@@ -22,13 +22,34 @@ type Candidate struct {
SourceRefs []source.SourceRef
}
// Selector identifies one candidate through its canonical name and source-free
// evidence ranges. It is the complete model-facing candidate descriptor.
type Selector struct {
Name string `json:"name"`
SourceRefs []SourceRange `json:"source_refs"`
}
// Clone returns an owned copy of the selector.
func (s Selector) Clone() Selector {
s.SourceRefs = cloneSourceRanges(s.SourceRefs)
return s
}
// SourceRange is a source-free evidence coordinate used in a selector.
type SourceRange struct {
StartUnitID int `json:"start_unit_id"`
EndUnitID int `json:"end_unit_id"`
}
// Materials contains owned prompt inputs and opaque candidate-key mappings.
type Materials struct {
Candidates contracts.LLMInputMaterial
Transcript contracts.LLMInputMaterial
candidateKeys []string
eligible map[string]struct{}
candidateKeys []string
eligible map[string]struct{}
keyBySelector map[string]string
collidedSelectors map[string]struct{}
}
// CandidateKeys returns all deterministic keys in candidate input order.
@@ -49,18 +70,7 @@ func (m Materials) EligibleCandidateKeys() []string {
}
type candidateInput struct {
Candidates []candidateView `json:"candidates"`
}
type candidateView struct {
Key string `json:"key"`
Name string `json:"name"`
SourceRefs []candidateSourceRef `json:"source_refs"`
}
type candidateSourceRef struct {
StartUnitID int `json:"start_unit_id"`
EndUnitID int `json:"end_unit_id"`
Candidates []Selector `json:"candidates"`
}
type transcriptInput struct {
@@ -91,8 +101,10 @@ func BuildContext(doc *source.SourceDocument, candidates []Candidate, radius int
return Materials{}, false, fmt.Errorf("build entity reconciliation context: radius must not be negative")
}
materials := Materials{
candidateKeys: make([]string, len(candidates)),
eligible: make(map[string]struct{}),
candidateKeys: make([]string, len(candidates)),
eligible: make(map[string]struct{}),
keyBySelector: make(map[string]string),
collidedSelectors: make(map[string]struct{}),
}
for index := range candidates {
key := fmt.Sprintf(candidateKeyFormat, index+1)
@@ -103,7 +115,15 @@ func BuildContext(doc *source.SourceDocument, candidates []Candidate, radius int
}
index := source.NewDocumentIndex(doc)
views := make([]candidateView, 0, len(candidates))
type preparedCandidate struct {
key string
selector Selector
intervals []sourceInterval
lookupKey string
}
prepared := make([]preparedCandidate, 0, len(candidates))
selectorCounts := make(map[string]int, len(candidates))
views := make([]Selector, 0, len(candidates))
intervals := make([]sourceInterval, 0)
cited := make([]bool, len(doc.Units))
for candidateIndex, candidate := range candidates {
@@ -112,9 +132,23 @@ func BuildContext(doc *source.SourceDocument, candidates []Candidate, radius int
continue
}
key := materials.candidateKeys[candidateIndex]
materials.eligible[key] = struct{}{}
views = append(views, candidateView{Key: key, Name: candidate.Name, SourceRefs: references})
for _, interval := range candidateIntervals {
selector := Selector{Name: candidate.Name, SourceRefs: references}
lookupKey, err := selectorLookupKey(selector)
if err != nil {
return Materials{}, false, fmt.Errorf("build entity reconciliation context: invalid candidate material")
}
prepared = append(prepared, preparedCandidate{key: key, selector: selector, intervals: candidateIntervals, lookupKey: lookupKey})
selectorCounts[lookupKey]++
}
for _, candidate := range prepared {
if selectorCounts[candidate.lookupKey] != 1 {
materials.collidedSelectors[candidate.lookupKey] = struct{}{}
continue
}
materials.eligible[candidate.key] = struct{}{}
materials.keyBySelector[candidate.lookupKey] = candidate.key
views = append(views, candidate.selector.Clone())
for _, interval := range candidate.intervals {
for position := interval.start; position <= interval.end; position++ {
cited[position] = true
}
@@ -145,24 +179,56 @@ func BuildContext(doc *source.SourceDocument, candidates []Candidate, radius int
return materials, true, nil
}
func candidateReferences(index source.DocumentIndex, refs []source.SourceRef) ([]candidateSourceRef, []sourceInterval, bool) {
func candidateReferences(index source.DocumentIndex, refs []source.SourceRef) ([]SourceRange, []sourceInterval, bool) {
if len(refs) == 0 {
return nil, nil, false
}
references := make([]candidateSourceRef, 0, len(refs))
intervals := make([]sourceInterval, 0, len(refs))
type referencedInterval struct {
reference SourceRange
interval sourceInterval
}
prepared := make([]referencedInterval, 0, len(refs))
for _, ref := range refs {
if err := index.ValidateRef(ref); err != nil {
return nil, nil, false
}
start, _ := index.Position(ref.StartUnitID)
end, _ := index.Position(ref.EndUnitID)
references = append(references, candidateSourceRef{StartUnitID: ref.StartUnitID, EndUnitID: ref.EndUnitID})
intervals = append(intervals, sourceInterval{start: start, end: end})
prepared = append(prepared, referencedInterval{reference: SourceRange{StartUnitID: ref.StartUnitID, EndUnitID: ref.EndUnitID}, interval: sourceInterval{start: start, end: end}})
}
sort.Slice(prepared, func(left, right int) bool {
if prepared[left].interval.start != prepared[right].interval.start {
return prepared[left].interval.start < prepared[right].interval.start
}
return prepared[left].interval.end < prepared[right].interval.end
})
references := make([]SourceRange, 0, len(prepared))
intervals := make([]sourceInterval, 0, len(prepared))
for _, item := range prepared {
if len(references) > 0 && references[len(references)-1] == item.reference {
continue
}
references = append(references, item.reference)
intervals = append(intervals, item.interval)
}
return references, intervals, true
}
func selectorLookupKey(selector Selector) (string, error) {
content, err := json.Marshal(selector.Clone())
if err != nil {
return "", err
}
return string(content), nil
}
func cloneSourceRanges(values []SourceRange) []SourceRange {
if len(values) == 0 {
return []SourceRange{}
}
return append([]SourceRange(nil), values...)
}
func coalesceIntervals(intervals []sourceInterval) []sourceInterval {
if len(intervals) == 0 {
return nil

93
internal/modules/dnd/shared/entityreconcile/entityreconcile_test.go
View File

@@ -13,7 +13,7 @@ import (
"github.com/santhosh-tekuri/jsonschema/v6"
)
func TestBuildContextUsesOpaqueKeysSourceOrderAndOwnedData(t *testing.T) {
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"}}},
@@ -52,10 +52,10 @@ func TestBuildContextUsesOpaqueKeysSourceOrderAndOwnedData(t *testing.T) {
if err := json.Unmarshal(materials.Candidates.Content, &candidatePayload); err != nil {
t.Fatal(err)
}
if len(candidatePayload.Candidates) != 2 || candidatePayload.Candidates[0].Key != "candidate-000001" || candidatePayload.Candidates[1].Key != "candidate-000002" || candidatePayload.Candidates[0].Name != candidatePayload.Candidates[1].Name {
t.Fatalf("candidate payload = %#v, want distinct opaque keys for equal names", candidatePayload)
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 != (candidateSourceRef{StartUnitID: 10, EndUnitID: 20}) {
if got := candidatePayload.Candidates[0].SourceRefs[0]; got != (SourceRange{StartUnitID: 10, EndUnitID: 20}) {
t.Fatalf("candidate reference = %#v", got)
}
@@ -86,6 +86,50 @@ func TestBuildContextUsesOpaqueKeysSourceOrderAndOwnedData(t *testing.T) {
}
}
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{
@@ -116,22 +160,21 @@ func TestBuildContextExcludesUnsafeReferencesAndCoalescesAdjacentWindows(t *test
func TestAssessmentRejectsEveryUnsafeProposalCategory(t *testing.T) {
materials := preparedMaterials(t, 4, true)
keys := materials.CandidateKeys()
selectors := materialSelectors(t, materials)
unsafe := []struct {
name string
response ProposalResponse
category string
}{
{"blank member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []string{"", keys[1]}, Canonical: keys[1]}}}, "member_blank"},
{"unknown member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []string{"candidate-999999", keys[1]}, Canonical: keys[1]}}}, "member_unknown"},
{"ineligible member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []string{keys[0], keys[3]}, Canonical: keys[0]}}}, "member_ineligible"},
{"repeated member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []string{keys[0], keys[0]}, Canonical: keys[0]}}}, "repeated_member"},
{"too small", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []string{keys[0]}, Canonical: keys[0]}}}, "fewer_than_two_members"},
{"canonical blank", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []string{keys[0], keys[1]}, Canonical: ""}}}, "canonical_blank"},
{"canonical not member", ProposalResponse{DuplicateGroups: []DuplicateGroup{{Members: []string{keys[0], keys[1]}, Canonical: keys[2]}}}, "canonical_not_member"},
{"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: []string{keys[0], keys[1]}, Canonical: keys[0]},
{Members: []string{keys[1], keys[2]}, Canonical: keys[2]},
{Members: []Selector{selectors[0], selectors[1]}, Canonical: selectors[0]},
{Members: []Selector{selectors[1], selectors[2]}, Canonical: selectors[2]},
}}, "overlapping_member"},
}
for _, test := range unsafe {
@@ -147,9 +190,10 @@ func TestAssessmentRejectsEveryUnsafeProposalCategory(t *testing.T) {
func TestAssessmentReturnsNonOverlappingSafeGroupsAndDefensiveCopies(t *testing.T) {
materials := preparedMaterials(t, 4, false)
keys := materials.CandidateKeys()
selectors := materialSelectors(t, materials)
assessment := materials.Assess(ProposalResponse{DuplicateGroups: []DuplicateGroup{
{Members: []string{keys[1], keys[0]}, Canonical: keys[1]},
{Members: []string{keys[3], keys[2]}, Canonical: keys[2]},
{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 {
@@ -183,12 +227,12 @@ func TestSharedResponseSchemaIsPrivateStrictAndRegisterableOnce(t *testing.T) {
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{""}, "canonical": ""}}}, 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": "candidate-000001", "name": "replacement"}}}, false},
{"replacement evidence", map[string]any{"duplicate_groups": []any{map[string]any{"members": []any{}, "canonical": "candidate-000001", "source_refs": []any{}}}}, false},
{"wrong key type", map[string]any{"duplicate_groups": []any{map[string]any{"members": []any{1}, "canonical": "candidate-000001"}}}, 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)
@@ -236,6 +280,15 @@ func preparedMaterials(t *testing.T, count int, includeIneligible bool) Material
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)

55
internal/modules/dnd/shared/entityreconcile/proposal.go
View File

@@ -6,15 +6,16 @@ import (
)
// ProposalResponse is the private structured response exchanged with the
// reconciliation prompt. It identifies candidates only by opaque keys.
// reconciliation prompt. It identifies candidates by contextual selectors.
type ProposalResponse struct {
DuplicateGroups []DuplicateGroup `json:"duplicate_groups"`
}
// DuplicateGroup proposes candidate keys that might denote one entity.
// DuplicateGroup proposes contextual candidate descriptors that might denote
// one entity.
type DuplicateGroup struct {
Members []string `json:"members"`
Canonical string `json:"canonical"`
Members []Selector `json:"members"`
Canonical Selector `json:"canonical"`
}
// Issue identifies one unsafe proposal category without prescribing a warning
@@ -59,16 +60,13 @@ func (a Assessment) DiscardedGroups() int { return a.discardedGroups }
// Issues returns the deterministic rejection categories in proposal order.
func (a Assessment) Issues() []Issue { return append([]Issue(nil), a.issues...) }
// Assess validates a proposal against the opaque keys created by BuildContext.
// Assess resolves contextual descriptors to internal candidate keys, then
// validates the proposal without exposing those keys to the model.
func (m Materials) Assess(response ProposalResponse) Assessment {
all := make(map[string]struct{}, len(m.candidateKeys))
for _, key := range m.candidateKeys {
all[key] = struct{}{}
}
groups := make([]assessedGroup, len(response.DuplicateGroups))
issues := make([]Issue, 0)
for groupIndex, proposal := range response.DuplicateGroups {
groups[groupIndex] = assessGroup(proposal, all, m.eligible)
groups[groupIndex] = m.assessGroup(proposal)
for _, category := range groups[groupIndex].issues {
issues = append(issues, Issue{GroupIndex: groupIndex, Category: category})
}
@@ -115,12 +113,13 @@ type assessedGroup struct {
conflicting bool
}
func assessGroup(proposal DuplicateGroup, all, eligible map[string]struct{}) assessedGroup {
func (m Materials) assessGroup(proposal DuplicateGroup) assessedGroup {
issues := make([]string, 0)
members := make([]string, 0, len(proposal.Members))
seen := make(map[string]struct{}, len(proposal.Members))
for _, key := range proposal.Members {
if category := keyCategory(key, all, eligible); category != "" {
for _, selector := range proposal.Members {
key, category := m.selectorKey(selector)
if category != "" {
issues = append(issues, "member_"+category)
continue
}
@@ -131,31 +130,39 @@ func assessGroup(proposal DuplicateGroup, all, eligible map[string]struct{}) ass
seen[key] = struct{}{}
members = append(members, key)
}
canonicalCategory := keyCategory(proposal.Canonical, all, eligible)
canonical, canonicalCategory := m.selectorKey(proposal.Canonical)
if canonicalCategory != "" {
issues = append(issues, "canonical_"+canonicalCategory)
}
if len(members) < 2 {
issues = append(issues, "fewer_than_two_members")
}
if canonicalCategory == "" && !contains(members, proposal.Canonical) {
if canonicalCategory == "" && !contains(members, canonical) {
issues = append(issues, "canonical_not_member")
}
sort.Strings(members)
return assessedGroup{members: members, canonical: proposal.Canonical, issues: issues, locallyValid: len(issues) == 0}
return assessedGroup{members: members, canonical: canonical, issues: issues, locallyValid: len(issues) == 0}
}
func keyCategory(key string, all, eligible map[string]struct{}) string {
if strings.TrimSpace(key) == "" {
return "blank"
func (m Materials) selectorKey(selector Selector) (string, string) {
if strings.TrimSpace(selector.Name) == "" {
return "", "blank"
}
if _, ok := all[key]; !ok {
return "unknown"
lookupKey, err := selectorLookupKey(selector)
if err != nil {
return "", "unknown"
}
if _, ok := eligible[key]; !ok {
return "ineligible"
if _, collided := m.collidedSelectors[lookupKey]; collided {
return "", "ineligible"
}
return ""
key, ok := m.keyBySelector[lookupKey]
if !ok {
return "", "unknown"
}
if _, eligible := m.eligible[key]; !eligible {
return "", "ineligible"
}
return key, ""
}
func contains(values []string, want string) bool {