audita/internal/framework/validators/llm_batching.go

package validators

import (
	"encoding/json"
	"fmt"
	"slices"

	"gitea.maximumdirect.net/eric/audita/v3/internal/core/chunking"
)

type LLMValidationBatch struct {
	BatchIndex int                 `json:"batch_index"`
	Items      []LLMValidationItem `json:"items"`
	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")
	}
	if len(items) == 0 {
		return nil, fmt.Errorf("validation input must contain at least one proposal")
	}
	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)
	currentTokens := 0

	for _, item := range items {
		singleTokens, err := estimateBatchTokens(estimator, []LLMValidationItem{item})
		if err != nil {
			return nil, err
		}
		if singleTokens > maxPromptTokens {
			return nil, fmt.Errorf("single validation proposal exceeds max prompt tokens")
		}

		candidate := append(append([]LLMValidationItem(nil), current...), item)
		candidateTokens, err := estimateBatchTokens(estimator, candidate)
		if err != nil {
			return nil, err
		}

		if len(current) > 0 && candidateTokens > maxPromptTokens {
			batches = append(batches, LLMValidationBatch{
				BatchIndex: len(batches),
				Items:      append([]LLMValidationItem(nil), current...),
				TokenCount: currentTokens,
			})
			current = []LLMValidationItem{item}
			currentTokens = singleTokens
			continue
		}

		current = candidate
		currentTokens = candidateTokens
	}

	if len(current) > 0 {
		batches = append(batches, LLMValidationBatch{
			BatchIndex: len(batches),
			Items:      append([]LLMValidationItem(nil), current...),
			TokenCount: currentTokens,
		})
	}

	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 {
		return 0, fmt.Errorf("marshal batch payload: %w", err)
	}
	return estimator.EstimateTokens(string(payload)), nil
}