From ff82d36c2fe5634358774db059b7a6320530e1f1 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 5 May 2026 12:02:46 +0300 Subject: [PATCH] add rebuild --descriptions: run VLM over stored picture bytes only --- CHANGELOG.md | 2 + docs/cli.md | 7 + docs/configuration/processing.md | 5 +- docs/python.md | 8 + haiku_rag_slim/haiku/rag/app.py | 1 + haiku_rag_slim/haiku/rag/cli.py | 16 +- haiku_rag_slim/haiku/rag/client/__init__.py | 2 + haiku_rag_slim/haiku/rag/client/rebuild.py | 148 ++++++++++++++++++ .../rag/providers/picture_description.py | 52 ++++++ tests/test_rebuild.py | 130 +++++++++++++++ 10 files changed, 368 insertions(+), 3 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/providers/picture_description.py diff --git a/CHANGELOG.md b/CHANGELOG.md index aee4c129..79da7e65 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,6 +3,8 @@ ### Added +- **`rebuild --descriptions` mode.** Adds VLM picture descriptions to an existing database without re-converting from source. Loads each document's stored docling blob, identifies pictures lacking `meta.description.text`, drives the configured VLM (via pydantic-ai `BinaryContent`) over the picture bytes already stored in `document_items.picture_data`, patches descriptions into the blob, then re-chunks + re-embeds so chunk text reflects them. Skips the docling parse entirely; only the VLM time is paid. Idempotent — pictures that already carry a description are not re-described, so the operation is safe to re-run after a partial failure. Errors clearly when `picture_description.enabled` is false. Exposed as `haiku-rag rebuild --descriptions` and `RebuildMode.DESCRIPTIONS`. +- **Silent-failure guard for picture descriptions.** When `picture_description.enabled=True` and a converted document has at least one picture but zero of them came back with a description, `client.processing.convert()` now logs a clear warning naming the source path, picture count, configured VLM model, and base URL. docling-serve swallows VLM errors (network failures, missing models, unreachable hosts) and returns a "successful" conversion with empty descriptions; this guard surfaces the failure before a long ingest produces a corpus with no descriptions in it. - **MCP image-query tool + CLI `--image PATH`.** New MCP tool `search_documents_by_image(image_base64, limit, include_images)` routes a base64-encoded image through `client.search()`. Registered only when the configured embedder supports images, so non-multimodal MCP servers don't expose a tool that would always fail. The `haiku-rag search` CLI gains an `--image PATH` flag that reads the file and runs the same image-as-query path. - **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`, ...). diff --git a/docs/cli.md b/docs/cli.md index 953815c3..668dd451 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -470,6 +470,10 @@ haiku-rag rebuild --embed-only # Only generate titles for untitled documents haiku-rag rebuild --title-only + +# Run the VLM over already-stored picture bytes and patch descriptions +# into the docling blob. Skips the docling parse entirely. +haiku-rag rebuild --descriptions ``` **Rebuild modes:** @@ -480,6 +484,9 @@ haiku-rag rebuild --title-only | Rechunk | `--rechunk` | Changed chunking strategy or chunk size | | Embed only | `--embed-only` | Changed embedding model or vector dimensions | | Title only | `--title-only` | Generate titles for documents without one | +| Descriptions | `--descriptions` | Add VLM picture descriptions to an existing database | + +**`--descriptions` mode** runs the configured VLM (`processing.conversion_options.picture_description.model`) over the picture bytes already stored in `document_items.picture_data`, patches each description into the stored docling blob's `pictures[i].meta.description.text`, and re-chunks + re-embeds so chunk text reflects the new descriptions. Requires `picture_description.enabled: true` in the config. Idempotent — pictures that already carry a description are skipped, so the operation is safe to re-run after a partial failure. The docling parse is skipped entirely; only the VLM time is paid. ### Download Models diff --git a/docs/configuration/processing.md b/docs/configuration/processing.md index 8c272c5e..941573a2 100644 --- a/docs/configuration/processing.md +++ b/docs/configuration/processing.md @@ -123,7 +123,10 @@ processing: When `enabled: false` (default), the VLM doesn't run; chunks contain only their natural text (captions, surrounding paragraphs). When `enabled: true`, each picture's description is woven into the chunk text and is searchable via FTS. -**Switching the VLM on or off on an existing database.** Picture bytes are already stored, so no reingest is required. Run `haiku-rag rebuild --rechunk` after flipping `enabled` so the chunk-text composition reflects the new setting. +**Switching the VLM on or off on an existing database.** Picture bytes are already stored, so no reingest is required. + +- To turn the VLM **off** (descriptions already exist, you want to drop them): flip `enabled: false` and run `haiku-rag rebuild --rechunk`. Chunk text recomposes from the stripped docling blob. +- To turn the VLM **on** (descriptions don't exist yet, you want them now): flip `enabled: true` and run `haiku-rag rebuild --descriptions`. The VLM is driven over the picture bytes already in `document_items.picture_data`, descriptions are patched into the docling blob, and chunks are recomposed. The docling parse is skipped entirely. See [Rebuild Database](../cli.md#rebuild-database) for full details. #### Picture descriptions × embedder × QA model: how the pieces compose diff --git a/docs/python.md b/docs/python.md index ca98b6df..813a17d6 100644 --- a/docs/python.md +++ b/docs/python.md @@ -222,6 +222,13 @@ async for doc_id in client.rebuild_database(mode=RebuildMode.RECHUNK): # Only regenerate embeddings (fastest, keeps existing chunks) async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY): print(f"Processed document {doc_id}") + +# Add VLM picture descriptions to an existing database — runs the VLM +# over already-stored picture bytes, patches descriptions into the +# docling blob, then re-chunks + re-embeds. Requires +# picture_description.enabled=true in the config. +async for doc_id in client.rebuild_database(mode=RebuildMode.DESCRIPTIONS): + print(f"Described pictures in {doc_id}") ``` **Rebuild modes:** @@ -230,6 +237,7 @@ async for doc_id in client.rebuild_database(mode=RebuildMode.EMBED_ONLY): - `RebuildMode.RECHUNK` - Re-chunk from existing document content, re-embed - `RebuildMode.EMBED_ONLY` - Keep existing chunks, only regenerate embeddings - `RebuildMode.TITLE_ONLY` - Generate titles for untitled documents (no re-chunking or re-embedding) +- `RebuildMode.DESCRIPTIONS` - Run the VLM over picture bytes already stored on `document_items.picture_data`, patch descriptions into the docling blob, re-chunk + re-embed. Skips the docling parse entirely. Idempotent — pictures already carrying `meta.description.text` are not re-described, so the operation is safe to re-run. ### Generating Titles diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 7c83f557..65281277 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -587,6 +587,7 @@ class HaikuRAGApp: # pragma: no cover RebuildMode.RECHUNK: "rechunk", RebuildMode.EMBED_ONLY: "embed only", RebuildMode.TITLE_ONLY: "title only", + RebuildMode.DESCRIPTIONS: "picture descriptions", }[mode] self.console.print( diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 29bb874e..50649ec0 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -490,13 +490,23 @@ def rebuild( "--title-only", help="Only generate titles for documents without one", ), + descriptions: bool = typer.Option( + False, + "--descriptions", + help=( + "Run the VLM over already-stored picture bytes, patch descriptions " + "into the docling blob, then re-chunk + re-embed. Skips the docling " + "parse entirely. Requires picture_description.enabled=true." + ), + ), ): from haiku.rag.client import RebuildMode - exclusive = sum([embed_only, rechunk, title_only]) + exclusive = sum([embed_only, rechunk, title_only, descriptions]) if exclusive > 1: typer.echo( - "Error: --embed-only, --rechunk, and --title-only are mutually exclusive" + "Error: --embed-only, --rechunk, --title-only, and --descriptions " + "are mutually exclusive" ) raise typer.Exit(1) @@ -506,6 +516,8 @@ def rebuild( mode = RebuildMode.RECHUNK elif title_only: # pragma: no cover mode = RebuildMode.TITLE_ONLY + elif descriptions: # pragma: no cover + mode = RebuildMode.DESCRIPTIONS else: # pragma: no cover mode = RebuildMode.FULL diff --git a/haiku_rag_slim/haiku/rag/client/__init__.py b/haiku_rag_slim/haiku/rag/client/__init__.py index c336ff7b..caca877f 100644 --- a/haiku_rag_slim/haiku/rag/client/__init__.py +++ b/haiku_rag_slim/haiku/rag/client/__init__.py @@ -46,6 +46,8 @@ class RebuildMode(Enum): RECHUNK = "rechunk" # Re-chunk from existing content, re-embed EMBED_ONLY = "embed_only" # Keep chunks, only regenerate embeddings TITLE_ONLY = "title_only" # Only generate titles for untitled documents + DESCRIPTIONS = "descriptions" # Run the VLM over already-stored picture + # bytes, patch descriptions into the docling blob, then re-chunk + re-embed. class HaikuRAG: diff --git a/haiku_rag_slim/haiku/rag/client/rebuild.py b/haiku_rag_slim/haiku/rag/client/rebuild.py index 1364930c..cf796118 100644 --- a/haiku_rag_slim/haiku/rag/client/rebuild.py +++ b/haiku_rag_slim/haiku/rag/client/rebuild.py @@ -57,6 +57,11 @@ async def rebuild_database( await client.store.recreate_embeddings_table() async for doc_id in _rebuild_rechunk(client, documents): yield doc_id + elif mode == RebuildMode.DESCRIPTIONS: + await client.chunk_repository.delete_all() + await client.store.recreate_embeddings_table() + async for doc_id in _rebuild_descriptions(client, documents): + yield doc_id else: # FULL await client.chunk_repository.delete_all() await client.store.recreate_embeddings_table() @@ -275,6 +280,149 @@ async def _rebuild_rechunk( yield doc_id +async def _patch_picture_descriptions(client: "HaikuRAG", doc: Document) -> int: + """Run the VLM against pictures lacking a description, patch the docling + blob in-place. Returns the number of newly described pictures. + Pictures that already carry ``meta.description.text`` are skipped, so the + operation is safe to re-run after a partial failure. + """ + from haiku.rag.providers.picture_description import describe_pictures + + assert doc.id is not None + docling_doc = doc.get_docling_document() + if docling_doc is None or not docling_doc.pictures: + return 0 + + needs_description: list[str] = [] + for pic in docling_doc.pictures: + meta = getattr(pic, "meta", None) + existing = ( + getattr(getattr(meta, "description", None), "text", None) if meta else None + ) + if not (isinstance(existing, str) and existing.strip()): + needs_description.append(pic.self_ref) + + if not needs_description: + return 0 + + bytes_by_ref = await client.document_item_repository.get_pictures_for_chunk( + doc.id, needs_description + ) + if not bytes_by_ref: + logger.warning( + "Document %s has %d pictures missing descriptions but no stored " + "picture bytes — skipping. Run a full rebuild from source to " + "recover the bytes.", + doc.id, + len(needs_description), + ) + return 0 + + descriptions = await describe_pictures(bytes_by_ref, config=client._config) + + if not descriptions: + return 0 + + # Patch the docling document in-place. PictureMeta + DescriptionMetaField + # are pydantic models; build them and assign. + from docling_core.types.doc.document import ( + DescriptionMetaField, + PictureMeta, + ) + + for pic in docling_doc.pictures: + text = descriptions.get(pic.self_ref) + if not text: + continue + if pic.meta is None: + pic.meta = PictureMeta() + pic.meta.description = DescriptionMetaField(text=text) + + doc.set_docling(docling_doc) + return len(descriptions) + + +async def _rebuild_descriptions( + client: "HaikuRAG", documents: list[Document] +) -> AsyncGenerator[str, None]: + """Run the VLM over already-stored picture bytes, patch descriptions into + the docling blob, then re-chunk + re-embed. + + Skips the docling parse entirely (the blob is already there); only the VLM + cost remains. Idempotent: pictures whose ``meta.description.text`` is + already populated are not re-described. + """ + from haiku.rag.embeddings import embed_chunks, get_embedder + + if not client._config.processing.conversion_options.picture_description.enabled: + raise ValueError( + "rebuild --descriptions requires " + "processing.conversion_options.picture_description.enabled = true " + "in your config." + ) + + pending_chunks: list[Chunk] = [] + pending_docs: list[Document] = [] + pending_doc_ids: list[str] = [] + embedder = get_embedder(client._config) + + described_total = 0 + for doc in documents: + assert doc.id is not None + + docling_document = doc.get_docling_document() + if docling_document is None: + raise ValueError( + f"Document {doc.id} has no stored docling document; " + "rebuild --descriptions requires it. Run a full rebuild instead." + ) + + n = await _patch_picture_descriptions(client, doc) + described_total += n + # Use the (possibly patched) docling document for chunking. + docling_document = doc.get_docling_document() + assert docling_document is not None + + existing_picture_data = ( + await client.document_item_repository.get_all_picture_data(doc.id) + if embedder.supports_images + else None + ) + chunks = await client.chunk( + docling_document, + existing_picture_data=existing_picture_data, + document_id=doc.id, + ) + embedded_chunks = await embed_chunks(chunks, client._config) + + for order, chunk in enumerate(embedded_chunks): + chunk.document_id = doc.id + chunk.order = order + + pending_chunks.extend(embedded_chunks) + pending_docs.append(doc) + pending_doc_ids.append(doc.id) + + if len(pending_docs) >= _REBUILD_BATCH_SIZE: + await _flush_rebuild_batch(client, pending_docs, pending_chunks) + for doc_id in pending_doc_ids: + yield doc_id + pending_chunks = [] + pending_docs = [] + pending_doc_ids = [] + + if pending_docs: + await _flush_rebuild_batch(client, pending_docs, pending_chunks) + for doc_id in pending_doc_ids: + yield doc_id + + logger.info( + "rebuild --descriptions: %d new picture descriptions added across %d documents", + described_total, + len(documents), + ) + + async def _rebuild_full( client: "HaikuRAG", documents: list[Document] ) -> AsyncGenerator[str, None]: diff --git a/haiku_rag_slim/haiku/rag/providers/picture_description.py b/haiku_rag_slim/haiku/rag/providers/picture_description.py new file mode 100644 index 00000000..e38b04f3 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/providers/picture_description.py @@ -0,0 +1,52 @@ +"""Direct VLM client for picture description, used by ``rebuild --descriptions``. + +The docling-serve converter normally drives picture description as a side-effect +of conversion. When we need to run the VLM against pictures already stored in +the DB (skipping the docling parse entirely), we drive the VLM through +pydantic-ai with ``BinaryContent`` parts so model construction goes through the +same ``get_model`` plumbing as every other agent in the codebase. +""" + +import logging + +from pydantic_ai import Agent +from pydantic_ai.messages import BinaryContent + +from haiku.rag.config import AppConfig +from haiku.rag.utils import get_model + +logger = logging.getLogger(__name__) + + +async def describe_pictures( + image_bytes_by_ref: dict[str, bytes], + *, + config: AppConfig, +) -> dict[str, str]: + """Describe pictures sequentially; returns ``{self_ref: text}``. + + Pictures whose VLM call fails or returns empty content are silently + dropped from the returned map so the caller can decide whether the + partial result is acceptable. + """ + pic_desc = config.processing.conversion_options.picture_description + model = get_model(pic_desc.model, config) + prompt = config.prompts.picture_description + agent: Agent[None, str] = Agent( + model=model, + output_type=str, + instructions=prompt, + ) + + out: dict[str, str] = {} + for ref, blob in image_bytes_by_ref.items(): + try: + result = await agent.run([BinaryContent(data=blob, media_type="image/png")]) + except Exception as e: + logger.warning("VLM call failed for %s: %s", ref, e) + continue + text = (result.output or "").strip() + if text: + out[ref] = text + + return out diff --git a/tests/test_rebuild.py b/tests/test_rebuild.py index d283c06f..7b7d0e57 100644 --- a/tests/test_rebuild.py +++ b/tests/test_rebuild.py @@ -425,3 +425,133 @@ async def test_rebuild_batch_size_flush(temp_db_path, monkeypatch): for doc_id in ids: chunks = await client.chunk_repository.get_by_document_id(doc_id) assert len(chunks) > 0 + + +async def test_rebuild_descriptions_requires_enabled(temp_db_path): + """Calling rebuild --descriptions without picture_description.enabled is + a config error: the user has nothing to gain and the resulting state is + indistinguishable from a plain --rechunk.""" + async with HaikuRAG(temp_db_path, create=True) as client: + with pytest.raises(ValueError, match="picture_description.enabled"): + async for _ in client.rebuild_database(mode=RebuildMode.DESCRIPTIONS): + pass + + +async def test_rebuild_descriptions_patches_blob_and_chunks(temp_db_path, monkeypatch): + """End-to-end: ingest a doc with a picture (no VLM at ingest), then run + rebuild --descriptions with the VLM mocked. The docling blob should gain + the description in meta, and the chunk text should pick it up.""" + from haiku.rag.client.documents import _store_document_with_chunks + from haiku.rag.config import AppConfig + from haiku.rag.store.models.document import Document + from tests.store.test_document_items import _docling_doc_with_picture + + docling_doc = _docling_doc_with_picture() + + config = AppConfig() + config.processing.conversion_options.picture_description.enabled = True + + async with HaikuRAG(temp_db_path, config=config, create=True) as rag: + document = Document(content="x", uri="test://doc") + document.set_docling(docling_doc) + created = await _store_document_with_chunks(rag, document, [], docling_doc) + assert created.id is not None + + from_blob = ( + await rag.document_repository.get_by_id(created.id) + ).get_docling_document() # type: ignore[union-attr] + assert from_blob is not None and from_blob.pictures + # No description in the freshly-ingested doc + meta = from_blob.pictures[0].meta + existing = ( + getattr(getattr(meta, "description", None), "text", None) if meta else None + ) + assert not existing + + async def fake_describe(image_bytes_by_ref, *, config): + return {ref: "A red square (mocked)." for ref in image_bytes_by_ref} + + monkeypatch.setattr( + "haiku.rag.client.rebuild.describe_pictures", fake_describe, raising=False + ) + # The function is imported lazily inside _patch_picture_descriptions, so + # patch the module-of-origin too. + monkeypatch.setattr( + "haiku.rag.providers.picture_description.describe_pictures", + fake_describe, + ) + + processed = [ + doc_id + async for doc_id in rag.rebuild_database(mode=RebuildMode.DESCRIPTIONS) + ] + assert created.id in processed + + # The stored docling blob now has the description + after = await rag.document_repository.get_by_id(created.id) + assert after is not None + after_doc = after.get_docling_document() + assert after_doc is not None and after_doc.pictures + meta = after_doc.pictures[0].meta + text = ( + getattr(getattr(meta, "description", None), "text", None) if meta else None + ) + assert text == "A red square (mocked)." + + # And the description reaches chunk text + chunks = await rag.chunk_repository.get_by_document_id(created.id) + assert any("A red square (mocked)." in (c.content or "") for c in chunks) + + +async def test_rebuild_descriptions_skips_already_described(temp_db_path, monkeypatch): + """Pictures that already carry a description must not be re-sent to the + VLM, so the operation is safe to re-run after a partial failure.""" + from haiku.rag.client.documents import _store_document_with_chunks + from haiku.rag.config import AppConfig + from haiku.rag.store.models.document import Document + from tests.store.test_document_items import _docling_doc_with_picture + + docling_doc = _docling_doc_with_picture() + + # Pre-populate the description directly on the docling document + from docling_core.types.doc.document import DescriptionMetaField, PictureMeta + + docling_doc.pictures[0].meta = PictureMeta( + description=DescriptionMetaField(text="Pre-existing description.") + ) + + config = AppConfig() + config.processing.conversion_options.picture_description.enabled = True + + async with HaikuRAG(temp_db_path, config=config, create=True) as rag: + document = Document(content="x", uri="test://doc") + document.set_docling(docling_doc) + created = await _store_document_with_chunks(rag, document, [], docling_doc) + assert created.id is not None + + called_with: list[dict[str, bytes]] = [] + + async def fake_describe(image_bytes_by_ref, *, config): + called_with.append(image_bytes_by_ref) + return {ref: "Should not be used." for ref in image_bytes_by_ref} + + monkeypatch.setattr( + "haiku.rag.providers.picture_description.describe_pictures", + fake_describe, + ) + + async for _ in rag.rebuild_database(mode=RebuildMode.DESCRIPTIONS): + pass + + # VLM was never called for this picture (it already had a description) + assert called_with == [] or all(not d for d in called_with) + + after = await rag.document_repository.get_by_id(created.id) + assert after is not None + after_doc = after.get_docling_document() + assert after_doc is not None + meta = after_doc.pictures[0].meta + text = ( + getattr(getattr(meta, "description", None), "text", None) if meta else None + ) + assert text == "Pre-existing description."