Merge pull request #462 from ggozad/fix/multi-modal-embedder-config

Decouple multimodal embedding from the provider name; add VoyageAI & Cohere multimodal embedders
This commit is contained in:
Yiorgis Gozadinos 2026-06-23 15:33:57 +03:00 committed by GitHub
commit 195dce7511
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
16 changed files with 7084 additions and 51 deletions

View file

@ -4,6 +4,13 @@
### Added
- `haiku-rag doctor` checks a database for consistency (orphaned chunks/items, chunk-less documents classified by content and embedder modality, dangling `doc_item_refs`, vector-dimension mismatch, unembedded chunks, missing picture data, settings/embedding drift, pending migrations, vector-index coverage, provider API keys) and probes configured provider endpoints (Ollama `/api/tags` with model presence, docling-serve `/health`, OpenAI-compatible/vLLM `/models`); exits 1 when any check fails.
- `embeddings.model.multimodal` (bool, default false) gates image embedding; `supports_images` derives from it instead of the provider name.
- VoyageAI multimodal embedder (`provider: voyageai`, `multimodal: true`, e.g. `voyage-multimodal-3`) embedding text and pictures into a shared vector space.
- Cohere multimodal embedder (`provider: cohere`, `multimodal: true`, e.g. `embed-v4.0`) embedding text and pictures into a shared vector space.
### Changed
- `provider: vllm` is text-only unless `embeddings.model.multimodal: true` is set. Existing multimodal vLLM configs must add the flag; `multimodal` is not part of the stored embedding identity, so changing it raises no drift error — re-ingest or `rebuild` to add or drop picture chunks.
## [0.60.0] - 2026-06-22

View file

@ -10,14 +10,14 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p
## Features
- **Hybrid search** — Vector + full-text with Reciprocal Rank Fusion
- **Multimodal & cross-modal search** — Multimodal embedders (vLLM) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query
- **Multimodal & cross-modal search** — Multimodal embedders (vLLM, VoyageAI, Cohere) put picture vectors in the same space as text; supports text-as-query → figure hits and image-as-query
- **Question answering** — RAG skill with citations (page numbers, section headings)
- **Vision QA** — Vision-capable models receive figure bytes alongside chunk text
- **Reranking** — MxBAI, Cohere, Zero Entropy, or vLLM
- **Analysis skill** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis)
- **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory
- **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, LM Studio, vLLM (multimodal). QA: any model supported by Pydantic AI
- **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, Cohere, LM Studio, vLLM (multimodal via `multimodal: true` on vLLM/VoyageAI/Cohere). QA: any model supported by Pydantic AI
- **Local-first** — Embedded LanceDB, no servers required. Also supports S3, GCS, Azure, and LanceDB Cloud
- **CLI & Python API** — Full functionality from command line or code
- **MCP server** — Expose as tools for AI assistants (Claude Desktop, etc.)

View file

