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:
commit
195dce7511
16 changed files with 7084 additions and 51 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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.)
|
||||
|
|
|
|||
|
|
@ -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 |
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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)),
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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."
|
||||
)
|
||||
|
|
|
|||
66
haiku_rag_slim/haiku/rag/embeddings/cohere.py
Normal file
66
haiku_rag_slim/haiku/rag/embeddings/cohere.py
Normal 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]
|
||||
|
|
@ -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}")
|
||||
|
|
|
|||
68
haiku_rag_slim/haiku/rag/embeddings/voyageai.py
Normal file
68
haiku_rag_slim/haiku/rag/embeddings/voyageai.py
Normal 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
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Reference in a new issue