danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

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

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-05-04 11:22:22 +03:00
parent 21f12cd770
commit ff656504d3
No known key found for this signature in database
9 changed files with 178 additions and 33 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
CHANGELOG.md
View file

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

3
haiku_rag_slim/haiku/rag/client/__init__.py
View file

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

50
haiku_rag_slim/haiku/rag/client/search.py
View file

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

20
haiku_rag_slim/haiku/rag/store/repositories/chunk.py
View file

@ -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()

10
tests/store/test_document_items.py
View file

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

3
tests/test_config.py
View file

@ -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):

2
tests/test_embedder.py
View file

@ -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():

9
tests/test_picture_in_context.py
View file

@ -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():

113
tests/test_search.py
View file

@ -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")