79 lines
2.1 KiB
Go
79 lines
2.1 KiB
Go
package validators
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
|
|
"gitea.maximumdirect.net/eric/audita/internal/core/chunking"
|
|
)
|
|
|
|
type LLMValidationBatch struct {
|
|
BatchIndex int `json:"batch_index"`
|
|
Items []LLMValidationItem `json:"items"`
|
|
TokenCount int `json:"token_count"`
|
|
}
|
|
|
|
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()
|
|
}
|
|
|
|
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 estimateBatchTokens(estimator chunking.TokenEstimator, batch []LLMValidationItem) (int, error) {
|
|
payload, err := json.Marshal(batch)
|
|
if err != nil {
|
|
return 0, fmt.Errorf("marshal batch payload: %w", err)
|
|
}
|
|
return estimator.EstimateTokens(string(payload)), nil
|
|
}
|