Add contextual location grounding

This commit is contained in:
2026-08-08 14:49:08 +00:00
parent 8e680cf96e
commit fc76805075
5 changed files with 399 additions and 10 deletions

View File

@@ -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...)
}

View File

@@ -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)
}

View File

@@ -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}
}

View File

@@ -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()},
}
}

View File

@@ -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) {