Add Cohere multimodal embedder

CHANGELOG.md

@@ -6,6 +6,7 @@
 - `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
 

haiku_rag_slim/haiku/rag/embeddings/__init__.py

@@ -1,3 +1,5 @@
+import base64
+import io
 from typing import TYPE_CHECKING, Any
 
 from pydantic_ai.embeddings import Embedder
@@ -69,6 +71,23 @@ class EmbedderWrapper:
         )
 
 
+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.
 
@@ -246,6 +265,11 @@ def _get_multimodal_embedder(
 
         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."

haiku_rag_slim/haiku/rag/embeddings/cohere.py

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

haiku_rag_slim/haiku/rag/embeddings/vllm.py

@@ -11,13 +11,11 @@ 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
@@ -123,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}")

tests/test_embedder.py

@@ -854,3 +854,135 @@ async def test_voyage_embed_text_and_image_end_to_end():
     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"