From 617b2df75e9eb3b66601c6a9e1276afc59ee7149 Mon Sep 17 00:00:00 2001
From: Sharang Parnerkar <30073382+mighty840@users.noreply.github.com>
Date: Tue, 21 Jul 2026 17:34:17 +0200
Subject: [PATCH] fix(llm): chunk embed() requests under the backend batch cap
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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>
---
 compliance-agent/src/llm/embedding.rs | 50 ++++++++++++++++++++++++++-
 1 file changed, 49 insertions(+), 1 deletion(-)

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<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"
+        );
+    }
+}