Reuse validation costs when fitting and batching

docs/internal/validators.md

@@ -142,7 +142,8 @@ payload token estimate. If one enriched item exceeds the configured limit, optio
 neighbors are removed from the end first, then glossary entries from the end,
 until it fits or only required candidate data remains. Required before/after text
 is never truncated. Existing oversized-candidate handling still applies if that
-required data does not fit. Context is supplemental, not a promise that every
+required data does not fit. The built-in estimator reuses encoded fragment
+counts for fitting and batching. Context is supplemental, not a promise that every
 relevant definition or neighboring utterance always fits. This uses the existing
 payload estimator, not a new guarantee about the provider's full context window.
 

docs/roadmap/remediation.md

@@ -601,6 +601,8 @@ immutable render inputs are shared safely across sections of the same stage.
 
 ## R17 — P3: Reuse Validation Payload Costs While Fitting and Batching
 
+**Status:** Completed. The built-in estimator composes compact JSON costs across context fitting and greedy batching. Other estimators retain exact serialization. Equivalence tests cover trimmed context, oversized items, token counts, batch splits, escaped text, and local references.
+
 **Category**
 
 Quadratic prefix serialization in the existing budgeting path now used with

internal/core/chunking/tokens.go

@@ -1,7 +1,6 @@
 package chunking
 
 import (
-	"strings"
 	"unicode"
 )
 
@@ -24,24 +23,24 @@ func NewSimpleTokenEstimator() *SimpleTokenEstimator {
 // This implementation uses a simple heuristic: count words and punctuation as tokens.
 // The estimate is deterministic and stable for the same input text.
 func (e *SimpleTokenEstimator) EstimateTokens(text string) int {
-	if text == "" {
-		return 0
-	}
+	words, punctuation := e.CountComponents(text)
+	return words + punctuation/2
+}
 
-	// Simple heuristic: split on whitespace and count non-empty segments
-	words := strings.Fields(text)
-	tokenCount := len(words)
-
-	// Add some estimate for punctuation that might be separate tokens
-	punctuationCount := 0
+// CountComponents returns unrounded counts for budgets composed from fragments.
+func (e *SimpleTokenEstimator) CountComponents(text string) (words, punctuation int) {
+	inWord := false
 	for _, r := range text {
+		nonspace := !unicode.IsSpace(r)
+		if nonspace && !inWord {
+			words++
+		}
+		inWord = nonspace
 		if unicode.IsPunct(r) && r != '\'' && r != '-' && r != '_' {
-			punctuationCount++
+			punctuation++
 		}
 	}
-
-	// Rough estimate: each word is a token, plus half the punctuation as separate tokens
-	return tokenCount + (punctuationCount / 2)
+	return words, punctuation
 }
 
 // ConstTokenEstimator returns a constant token count for testing purposes.

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...)