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

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

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