From fc76805075df861d6e769647cae7eeab1f045ac1 Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Sat, 8 Aug 2026 14:49:08 +0000 Subject: [PATCH] Add contextual location grounding --- .../dnd/locations/registry/grounding.go | 197 ++++++++++++++++++ .../dnd/locations/registry/registry.go | 29 ++- .../dnd/locations/registry/registry_test.go | 179 ++++++++++++++++ .../locationoccurrences/normalizer.go | 2 +- .../locationoccurrences/registry/validator.go | 2 +- 5 files changed, 399 insertions(+), 10 deletions(-) create mode 100644 internal/modules/dnd/locations/registry/grounding.go diff --git a/internal/modules/dnd/locations/registry/grounding.go b/internal/modules/dnd/locations/registry/grounding.go new file mode 100644 index 0000000..56803e5 --- /dev/null +++ b/internal/modules/dnd/locations/registry/grounding.go @@ -0,0 +1,197 @@ +package registry + +import ( + "encoding/json" + "fmt" + + "gitea.maximumdirect.net/eric/notarius/internal/core/source" + "gitea.maximumdirect.net/eric/notarius/internal/framework/contracts" + "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd" + locationcodec "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/codec/locationregistry" + "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/identity" + "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared" +) + +// Selector identifies one location using its canonical name and, when needed, +// source-free registry ranges. +type Selector struct { + Name string `json:"name"` + RegistryRefs []RegistryRef `json:"registry_refs"` +} + +// RegistryRef is a source-free range used only for location identity grounding. +type RegistryRef struct { + StartUnitID int `json:"start_unit_id"` + EndUnitID int `json:"end_unit_id"` +} + +// ContextUnit supplies the bounded transcript context for an ambiguous +// location's registry ranges. +type ContextUnit struct { + UnitID int `json:"unit_id"` + Text string `json:"text"` +} + +// Grounding is an immutable, operation-scoped location selector projection. +// Its prompt input contains no durable IDs or source IDs. +type Grounding struct { + promptInput contracts.LLMInputMaterial + projectionDigest string + locationsBySelector map[string]dnd.Location +} + +type groundingProjection struct { + Locations []groundingProjectionLocation `json:"locations"` +} + +type groundingProjectionLocation struct { + Name string `json:"name"` + RegistryRefs []RegistryRef `json:"registry_refs"` + Context []ContextUnit `json:"context"` +} + +// NewGrounding creates the contextual location projection for one source +// document and resolved registry. +func NewGrounding(registry *Registry, doc *source.SourceDocument) (*Grounding, error) { + if registry == nil { + return nil, fmt.Errorf("location registry is required") + } + if doc == nil { + return nil, fmt.Errorf("source document is required") + } + + index := source.NewDocumentIndex(doc) + order := shared.NewSourceRefOrderFromIndex(index) + locations := registry.Locations() + grouped := make(map[string][]int, len(locations)) + for locationIndex, location := range locations { + key := identity.ComparisonKey(location.Name) + if key == "" { + return nil, fmt.Errorf("location registry record %d has an empty comparison name", locationIndex) + } + grouped[key] = append(grouped[key], locationIndex) + } + + projection := groundingProjection{Locations: make([]groundingProjectionLocation, len(locations))} + locationsBySelector := make(map[string]dnd.Location, len(locations)) + for locationIndex, location := range locations { + key := identity.ComparisonKey(location.Name) + registryRefs := make([]RegistryRef, 0) + context := make([]ContextUnit, 0) + if len(grouped[key]) > 1 { + var err error + registryRefs, context, err = contextualRanges(locationIndex, location.SourceRefs, doc, index, order) + if err != nil { + return nil, err + } + } + projection.Locations[locationIndex] = groundingProjectionLocation{ + Name: location.Name, + RegistryRefs: cloneRegistryRefs(registryRefs), + Context: cloneContextUnits(context), + } + selectorKey := contextualSelectorKey(key, registryRefs) + if _, exists := locationsBySelector[selectorKey]; exists { + return nil, fmt.Errorf("location registry records produce the same contextual selector") + } + locationsBySelector[selectorKey] = cloneLocation(location) + } + + content, err := json.Marshal(projection) + if err != nil { + return nil, fmt.Errorf("encode location grounding projection: %w", err) + } + digest := semanticDigest(content) + return &Grounding{ + promptInput: contracts.NewLLMInputMaterial(ReferenceSlot, locationcodec.MediaType, content, digest, ""), + projectionDigest: digest, + locationsBySelector: locationsBySelector, + }, nil +} + +// PromptInput returns a defensive copy of the contextual projection for the +// location_registry prompt slot. +func (g *Grounding) PromptInput() contracts.LLMInputMaterial { + if g == nil { + return contracts.LLMInputMaterial{} + } + return g.promptInput.Clone() +} + +// ProjectionDigest returns the digest of the exact contextual prompt input. +func (g *Grounding) ProjectionDigest() string { + if g == nil { + return "" + } + return g.projectionDigest +} + +// Resolve maps a contextual selector to one canonical location. +func (g *Grounding) Resolve(selector Selector) (dnd.Location, bool) { + if g == nil { + return dnd.Location{}, false + } + key := identity.ComparisonKey(selector.Name) + if key == "" { + return dnd.Location{}, false + } + location, ok := g.locationsBySelector[contextualSelectorKey(key, selector.RegistryRefs)] + if !ok { + return dnd.Location{}, false + } + return cloneLocation(location), true +} + +func contextualRanges(locationIndex int, refs []source.SourceRef, doc *source.SourceDocument, index source.DocumentIndex, order shared.SourceRefOrder) ([]RegistryRef, []ContextUnit, error) { + for _, ref := range refs { + if ref.SourceID != doc.ID || index.ValidateRef(ref) != nil { + return nil, nil, fmt.Errorf("location registry record %d has an invalid source reference", locationIndex) + } + } + canonical := order.Canonicalize(refs) + ranges := make([]RegistryRef, len(canonical)) + for refIndex, ref := range canonical { + ranges[refIndex] = RegistryRef{StartUnitID: ref.StartUnitID, EndUnitID: ref.EndUnitID} + } + return ranges, contextUnits(doc, index, canonical), nil +} + +func contextUnits(doc *source.SourceDocument, index source.DocumentIndex, refs []source.SourceRef) []ContextUnit { + units := make([]ContextUnit, 0) + seen := make(map[int]struct{}) + for _, ref := range refs { + start, _ := index.Position(ref.StartUnitID) + end, _ := index.Position(ref.EndUnitID) + for position := start; position <= end; position++ { + unit := doc.Units[position] + if _, exists := seen[unit.ID]; exists { + continue + } + seen[unit.ID] = struct{}{} + units = append(units, ContextUnit{UnitID: unit.ID, Text: unit.Text}) + } + } + return units +} + +func contextualSelectorKey(name string, refs []RegistryRef) string { + key, _ := json.Marshal(struct { + Name string `json:"name"` + RegistryRefs []RegistryRef `json:"registry_refs"` + }{Name: name, RegistryRefs: cloneRegistryRefs(refs)}) + return string(key) +} + +func cloneRegistryRefs(values []RegistryRef) []RegistryRef { + if len(values) == 0 { + return []RegistryRef{} + } + return append([]RegistryRef(nil), values...) +} + +func cloneContextUnits(values []ContextUnit) []ContextUnit { + if len(values) == 0 { + return []ContextUnit{} + } + return append([]ContextUnit(nil), values...) +} diff --git a/internal/modules/dnd/locations/registry/registry.go b/internal/modules/dnd/locations/registry/registry.go index ba9a9b3..737da44 100644 --- a/internal/modules/dnd/locations/registry/registry.go +++ b/internal/modules/dnd/locations/registry/registry.go @@ -31,6 +31,7 @@ type Registry struct { canonical []byte digest string projectionDigest string + identityDigest string promptInput contracts.LLMInputMaterial lookupByID map[string]int } @@ -105,6 +106,7 @@ func emptyRegistry() *Registry { list: dnd.LocationRegistry{Locations: []dnd.Location{}}, canonical: append([]byte(nil), content...), projectionDigest: projectionDigest, + identityDigest: projectionDigest, promptInput: contracts.NewLLMInputMaterial(ReferenceSlot, locationcodec.MediaType, content, projectionDigest, ""), lookupByID: map[string]int{}, } @@ -129,7 +131,7 @@ func loadRegistry(referenceContent []byte) (*Registry, error) { for index, location := range list.Locations { lookupByID[location.ID] = index } - projection, err := promptProjection(list) + projection, err := identityProjection(list) if err != nil { return nil, fmt.Errorf("encode location prompt projection: %w", err) } @@ -141,6 +143,7 @@ func loadRegistry(referenceContent []byte) (*Registry, error) { canonical: append([]byte(nil), content...), digest: digest, projectionDigest: projectionDigest, + identityDigest: projectionDigest, promptInput: contracts.NewLLMInputMaterial(ReferenceSlot, locationcodec.MediaType, projection, projectionDigest, ""), lookupByID: lookupByID, }, nil @@ -191,6 +194,15 @@ func (r *Registry) ProjectionDigest() string { return r.projectionDigest } +// IdentityDigest returns the digest of the ordered ID/name identity projection +// used by deterministic consumers. +func (r *Registry) IdentityDigest() string { + if r == nil { + return "" + } + return r.identityDigest +} + // Count returns the number of validated location records. func (r *Registry) Count() int { if r == nil { @@ -200,7 +212,8 @@ func (r *Registry) Count() int { } // PromptInput returns the ordered ID-and-name projection without evidence or -// reference provenance. +// reference provenance. It remains available only while the location occurrence +// extractor transitions to contextual grounding. func (r *Registry) PromptInput() contracts.LLMInputMaterial { if r == nil { return contracts.LLMInputMaterial{} @@ -231,19 +244,19 @@ func semanticDigest(content []byte) string { return "sha256:" + hex.EncodeToString(sum[:]) } -type projectedLocation struct { +type identityProjectedLocation struct { ID string `json:"id"` Name string `json:"name"` } -type projectedLocationRegistry struct { - Locations []projectedLocation `json:"locations"` +type identityProjectedLocationRegistry struct { + Locations []identityProjectedLocation `json:"locations"` } -func promptProjection(list dnd.LocationRegistry) ([]byte, error) { - projection := projectedLocationRegistry{Locations: make([]projectedLocation, len(list.Locations))} +func identityProjection(list dnd.LocationRegistry) ([]byte, error) { + projection := identityProjectedLocationRegistry{Locations: make([]identityProjectedLocation, len(list.Locations))} for index, location := range list.Locations { - projection.Locations[index] = projectedLocation{ID: location.ID, Name: location.Name} + projection.Locations[index] = identityProjectedLocation{ID: location.ID, Name: location.Name} } return json.Marshal(projection) } diff --git a/internal/modules/dnd/locations/registry/registry_test.go b/internal/modules/dnd/locations/registry/registry_test.go index 3986398..ea185ce 100644 --- a/internal/modules/dnd/locations/registry/registry_test.go +++ b/internal/modules/dnd/locations/registry/registry_test.go @@ -2,6 +2,8 @@ package registry import ( "bytes" + "encoding/json" + "reflect" "strings" "testing" @@ -189,3 +191,180 @@ func item(content []byte) contracts.ReferenceItem { func referenceSet(items ...contracts.ReferenceItem) contracts.ReferenceSet { return contracts.ReferenceSet{Slots: map[string]contracts.ResolvedReferenceSlot{ReferenceSlot: {Items: items}}} } + +func TestGroundingProjectsAndResolvesUniqueAndSameNameLocations(t *testing.T) { + doc := groundingDocument() + sharedName := "The Tavern" + firstRefs := []source.SourceRef{{SourceID: doc.ID, StartUnitID: 20, EndUnitID: 20}, {SourceID: doc.ID, StartUnitID: 30, EndUnitID: 30}} + secondRefs := []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}} + marketRefs := []source.SourceRef{{SourceID: doc.ID, StartUnitID: 40, EndUnitID: 40}} + list := dnd.LocationRegistry{Locations: []dnd.Location{ + location(sharedName, firstRefs), + location(sharedName, secondRefs), + location("Market", marketRefs), + }} + registry := resolveList(t, list) + grounding, err := NewGrounding(registry, doc) + if err != nil { + t.Fatal(err) + } + + projection := decodeGroundingProjection(t, grounding.PromptInput().Content) + if len(projection.Locations) != 3 { + t.Fatalf("projection = %#v", projection) + } + first := projection.Locations[0] + if first.Name != sharedName || !reflect.DeepEqual(first.RegistryRefs, []RegistryRef{{StartUnitID: 30, EndUnitID: 30}, {StartUnitID: 20, EndUnitID: 20}}) || !reflect.DeepEqual(first.Context, []ContextUnit{{UnitID: 30, Text: "Tavern entrance"}, {UnitID: 20, Text: "Tavern cellar"}}) { + t.Fatalf("first contextual projection = %#v", first) + } + second := projection.Locations[1] + if !reflect.DeepEqual(second.RegistryRefs, []RegistryRef{{StartUnitID: 10, EndUnitID: 10}}) || !reflect.DeepEqual(second.Context, []ContextUnit{{UnitID: 10, Text: "Tavern common room"}}) { + t.Fatalf("second contextual projection = %#v", second) + } + market := projection.Locations[2] + if market.Name != "Market" || len(market.RegistryRefs) != 0 || len(market.Context) != 0 { + t.Fatalf("unique-name projection = %#v", market) + } + if bytes.Contains(grounding.PromptInput().Content, []byte("location:sha256:")) || bytes.Contains(grounding.PromptInput().Content, []byte(`"source_id"`)) { + t.Fatalf("grounding projection exposed durable identity: %s", grounding.PromptInput().Content) + } + if grounding.ProjectionDigest() == "" || grounding.ProjectionDigest() == registry.IdentityDigest() || grounding.PromptInput().Digest != grounding.ProjectionDigest() { + t.Fatalf("grounding/identity digests = %q/%q", grounding.ProjectionDigest(), registry.IdentityDigest()) + } + + resolved, ok := grounding.Resolve(Selector{Name: " the tavern ", RegistryRefs: []RegistryRef{{StartUnitID: 30, EndUnitID: 30}, {StartUnitID: 20, EndUnitID: 20}}}) + if !ok || resolved.ID != list.Locations[0].ID || resolved.Name != sharedName { + t.Fatalf("Resolve(same name) = %#v, %t", resolved, ok) + } + if resolved, ok = grounding.Resolve(Selector{Name: "MARKET", RegistryRefs: []RegistryRef{}}); !ok || resolved.ID != list.Locations[2].ID { + t.Fatalf("Resolve(unique name) = %#v, %t", resolved, ok) + } + for _, selector := range []Selector{ + {Name: "Market", RegistryRefs: []RegistryRef{{StartUnitID: 40, EndUnitID: 40}}}, + {Name: sharedName, RegistryRefs: []RegistryRef{}}, + {Name: sharedName, RegistryRefs: []RegistryRef{{StartUnitID: 30, EndUnitID: 30}}}, + {Name: sharedName, RegistryRefs: []RegistryRef{{StartUnitID: 20, EndUnitID: 20}, {StartUnitID: 30, EndUnitID: 30}}}, + {Name: "Unknown", RegistryRefs: []RegistryRef{}}, + } { + if _, ok := grounding.Resolve(selector); ok { + t.Fatalf("Resolve(%#v) accepted an unsupported selector", selector) + } + } + + input := grounding.PromptInput() + input.Content[0] = '[' + resolved.SourceRefs[0].SourceID = "changed" + if grounding.PromptInput().Content[0] != '{' { + t.Fatal("grounding prompt input was mutable") + } + if next, ok := grounding.Resolve(Selector{Name: sharedName, RegistryRefs: []RegistryRef{{StartUnitID: 30, EndUnitID: 30}, {StartUnitID: 20, EndUnitID: 20}}}); !ok || next.SourceRefs[0].SourceID != doc.ID { + t.Fatalf("grounding resolved location was mutable: %#v, %t", next, ok) + } +} + +func TestGroundingHandlesEmptyRegistriesAndIdentityFingerprints(t *testing.T) { + doc := groundingDocument() + empty := resolveList(t, dnd.LocationRegistry{Locations: []dnd.Location{}}) + grounding, err := NewGrounding(empty, doc) + if err != nil { + t.Fatal(err) + } + if projection := decodeGroundingProjection(t, grounding.PromptInput().Content); projection.Locations == nil || len(projection.Locations) != 0 { + t.Fatalf("empty projection = %#v", projection) + } + if _, ok := grounding.Resolve(Selector{Name: "Market", RegistryRefs: []RegistryRef{}}); ok { + t.Fatal("empty grounding resolved a location") + } + + baseRefs := []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}} + base := resolveList(t, dnd.LocationRegistry{Locations: []dnd.Location{location("Market", baseRefs)}}) + expandedRefs := append(append([]source.SourceRef(nil), baseRefs...), source.SourceRef{SourceID: doc.ID, StartUnitID: 40, EndUnitID: 40}) + expanded := resolveList(t, dnd.LocationRegistry{Locations: []dnd.Location{location("Market", expandedRefs)}}) + baseGrounding, err := NewGrounding(base, doc) + if err != nil { + t.Fatal(err) + } + expandedGrounding, err := NewGrounding(expanded, doc) + if err != nil { + t.Fatal(err) + } + if base.IdentityDigest() != expanded.IdentityDigest() || baseGrounding.ProjectionDigest() != expandedGrounding.ProjectionDigest() { + t.Fatalf("identity/model fingerprints = %q/%q and %q/%q", base.IdentityDigest(), expanded.IdentityDigest(), baseGrounding.ProjectionDigest(), expandedGrounding.ProjectionDigest()) + } +} + +func TestGroundingRejectsUnsafeContextualConstruction(t *testing.T) { + doc := groundingDocument() + validRefs := []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}} + valid := resolveList(t, dnd.LocationRegistry{Locations: []dnd.Location{location("Market", validRefs)}}) + if _, err := NewGrounding(valid, nil); err == nil { + t.Fatal("NewGrounding(nil source) error = nil") + } + + for _, test := range []struct { + name string + registry *Registry + }{ + { + name: "foreign reference", + registry: resolveList(t, dnd.LocationRegistry{Locations: []dnd.Location{ + location("The Tavern", []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}}), + location("The Tavern", []source.SourceRef{{SourceID: "other-source", StartUnitID: 10, EndUnitID: 10}}), + }}), + }, + { + name: "invalid reference", + registry: resolveList(t, dnd.LocationRegistry{Locations: []dnd.Location{ + location("The Tavern", []source.SourceRef{{SourceID: doc.ID, StartUnitID: 10, EndUnitID: 10}}), + location("The Tavern", []source.SourceRef{{SourceID: doc.ID, StartUnitID: 999, EndUnitID: 999}}), + }}), + }, + { + name: "selector collision", + registry: &Registry{list: dnd.LocationRegistry{Locations: []dnd.Location{ + {ID: "first", Name: "The Tavern", SourceRefs: validRefs}, + {ID: "second", Name: "The Tavern", SourceRefs: validRefs}, + }}}, + }, + { + name: "empty comparison name", + registry: &Registry{list: dnd.LocationRegistry{Locations: []dnd.Location{{ID: "first", Name: " ", SourceRefs: validRefs}}}}, + }, + } { + t.Run(test.name, func(t *testing.T) { + if _, err := NewGrounding(test.registry, doc); err == nil || strings.Contains(err.Error(), "Tavern entrance") { + t.Fatalf("NewGrounding() error = %v, want content-safe failure", err) + } + }) + } +} + +type decodedGroundingProjection struct { + Locations []struct { + Name string `json:"name"` + RegistryRefs []RegistryRef `json:"registry_refs"` + Context []ContextUnit `json:"context"` + } `json:"locations"` +} + +func decodeGroundingProjection(t *testing.T, content []byte) decodedGroundingProjection { + t.Helper() + var projection decodedGroundingProjection + if err := json.Unmarshal(content, &projection); err != nil { + t.Fatal(err) + } + return projection +} + +func groundingDocument() *source.SourceDocument { + return &source.SourceDocument{ID: "session-alpha", Units: []source.SourceUnit{ + {ID: 30, Text: "Tavern entrance"}, + {ID: 10, Text: "Tavern common room"}, + {ID: 20, Text: "Tavern cellar"}, + {ID: 40, Text: "Market square"}, + }} +} + +func location(name string, refs []source.SourceRef) dnd.Location { + return dnd.Location{ID: identity.DeriveID(name, refs), Name: name, SourceRefs: refs} +} diff --git a/internal/modules/dnd/normalize/locationoccurrences/normalizer.go b/internal/modules/dnd/normalize/locationoccurrences/normalizer.go index 665eefa..e6ffa25 100644 --- a/internal/modules/dnd/normalize/locationoccurrences/normalizer.go +++ b/internal/modules/dnd/normalize/locationoccurrences/normalizer.go @@ -86,7 +86,7 @@ func (n *Normalizer) CheckpointFingerprints() []pipeline.CheckpointFingerprint { } return []pipeline.CheckpointFingerprint{ {Name: "normalization_policy", Value: normalizationPolicy}, - {Name: "location_registry", Value: n.locationResolver.Seeded().ProjectionDigest()}, + {Name: "location_registry", Value: n.locationResolver.Seeded().IdentityDigest()}, } } diff --git a/internal/modules/dnd/validate/locationoccurrences/registry/validator.go b/internal/modules/dnd/validate/locationoccurrences/registry/validator.go index dc4a3ae..1c81f13 100644 --- a/internal/modules/dnd/validate/locationoccurrences/registry/validator.go +++ b/internal/modules/dnd/validate/locationoccurrences/registry/validator.go @@ -61,7 +61,7 @@ func (v *Validator) CheckpointFingerprints() []pipeline.CheckpointFingerprint { if v == nil || v.locationResolver == nil { return nil } - return []pipeline.CheckpointFingerprint{{Name: "policy", Value: policy}, {Name: "location_registry", Value: v.locationResolver.Seeded().ProjectionDigest()}} + return []pipeline.CheckpointFingerprint{{Name: "policy", Value: policy}, {Name: "location_registry", Value: v.locationResolver.Seeded().IdentityDigest()}} } func (v *Validator) Validate(_ context.Context, req contracts.TypedValidationRequest[dnd.LocationOccurrenceList]) (contracts.ValidationResult, error) {