diff --git a/compliance-agent/src/llm/embedding.rs b/compliance-agent/src/llm/embedding.rs index 3e8e0d9..cf6f9a5 100644 --- a/compliance-agent/src/llm/embedding.rs +++ b/compliance-agent/src/llm/embedding.rs @@ -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) -> Result>, 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) -> Result>, 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" + ); + } +}