Update megumin_agent/retrieval.py

megumin_agent/retrieval.py

@@ -36,6 +36,7 @@ ANSWER_KEYS = (
 COLLECTION_KEYS = ("items", "data", "examples", "dataset", "records")
 EMBEDDING_MODEL_NAME = os.getenv("MEGUMIN_EMBEDDING_MODEL", "gemini-embedding-001")
 EMBEDDING_DIMENSION = int(os.getenv("MEGUMIN_EMBEDDING_DIM", "768"))
+EMBEDDING_BATCH_SIZE = int(os.getenv("MEGUMIN_EMBEDDING_BATCH_SIZE", "100"))
 
 
 def _normalize_text(value: Any) -> str:
@@ -180,22 +181,30 @@ def _embed_texts(
     if not texts:
         return np.zeros((0, output_dimensionality), dtype="float32")
 
-    response = _get_genai_client().models.embed_content(
-        model=embedding_model,
-        contents=texts,
-        config=types.EmbedContentConfig(
-            task_type=task_type,
-            output_dimensionality=output_dimensionality,
-        ),
-    )
-    vectors = np.array(
-        [embedding.values for embedding in response.embeddings],
-        dtype="float32",
-    )
-    if vectors.size == 0:
+    batches: list[np.ndarray] = []
+    batch_size = max(1, min(EMBEDDING_BATCH_SIZE, 100))
+    for start in range(0, len(texts), batch_size):
+        chunk = texts[start : start + batch_size]
+        response = _get_genai_client().models.embed_content(
+            model=embedding_model,
+            contents=chunk,
+            config=types.EmbedContentConfig(
+                task_type=task_type,
+                output_dimensionality=output_dimensionality,
+            ),
+        )
+        vectors = np.array(
+            [embedding.values for embedding in response.embeddings],
+            dtype="float32",
+        )
+        if vectors.size == 0:
+            continue
+        faiss.normalize_L2(vectors)
+        batches.append(vectors)
+
+    if not batches:
         return np.zeros((0, output_dimensionality), dtype="float32")
-    faiss.normalize_L2(vectors)
-    return vectors
+    return np.vstack(batches)
 
 
 @lru_cache(maxsize=8)