91 lines
3.1 KiB
Go
91 lines
3.1 KiB
Go
package validators
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"gitea.maximumdirect.net/eric/audita/v3/internal/core/chunking"
|
|
)
|
|
|
|
// Compact JSON fragments in this budget begin and end with non-whitespace.
|
|
// When fragments join, their touching word runs merge. Counting only extra
|
|
// runs inside each fragment makes costs additive, with one run for the whole
|
|
// payload and punctuation rounded only after assembling the array.
|
|
type validationCost struct {
|
|
extraRuns int
|
|
punctuation int
|
|
}
|
|
|
|
func (c validationCost) add(other validationCost) validationCost {
|
|
return validationCost{c.extraRuns + other.extraRuns, c.punctuation + other.punctuation}
|
|
}
|
|
|
|
func (c validationCost) arrayTokens() int {
|
|
return 1 + c.extraRuns + (c.punctuation+2)/2 // opening and closing brackets
|
|
}
|
|
|
|
func measureJSON(estimator *chunking.SimpleTokenEstimator, fragment []byte) validationCost {
|
|
words, punctuation := estimator.CountComponents(string(fragment))
|
|
return validationCost{words - 1, punctuation}
|
|
}
|
|
|
|
func measureValidationItem(estimator *chunking.SimpleTokenEstimator, item LLMValidationItem) (validationCost, error) {
|
|
// Every positive local reference has the same cost, regardless of digit count.
|
|
item.CorrectionIndex = 1
|
|
payload, err := json.Marshal(item)
|
|
if err != nil {
|
|
return validationCost{}, fmt.Errorf("marshal batch payload: %w", err)
|
|
}
|
|
return measureJSON(estimator, payload), nil
|
|
}
|
|
|
|
func measureContextPrefixes[T any](estimator *chunking.SimpleTokenEstimator, field string, entries []T) ([]validationCost, error) {
|
|
prefixes := make([]validationCost, len(entries)+1)
|
|
cost := measureJSON(estimator, []byte(`,"`+field+`":[]`))
|
|
for i, entry := range entries {
|
|
payload, err := json.Marshal(entry)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshal validation %s: %w", field, err)
|
|
}
|
|
if i > 0 {
|
|
cost.punctuation++ // separator between entries
|
|
}
|
|
cost = cost.add(measureJSON(estimator, payload))
|
|
prefixes[i+1] = cost
|
|
}
|
|
return prefixes, nil
|
|
}
|
|
|
|
func fitSimpleValidationContext(item LLMValidationItem, budget int, estimator *chunking.SimpleTokenEstimator) (LLMValidationItem, validationCost, error) {
|
|
required := item
|
|
required.NeighboringSegments, required.GlossaryEntries = nil, nil
|
|
base, err := measureValidationItem(estimator, required)
|
|
if err != nil {
|
|
return item, validationCost{}, err
|
|
}
|
|
if len(item.NeighboringSegments)+len(item.GlossaryEntries) == 0 {
|
|
return item, base, nil
|
|
}
|
|
neighbors, err := measureContextPrefixes(estimator, "neighboring_segments", item.NeighboringSegments)
|
|
if err != nil {
|
|
return item, validationCost{}, err
|
|
}
|
|
glossary, err := measureContextPrefixes(estimator, "glossary_entries", item.GlossaryEntries)
|
|
if err != nil {
|
|
return item, validationCost{}, err
|
|
}
|
|
n, g := len(item.NeighboringSegments), len(item.GlossaryEntries)
|
|
cost := base.add(neighbors[n]).add(glossary[g])
|
|
for n > 0 && cost.arrayTokens() > budget {
|
|
n--
|
|
cost = base.add(neighbors[n]).add(glossary[g])
|
|
}
|
|
for g > 0 && cost.arrayTokens() > budget {
|
|
g--
|
|
cost = base.add(glossary[g])
|
|
}
|
|
item.NeighboringSegments = item.NeighboringSegments[:n]
|
|
item.GlossaryEntries = item.GlossaryEntries[:g]
|
|
return item, cost, nil
|
|
}
|