Files
notarius/internal/framework/semanticreconcile/preparation.go

339 lines
10 KiB
Go

package semanticreconcile
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 (
candidateInputName = "candidates"
transcriptInputName = "transcript"
jsonMediaType = "application/json"
)
var defaultLimits = Limits{
ContextRadius: 2,
MaximumCandidates: 128,
MaximumMaterialBytes: 262144,
}
// Candidate is contextual source-backed input supplied by a typed consumer.
// Prepare does not retain or mutate Label or SourceRefs.
type Candidate struct {
Label string
SourceRefs []source.SourceRef
}
// Limits bounds source context and serialized model input.
type Limits struct {
ContextRadius int
MaximumCandidates int
MaximumMaterialBytes int
}
// DefaultLimits returns the core-owned production limits.
func DefaultLimits() Limits {
return defaultLimits
}
// Validate rejects limits that cannot safely bound preparation.
func (limits Limits) Validate() error {
if limits.ContextRadius < 0 {
return fmt.Errorf("semantic reconciliation limits: context radius must not be negative")
}
if limits.MaximumCandidates <= 0 {
return fmt.Errorf("semantic reconciliation limits: maximum candidates must be positive")
}
if limits.MaximumMaterialBytes <= 0 {
return fmt.Errorf("semantic reconciliation limits: maximum bytes must be positive")
}
return nil
}
// Disposition describes whether prepared materials may be sent to a model.
type Disposition uint8
const (
// Ready indicates that the result contains complete bounded materials.
Ready Disposition = iota + 1
// InsufficientCandidates indicates that fewer than two candidates were
// eligible after source-reference validation.
InsufficientCandidates
// LimitExceeded indicates that a candidate or serialized-material bound was
// exceeded and no request should be split or sent.
LimitExceeded
)
// CandidateMapping relates one model-visible request-local ID to the
// corresponding zero-based position in the caller's candidate slice.
type CandidateMapping struct {
CandidateID int
CandidatePosition int
}
// Preparation owns the visible candidate mapping and prompt materials.
type Preparation struct {
disposition Disposition
mappings []CandidateMapping
materials contracts.LLMInputSet
}
// Disposition returns the preparation outcome.
func (preparation Preparation) Disposition() Disposition {
return preparation.disposition
}
// CandidateMappings returns an owned copy in model-visible candidate order.
func (preparation Preparation) CandidateMappings() []CandidateMapping {
return append([]CandidateMapping(nil), preparation.mappings...)
}
// Materials returns independently owned candidate and transcript materials.
// It is empty unless Disposition returns Ready.
func (preparation Preparation) Materials() contracts.LLMInputSet {
return preparation.materials.Clone()
}
type sourceRange struct {
StartUnitID int `json:"start_unit_id"`
EndUnitID int `json:"end_unit_id"`
}
type visibleCandidate struct {
CandidateID int `json:"candidate_id"`
Label string `json:"label"`
SourceRefs []sourceRange `json:"source_refs"`
}
type candidateInput struct {
Candidates []visibleCandidate `json:"candidates"`
}
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
}
type preparedCandidate struct {
position int
references []sourceRange
intervals []sourceInterval
}
// Prepare validates candidates and constructs bounded, source-ordered model
// inputs. Deterministic skip conditions are represented by the returned
// disposition rather than an error.
func Prepare(document *source.SourceDocument, candidates []Candidate, limits Limits) (Preparation, error) {
if err := limits.Validate(); err != nil {
return Preparation{}, err
}
documentIndex := source.NewDocumentIndex(document)
prepared := make([]preparedCandidate, 0, len(candidates))
for candidatePosition, candidate := range candidates {
references, intervals, valid := prepareReferences(documentIndex, candidate.SourceRefs)
if !valid {
continue
}
prepared = append(prepared, preparedCandidate{
position: candidatePosition,
references: references,
intervals: intervals,
})
}
result := Preparation{
disposition: InsufficientCandidates,
mappings: make([]CandidateMapping, len(prepared)),
}
views := make([]visibleCandidate, len(prepared))
for index, candidate := range prepared {
candidateID := index + 1
result.mappings[index] = CandidateMapping{
CandidateID: candidateID,
CandidatePosition: candidate.position,
}
views[index] = visibleCandidate{
CandidateID: candidateID,
Label: candidates[candidate.position].Label,
SourceRefs: cloneSourceRanges(candidate.references),
}
}
if len(prepared) < 2 {
return result, nil
}
if len(prepared) > limits.MaximumCandidates {
result.disposition = LimitExceeded
return result, nil
}
candidateContent, err := json.Marshal(candidateInput{Candidates: views})
if err != nil {
return Preparation{}, fmt.Errorf("prepare semantic reconciliation: encode candidate material: %w", err)
}
if len(candidateContent) > limits.MaximumMaterialBytes {
result.disposition = LimitExceeded
return result, nil
}
intervals := make([]sourceInterval, 0)
cited := make([]bool, len(document.Units))
for _, candidate := range prepared {
for _, interval := range candidate.intervals {
for position := interval.start; position <= interval.end; position++ {
cited[position] = true
}
intervals = append(intervals, sourceInterval{
start: max(0, interval.start-limits.ContextRadius),
end: min(len(document.Units)-1, interval.end+limits.ContextRadius),
})
}
}
windows, err := buildContextWindows(document.Units, coalesceIntervals(intervals), cited)
if err != nil {
return Preparation{}, fmt.Errorf("prepare semantic reconciliation: build transcript material: %w", err)
}
transcriptContent, err := json.Marshal(transcriptInput{Windows: windows})
if err != nil {
return Preparation{}, fmt.Errorf("prepare semantic reconciliation: encode transcript material: %w", err)
}
if len(transcriptContent) > limits.MaximumMaterialBytes-len(candidateContent) {
result.disposition = LimitExceeded
return result, nil
}
result.disposition = Ready
result.materials = contracts.LLMInputSet{
candidateInputName: newInputMaterial(candidateInputName, candidateContent),
transcriptInputName: newInputMaterial(transcriptInputName, transcriptContent),
}
return result, nil
}
func prepareReferences(index source.DocumentIndex, references []source.SourceRef) ([]sourceRange, []sourceInterval, bool) {
if len(references) == 0 {
return nil, nil, false
}
type referencedInterval struct {
reference sourceRange
interval sourceInterval
}
prepared := make([]referencedInterval, 0, len(references))
for _, reference := range references {
if err := index.ValidateRef(reference); err != nil {
return nil, nil, false
}
start, _ := index.Position(reference.StartUnitID)
end, _ := index.Position(reference.EndUnitID)
prepared = append(prepared, referencedInterval{
reference: sourceRange{StartUnitID: reference.StartUnitID, EndUnitID: reference.EndUnitID},
interval: sourceInterval{start: start, end: end},
})
}
sort.Slice(prepared, func(left, right int) bool {
if prepared[left].interval.start != prepared[right].interval.start {
return prepared[left].interval.start < prepared[right].interval.start
}
return prepared[left].interval.end < prepared[right].interval.end
})
canonicalReferences := make([]sourceRange, 0, len(prepared))
intervals := make([]sourceInterval, 0, len(prepared))
for _, item := range prepared {
if len(canonicalReferences) > 0 && canonicalReferences[len(canonicalReferences)-1] == item.reference {
continue
}
canonicalReferences = append(canonicalReferences, item.reference)
intervals = append(intervals, item.interval)
}
return canonicalReferences, intervals, true
}
func cloneSourceRanges(ranges []sourceRange) []sourceRange {
if len(ranges) == 0 {
return []sourceRange{}
}
return append([]sourceRange(nil), ranges...)
}
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 buildContextWindows(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,
jsonMediaType,
content,
"sha256:"+hex.EncodeToString(digest[:]),
"",
)
}