notarius/internal/modules/dnd/shared/entityreconcile/context.go

// Package entityreconcile provides safe, D&D-specific duplicate proposal
// materials shared by entity normalizers.
package entityreconcile

import (
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"sort"

	"gitea.maximumdirect.net/eric/notarius/internal/core/source"
	"gitea.maximumdirect.net/eric/notarius/internal/framework/contracts"
)

const candidateKeyFormat = "candidate-%06d"

// Candidate is one domain-neutral entity candidate supplied by a normalizer.
// BuildContext never retains or mutates its source references.
type Candidate struct {
	Name       string
	SourceRefs []source.SourceRef
}

// Materials contains owned prompt inputs and opaque candidate-key mappings.
type Materials struct {
	Candidates contracts.LLMInputMaterial
	Transcript contracts.LLMInputMaterial

	candidateKeys []string
	eligible      map[string]struct{}
}

// CandidateKeys returns all deterministic keys in candidate input order.
func (m Materials) CandidateKeys() []string {
	return append([]string(nil), m.candidateKeys...)
}

// EligibleCandidateKeys returns only candidates whose evidence safely produced
// transcript context, preserving candidate input order.
func (m Materials) EligibleCandidateKeys() []string {
	keys := make([]string, 0, len(m.eligible))
	for _, key := range m.candidateKeys {
		if _, ok := m.eligible[key]; ok {
			keys = append(keys, key)
		}
	}
	return keys
}

type candidateInput struct {
	Candidates []candidateView `json:"candidates"`
}

type candidateView struct {
	Key        string               `json:"key"`
	Name       string               `json:"name"`
	SourceRefs []candidateSourceRef `json:"source_refs"`
}

type candidateSourceRef struct {
	StartUnitID int `json:"start_unit_id"`
	EndUnitID   int `json:"end_unit_id"`
}

type transcriptInput struct {
	Windows []transcriptWindow `json:"windows"`
}

type transcriptWindow struct {
	Units []transcriptUnit `json:"units"`
}

type transcriptUnit struct {
	ID       int            `json:"id"`
	Kind     string         `json:"kind"`
	Text     string         `json:"text"`
	Metadata map[string]any `json:"metadata,omitempty"`
	Cited    bool           `json:"cited"`
}

type sourceInterval struct {
	start int
	end   int
}

// BuildContext constructs bounded, source-ordered prompt inputs. It returns
// ready=false when fewer than two candidates have safe context.
func BuildContext(doc *source.SourceDocument, candidates []Candidate, radius int) (Materials, bool, error) {
	if radius < 0 {
		return Materials{}, false, fmt.Errorf("build entity reconciliation context: radius must not be negative")
	}
	materials := Materials{
		candidateKeys: make([]string, len(candidates)),
		eligible:      make(map[string]struct{}),
	}
	for index := range candidates {
		key := fmt.Sprintf(candidateKeyFormat, index+1)
		materials.candidateKeys[index] = key
	}
	if doc == nil {
		return materials, false, nil
	}

	index := source.NewDocumentIndex(doc)
	views := make([]candidateView, 0, len(candidates))
	intervals := make([]sourceInterval, 0)
	cited := make([]bool, len(doc.Units))
	for candidateIndex, candidate := range candidates {
		references, candidateIntervals, valid := candidateReferences(index, candidate.SourceRefs)
		if !valid {
			continue
		}
		key := materials.candidateKeys[candidateIndex]
		materials.eligible[key] = struct{}{}
		views = append(views, candidateView{Key: key, Name: candidate.Name, SourceRefs: references})
		for _, interval := range candidateIntervals {
			for position := interval.start; position <= interval.end; position++ {
				cited[position] = true
			}
			intervals = append(intervals, sourceInterval{
				start: maxInt(0, interval.start-radius),
				end:   minInt(len(doc.Units)-1, interval.end+radius),
			})
		}
	}
	if len(views) < 2 {
		return materials, false, nil
	}

	windows, err := contextWindows(doc.Units, coalesceIntervals(intervals), cited)
	if err != nil {
		return Materials{}, false, fmt.Errorf("build entity reconciliation context: invalid source metadata")
	}
	candidateContent, err := json.Marshal(candidateInput{Candidates: views})
	if err != nil {
		return Materials{}, false, fmt.Errorf("build entity reconciliation context: invalid candidate material")
	}
	transcriptContent, err := json.Marshal(transcriptInput{Windows: windows})
	if err != nil {
		return Materials{}, false, fmt.Errorf("build entity reconciliation context: invalid transcript material")
	}
	materials.Candidates = newInputMaterial("candidates", candidateContent)
	materials.Transcript = newInputMaterial("transcript", transcriptContent)
	return materials, true, nil
}

func candidateReferences(index source.DocumentIndex, refs []source.SourceRef) ([]candidateSourceRef, []sourceInterval, bool) {
	if len(refs) == 0 {
		return nil, nil, false
	}
	references := make([]candidateSourceRef, 0, len(refs))
	intervals := make([]sourceInterval, 0, len(refs))
	for _, ref := range refs {
		if err := index.ValidateRef(ref); err != nil {
			return nil, nil, false
		}
		start, _ := index.Position(ref.StartUnitID)
		end, _ := index.Position(ref.EndUnitID)
		references = append(references, candidateSourceRef{StartUnitID: ref.StartUnitID, EndUnitID: ref.EndUnitID})
		intervals = append(intervals, sourceInterval{start: start, end: end})
	}
	return references, intervals, true
}

func coalesceIntervals(intervals []sourceInterval) []sourceInterval {
	if len(intervals) == 0 {
		return nil
	}
	ordered := append([]sourceInterval(nil), intervals...)
	sort.Slice(ordered, func(left, right int) bool {
		if ordered[left].start != ordered[right].start {
			return ordered[left].start < ordered[right].start
		}
		return ordered[left].end < ordered[right].end
	})
	coalesced := make([]sourceInterval, 0, len(ordered))
	for _, interval := range ordered {
		if len(coalesced) == 0 || interval.start > coalesced[len(coalesced)-1].end+1 {
			coalesced = append(coalesced, interval)
			continue
		}
		if interval.end > coalesced[len(coalesced)-1].end {
			coalesced[len(coalesced)-1].end = interval.end
		}
	}
	return coalesced
}

func contextWindows(units []source.SourceUnit, intervals []sourceInterval, cited []bool) ([]transcriptWindow, error) {
	windows := make([]transcriptWindow, 0, len(intervals))
	for _, interval := range intervals {
		window := transcriptWindow{Units: make([]transcriptUnit, 0, interval.end-interval.start+1)}
		for position := interval.start; position <= interval.end; position++ {
			unit := units[position]
			metadata, err := source.CloneMetadata(unit.Metadata)
			if err != nil {
				return nil, err
			}
			window.Units = append(window.Units, transcriptUnit{
				ID: unit.ID, Kind: unit.Kind, Text: unit.Text, Metadata: metadata, Cited: cited[position],
			})
		}
		windows = append(windows, window)
	}
	return windows, nil
}

func newInputMaterial(name string, content []byte) contracts.LLMInputMaterial {
	digest := sha256.Sum256(content)
	return contracts.NewLLMInputMaterial(name, "application/json", content, "sha256:"+hex.EncodeToString(digest[:]), "")
}

func minInt(left, right int) int {
	if left < right {
		return left
	}
	return right
}

func maxInt(left, right int) int {
	if left > right {
		return left
	}
	return right
}