notarius/internal/modules/dnd/normalize/npcs/context_material.go

package npcs

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"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd"
	"gitea.maximumdirect.net/eric/notarius/internal/modules/dnd/npcs/identity"
)

const semanticContextRadius = 2

type normalizeContextMaterials struct {
	Candidates         contracts.LLMInputMaterial
	Transcript         contracts.LLMInputMaterial
	candidatePositions []int
}

type normalizeCandidateInput struct {
	NPCs []normalizeCandidate `json:"npcs"`
}

type normalizeCandidate struct {
	Name       string                        `json:"name"`
	SourceRefs []normalizeCandidateSourceRef `json:"source_refs"`
}

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

type normalizeTranscriptInput struct {
	Windows []normalizeTranscriptWindow `json:"windows"`
}

type normalizeTranscriptWindow struct {
	Units []normalizeTranscriptUnit `json:"units"`
}

type normalizeTranscriptUnit 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 normalizeProposalResponse struct {
	DuplicateGroups []normalizeProposalGroup `json:"duplicate_groups"`
}

type normalizeProposalGroup struct {
	Members       []string `json:"members"`
	CanonicalName string   `json:"canonical_name"`
}

type sourceInterval struct {
	start int
	end   int
}

func buildDefaultNormalizeContextMaterials(doc *source.SourceDocument, records []dnd.NPC) (normalizeContextMaterials, bool, error) {
	return buildNormalizeContextMaterials(doc, records, semanticContextRadius)
}

// buildNormalizeContextMaterials prepares the owned prompt inputs for a
// document-level normalization proposal. A false ready value means semantic
// normalization has no comparison-distinct eligible candidates to consider.
func buildNormalizeContextMaterials(doc *source.SourceDocument, records []dnd.NPC, radius int) (materials normalizeContextMaterials, ready bool, err error) {
	if doc == nil {
		return normalizeContextMaterials{}, false, nil
	}
	if radius < 0 {
		return normalizeContextMaterials{}, false, fmt.Errorf("build NPC normalization context: radius must not be negative")
	}

	index := source.NewDocumentIndex(doc)
	candidates := make([]normalizeCandidate, 0, len(records))
	candidatePositions := make([]int, 0, len(records))
	intervals := make([]sourceInterval, 0)
	cited := make([]bool, len(doc.Units))
	seenKeys := make(map[string]struct{}, len(records))

	for position, record := range records {
		key := identity.ComparisonKey(record.Name)
		if key == "" || len(record.SourceRefs) == 0 {
			continue
		}
		if _, exists := seenKeys[key]; exists {
			continue
		}

		references, recordIntervals, valid := normalizeRecordReferences(index, record.SourceRefs)
		if !valid {
			continue
		}
		seenKeys[key] = struct{}{}
		candidates = append(candidates, normalizeCandidate{Name: record.Name, SourceRefs: references})
		candidatePositions = append(candidatePositions, position)
		for _, interval := range recordIntervals {
			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(candidates) < 2 {
		return normalizeContextMaterials{}, false, nil
	}

	windows, err := normalizeContextWindows(doc.Units, coalesceIntervals(intervals), cited)
	if err != nil {
		return normalizeContextMaterials{}, false, fmt.Errorf("build NPC normalization context: invalid source metadata")
	}
	candidateContent, err := json.Marshal(normalizeCandidateInput{NPCs: candidates})
	if err != nil {
		return normalizeContextMaterials{}, false, fmt.Errorf("build NPC normalization context: invalid candidate material")
	}
	transcriptContent, err := json.Marshal(normalizeTranscriptInput{Windows: windows})
	if err != nil {
		return normalizeContextMaterials{}, false, fmt.Errorf("build NPC normalization context: invalid transcript material")
	}
	return normalizeContextMaterials{
		Candidates:         newNormalizeInputMaterial("candidates", candidateContent),
		Transcript:         newNormalizeInputMaterial("transcript", transcriptContent),
		candidatePositions: candidatePositions,
	}, true, nil
}

func normalizeRecordReferences(index source.DocumentIndex, refs []source.SourceRef) ([]normalizeCandidateSourceRef, []sourceInterval, bool) {
	references := make([]normalizeCandidateSourceRef, 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, normalizeCandidateSourceRef{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(i, j int) bool {
		if ordered[i].start != ordered[j].start {
			return ordered[i].start < ordered[j].start
		}
		return ordered[i].end < ordered[j].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 normalizeContextWindows(units []source.SourceUnit, intervals []sourceInterval, cited []bool) ([]normalizeTranscriptWindow, error) {
	windows := make([]normalizeTranscriptWindow, 0, len(intervals))
	for _, interval := range intervals {
		window := normalizeTranscriptWindow{Units: make([]normalizeTranscriptUnit, 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, normalizeTranscriptUnit{
				ID: unit.ID, Kind: unit.Kind, Text: unit.Text, Metadata: metadata, Cited: cited[position],
			})
		}
		windows = append(windows, window)
	}
	return windows, nil
}

func newNormalizeInputMaterial(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
}