@ -283,9 +283,11 @@ Three independent settings drive ingest, retrieval, and QA:
| Setting | Question it answers | Values |
|---|---|---|
| `processing.pictures` | Generate and/or describe pictures at ingest? | `none` / `description` / `image` (default) |
| `embeddings.model.provider` | Can the embedder index image content? | text-only (`ollama`, `openai`, `cohere`, `sentence-transformers`) vs `vllm` (multimodal) |
| `embeddings.model.multimodal` | Can the embedder index image content? | `false` (default, text-only) / `true` (supported on `vllm`, `voyageai`, `cohere`) |
| `qa.model.vision` | Can the QA model interpret images? | `false` (default) / `true` |
The Embedder column below is driven by `embeddings.model.multimodal`, not the provider name — a vision-capable model under a text-only configuration still indexes no images, and an image-only document then produces zero chunks. See [Multimodal embedders](providers.md#multimodal-embedders).
**What gets stored** by `pictures` × embedder:
| `pictures` | Embedder | Text chunks | Synthetic picture chunks |

View file

@ -228,11 +228,15 @@ embeddings:
base_url: http://localhost:1234/v1
```
**Note:** The `base_url` must include the `/v1` path for OpenAI-compatible endpoints.
**Note:** The `base_url` must include the `/v1` path for OpenAI-compatible endpoints. This path is text-only. For a vision-language model served by vLLM, use `provider: vllm` with `multimodal: true` (below), not `provider: openai`.
### vLLM (multimodal)
### Multimodal embedders
For cross-modal retrieval (text and pictures share a single vector space), use the dedicated `vllm` provider against a vLLM server hosting a multimodal embedding model:
For cross-modal retrieval (text and pictures share a single vector space), set `embeddings.model.multimodal: true`. Capability is decided by this flag, not the provider name: each provider passes images in its own wire format, so multimodal is supported only on `vllm`, `voyageai`, and `cohere`. Setting it on any other provider raises at startup.
A model produces picture chunks at ingest only when its embedder is multimodal. Without the flag, an image-only document produces zero chunks and is not retrievable. Switching `multimodal` on or off does not change the stored embedding identity, so it raises no drift error; re-ingest or `rebuild` to add or drop picture chunks.
**vLLM** — a vLLM server hosting a multimodal embedding model. Text inputs use the standard OpenAI `input` field; image inputs use vLLM's `messages`-with-`image_url` superset. Tested with `Qwen/Qwen3-VL-Embedding-8B` (4096-dim) and `jinaai/jina-embeddings-v4` (2048-dim). Run vLLM separately; haiku.rag adds no Python ML dependencies for this path.
```yaml
embeddings:
@ -241,11 +245,34 @@ embeddings:
name: Qwen/Qwen3-VL-Embedding-8B
vector_dim: 4096
base_url: http://localhost:8000/v1
multimodal: true
```
Tested with `Qwen/Qwen3-VL-Embedding-8B` (4096-dim) and `jinaai/jina-embeddings-v4` (2048-dim). Run vLLM separately. haiku.rag adds no Python ML dependencies for this path. Text inputs use the standard OpenAI `input` field. Image inputs use vLLM's `messages`-with-`image_url` superset, transparently to the caller.
**VoyageAI** — `voyage-multimodal-3` (1024-dim) via the `voyageai` extra. Reads `VOYAGE_API_KEY` from the environment.
Picture chunks for retrieval are emitted at ingest under any embedder reporting `supports_images=True`. See [Picture Handling](processing.md#picture-handling).
```yaml
embeddings:
model:
provider: voyageai
name: voyage-multimodal-3
vector_dim: 1024
multimodal: true
```
**Cohere** — `embed-v4.0` (configurable `vector_dim`, e.g. 1536) via the `cohere` extra. Reads `CO_API_KEY` from the environment.
```yaml
embeddings:
model:
provider: cohere
name: embed-v4.0
vector_dim: 1536
multimodal: true
```
A text-only model served by vLLM uses `provider: vllm` without the flag (or `provider: openai` with a `base_url`).
Picture chunks for retrieval are emitted at ingest under any multimodal embedder. See [Picture Handling](processing.md#picture-handling).
## Question Answering Providers

View file

@ -260,7 +260,7 @@ results = await client.search(
### Image queries
`client.search()` accepts an image instead of a text query when the configured embedder is multimodal (e.g. `provider: vllm` against a vision-language embedding model). The image is embedded once and the chunks table is searched vector-only. Full-text search and reranking don't apply without a text query.
`client.search()` accepts an image instead of a text query when the configured embedder is multimodal (`embeddings.model.multimodal: true` on a vLLM, VoyageAI, or Cohere model). The image is embedded once and the chunks table is searched vector-only. Full-text search and reranking don't apply without a text query.
```python
from PIL import Image

View file

@ -58,8 +58,9 @@ async def search(
embedder = client.embedder
if not embedder.supports_images:
raise ValueError(
"Image queries require a multimodal embedder. Configure "
"provider='vllm' (or another image-capable provider)."
"Image queries require a multimodal embedder. Set "
"embeddings.model.multimodal: true on a vllm, voyageai, or cohere "
"model."
)
query_vector = await embedder.embed_image(query)
chunk_results = await client.chunk_repository.search(

View file

@ -39,16 +39,20 @@ class EmbeddingModelConfig(BaseModel):
"""Configuration for an embedding model.
Attributes:
provider: Model provider (ollama, openai, voyageai, cohere, sentence-transformers)
provider: Model provider (ollama, openai, voyageai, cohere, sentence-transformers, vllm)
name: Model name/identifier
vector_dim: Vector dimensions produced by the model
base_url: Optional base URL for OpenAI-compatible servers (vLLM, LM Studio, etc.)
multimodal: Whether the model embeds images into the same vector space as
text. Supported on the vllm, voyageai, and cohere providers; other
providers raise when this is set.
"""
provider: str = "ollama"
name: str = "qwen3-embedding:4b"
vector_dim: int = 2560
base_url: str | None = None
multimodal: bool = False
class StorageConfig(BaseModel):

View file

@ -209,7 +209,10 @@ def _classify_unchunked(
f"{len(picture_docs)} image-only document(s) have no chunks; "
"a text-only embedder cannot index images."
),
remediation="Configure a multimodal embedder and rebuild to index images.",
remediation=(
"Set embeddings.model.multimodal: true on a vllm, voyageai, or "
"cohere model and rebuild to index images."
),
details=_sample(sorted(picture_docs)),
)
)

View file

@ -1,3 +1,5 @@
import base64
import io
from typing import TYPE_CHECKING, Any
from pydantic_ai.embeddings import Embedder
@ -10,6 +12,7 @@ from haiku.rag.config import AppConfig, Config
if TYPE_CHECKING:
from PIL import Image as PILImage
from haiku.rag.config.models import EmbeddingModelConfig
from haiku.rag.store.models.chunk import Chunk
@ -19,15 +22,23 @@ ImageInput = "bytes | PILImage.Image"
class EmbedderWrapper:
"""Wrapper around pydantic-ai Embedder with explicit query/document methods.
Subclasses set ``supports_images = True`` and override the image methods
when the underlying model can encode pictures into the same vector space.
Subclasses that can encode pictures into the same vector space as text either
set the ``supports_images`` class attribute or pass ``supports_images=True``,
and override the image methods.
"""
supports_images: bool = False
def __init__(self, embedder: Embedder | None, vector_dim: int):
def __init__(
self,
embedder: Embedder | None,
vector_dim: int,
supports_images: bool | None = None,
):
self._embedder = embedder
self._vector_dim = vector_dim
if supports_images is not None:
self.supports_images = supports_images
@property
def vector_dim(self) -> int:
@ -55,11 +66,28 @@ class EmbedderWrapper:
``messages`` superset. Callers loop when they need many.
"""
raise NotImplementedError(
f"{type(self).__name__} does not support image embedding. "
"Configure a multimodal provider (e.g. provider='vllm')."
f"{type(self).__name__} does not support image embedding. Set "
"embeddings.model.multimodal: true on a vllm, voyageai, or cohere model."
)
def _to_data_uri(image: "bytes | PILImage.Image") -> str:
"""Render an image as a ``data:image/png;base64,...`` URI."""
if isinstance(image, bytes):
return f"data:image/png;base64,{base64.b64encode(image).decode('ascii')}"
from PIL import Image as PILImageModule
if isinstance(image, PILImageModule.Image):
buf = io.BytesIO()
image.save(buf, format="PNG")
return (
f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('ascii')}"
)
raise TypeError(f"Unsupported image type: {type(image)!r}")
def contextualize(chunks: list["Chunk"]) -> list[str]:
"""Prepare chunk content for embedding/FTS by adding context.
@ -117,8 +145,9 @@ async def embed_chunks(
if picture_chunks:
if not embedder.supports_images:
raise ValueError(
"Picture chunks require a multimodal embedder. Configure "
"provider='vllm', or omit picture chunks."
"Picture chunks require a multimodal embedder. Set "
"embeddings.model.multimodal: true on a vllm, voyageai, or cohere "
"model, or omit picture chunks."
)
for chunk in picture_chunks:
picture_embeddings.append(await embedder.embed_image(chunk._picture_data))
@ -159,6 +188,9 @@ def get_embedder(config: AppConfig = Config) -> EmbedderWrapper:
model_name = embedding_model.name
vector_dim = embedding_model.vector_dim
if embedding_model.multimodal:
return _get_multimodal_embedder(embedding_model)
if provider == "ollama":
# Use model-level base_url if set, otherwise fall back to providers config
base_url = embedding_model.base_url or config.providers.ollama.base_url
@ -193,9 +225,52 @@ def get_embedder(config: AppConfig = Config) -> EmbedderWrapper:
if provider == "vllm":
from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder
base_url = embedding_model.base_url or "http://localhost:8000/v1"
if not base_url.rstrip("/").endswith("/v1"):
base_url = base_url.rstrip("/") + "/v1"
return VLLMMultimodalEmbedder(model_name, vector_dim, base_url=base_url)
base_url = _vllm_base_url(embedding_model.base_url)
return VLLMMultimodalEmbedder(
model_name, vector_dim, base_url=base_url, supports_images=False
)
raise ValueError(f"Unsupported embedding provider: {provider}")
def _vllm_base_url(base_url: str | None) -> str:
base_url = base_url or "http://localhost:8000/v1"
if not base_url.rstrip("/").endswith("/v1"):
base_url = base_url.rstrip("/") + "/v1"
return base_url
def _get_multimodal_embedder(
embedding_model: "EmbeddingModelConfig",
) -> EmbedderWrapper:
"""Build an image-capable embedder for providers that support multimodal.
Each provider passes images in its own wire format, so the capability lives
in a per-provider embedder rather than a generic flag.
"""
provider = embedding_model.provider
model_name = embedding_model.name
vector_dim = embedding_model.vector_dim
if provider == "vllm":
from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder
base_url = _vllm_base_url(embedding_model.base_url)
return VLLMMultimodalEmbedder(
model_name, vector_dim, base_url=base_url, supports_images=True
)
if provider == "voyageai":
from haiku.rag.embeddings.voyageai import VoyageMultimodalEmbedder
return VoyageMultimodalEmbedder(model_name, vector_dim)
if provider == "cohere":
from haiku.rag.embeddings.cohere import CohereMultimodalEmbedder
return CohereMultimodalEmbedder(model_name, vector_dim)
raise ValueError(
f"Provider '{provider}' does not support multimodal embedding. Set "
"embeddings.model.multimodal: true on a vllm, voyageai, or cohere model."
)

View file

@ -0,0 +1,66 @@
"""Multimodal embedder backed by Cohere's ``embed`` API (``embed-v4.0``).
``embed-v4.0`` maps text and images into a shared vector space. Text uses the
``search_document``/``search_query`` input types; images are passed as base64
data URIs with the ``image`` input type. The API key is read from the
environment (``CO_API_KEY``) like the text-only Cohere path.
"""
from typing import TYPE_CHECKING
from haiku.rag.embeddings import EmbedderWrapper, _to_data_uri
if TYPE_CHECKING:
from PIL import Image as PILImage
class CohereMultimodalEmbedder(EmbedderWrapper):
def __init__(
self,
model_name: str,
vector_dim: int,
api_key: str | None = None,
):
super().__init__(embedder=None, vector_dim=vector_dim, supports_images=True)
import cohere
self._model_name = model_name
self._client = cohere.AsyncClientV2(api_key=api_key)
async def _embed_texts(
self, texts: list[str], input_type: str
) -> list[list[float]]:
result = await self._client.embed(
model=self._model_name,
input_type=input_type,
texts=texts,
output_dimension=self._vector_dim,
embedding_types=["float"],
)
return _floats(result)
async def embed_query(self, text: str) -> list[float]:
rows = await self._embed_texts([text], "search_query")
return rows[0]
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
if not texts:
return []
return await self._embed_texts(texts, "search_document")
async def embed_image(self, image: "bytes | PILImage.Image") -> list[float]:
result = await self._client.embed(
model=self._model_name,
input_type="image",
images=[_to_data_uri(image)],
output_dimension=self._vector_dim,
embedding_types=["float"],
)
return _floats(result)[0]
def _floats(result: object) -> list[list[float]]:
floats = result.embeddings.float_ # type: ignore[attr-defined] # ty: ignore[unresolved-attribute]
if floats is None:
raise ValueError("Cohere returned no float embeddings.")
return [list(e) for e in floats]

View file

@ -11,21 +11,17 @@ Models like ``Qwen/Qwen3-VL-Embedding-8B`` and ``jinaai/jina-embeddings-v4``
ship with chat templates that map both shapes into a shared vector space.
"""
import base64
import io
from typing import TYPE_CHECKING, Any
import httpx
from haiku.rag.embeddings import EmbedderWrapper
from haiku.rag.embeddings import EmbedderWrapper, _to_data_uri
if TYPE_CHECKING:
from PIL import Image as PILImage
class VLLMMultimodalEmbedder(EmbedderWrapper):
supports_images = True
def __init__(
self,
model_name: str,
@ -33,8 +29,11 @@ class VLLMMultimodalEmbedder(EmbedderWrapper):
base_url: str,
api_key: str | None = None,
timeout: float = 60.0,
supports_images: bool = True,
):
super().__init__(embedder=None, vector_dim=vector_dim)
super().__init__(
embedder=None, vector_dim=vector_dim, supports_images=supports_images
)
self._model_name = model_name
self._base_url = base_url.rstrip("/")
self._api_key = api_key
@ -99,6 +98,11 @@ class VLLMMultimodalEmbedder(EmbedderWrapper):
)
async def embed_image(self, image: "bytes | PILImage.Image") -> list[float]:
if not self.supports_images:
raise NotImplementedError(
"This vLLM embedder is text-only. Set "
"embeddings.model.multimodal: true to embed images."
)
rows = await self._post(
{
"model": self._model_name,
@ -117,20 +121,3 @@ class VLLMMultimodalEmbedder(EmbedderWrapper):
}
)
return rows[0]
def _to_data_uri(image: "bytes | PILImage.Image") -> str:
"""Render an image as a ``data:image/png;base64,...`` URI."""
if isinstance(image, bytes):
return f"data:image/png;base64,{base64.b64encode(image).decode('ascii')}"
from PIL import Image as PILImageModule
if isinstance(image, PILImageModule.Image):
buf = io.BytesIO()
image.save(buf, format="PNG")
return (
f"data:image/png;base64,{base64.b64encode(buf.getvalue()).decode('ascii')}"
)
raise TypeError(f"Unsupported image type: {type(image)!r}")

View file

@ -0,0 +1,68 @@
"""Multimodal embedder backed by VoyageAI's ``multimodal_embed`` endpoint.
``voyage-multimodal-3`` maps text and images into a shared vector space. Text is
embedded as single-element content lists; images are passed as ``PIL.Image``
objects (the SDK accepts them directly). The API key is read from the
environment (``VOYAGE_API_KEY``) like the text-only Voyage path.
"""
import io
from typing import TYPE_CHECKING
from haiku.rag.embeddings import EmbedderWrapper
if TYPE_CHECKING:
from PIL import Image as PILImage
class VoyageMultimodalEmbedder(EmbedderWrapper):
def __init__(
self,
model_name: str,
vector_dim: int,
api_key: str | None = None,
):
super().__init__(embedder=None, vector_dim=vector_dim, supports_images=True)
import voyageai
self._model_name = model_name
self._client = voyageai.AsyncClient(api_key=api_key)
async def embed_query(self, text: str) -> list[float]:
result = await self._client.multimodal_embed(
inputs=[[text]],
model=self._model_name,
input_type="query",
output_dimension=self._vector_dim,
)
return list(result.embeddings[0])
async def embed_documents(self, texts: list[str]) -> list[list[float]]:
if not texts:
return []
result = await self._client.multimodal_embed(
inputs=[[text] for text in texts],
model=self._model_name,
input_type="document",
output_dimension=self._vector_dim,
)
return [list(e) for e in result.embeddings]
async def embed_image(self, image: "bytes | PILImage.Image") -> list[float]:
result = await self._client.multimodal_embed(
inputs=[[_to_pil(image)]],
model=self._model_name,
input_type="document",
output_dimension=self._vector_dim,
)
return list(result.embeddings[0])
def _to_pil(image: "bytes | PILImage.Image") -> "PILImage.Image":
from PIL import Image as PILImageModule
if isinstance(image, bytes):
return PILImageModule.open(io.BytesIO(image))
if isinstance(image, PILImageModule.Image):
return image
raise TypeError(f"Unsupported image type: {type(image)!r}")

File diff suppressed because it is too large Load diff

File diff suppressed because one or more lines are too long

View file

@ -49,11 +49,19 @@ VECTOR_DIM = 4
ChunkRecord = create_chunk_model(VECTOR_DIM)
def _config(provider: str = "ollama", name: str = "test", vector_dim: int = VECTOR_DIM):
def _config(
provider: str = "ollama",
name: str = "test",
vector_dim: int = VECTOR_DIM,
multimodal: bool = False,
):
return AppConfig(
embeddings=EmbeddingsConfig(
model=EmbeddingModelConfig(
provider=provider, name=name, vector_dim=vector_dim
provider=provider,
name=name,
vector_dim=vector_dim,
multimodal=multimodal,
)
)
)
@ -312,7 +320,7 @@ async def test_image_only_document_multimodal_embedder_warns(temp_db_path):
],
)
report = await run_doctor(
_config(provider="vllm", name="qwen-vl"), temp_db_path, {}
_config(provider="vllm", name="qwen-vl", multimodal=True), temp_db_path, {}
)
result = _result(report, "documents_pictures_no_chunks")
assert result.severity is Severity.WARN

View file

@ -308,7 +308,7 @@ def _ollama_text_only_config():
async def test_text_only_embedder_does_not_support_images():
embedder = get_embedder(_ollama_text_only_config())
assert embedder.supports_images is False
with pytest.raises(NotImplementedError, match="multimodal provider"):
with pytest.raises(NotImplementedError, match="multimodal"):
await embedder.embed_image(b"\x89PNG\r\n\x1a\n")
@ -417,6 +417,20 @@ async def test_vllm_supports_images_flag():
assert embedder.supports_images is True
async def test_vllm_text_only_embed_image_raises():
from haiku.rag.embeddings.vllm import VLLMMultimodalEmbedder
embedder = VLLMMultimodalEmbedder(
model_name="x",
vector_dim=2,
base_url="http://localhost:8000/v1",
supports_images=False,
)
assert embedder.supports_images is False
with pytest.raises(NotImplementedError, match="text-only"):
await embedder.embed_image(b"\x89PNG\r\n\x1a\n")
async def test_vllm_connect_error_surfaces_helpful_message(monkeypatch):
import httpx
@ -630,6 +644,7 @@ async def test_vllm_get_embedder_routes_to_multimodal():
name="Qwen/Qwen3-VL-Embedding-2B",
vector_dim=2048,
base_url="http://my-vllm:8000/v1",
multimodal=True,
)
)
)
@ -638,6 +653,40 @@ async def test_vllm_get_embedder_routes_to_multimodal():
assert embedder._base_url == "http://my-vllm:8000/v1" # type: ignore[attr-defined] # ty: ignore[unresolved-attribute]
def test_multimodal_defaults_to_false():
model = EmbeddingModelConfig(provider="vllm", name="x", vector_dim=2)
assert model.multimodal is False
assert "multimodal" in model.model_dump()
async def test_vllm_text_only_when_multimodal_unset():
config = AppConfig(
embeddings=EmbeddingsConfig(
model=EmbeddingModelConfig(
provider="vllm",
name="qwen3-embedding:4b",
vector_dim=2560,
base_url="http://my-vllm:8000/v1",
)
)
)
embedder = get_embedder(config)
assert embedder.supports_images is False
@pytest.mark.parametrize("provider", ["ollama", "openai", "sentence-transformers"])
async def test_multimodal_unsupported_provider_raises(provider):
config = AppConfig(
embeddings=EmbeddingsConfig(
model=EmbeddingModelConfig(
provider=provider, name="x", vector_dim=2, multimodal=True
)
)
)
with pytest.raises(ValueError, match="does not support multimodal"):
get_embedder(config)
@pytest.mark.vcr()
async def test_vllm_embed_text_and_image_end_to_end():
"""End-to-end against a real vLLM ``/v1/embeddings`` server: confirm
@ -672,3 +721,268 @@ async def test_vllm_embed_text_and_image_end_to_end():
image_vec = await embedder.embed_image(image)
assert len(image_vec) == 4096
assert any(abs(x) > 1e-6 for x in image_vec), "image embedding is all zeros"
class _FakeVoyageResult:
def __init__(self, embeddings):
self.embeddings = embeddings
def _fake_voyage_client(captured, embeddings):
class FakeAsyncClient:
def __init__(self, *args, **kwargs):
captured["init"] = kwargs
async def multimodal_embed(self, **kwargs):
captured.update(kwargs)
return _FakeVoyageResult(embeddings)
return FakeAsyncClient
async def test_voyage_embed_documents_request_shape(monkeypatch):
from haiku.rag.embeddings.voyageai import VoyageMultimodalEmbedder
captured: dict = {}
monkeypatch.setattr(
"voyageai.AsyncClient",
_fake_voyage_client(captured, [[0.1, 0.2], [0.3, 0.4]]),
)
embedder = VoyageMultimodalEmbedder("voyage-multimodal-3", vector_dim=2)
vecs = await embedder.embed_documents(["a cat", "a dog"])
assert vecs == [[0.1, 0.2], [0.3, 0.4]]
assert captured["model"] == "voyage-multimodal-3"
assert captured["input_type"] == "document"
assert captured["inputs"] == [["a cat"], ["a dog"]]
assert captured["output_dimension"] == 2
async def test_voyage_embed_query_request_shape(monkeypatch):
from haiku.rag.embeddings.voyageai import VoyageMultimodalEmbedder
captured: dict = {}
monkeypatch.setattr(
"voyageai.AsyncClient", _fake_voyage_client(captured, [[0.5, 0.6]])
)
embedder = VoyageMultimodalEmbedder("voyage-multimodal-3", vector_dim=2)
vec = await embedder.embed_query("find the cat")
assert vec == [0.5, 0.6]
assert captured["model"] == "voyage-multimodal-3"
assert captured["input_type"] == "query"
assert captured["inputs"] == [["find the cat"]]
async def test_voyage_embed_image_passes_pil(monkeypatch):
from PIL import Image
from haiku.rag.embeddings.voyageai import VoyageMultimodalEmbedder
captured: dict = {}
monkeypatch.setattr(
"voyageai.AsyncClient", _fake_voyage_client(captured, [[0.7, 0.8]])
)
embedder = VoyageMultimodalEmbedder("voyage-multimodal-3", vector_dim=2)
import io
buf = io.BytesIO()
Image.new("RGB", (4, 4), "red").save(buf, format="PNG")
vec = await embedder.embed_image(buf.getvalue())
assert vec == [0.7, 0.8]
assert captured["model"] == "voyage-multimodal-3"
inputs = captured["inputs"]
assert len(inputs) == 1 and len(inputs[0]) == 1
assert isinstance(inputs[0][0], Image.Image)
assert captured["input_type"] == "document"
async def test_voyage_embed_documents_empty_list_skips_request(monkeypatch):
from haiku.rag.embeddings.voyageai import VoyageMultimodalEmbedder
captured: dict = {}
monkeypatch.setattr("voyageai.AsyncClient", _fake_voyage_client(captured, []))
embedder = VoyageMultimodalEmbedder("voyage-multimodal-3", vector_dim=2)
assert await embedder.embed_documents([]) == []
assert "inputs" not in captured
async def test_voyage_get_embedder_routes_to_multimodal(monkeypatch):
monkeypatch.setattr("voyageai.AsyncClient", _fake_voyage_client({}, []))
from haiku.rag.embeddings.voyageai import VoyageMultimodalEmbedder
config = AppConfig(
embeddings=EmbeddingsConfig(
model=EmbeddingModelConfig(
provider="voyageai",
name="voyage-multimodal-3",
vector_dim=1024,
multimodal=True,
)
)
)
embedder = get_embedder(config)
assert isinstance(embedder, VoyageMultimodalEmbedder)
assert embedder.supports_images is True
@pytest.mark.vcr()
async def test_voyage_embed_text_and_image_end_to_end():
"""End-to-end against the real VoyageAI ``multimodal_embed`` API: text and
image inputs return embeddings of the configured dimension in a shared
vector space. To re-record, set ``VOYAGE_API_KEY`` and run with
``--record-mode=rewrite``."""
from PIL import Image
from haiku.rag.embeddings.voyageai import VoyageMultimodalEmbedder
embedder = VoyageMultimodalEmbedder("voyage-multimodal-3", vector_dim=1024)
text_vec = await embedder.embed_query("a photo of a red square")
assert len(text_vec) == 1024
assert any(abs(x) > 1e-6 for x in text_vec), "text embedding is all zeros"
text_batch = await embedder.embed_documents(["hello world", "another doc"])
assert len(text_batch) == 2
assert all(len(v) == 1024 for v in text_batch)
image_vec = await embedder.embed_image(Image.new("RGB", (64, 64), (255, 0, 0)))
assert len(image_vec) == 1024
assert any(abs(x) > 1e-6 for x in image_vec), "image embedding is all zeros"
class _FakeCohereEmbeddings:
def __init__(self, float_):
self.float_ = float_
class _FakeCohereResult:
def __init__(self, float_):
self.embeddings = _FakeCohereEmbeddings(float_)
def _fake_cohere_client(captured, float_):
class FakeAsyncClientV2:
def __init__(self, *args, **kwargs):
captured["init"] = kwargs
async def embed(self, **kwargs):
captured.update(kwargs)
return _FakeCohereResult(float_)
return FakeAsyncClientV2
async def test_cohere_embed_documents_request_shape(monkeypatch):
from haiku.rag.embeddings.cohere import CohereMultimodalEmbedder
captured: dict = {}
monkeypatch.setattr(
"cohere.AsyncClientV2", _fake_cohere_client(captured, [[0.1, 0.2], [0.3, 0.4]])
)
embedder = CohereMultimodalEmbedder("embed-v4.0", vector_dim=2)
vecs = await embedder.embed_documents(["a cat", "a dog"])
assert vecs == [[0.1, 0.2], [0.3, 0.4]]
assert captured["model"] == "embed-v4.0"
assert captured["input_type"] == "search_document"
assert captured["texts"] == ["a cat", "a dog"]
assert captured["output_dimension"] == 2
assert captured["embedding_types"] == ["float"]
async def test_cohere_embed_query_request_shape(monkeypatch):
from haiku.rag.embeddings.cohere import CohereMultimodalEmbedder
captured: dict = {}
monkeypatch.setattr(
"cohere.AsyncClientV2", _fake_cohere_client(captured, [[0.5, 0.6]])
)
embedder = CohereMultimodalEmbedder("embed-v4.0", vector_dim=2)
vec = await embedder.embed_query("find the cat")
assert vec == [0.5, 0.6]
assert captured["model"] == "embed-v4.0"
assert captured["input_type"] == "search_query"
assert captured["texts"] == ["find the cat"]
async def test_cohere_embed_image_uses_image_input_type(monkeypatch):
from haiku.rag.embeddings.cohere import CohereMultimodalEmbedder
captured: dict = {}
monkeypatch.setattr(
"cohere.AsyncClientV2", _fake_cohere_client(captured, [[0.7, 0.8]])
)
embedder = CohereMultimodalEmbedder("embed-v4.0", vector_dim=2)
vec = await embedder.embed_image(b"\x89PNG\r\n\x1a\nfake")
assert vec == [0.7, 0.8]
assert captured["model"] == "embed-v4.0"
assert captured["input_type"] == "image"
images = captured["images"]
assert len(images) == 1
assert images[0].startswith("data:image/png;base64,")
async def test_cohere_embed_documents_empty_list_skips_request(monkeypatch):
from haiku.rag.embeddings.cohere import CohereMultimodalEmbedder
captured: dict = {}
monkeypatch.setattr("cohere.AsyncClientV2", _fake_cohere_client(captured, []))
embedder = CohereMultimodalEmbedder("embed-v4.0", vector_dim=2)
assert await embedder.embed_documents([]) == []
assert "texts" not in captured
async def test_cohere_get_embedder_routes_to_multimodal(monkeypatch):
monkeypatch.setattr("cohere.AsyncClientV2", _fake_cohere_client({}, []))
from haiku.rag.embeddings.cohere import CohereMultimodalEmbedder
config = AppConfig(
embeddings=EmbeddingsConfig(
model=EmbeddingModelConfig(
provider="cohere",
name="embed-v4.0",
vector_dim=1536,
multimodal=True,
)
)
)
embedder = get_embedder(config)
assert isinstance(embedder, CohereMultimodalEmbedder)
assert embedder.supports_images is True
@pytest.mark.vcr()
async def test_cohere_embed_text_and_image_end_to_end():
"""End-to-end against the real Cohere ``embed`` API (``embed-v4.0``): text
and image inputs return embeddings of the configured dimension in a shared
vector space. To re-record, set ``CO_API_KEY`` and run with
``--record-mode=rewrite``."""
from PIL import Image
from haiku.rag.embeddings.cohere import CohereMultimodalEmbedder
embedder = CohereMultimodalEmbedder("embed-v4.0", vector_dim=1536)
text_vec = await embedder.embed_query("a photo of a red square")
assert len(text_vec) == 1536
assert any(abs(x) > 1e-6 for x in text_vec), "text embedding is all zeros"
text_batch = await embedder.embed_documents(["hello world", "another doc"])
assert len(text_batch) == 2
assert all(len(v) == 1536 for v in text_batch)
image_vec = await embedder.embed_image(Image.new("RGB", (64, 64), (255, 0, 0)))
assert len(image_vec) == 1536
assert any(abs(x) > 1e-6 for x in image_vec), "image embedding is all zeros"