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 {