Expose embedding batch size as config

This commit is contained in:
Yiorgis Gozadinos 2026-05-13 14:09:47 +03:00
parent d11b8a4e6e
commit 8b135c4d4d
No known key found for this signature in database
5 changed files with 25 additions and 15 deletions

View file

@ -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

View file

@ -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

View file

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

View file

@ -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]] = []

View file

@ -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)