diff --git a/CHANGELOG.md b/CHANGELOG.md index 44f3a420..2da633b2 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/configuration/providers.md b/docs/configuration/providers.md index 0e17e05e..176f0dda 100644 --- a/docs/configuration/providers.md +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 5a12f384..39060348 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/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): diff --git a/haiku_rag_slim/haiku/rag/embeddings/__init__.py b/haiku_rag_slim/haiku/rag/embeddings/__init__.py index 3cd77a4f..8b8cb563 100644 --- a/haiku_rag_slim/haiku/rag/embeddings/__init__.py +++ b/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]] = [] diff --git a/tests/test_embedder.py b/tests/test_embedder.py index 61bdb087..49552e98 100644 --- a/tests/test_embedder.py +++ b/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)