fix(llm): chunk embed() requests under the backend batch cap
CI / Check (push) Skipped
CI / Check (pull_request) Successful in 5m42s
CI / Detect Changes (pull_request) Skipped
CI / Deploy Agent (pull_request) Skipped
CI / Deploy Dashboard (pull_request) Skipped
CI / Deploy Docs (pull_request) Skipped
CI / Deploy MCP (pull_request) Skipped
CI / Check (push) Skipped
CI / Check (pull_request) Successful in 5m42s
CI / Detect Changes (pull_request) Skipped
CI / Deploy Agent (pull_request) Skipped
CI / Deploy Dashboard (pull_request) Skipped
CI / Deploy Docs (pull_request) Skipped
CI / Deploy MCP (pull_request) Skipped
Found while prepping the C5 live test: the embeddings backend
(bge-multilingual-gemma2 via LiteLLM) caps input arrays at 25 per request
("given batch size overflow maximal one", max: 25), but embed() sent the whole
input in a single request. ControlIndex::build embeds the entire master-controls
corpus (~1.8k texts) in one embed() call, and the RAG pipeline batches docs too —
both 500 at scale. Unit tests passed only because they embed <=3 texts.
Fix: embed() now chunks into EMBED_BATCH_SIZE (16) requests and concatenates in
order; empty input short-circuits. Validated live — 1,784 texts over 112 chunked
requests succeed (~70s, one-time + cached by the index).
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Fable 5
parent
a7ff36edf3
commit
617b2df75e
@@ -22,6 +22,11 @@ struct EmbeddingData {
|
||||
index: usize,
|
||||
}
|
||||
|
||||
/// Max inputs per embedding request. The bge/OpenAI-like backends cap the input
|
||||
/// array (bge-multilingual-gemma2 rejects >25 with "batch size overflow"), so we
|
||||
/// chunk larger corpora — a whole control catalog (~1.8k) would otherwise 500.
|
||||
const EMBED_BATCH_SIZE: usize = 16;
|
||||
|
||||
// ── Embedding implementation ───────────────────────────────────
|
||||
|
||||
impl LlmClient {
|
||||
@@ -29,8 +34,21 @@ impl LlmClient {
|
||||
&self.embed_model
|
||||
}
|
||||
|
||||
/// Generate embeddings for a batch of texts
|
||||
/// Generate embeddings for a batch of texts, chunking into backend-sized
|
||||
/// requests and preserving input order across chunks.
|
||||
pub async fn embed(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
|
||||
if texts.is_empty() {
|
||||
return Ok(Vec::new());
|
||||
}
|
||||
let mut out = Vec::with_capacity(texts.len());
|
||||
for chunk in texts.chunks(EMBED_BATCH_SIZE) {
|
||||
out.extend(self.embed_batch(chunk.to_vec()).await?);
|
||||
}
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
/// Embed one backend-sized batch (≤ [`EMBED_BATCH_SIZE`]) in a single request.
|
||||
async fn embed_batch(&self, texts: Vec<String>) -> Result<Vec<Vec<f64>>, AgentError> {
|
||||
let url = format!("{}/v1/embeddings", self.base_url.trim_end_matches('/'));
|
||||
|
||||
let request_body = EmbeddingRequest {
|
||||
@@ -72,3 +90,33 @@ impl LlmClient {
|
||||
Ok(data.into_iter().map(|d| d.embedding).collect())
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use secrecy::SecretString;
|
||||
|
||||
fn client() -> LlmClient {
|
||||
LlmClient::new(
|
||||
"http://unused".into(),
|
||||
SecretString::from(String::new()),
|
||||
"m".into(),
|
||||
"e".into(),
|
||||
)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn empty_input_makes_no_request() {
|
||||
// Must short-circuit before any HTTP call (base_url is unroutable).
|
||||
let out = client().embed(Vec::new()).await.unwrap();
|
||||
assert!(out.is_empty());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn batch_size_is_within_backend_cap() {
|
||||
assert!(
|
||||
EMBED_BATCH_SIZE <= 25,
|
||||
"must stay under the bge 25-input cap"
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user