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

Ground location occurrences with contextual selectors

Browse Source
This commit is contained in:
Eric Rakestraw
2026-08-08 14:56:43 +00:00
parent fc76805075
commit 51d62de1f3
13 changed files with 181 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
internal/modules/dnd/extract/locationoccurrences/canonicalize.go
View File

@@ -1,11 +1,13 @@
package locationoccurrences
import (
"fmt"
"reflect"
"sort"
"gitea.maximumdirect.net/eric/notarius/internal/core/source"
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
locationregistry "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/registry"
"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
)
@@ -15,17 +17,21 @@ type orderedOccurrence struct {
hasEvidence bool
}
func canonicalOccurrenceList(response extractionResponse, order shared.SourceRefOrder, sourceID string) dnd.LocationOccurrenceList {
func canonicalOccurrenceList(response extractionResponse, order shared.SourceRefOrder, sourceID string, grounding *locationregistry.Grounding) (dnd.LocationOccurrenceList, error) {
if response.Occurrences == nil {
return dnd.LocationOccurrenceList{}
return dnd.LocationOccurrenceList{}, nil
}
ordered := make([]orderedOccurrence, len(response.Occurrences))
for index, occurrence := range response.Occurrences {
location, ok := grounding.Resolve(locationregistry.Selector{Name: occurrence.Name, RegistryRefs: occurrence.RegistryRefs})
if !ok {
return dnd.LocationOccurrenceList{}, fmt.Errorf("occurrence %d does not match a supplied location selector", index)
}
refs := order.Canonicalize(canonicalSourceRefs(occurrence.SourceRefs, sourceID))
earliest, hasEvidence := order.EarliestValid(refs)
ordered[index] = orderedOccurrence{value: dnd.LocationOccurrence{
LocationID: occurrence.LocationID,
Name: occurrence.Name,
LocationID: location.ID,
Name: location.Name,
Kind: dnd.LocationOccurrenceKind(occurrence.Kind),
SourceRefs: refs,
}, earliest: earliest, hasEvidence: hasEvidence}
@@ -39,7 +45,7 @@ func canonicalOccurrenceList(response extractionResponse, order shared.SourceRef
occurrences = append(occurrences, occurrence.value)
}
}
return dnd.LocationOccurrenceList{Occurrences: occurrences}
return dnd.LocationOccurrenceList{Occurrences: occurrences}, nil
}
func lessOccurrence(left, right orderedOccurrence, order shared.SourceRefOrder) bool {

16
internal/modules/dnd/extract/locationoccurrences/extractor.go
View File

@@ -15,7 +15,7 @@ import (
const (
Key = "dnd/location-occurrences"
mappingPolicy = "dnd.location_occurrences.extract_mapping.v1"
mappingPolicy = "dnd.location_occurrences.extract_mapping.v2"
)
const (
@@ -121,7 +121,7 @@ func (e *Extractor) CheckpointFingerprints() []pipeline.CheckpointFingerprint {
{Name: "prompt", Value: e.promptSHA},
{Name: "response_schema", Value: e.responseSchemaSHA},
{Name: "mapping_policy", Value: mappingPolicy},
{Name: "location_registry", Value: e.locationResolver.Seeded().ProjectionDigest()},
{Name: "location_registry", Value: e.locationResolver.Seeded().IdentityDigest()},
}
}
@@ -143,17 +143,25 @@ func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRe
if !registry.Bound() {
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("location registry reference is required")
}
grounding, err := locationregistry.NewGrounding(registry, req.Source)
if err != nil {
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("prepare location grounding: %w", err)
}
var response extractionResponse
inputs := shared.PromptInputs(sourceInput, req.References)
inputs[LocationRegistryReferenceSlot] = registry.PromptInput()
inputs[LocationRegistryReferenceSlot] = grounding.PromptInput()
if _, err := e.llm.CompleteStructured(ctx, contracts.StructuredCompletionRequest{
StageName: Key, PromptID: PromptID, PromptVersion: SchemaVersion,
ProfileID: req.LLMProfile, SessionID: req.SessionID, Inputs: inputs,
}, &response); err != nil {
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("complete structured output: %w", err)
}
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{Value: canonicalOccurrenceList(response, shared.NewSourceRefOrder(req.Source), req.Source.ID)}, nil
occurrences, err := canonicalOccurrenceList(response, shared.NewSourceRefOrder(req.Source), req.Source.ID, grounding)
if err != nil {
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{}, extractorErrorf("resolve location grounding: %w", err)
}
return contracts.TypedExtractionResult[dnd.LocationOccurrenceList]{Value: occurrences}, nil
}
func ModuleSpec() pipeline.ModuleSpec {

92
internal/modules/dnd/extract/locationoccurrences/extractor_test.go
View File

@@ -14,19 +14,20 @@ import (
"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"
locationregistry "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/registry"
)
func TestExtractMapsKindsOrdersOccurrencesAndPreservesIndependentFacts(t *testing.T) {
locations := locationRegistry(t, "The Tavern", "The Tavern")
first, second := locations.Locations[0], locations.Locations[1]
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{
{LocationID: second.ID, Name: second.Name, Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30)},
{LocationID: first.ID, Name: first.Name, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)},
{LocationID: first.ID, Name: first.Name, Kind: "recalled", SourceRefs: occurrenceRefs(10, 10)},
{LocationID: first.ID, Name: first.Name, Kind: "planned", SourceRefs: occurrenceRefs(10, 10)},
{LocationID: first.ID, Name: first.Name, Kind: "visited", SourceRefs: append(occurrenceRefs(10, 10), occurrenceRefs(10, 10)...)},
{LocationID: first.ID, Name: first.Name, Kind: "visited", SourceRefs: occurrenceRefs(20, 20)},
{LocationID: first.ID, Name: first.Name, Kind: "visited", SourceRefs: occurrenceRefs(10, 10)},
{Name: second.Name, RegistryRefs: registryRefs(second), Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30)},
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)},
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "recalled", SourceRefs: occurrenceRefs(10, 10)},
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "planned", SourceRefs: occurrenceRefs(10, 10)},
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "visited", SourceRefs: append(occurrenceRefs(10, 10), occurrenceRefs(10, 10)...)},
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "visited", SourceRefs: occurrenceRefs(20, 20)},
{Name: first.Name, RegistryRefs: registryRefs(first), Kind: "visited", SourceRefs: occurrenceRefs(10, 10)},
}}}
references := registryReferences(t, locations)
req := extractionRequest()
@@ -50,11 +51,11 @@ func TestExtractMapsKindsOrdersOccurrencesAndPreservesIndependentFacts(t *testin
}
}
func TestExtractUsesIDsNamesAndCurrentTranscriptEvidenceOnly(t *testing.T) {
func TestExtractResolvesContextualSelectorsAndUsesCurrentTranscriptEvidenceOnly(t *testing.T) {
locations := locationRegistry(t, "The Tavern", "The Tavern")
first, second := locations.Locations[0], locations.Locations[1]
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{
LocationID: second.ID, Name: second.Name, Kind: "visited", SourceRefs: occurrenceRefs(10, 10),
Name: second.Name, RegistryRefs: registryRefs(second), Kind: "visited", SourceRefs: occurrenceRefs(10, 10),
}}}}
references := registryReferences(t, locations)
req := extractionRequest()
@@ -67,10 +68,10 @@ func TestExtractUsesIDsNamesAndCurrentTranscriptEvidenceOnly(t *testing.T) {
t.Fatalf("occurrence = %#v", occurrence)
}
input := client.requests[0].Inputs[LocationRegistryReferenceSlot]
if input.Name != LocationRegistryReferenceSlot || !strings.Contains(string(input.Content), first.ID) || !strings.Contains(string(input.Content), second.ID) {
if input.Name != LocationRegistryReferenceSlot || !strings.Contains(string(input.Content), `"registry_refs":[{"start_unit_id":20,"end_unit_id":20}]`) {
t.Fatalf("location prompt input = %#v", input)
}
for _, forbidden := range []string{"source_refs", "source_id", "other-session"} {
for _, forbidden := range []string{"source_refs", "source_id", first.ID, second.ID} {
if strings.Contains(string(input.Content), forbidden) {
t.Fatalf("location prompt leaked %q: %s", forbidden, input.Content)
}
@@ -84,22 +85,46 @@ func TestExtractUsesIDsNamesAndCurrentTranscriptEvidenceOnly(t *testing.T) {
}
}
func TestExtractPreservesUnknownOrMismatchedGroundingForValidators(t *testing.T) {
func TestExtractRejectsUnknownMalformedOrMismatchedSelectorsAtomically(t *testing.T) {
locations := locationRegistry(t, "The Mill")
known := locations.Locations[0]
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{
{LocationID: "location:sha256:unknown", Name: "The Mill", Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)},
{LocationID: known.ID, Name: "A Different Mill", Kind: "mentioned", SourceRefs: occurrenceRefs(20, 20)},
}}}
references := registryReferences(t, locations)
req := extractionRequest()
req.References = references
result, err := newExtractor(t, client, references).Extract(context.Background(), req)
if err != nil {
t.Fatal(err)
for _, occurrences := range [][]occurrenceResponse{
{{Name: "Unknown", RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}},
{{Name: known.Name, RegistryRefs: []locationregistry.RegistryRef{{StartUnitID: 10, EndUnitID: 0}}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}},
{{Name: known.Name, RegistryRefs: []locationregistry.RegistryRef{{StartUnitID: 10, EndUnitID: 10}}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}},
{{Name: known.Name, RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}, {Name: "Unknown", RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(20, 20)}},
} {
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: occurrences}}
if result, err := newExtractor(t, client, references).Extract(context.Background(), req); err == nil || result.Value.Occurrences != nil || !strings.Contains(err.Error(), "location selector") {
t.Fatalf("Extract() = %#v, %v", result, err)
}
}
if result.Value.Occurrences[0].LocationID != "location:sha256:unknown" || result.Value.Occurrences[1].Name != "A Different Mill" {
t.Fatalf("extractor repaired validator-owned grounding errors: %#v", result.Value.Occurrences)
sharedName := "The Tavern"
firstRefs := []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 10, EndUnitID: 10}, {SourceID: req.Source.ID, StartUnitID: 20, EndUnitID: 20}}
secondRefs := []source.SourceRef{{SourceID: req.Source.ID, StartUnitID: 30, EndUnitID: 30}}
duplicateLocations := dnd.LocationRegistry{Locations: []dnd.Location{
{ID: identity.DeriveID(sharedName, firstRefs), Name: sharedName, SourceRefs: firstRefs},
{ID: identity.DeriveID(sharedName, secondRefs), Name: sharedName, SourceRefs: secondRefs},
}}
duplicateReferences := registryReferences(t, duplicateLocations)
for _, selector := range []struct {
name string
refs []locationregistry.RegistryRef
}{
{name: sharedName, refs: []locationregistry.RegistryRef{}},
{name: sharedName, refs: []locationregistry.RegistryRef{{StartUnitID: 10, EndUnitID: 10}}},
{name: sharedName, refs: []locationregistry.RegistryRef{{StartUnitID: 20, EndUnitID: 20}, {StartUnitID: 10, EndUnitID: 10}}},
} {
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{Name: selector.name, RegistryRefs: selector.refs, Kind: "mentioned", SourceRefs: occurrenceRefs(10, 10)}}}}
duplicateRequest := extractionRequest()
duplicateRequest.References = duplicateReferences
if _, err := newExtractor(t, client, duplicateReferences).Extract(context.Background(), duplicateRequest); err == nil || !strings.Contains(err.Error(), "location selector") {
t.Fatalf("Extract(%#v) error = %v", selector, err)
}
}
}
@@ -125,7 +150,7 @@ func TestExtractResolvesGeneratedRegistryAtOperationTimeAndDoesNotMutateResponse
locations := locationRegistry(t, "The Mill")
location := locations.Locations[0]
client := &fakeOccurrencesLLMClient{response: extractionResponse{Occurrences: []occurrenceResponse{{
LocationID: location.ID, Name: location.Name, Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30),
Name: location.Name, RegistryRefs: []locationregistry.RegistryRef{}, Kind: "mentioned", SourceRefs: occurrenceRefs(30, 30),
}}}}
references := registryReferences(t, locations)
req := extractionRequest()
@@ -135,7 +160,7 @@ func TestExtractResolvesGeneratedRegistryAtOperationTimeAndDoesNotMutateResponse
if err != nil || result.Value.Occurrences[0].Name != "The Mill" {
t.Fatalf("Extract() = %#v, %v", result, err)
}
if input := client.requests[0].Inputs[LocationRegistryReferenceSlot]; !strings.Contains(string(input.Content), location.ID) || input.OriginURI != "" {
if input := client.requests[0].Inputs[LocationRegistryReferenceSlot]; strings.Contains(string(input.Content), location.ID) || input.OriginURI != "" {
t.Fatalf("generated registry prompt input = %#v", input)
}
if _, ok := extractor.ManifestMetadata()["location_registry_digest"]; ok {
@@ -215,21 +240,38 @@ func TestExtractorContractsMetadataAndFailures(t *testing.T) {
t.Fatalf("metadata[%q] = %#v", key, metadata[key])
}
}
if got := newExtractor(t, &fakeOccurrencesLLMClient{}, references).CheckpointFingerprints(); len(got) != 4 || got[3].Name != "location_registry" {
if got := newExtractor(t, &fakeOccurrencesLLMClient{}, references).CheckpointFingerprints(); len(got) != 4 || got[3].Name != "location_registry" || got[3].Value != locationRegistryIdentityDigest(t, references) {
t.Fatalf("fingerprints = %#v", got)
}
}
func locationRegistryIdentityDigest(t *testing.T, references contracts.ReferenceSet) string {
t.Helper()
resolver, err := locationregistry.NewResolver(references)
if err != nil {
t.Fatal(err)
}
return resolver.Seeded().IdentityDigest()
}
func locationRegistry(t *testing.T, names ...string) dnd.LocationRegistry {
t.Helper()
locations := make([]dnd.Location, len(names))
for index, name := range names {
refs := []source.SourceRef{{SourceID: "other-session", StartUnitID: index + 1, EndUnitID: index + 1}}
refs := []source.SourceRef{{SourceID: sourceDocument().ID, StartUnitID: (index + 1) * 10, EndUnitID: (index + 1) * 10}}
locations[index] = dnd.Location{ID: identity.DeriveID(name, refs), Name: name, SourceRefs: refs}
}
return dnd.LocationRegistry{Locations: locations}
}
func registryRefs(location dnd.Location) []locationregistry.RegistryRef {
refs := make([]locationregistry.RegistryRef, len(location.SourceRefs))
for index, ref := range location.SourceRefs {
refs[index] = locationregistry.RegistryRef{StartUnitID: ref.StartUnitID, EndUnitID: ref.EndUnitID}
}
return refs
}
func registryReferences(t *testing.T, locations dnd.LocationRegistry) contracts.ReferenceSet {
t.Helper()
content, err := locationcodec.New().Encode(locations)

10
internal/modules/dnd/extract/locationoccurrences/model.go
View File

@@ -1,14 +1,16 @@
package locationoccurrences
import locationregistry "gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/locations/registry"
type extractionResponse struct {
Occurrences []occurrenceResponse `json:"occurrences"`
}
type occurrenceResponse struct {
LocationID string `json:"location_id"`
Name string `json:"name"`
Kind string `json:"kind"`
SourceRefs []occurrenceSourceRefResponse `json:"source_refs"`
Name string `json:"name"`
RegistryRefs []locationregistry.RegistryRef `json:"registry_refs"`
Kind string `json:"kind"`
SourceRefs []occurrenceSourceRefResponse `json:"source_refs"`
}
type occurrenceSourceRefResponse struct {

8
internal/modules/dnd/extract/locationoccurrences/prompt_assets_test.go
View File

@@ -30,7 +30,7 @@ func TestRegisterPromptAssetsPreparesLocationOccurrencePrompt(t *testing.T) {
PromptID: PromptID, PromptVersion: SchemaVersion, ProfileID: "location-occurrences-test",
Inputs: map[string]promptkit.ArtifactRef{
"transcript": promptkit.Inline(`{"units":[{"sentinel":"location-occurrence-transcript"}]}`), "players": promptkit.Inline("location-occurrence-player"), "party": promptkit.Inline(" "), "glossary": promptkit.Inline(" "),
"location_registry": promptkit.Inline(`{"locations":[{"id":"location:sha256:test","name":"location-occurrence-registry"}]}`),
"location_registry": promptkit.Inline(`{"locations":[{"name":"location-occurrence-registry","registry_refs":[{"start_unit_id":7,"end_unit_id":7}],"context":[{"unit_id":7,"text":"registry context"}]}]}`),
},
})
if err != nil {
@@ -41,11 +41,11 @@ func TestRegisterPromptAssetsPreparesLocationOccurrencePrompt(t *testing.T) {
}
var registryMessage string
for _, message := range prepared.Messages {
if strings.Contains(message.Content, "normalized location registry") {
if strings.Contains(message.Content, "contextual location registry") {
registryMessage = message.Content
}
}
if !strings.Contains(registryMessage, "location:sha256:test") || !strings.Contains(registryMessage, "location-occurrence-registry") || strings.Contains(registryMessage, "source_refs") {
if !strings.Contains(registryMessage, "location-occurrence-registry") || !strings.Contains(registryMessage, `"registry_refs":[{"start_unit_id":7,"end_unit_id":7}]`) || strings.Contains(registryMessage, "source_id") || strings.Contains(registryMessage, "location:sha256") {
t.Fatalf("rendered prompt did not preserve source-free registry grounding: %s", registryMessage)
}
content := make([]string, len(prepared.Messages))
@@ -54,7 +54,7 @@ func TestRegisterPromptAssetsPreparesLocationOccurrencePrompt(t *testing.T) {
}
rendered := strings.Join(content, "\n")
policy := strings.ReplaceAll(rendered, "\n", " ")
if !strings.Contains(policy, "context supports that coreference") || !strings.Contains(policy, "must not create a registry location") || !strings.Contains(policy, "provenance must never replace current-chunk evidence") {
if !strings.Contains(policy, "context supports that coreference") || !strings.Contains(policy, "must not create a registry location") || !strings.Contains(policy, "provenance must never replace current-chunk evidence") || !strings.Contains(policy, "complete, ordered") {
t.Fatalf("rendered prompt = %q, want contextual coreference without registry-derived evidence", rendered)
}
}

5
internal/modules/dnd/extract/locationoccurrences/schema_test.go
View File

@@ -18,16 +18,17 @@ func TestResponseSchemaRestrictsPrivateOccurrenceStructureAndKinds(t *testing.T)
t.Fatalf("schema = %#v", schema)
}
valid := map[string]any{"occurrences": []any{map[string]any{
"location_id": "location:sha256:test", "name": "The Mill", "kind": "visited",
"name": "The Mill", "registry_refs": []any{map[string]any{"start_unit_id": 1, "end_unit_id": 2}}, "kind": "visited",
"source_refs": []any{map[string]any{"start_unit_id": 1, "end_unit_id": 2}},
}}}
if err := validateSchema(t, valid, schema.JSONSchema); err != nil {
t.Fatalf("valid response rejected: %v", err)
}
for _, mutate := range []func(map[string]any){
func(value map[string]any) { delete(value, "location_id") },
func(value map[string]any) { delete(value, "registry_refs") },
func(value map[string]any) { value["kind"] = "other" },
func(value map[string]any) { value["unexpected"] = true },
func(value map[string]any) { value["registry_refs"].([]any)[0].(map[string]any)["start_unit_id"] = 0 },
func(value map[string]any) {
value["source_refs"].([]any)[0].(map[string]any)["source_id"] = "assigned later"
},