Add contextual location grounding
This commit is contained in:
@@ -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}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user