Reuse validation costs when fitting and batching

internal/framework/validators/llm_batching.go

@@ -3,6 +3,7 @@ package validators
 import (
 	"encoding/json"
 	"fmt"
+	"slices"
 
 	"gitea.maximumdirect.net/eric/audita/v3/internal/core/chunking"
 )
@@ -13,6 +14,39 @@ type LLMValidationBatch struct {
 	TokenCount int                 `json:"token_count"`
 }
 
+func prepareValidationBatches(items []LLMValidationItem, budget int, estimator chunking.TokenEstimator) ([]LLMValidationBatch, []Decision, error) {
+	if simple, ok := estimator.(*chunking.SimpleTokenEstimator); ok {
+		usable := make([]LLMValidationItem, 0, len(items))
+		costs := make([]validationCost, 0, len(items))
+		var oversized []Decision
+		for _, item := range items {
+			fitted, cost, err := fitSimpleValidationContext(item, budget, simple)
+			if err != nil {
+				return nil, oversized, err
+			}
+			if cost.arrayTokens() > budget {
+				oversized = append(oversized, oversizedValidationDecision(item.CorrectionIndex))
+			} else {
+				usable = append(usable, fitted)
+				costs = append(costs, cost)
+			}
+		}
+		batches, err := chunkSimpleValidationItems(usable, costs, budget, simple)
+		return batches, oversized, err
+	}
+
+	for i := range items {
+		items[i] = fitValidationContext(items[i], budget, estimator)
+	}
+	oversized := oversizedValidationDecisions(items, budget, estimator)
+	usable := filterItemsByDecision(items, oversized)
+	if len(usable) == 0 {
+		return nil, oversized, nil
+	}
+	batches, err := ChunkLLMValidationItems(usable, budget, estimator)
+	return batches, oversized, err
+}
+
 func ChunkLLMValidationItems(items []LLMValidationItem, maxPromptTokens int, estimator chunking.TokenEstimator) ([]LLMValidationBatch, error) {
 	if maxPromptTokens <= 0 {
 		return nil, fmt.Errorf("validation max prompt tokens must be greater than zero")
@@ -23,6 +57,9 @@ func ChunkLLMValidationItems(items []LLMValidationItem, maxPromptTokens int, est
 	if estimator == nil {
 		estimator = chunking.NewSimpleTokenEstimator()
 	}
+	if simple, ok := estimator.(*chunking.SimpleTokenEstimator); ok {
+		return chunkSimpleValidationItems(items, nil, maxPromptTokens, simple)
+	}
 
 	batches := make([]LLMValidationBatch, 0)
 	current := make([]LLMValidationItem, 0)
@@ -69,6 +106,48 @@ func ChunkLLMValidationItems(items []LLMValidationItem, maxPromptTokens int, est
 	return batches, nil
 }
 
+func chunkSimpleValidationItems(items []LLMValidationItem, costs []validationCost, budget int, estimator *chunking.SimpleTokenEstimator) ([]LLMValidationBatch, error) {
+	var batches []LLMValidationBatch
+	start := 0
+	var batchCost validationCost
+	for i, item := range items {
+		var cost validationCost
+		if costs != nil {
+			cost = costs[i]
+		} else {
+			var err error
+			cost, err = measureValidationItem(estimator, item)
+			if err != nil {
+				return nil, err
+			}
+		}
+		if cost.arrayTokens() > budget {
+			return nil, fmt.Errorf("single validation proposal exceeds max prompt tokens")
+		}
+		candidateCost := batchCost.add(cost)
+		if i > start {
+			candidateCost.punctuation++ // separator between items
+			if candidateCost.arrayTokens() > budget {
+				batches = append(batches, LLMValidationBatch{
+					BatchIndex: len(batches),
+					Items:      slices.Clone(items[start:i]),
+					TokenCount: batchCost.arrayTokens(),
+				})
+				start, candidateCost = i, cost
+			}
+		}
+		batchCost = candidateCost
+	}
+	if start < len(items) {
+		batches = append(batches, LLMValidationBatch{
+			BatchIndex: len(batches),
+			Items:      slices.Clone(items[start:]),
+			TokenCount: batchCost.arrayTokens(),
+		})
+	}
+	return batches, nil
+}
+
 func estimateBatchTokens(estimator chunking.TokenEstimator, batch []LLMValidationItem) (int, error) {
 	payload, err := json.Marshal(localValidationItems(batch))
 	if err != nil {

internal/framework/validators/llm_budget.go

@@ -0,0 +1,90 @@
+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
+}

internal/framework/validators/llm_budget_test.go

@@ -0,0 +1,74 @@
+package validators
+
+import (
+	"encoding/json"
+	"reflect"
+	"slices"
+	"strings"
+	"testing"
+
+	"gitea.maximumdirect.net/eric/audita/v3/internal/core/chunking"
+	"gitea.maximumdirect.net/eric/audita/v3/internal/core/glossary"
+)
+
+// Wrapping the same estimator exercises budgeting through exact serialization.
+type serializedBudgetEstimator struct{ *chunking.SimpleTokenEstimator }
+
+func TestComposedValidationBudgetMatchesSerializedPayload(t *testing.T) {
+	items := make([]LLMValidationItem, 12)
+	for i := range items {
+		items[i] = LLMValidationItem{
+			CorrectionIndex:      100 + i*7,
+			OriginalText:         "one\u00a0two!",
+			CorrectedText:        "one two",
+			OriginalSegmentText:  "she said \"one\u00a0two!\"\n",
+			CorrectedSegmentText: "she said <one & two> — there.",
+			Categories:           []string{"say, then write", "repeat"},
+		}
+	}
+	items[0].NeighboringSegments = []NeighboringSegment{
+		{Position: "preceding", Text: "spaces\tand escapes\\", Categories: []string{"speaker line"}},
+		{Position: "following", Text: "‘quoted’\u2028words"},
+	}
+	items[0].GlossaryEntries = []glossary.Entry{
+		{ID: "name", Name: "Name?", Aliases: []string{"N. M.", "Name's"}},
+		{ID: "two", Name: "One Two", Description: "some\u2003meaning…"},
+	}
+	items[1].CorrectedSegmentText = strings.Repeat("longer candidate. ", 12)
+
+	simple := chunking.NewSimpleTokenEstimator()
+	serialized := serializedBudgetEstimator{simple}
+	full, err := estimateBatchTokens(serialized, items)
+	if err != nil {
+		t.Fatal(err)
+	}
+	for budget := 1; budget <= full; budget++ {
+		got, rejected, err := prepareValidationBatches(slices.Clone(items), budget, simple)
+		want, wantRejected, serializedErr := prepareValidationBatches(slices.Clone(items), budget, serialized)
+		if err != nil || serializedErr != nil || !reflect.DeepEqual(got, want) || !reflect.DeepEqual(rejected, wantRejected) {
+			t.Fatalf("budget %d: composed=%+v rejected=%+v err=%v; serialized=%+v rejected=%+v err=%v",
+				budget, got, rejected, err, want, wantRejected, serializedErr)
+		}
+	}
+
+	batches, err := ChunkLLMValidationItems(items, full, simple)
+	if err != nil || len(batches) != 1 || batches[0].TokenCount != full {
+		t.Fatalf("full batch: %+v, %v", batches, err)
+	}
+	variables, err := BuildValidationVariables(batches[0].Items, "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	var payload []LLMValidationItem
+	if err := json.Unmarshal([]byte(variables["PayloadJSON"]), &payload); err != nil {
+		t.Fatal(err)
+	}
+	if len(payload) != len(items) {
+		t.Fatalf("sent %d candidates, want %d", len(payload), len(items))
+	}
+	for i := range payload {
+		if payload[i].CorrectionIndex != i+1 || batches[0].Items[i].CorrectionIndex != items[i].CorrectionIndex {
+			t.Fatalf("candidate %d lost its local or run reference", i)
+		}
+	}
+}

internal/framework/validators/llm_validators.go

@@ -75,10 +75,7 @@ func (v *LLMBackedValidator) Validate(ctx context.Context, req Request) (Result,
 		maxTokens = req.Settings.Validation.MaxPromptTokens
 	}
 
-	for i := range validationReq.Items {
-		validationReq.Items[i] = fitValidationContext(validationReq.Items[i], maxTokens, v.estimator)
-	}
-	oversized := oversizedValidationDecisions(validationReq.Items, maxTokens, v.estimator)
+	batches, oversized, err := prepareValidationBatches(validationReq.Items, maxTokens, v.estimator)
 	operations := make([]outcomes.Operation, 0)
 	llmDecisions := make([]Decision, 0)
 	partialResult := func() Result {
@@ -110,18 +107,6 @@ func (v *LLMBackedValidator) Validate(ctx context.Context, req Request) (Result,
 			ProposalIndices: indices,
 		})
 	}
-	itemsForBatching := filterItemsByDecision(validationReq.Items, oversized)
-	if len(itemsForBatching) == 0 {
-		result := partialResult()
-		all := result.Decisions
-		if err := EnforceDecisionCardinality(req.CandidateProposal, all); err != nil {
-			return result, err
-		}
-		sort.SliceStable(all, func(i, j int) bool { return all[i].ProposalIndex < all[j].ProposalIndex })
-		return result, nil
-	}
-
-	batches, err := ChunkLLMValidationItems(itemsForBatching, maxTokens, v.estimator)
 	if err != nil {
 		return partialResult(), err
 	}
@@ -355,17 +340,21 @@ func reconcileLLMResponse(items []LLMValidationItem, response LLMValidationRespo
 }
 
 func oversizedValidationDecisions(items []LLMValidationItem, maxPromptTokens int, estimator chunking.TokenEstimator) []Decision {
-	out := make([]Decision, 0)
+	var out []Decision
 	for _, item := range items {
 		singleTokens, err := estimateBatchTokens(estimator, []LLMValidationItem{item})
 		if err != nil || singleTokens <= maxPromptTokens {
 			continue
 		}
-		out = append(out, unavailable(item.CorrectionIndex, ReasonValidatorInputTooLarge, "validation input exceeds max prompt tokens"))
+		out = append(out, oversizedValidationDecision(item.CorrectionIndex))
 	}
 	return out
 }
 
+func oversizedValidationDecision(index int) Decision {
+	return unavailable(index, ReasonValidatorInputTooLarge, "validation input exceeds max prompt tokens")
+}
+
 func filterItemsByDecision(items []LLMValidationItem, decisions []Decision) []LLMValidationItem {
 	if len(decisions) == 0 {
 		return append([]LLMValidationItem(nil), items...)