Simplify contextual entity grounding

internal/modules/dnd/extract/itemoccurrences/canonicalize.go

@@ -10,26 +10,34 @@ import (
 	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/shared"
 )
 
-type orderedItemOccurrenceResponse struct {
-	value       itemOccurrenceResponse
+type orderedItemOccurrence struct {
+	value       dnd.ItemOccurrence
 	earliest    int
 	hasEvidence bool
 }
 
-func canonicalizeResponse(response *extractionResponse, order shared.SourceRefOrder, sourceID string, registry *itemregistry.Registry) error {
-	if response == nil {
-		return nil
+func canonicalItemOccurrenceList(response extractionResponse, order shared.SourceRefOrder, sourceID string, registry *itemregistry.Registry) (dnd.ItemOccurrenceList, error) {
+	if response.Occurrences == nil {
+		return dnd.ItemOccurrenceList{}, nil
 	}
-	ordered := make([]orderedItemOccurrenceResponse, len(response.Occurrences))
-	for index := range response.Occurrences {
-		canonical, found := registry.Lookup(response.Occurrences[index].Name)
+	ordered := make([]orderedItemOccurrence, len(response.Occurrences))
+	for index, occurrence := range response.Occurrences {
+		item, found := registry.Lookup(occurrence.Name)
 		if !found {
-			return fmt.Errorf("occurrences[%d].name is not in the item registry", index)
+			return dnd.ItemOccurrenceList{}, fmt.Errorf("occurrences[%d].name is not in the item registry", index)
 		}
-		response.Occurrences[index].Name = canonical.Name
-		earliest, hasEvidence := canonicalizeItemOccurrence(&response.Occurrences[index], order, sourceID)
-		ordered[index] = orderedItemOccurrenceResponse{
-			value:       response.Occurrences[index],
+		refs := order.Canonicalize(itemOccurrenceSourceRefs(occurrence.SourceRefs, sourceID))
+		earliest, hasEvidence := order.EarliestValid(refs)
+		ordered[index] = orderedItemOccurrence{
+			value: dnd.ItemOccurrence{
+				ItemID:     item.ID,
+				Name:       item.Name,
+				Kind:       dnd.ItemOccurrenceKind(occurrence.Kind),
+				Quantity:   cloneQuantity(occurrence.Quantity),
+				From:       occurrence.From,
+				To:         occurrence.To,
+				SourceRefs: refs,
+			},
 			earliest:    earliest,
 			hasEvidence: hasEvidence,
 		}
@@ -43,40 +51,9 @@ func canonicalizeResponse(response *extractionResponse, order shared.SourceRefOr
 		}
 		return ordered[left].earliest < ordered[right].earliest
 	})
-	for index := range ordered {
-		response.Occurrences[index] = ordered[index].value
-	}
-	return nil
-}
-
-func canonicalizeItemOccurrence(occurrence *itemOccurrenceResponse, order shared.SourceRefOrder, sourceID string) (int, bool) {
-	if occurrence == nil {
-		return 0, false
-	}
-	refs := order.Canonicalize(itemOccurrenceSourceRefs(occurrence.SourceRefs, sourceID))
-	occurrence.SourceRefs = itemOccurrenceResponseRefs(refs)
-	return order.EarliestValid(refs)
-}
-
-func canonicalItemOccurrenceList(response extractionResponse, sourceID string, registry *itemregistry.Registry) (dnd.ItemOccurrenceList, error) {
-	if response.Occurrences == nil {
-		return dnd.ItemOccurrenceList{}, nil
-	}
-	occurrences := make([]dnd.ItemOccurrence, 0, len(response.Occurrences))
-	for index, occurrence := range response.Occurrences {
-		item, found := registry.Lookup(occurrence.Name)
-		if !found {
-			return dnd.ItemOccurrenceList{}, fmt.Errorf("occurrences[%d].name is not in the item registry", index)
-		}
-		occurrences = append(occurrences, dnd.ItemOccurrence{
-			ItemID:     item.ID,
-			Name:       item.Name,
-			Kind:       dnd.ItemOccurrenceKind(occurrence.Kind),
-			Quantity:   cloneQuantity(occurrence.Quantity),
-			From:       occurrence.From,
-			To:         occurrence.To,
-			SourceRefs: itemOccurrenceSourceRefs(occurrence.SourceRefs, sourceID),
-		})
+	occurrences := make([]dnd.ItemOccurrence, len(ordered))
+	for index, occurrence := range ordered {
+		occurrences[index] = occurrence.value
 	}
 	return dnd.ItemOccurrenceList{Occurrences: occurrences}, nil
 }
@@ -92,17 +69,6 @@ func itemOccurrenceSourceRefs(refs []itemOccurrenceSourceRefResponse, sourceID s
 	return values
 }
 
-func itemOccurrenceResponseRefs(refs []source.SourceRef) []itemOccurrenceSourceRefResponse {
-	if refs == nil {
-		return nil
-	}
-	values := make([]itemOccurrenceSourceRefResponse, len(refs))
-	for index, ref := range refs {
-		values[index] = itemOccurrenceSourceRefResponse{StartSegment: ref.StartUnitID, EndSegment: ref.EndUnitID}
-	}
-	return values
-}
-
 func cloneQuantity(value *int) *int {
 	if value == nil {
 		return nil

internal/modules/dnd/extract/itemoccurrences/extractor.go

@@ -161,12 +161,9 @@ func (e *Extractor) Extract(ctx context.Context, req contracts.TypedExtractionRe
 	}, &response); err != nil {
 		return contracts.TypedExtractionResult[dnd.ItemOccurrenceList]{}, extractorErrorf("complete structured output: %w", err)
 	}
-	if err := canonicalizeResponse(&response, order, req.Source.ID, registry); err != nil {
-		return contracts.TypedExtractionResult[dnd.ItemOccurrenceList]{}, extractorErrorf("map item occurrence response: %w", err)
-	}
-	value, err := canonicalItemOccurrenceList(response, req.Source.ID, registry)
+	value, err := canonicalItemOccurrenceList(response, order, req.Source.ID, registry)
 	if err != nil {
-		return contracts.TypedExtractionResult[dnd.ItemOccurrenceList]{}, extractorErrorf("resolve canonical item occurrence names: %w", err)
+		return contracts.TypedExtractionResult[dnd.ItemOccurrenceList]{}, extractorErrorf("map item occurrence response: %w", err)
 	}
 	return contracts.TypedExtractionResult[dnd.ItemOccurrenceList]{Value: value}, nil
 }

internal/modules/dnd/locations/registry/grounding.go

@@ -36,7 +36,6 @@ type ContextUnit struct {
 // Its prompt input contains no durable IDs or source IDs.
 type Grounding struct {
 	promptInput         contracts.LLMInputMaterial
-	projectionDigest    string
 	locationsBySelector map[string]dnd.Location
 }
 
@@ -104,7 +103,6 @@ func NewGrounding(registry *Registry, doc *source.SourceDocument) (*Grounding, e
 	digest := semanticDigest(content)
 	return &Grounding{
 		promptInput:         contracts.NewLLMInputMaterial(ReferenceSlot, locationcodec.MediaType, content, digest, ""),
-		projectionDigest:    digest,
 		locationsBySelector: locationsBySelector,
 	}, nil
 }
@@ -118,14 +116,6 @@ func (g *Grounding) PromptInput() 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 {

internal/modules/dnd/locations/registry/registry_test.go

@@ -216,8 +216,8 @@ func TestGroundingProjectsAndResolvesUniqueAndSameNameLocations(t *testing.T) {
 	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())
+	if grounding.PromptInput().Digest == "" || grounding.PromptInput().Digest == registry.IdentityDigest() {
+		t.Fatalf("grounding/identity digests = %q/%q", grounding.PromptInput().Digest, registry.IdentityDigest())
 	}
 
 	resolved, ok := grounding.Resolve(Selector{Name: " the tavern ", RegistryRefs: []RegistryRef{{StartUnitID: 30, EndUnitID: 30}, {StartUnitID: 20, EndUnitID: 20}}})
@@ -276,8 +276,8 @@ func TestGroundingHandlesEmptyRegistriesAndIdentityFingerprints(t *testing.T) {
 	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())
+	if base.IdentityDigest() != expanded.IdentityDigest() || baseGrounding.PromptInput().Digest != expandedGrounding.PromptInput().Digest {
+		t.Fatalf("identity/model fingerprints = %q/%q and %q/%q", base.IdentityDigest(), expanded.IdentityDigest(), baseGrounding.PromptInput().Digest, expandedGrounding.PromptInput().Digest)
 	}
 }