fix(llm): chunk embed() requests under the backend batch cap (#221)
This commit was merged in pull request #221.
This commit is contained in:
@@ -22,6 +22,11 @@ struct EmbeddingData {
|
|||||||
index: usize,
|
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 ───────────────────────────────────
|
// ── Embedding implementation ───────────────────────────────────
|
||||||
|
|
||||||
impl LlmClient {
|
impl LlmClient {
|
||||||
@@ -29,8 +34,21 @@ impl LlmClient {
|
|||||||
&self.embed_model
|
&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> {
|
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 url = format!("{}/v1/embeddings", self.base_url.trim_end_matches('/'));
|
||||||
|
|
||||||
let request_body = EmbeddingRequest {
|
let request_body = EmbeddingRequest {
|
||||||
@@ -72,3 +90,33 @@ impl LlmClient {
|
|||||||
Ok(data.into_iter().map(|d| d.embedding).collect())
|
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