Merge pull request #369 from ggozad/fix/embedding-batch-size

Expose embedding batch size as config

CHANGELOG.md

@@ -9,6 +9,7 @@
 - **`[s3]` optional extra** (`obstore>=0.9`). Required for `s3://` sources and the S3 watcher. Uses obstore — the Python binding to the same Rust `object_store` crate that LanceDB uses internally — so `monitor.s3[*].storage_options` accepts the same dict shape as `lancedb.storage_options`. Empty/missing options fall back to the AWS default credential chain.
 - **`scripts/run-integration-tests.sh`** — wraps `docker compose up --wait`, `pytest -m integration`, and tear-down so the SeaweedFS-backed integration suite is a one-liner.
 - **`ModelConfig.extra_body`**. Optional dict forwarded verbatim to `ModelSettings.extra_body`, the raw pass-through pydantic-ai exposes for openai/ollama/anthropic/groq. Lets configs reach provider-specific keys without haiku.rag modelling them — e.g. `extra_body: {chat_template_kwargs: {enable_thinking: false}}` to disable Qwen3 thinking on a vLLM endpoint, where the high-level `enable_thinking` flag is a no-op.
+- **`embeddings.batch_size`** (default `512`). Number of text chunks per `/v1/embeddings` call during ingest. Lower it when your provider caps total tokens per request. Closes #365.
 
 ### Changed
 

docs/configuration/providers.md

@@ -92,6 +92,10 @@ vLLM serves Qwen3 chat templates that read their thinking switch from `chat_temp
 
 Embedding models require three settings: `provider`, `name`, and `vector_dim`. Optionally, use `base_url` for OpenAI-compatible servers.
 
+### Batch Size
+
+`embeddings.batch_size` (default `512`) sets how many text chunks are sent per `/v1/embeddings` call during ingest. Lower it if your provider caps total tokens per request. Picture embeddings are always sent one image per call and are unaffected.
+
 ### Ollama (Default)
 
 ```yaml

haiku_rag_slim/haiku/rag/config/models.py

@@ -83,6 +83,7 @@ class LanceDBConfig(BaseModel):
 
 class EmbeddingsConfig(BaseModel):
     model: EmbeddingModelConfig = Field(default_factory=EmbeddingModelConfig)
+    batch_size: int = 512
 
 
 class RerankingConfig(BaseModel):

haiku_rag_slim/haiku/rag/embeddings/__init__.py

@@ -82,9 +82,6 @@ def contextualize(chunks: list["Chunk"]) -> list[str]:
     return texts
 
 
-EMBEDDING_BATCH_SIZE = 512
-
-
 async def embed_chunks(
     chunks: list["Chunk"], config: AppConfig = Config
 ) -> list["Chunk"]:
@@ -113,8 +110,9 @@ async def embed_chunks(
     text_embeddings: list[list[float]] = []
     if text_chunks:
         texts = contextualize(text_chunks)
-        for i in range(0, len(texts), EMBEDDING_BATCH_SIZE):
-            batch = texts[i : i + EMBEDDING_BATCH_SIZE]
+        batch_size = config.embeddings.batch_size
+        for i in range(0, len(texts), batch_size):
+            batch = texts[i : i + batch_size]
             text_embeddings.extend(await embedder.embed_documents(batch))
 
     picture_embeddings: list[list[float]] = []

tests/test_embedder.py

@@ -179,9 +179,15 @@ async def test_embed_chunks_picture_with_text_only_embedder_raises():
         await embed_chunks([chunk], config)
 
 
-async def test_embed_chunks_batches_large_inputs(monkeypatch):
-    """Test that embed_chunks batches calls when chunk count exceeds batch size."""
-    from haiku.rag.embeddings import EMBEDDING_BATCH_SIZE, EmbedderWrapper
+async def test_embed_chunks_respects_configured_batch_size(monkeypatch):
+    """`embeddings.batch_size` controls how `embed_chunks` slices its input.
+
+    Voyage models cap total tokens per /embeddings call (120K for
+    voyage-3-large and friends). Exposing the slice size lets users tune
+    it down without dropping `chunk_size` and harming retrieval quality.
+    """
+    from haiku.rag.config import AppConfig
+    from haiku.rag.embeddings import EmbedderWrapper
 
     call_sizes: list[int] = []
 
@@ -191,20 +197,20 @@ async def test_embed_chunks_batches_large_inputs(monkeypatch):
 
     monkeypatch.setattr(EmbedderWrapper, "embed_documents", tracking_embed)
 
-    # Create more chunks than one batch
-    num_chunks = EMBEDDING_BATCH_SIZE + 100
+    config = AppConfig()
+    config.embeddings.batch_size = 7
+
+    num_chunks = 20
     chunks = [
         Chunk(id=f"chunk-{i}", content=f"Content {i}", order=i)
         for i in range(num_chunks)
     ]
 
-    result = await embed_chunks(chunks)
+    result = await embed_chunks(chunks, config)
 
     assert len(result) == num_chunks
-    assert len(call_sizes) == 2
-    assert call_sizes[0] == EMBEDDING_BATCH_SIZE
-    assert call_sizes[1] == 100
-    # Verify order is preserved
+    # 20 chunks / 7 per batch -> 7, 7, 6
+    assert call_sizes == [7, 7, 6]
     assert result[0].id == "chunk-0"
     assert result[-1].id == f"chunk-{num_chunks - 1}"
     assert all(r.embedding == [0.1] * 10 for r in result)