From a9b7c37a7edd377ece775650416f969300b78bbf Mon Sep 17 00:00:00 2001 From: Eric Rakestraw Date: Tue, 10 Nov 2026 20:45:38 +0000 Subject: [PATCH] Reuse validation costs when fitting and batching --- docs/internal/validators.md | 3 +- docs/roadmap/remediation.md | 2 + internal/core/chunking/tokens.go | 27 +++--- internal/framework/validators/llm_batching.go | 79 ++++++++++++++++ internal/framework/validators/llm_budget.go | 90 +++++++++++++++++++ .../framework/validators/llm_budget_test.go | 74 +++++++++++++++ .../framework/validators/llm_validators.go | 25 ++---- 7 files changed, 267 insertions(+), 33 deletions(-) create mode 100644 internal/framework/validators/llm_budget.go create mode 100644 internal/framework/validators/llm_budget_test.go diff --git a/docs/internal/validators.md b/docs/internal/validators.md index 9e234ab..1131bc1 100644 --- a/docs/internal/validators.md +++ b/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. diff --git a/docs/roadmap/remediation.md b/docs/roadmap/remediation.md index 89d0047..ff1006e 100644 --- a/docs/roadmap/remediation.md +++ b/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 diff --git a/internal/core/chunking/tokens.go b/internal/core/chunking/tokens.go index 4666821..b65601b 100644 --- a/internal/core/chunking/tokens.go +++ b/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. diff --git a/internal/framework/validators/llm_batching.go b/internal/framework/validators/llm_batching.go index a27812c..fb015db 100644 --- a/internal/framework/validators/llm_batching.go +++ b/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 { diff --git a/internal/framework/validators/llm_budget.go b/internal/framework/validators/llm_budget.go new file mode 100644 index 0000000..b90ca1f --- /dev/null +++ b/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 +} diff --git a/internal/framework/validators/llm_budget_test.go b/internal/framework/validators/llm_budget_test.go new file mode 100644 index 0000000..491f5cf --- /dev/null +++ b/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 — 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) + } + } +} diff --git a/internal/framework/validators/llm_validators.go b/internal/framework/validators/llm_validators.go index 66cb22e..bf2a2d4 100644 --- a/internal/framework/validators/llm_validators.go +++ b/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...)