From ff656504d3c860c0162ef6feababf5ed4fdd74b2 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 4 May 2026 11:22:22 +0300 Subject: [PATCH] make client.search() polymorphic on query type: str | bytes | PIL.Image.Image. Bytes/PIL queries embed via embed_image_query and run vector-only against the chunks table --- CHANGELOG.md | 1 + haiku_rag_slim/haiku/rag/client/__init__.py | 3 +- haiku_rag_slim/haiku/rag/client/search.py | 50 +++++--- .../haiku/rag/store/repositories/chunk.py | 20 +++- tests/store/test_document_items.py | 10 +- tests/test_config.py | 3 +- tests/test_embedder.py | 2 +- tests/test_picture_in_context.py | 9 +- tests/test_search.py | 113 ++++++++++++++++++ 9 files changed, 178 insertions(+), 33 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index df0d51bd..f7296518 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,7 @@ ### Added +- **Image-as-query search.** `client.search()` now accepts `str | bytes | PIL.Image.Image`. Bytes/PIL queries embed via the multimodal embedder's `embed_image_query` and dispatch to vector-only chunk search (FTS doesn't apply to non-text queries; reranking is also skipped). Raises a clear error if the configured embedder is text-only. `ChunkRepository.search()` gains an optional `query_vector` parameter that bypasses `embed_query` and forces the vector-only path. - **`vision: bool` flag on `ModelConfig`.** Tracks whether a configured language model can interpret images. Default `False`. The agent's `search` tool only attaches picture bytes (as `BinaryContent`) to the `ToolReturn` when `qa.model.vision = True`. Without the gate, sending image content to a text-only model behaves inconsistently across providers — Ollama silently accepts and the model hallucinates a confident wrong answer; OpenAI returns 400; others vary. Capability detection from a probe or a model-name whitelist is unreliable, so `vision` is an explicit user-set capability declaration. Set it to `True` for vision-capable QA models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, ...). - **Synthetic picture chunks at ingest under multimodal embedders.** `build_picture_chunks` (in `client/processing.py`) walks a `DoclingDocument`'s `pictures` and emits one synthetic `Chunk` per `PictureItem` with available bytes. Bytes ride on a `Chunk._picture_data` PrivateAttr (not serialized) so `embed_chunks` can route them through `embed_images` while text chunks keep going through `embed_documents`. Wired into the three ingest paths (`create_document`, `_create_document_from_file`, `_create_or_update_document_from_url`, `_update_document_with_chunks`, and `_rebuild_rechunk`) — guarded by `embedder.supports_images` so text-only configurations are unaffected. Snapshot/merge with `existing_picture_data` keeps picture chunks alive across rebuild round-trips. Picture chunks land in the same `chunks` table with the same vector dim as text chunks, so cross-modal search reuses the existing hybrid+RRF pipeline. - **Multimodal embedder support (`provider="vllm"`).** `EmbedderWrapper` gains `supports_images: bool` and `embed_image_query`. The `vllm` provider talks HTTP to a vLLM server's OpenAI-compatible `/v1/embeddings` endpoint — text inputs use the standard `input` field for true server-side batching; image inputs use vLLM's `messages` superset with `image_url` content parts carrying base64 data URIs. Works with `Qwen/Qwen3-VL-Embedding-8B` and `jinaai/jina-embeddings-v4`. No Python ML deps added — uses `httpx`. Text-only providers (`ollama`, `openai`, `cohere`, `sentence-transformers`) report `supports_images=False` and raise a clear error if image methods are called. diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index 3987bc33..c336ff7b 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -28,6 +28,7 @@ from haiku.rag.utils import escape_sql_string if TYPE_CHECKING: from docling_core.types.doc.document import DoclingDocument + from PIL import Image as PILImage from haiku.rag.agents.analysis.models import AnalysisResult from haiku.rag.agents.research.models import ( @@ -329,7 +330,7 @@ class HaikuRAG: async def search( self, - query: str, + query: "str | bytes | PILImage.Image", limit: int | None = None, search_type: str = "hybrid", filter: str | None = None, diff --git a/haiku_rag_slim/haiku/rag/client/search.py b/haiku_rag_slim/haiku/rag/client/search.py index 9cdb8a3c..5dedc596 100644 --- a/haiku_rag_slim/haiku/rag/client/search.py +++ b/haiku_rag_slim/haiku/rag/client/search.py @@ -5,12 +5,14 @@ from haiku.rag.reranking import get_reranker from haiku.rag.store.models.chunk import Chunk, SearchResult if TYPE_CHECKING: + from PIL import Image as PILImage + from haiku.rag.client import HaikuRAG async def search( client: "HaikuRAG", - query: str, + query: "str | bytes | PILImage.Image", limit: int | None = None, search_type: str = "hybrid", filter: str | None = None, @@ -20,14 +22,13 @@ async def search( Args: client: The HaikuRAG client (provides config + chunk repository). - query: The search query string. + query: Text (``str``) or image (``bytes`` / ``PIL.Image.Image``). + Image queries require a multimodal embedder and run vector-only. limit: Maximum number of results to return. Defaults to config.search.limit. - search_type: Type of search - "vector", "fts", or "hybrid" (default). + search_type: "vector", "fts", or "hybrid" (default). Text queries only. filter: Optional SQL WHERE clause to filter documents before searching chunks. include_images: When True, populate ``SearchResult.image_data`` with - base64-encoded picture bytes for picture-labeled chunks. Set to - False to skip the lookup (e.g. for plain-text MCP consumers that - don't want the JSON bloat). + base64 picture bytes for picture-labeled chunks. Returns: List of SearchResult objects ordered by relevance. @@ -35,19 +36,36 @@ async def search( if limit is None: limit = client._config.search.limit - reranker = get_reranker(config=client._config) + if isinstance(query, str): + reranker = get_reranker(config=client._config) - if reranker is None: - chunk_results = await client.chunk_repository.search( - query, limit, search_type, filter - ) + if reranker is None: + chunk_results = await client.chunk_repository.search( + query, limit, search_type, filter + ) + else: + search_limit = limit * 10 + raw_results = await client.chunk_repository.search( + query, search_limit, search_type, filter + ) + chunks = [chunk for chunk, _ in raw_results] + chunk_results = await reranker.rerank(query, chunks, top_n=limit) else: - search_limit = limit * 10 - raw_results = await client.chunk_repository.search( - query, search_limit, search_type, filter + from haiku.rag.embeddings import get_embedder + + embedder = get_embedder(client._config) + if not embedder.supports_images: + raise ValueError( + "Image queries require a multimodal embedder. Configure " + "provider='vllm' (or another image-capable provider)." + ) + query_vector = await embedder.embed_image_query(query) + chunk_results = await client.chunk_repository.search( + query="", + limit=limit, + filter=filter, + query_vector=query_vector, ) - chunks = [chunk for chunk, _ in raw_results] - chunk_results = await reranker.rerank(query, chunks, top_n=limit) results = [SearchResult.from_chunk(chunk, score) for chunk, score in chunk_results] diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index ed6286ca..f32121c8 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -218,23 +218,25 @@ class ChunkRepository: async def search( self, - query: str, + query: str = "", limit: int = 5, search_type: str = "hybrid", filter: str | None = None, + query_vector: list[float] | None = None, ) -> list[tuple[Chunk, float]]: """Search for relevant chunks using the specified search method. Args: - query: The search query string. + query: Text query. Empty when ``query_vector`` is supplied. limit: Maximum number of results to return. - search_type: Type of search - "vector", "fts", or "hybrid" (default). + search_type: "vector", "fts", or "hybrid" (default). filter: Optional SQL WHERE clause to filter documents before searching chunks. + query_vector: Pre-computed query embedding; forces vector-only search. Returns: List of (chunk, score) tuples ordered by relevance. """ - if not query.strip(): + if query_vector is None and not query.strip(): return [] chunk_filter: str | None = None @@ -255,7 +257,15 @@ class ChunkRepository: id_list = ", ".join(f"'{d}'" for d in docs_df["id"]) chunk_filter = f"document_id IN ({id_list})" - if search_type == "vector": + if query_vector is not None: + # Image-as-query: vector-only against the pre-computed embedding. + results = ( + self.store.chunks_table.query() + .nearest_to(query_vector) + .column("vector") + .refine_factor(self.store._config.search.vector_refine_factor) + ) + elif search_type == "vector": query_embedding = await self.embedder.embed_query(query) results = ( self.store.chunks_table.query() diff --git a/tests/store/test_document_items.py b/tests/store/test_document_items.py index f55e6c74..24ead8b7 100644 --- a/tests/store/test_document_items.py +++ b/tests/store/test_document_items.py @@ -448,7 +448,7 @@ def _docling_doc_with_picture(): class TestExtractItemsPictureBytes: - """A2b: extract_items decodes picture image bytes from data URIs.""" + """extract_items decodes picture image bytes from data URIs.""" def test_decodes_picture_bytes_from_live_doc(self): doc = _docling_doc_with_picture() @@ -483,7 +483,7 @@ class TestExtractItemsPictureBytes: class TestExtractItemTextDescription: - """A2b: extract_item_text returns VLM description text for PictureItems.""" + """extract_item_text returns VLM description text for PictureItems.""" def test_returns_description_text_when_present(self): from docling_core.types.doc.document import ( @@ -503,7 +503,7 @@ class TestExtractItemTextDescription: class TestCompressDoclingSplitStripsPictureUris: - """A2b: compress_docling_split removes inline picture URIs from the structure.""" + """compress_docling_split removes inline picture URIs from the structure.""" def test_picture_image_set_to_none_in_structure(self): import json @@ -540,7 +540,7 @@ class TestCompressDoclingSplitStripsPictureUris: @pytest.mark.asyncio class TestPictureDataMigrationBackfill: - """A2b: v0.45.0 migration backfills picture_data and strips URIs from blobs.""" + """0.45.0 migration backfills picture_data and strips URIs from blobs.""" async def test_backfill_populates_column_and_strips_blob(self, temp_db_path): import base64 @@ -630,7 +630,7 @@ class TestPictureDataMigrationBackfill: @pytest.mark.asyncio class TestPictureDataPreservedThroughRoundTrip: - """A2b: snapshot/merge keeps picture bytes through update / rebuild cycles.""" + """Snapshot/merge keeps picture bytes through update / rebuild cycles.""" async def test_update_preserves_picture_data_when_blob_round_tripped( self, temp_db_path diff --git a/tests/test_config.py b/tests/test_config.py index 2e7dc804..87a50dcf 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -227,7 +227,8 @@ def test_init_config_creates_valid_yaml(tmp_path): assert config.environment == "production" -# A4: legacy `generate_picture_images` + `picture_description.enabled` translation +# Legacy picture-handling field translation (`generate_picture_images` and +# `picture_description.enabled` → `processing.pictures` enum) def _write(tmp_path, body: str): diff --git a/tests/test_embedder.py b/tests/test_embedder.py index c15b8d87..2f1f1541 100644 --- a/tests/test_embedder.py +++ b/tests/test_embedder.py @@ -218,7 +218,7 @@ async def test_embed_chunks_preserves_all_fields(allow_model_requests): assert embedded[0].embedding is not None -# B1: multimodal embedder support +# Multimodal embedder support def _ollama_text_only_config(): diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index 7123545a..8ccf9229 100644 --- a/tests/test_picture_in_context.py +++ b/tests/test_picture_in_context.py @@ -1,4 +1,4 @@ -"""A3: SearchResult.image_data, expand_context preservation, multimodal ToolReturn.""" +"""Picture-bearing search results: image_data attachment, expansion, multimodal ToolReturn.""" import base64 from dataclasses import dataclass @@ -107,8 +107,9 @@ async def test_expand_context_preserves_picture_refs_with_empty_text(temp_db_pat """A picture item with empty text must keep its self_ref through expansion.""" async with HaikuRAG(temp_db_path, create=True) as rag: # Build an items table with a section header + a paragraph match + an - # adjacent picture row that has no text. expand_with_items used to - # filter the picture out via the `if item.text:` guard; A3 keeps it. + # adjacent picture row that has no text. The expansion must keep + # picture self_refs even when item.text is empty so picture bytes + # are still attached downstream. await rag.document_item_repository.create_items( "doc-1", [ @@ -326,7 +327,7 @@ async def test_search_tool_returns_multimodal_when_picture_present(): assert part.data == PICTURE_BYTES -# B2: synthetic picture chunks at ingestion +# Synthetic picture chunks at ingest def test_build_picture_chunks_uses_live_uri(): diff --git a/tests/test_search.py b/tests/test_search.py index d235bfb1..5ced315e 100644 --- a/tests/test_search.py +++ b/tests/test_search.py @@ -327,3 +327,116 @@ def test_search_result_primary_label_prioritizes_structural_types(): labels=[], ) assert result._get_primary_label() is None + + +# Image queries (bytes / PIL.Image) + + +@pytest.mark.asyncio +async def test_search_with_bytes_query_uses_multimodal_embedder( + temp_db_path, monkeypatch +): + """``client.search(bytes)`` embeds via ``embed_image_query`` and dispatches + to vector-only chunk search (skipping FTS and reranker).""" + from haiku.rag.embeddings import EmbedderWrapper + from haiku.rag.store.models.chunk import Chunk + + image_calls: list[bytes] = [] + + class StubMultimodal(EmbedderWrapper): + supports_images = True + + def __init__(self): + super().__init__(embedder=None, vector_dim=4) + + async def embed_image_query(self, image): + image_calls.append(image) + return [0.5, 0.5, 0.5, 0.5] + + monkeypatch.setattr( + "haiku.rag.embeddings.get_embedder", + lambda *a, **kw: StubMultimodal(), + ) + + received_kwargs: dict = {} + + async def fake_chunk_search( + query="", limit=5, search_type="hybrid", filter=None, query_vector=None + ): + received_kwargs.update( + { + "query": query, + "limit": limit, + "search_type": search_type, + "filter": filter, + "query_vector": query_vector, + } + ) + return [ + ( + Chunk( + content="figure 1", + metadata={"labels": ["picture"], "doc_item_refs": ["#/pictures/0"]}, + ), + 0.91, + ) + ] + + async with HaikuRAG(temp_db_path, create=True) as rag: + rag.chunk_repository.search = fake_chunk_search # type: ignore[method-assign] + results = await rag.search(b"\x89PNG\r\n\x1a\n", limit=3, include_images=False) + + assert len(results) == 1 + assert results[0].score == 0.91 + # The bytes were sent through the image embedder once. + assert image_calls == [b"\x89PNG\r\n\x1a\n"] + # The chunk repo received a pre-computed vector and an empty text query. + assert received_kwargs["query_vector"] == [0.5, 0.5, 0.5, 0.5] + assert received_kwargs["query"] == "" + + +@pytest.mark.asyncio +async def test_search_with_pil_image_works_like_bytes(temp_db_path, monkeypatch): + from PIL import Image as PILImageModule + + from haiku.rag.embeddings import EmbedderWrapper + from haiku.rag.store.models.chunk import Chunk + + seen_types: list[type] = [] + + class StubMultimodal(EmbedderWrapper): + supports_images = True + + def __init__(self): + super().__init__(embedder=None, vector_dim=4) + + async def embed_image_query(self, image): + seen_types.append(type(image)) + return [0.1] * 4 + + monkeypatch.setattr( + "haiku.rag.embeddings.get_embedder", + lambda *a, **kw: StubMultimodal(), + ) + + async def fake_chunk_search(**kwargs): + return [(Chunk(content="x", metadata={}), 1.0)] + + async with HaikuRAG(temp_db_path, create=True) as rag: + rag.chunk_repository.search = fake_chunk_search # type: ignore[method-assign] + img = PILImageModule.new("RGB", (8, 8), "red") + results = await rag.search(img, include_images=False) + + assert len(results) == 1 + assert seen_types == [PILImageModule.Image] + + +@pytest.mark.asyncio +async def test_search_with_bytes_query_raises_for_text_only_embedder( + temp_db_path, +): + """A text-only embedder configured for QA must reject image queries + with a clear error rather than silently degrading.""" + async with HaikuRAG(temp_db_path, create=True) as rag: + with pytest.raises(ValueError, match="multimodal embedder"): + await rag.search(b"\x89PNG\r\n\x1a\n")