feat(multimodal): add mm_assets image retrieval via vLLM

Phase 1: orthogonal multimodal index (mm_assets) for text→image and image→image search.

- Add multimodal config + vLLM embedder client
- Index Docling PictureItems (bbox crops) into mm_assets with rollback support
- Add CLI commands: search-image-text, search-image, visualize-asset
- Extend inspector with mm_assets browser + doc index overlay
- Add multimodal eval commands (mm-build/mm) and docs
This commit is contained in:
jvdrag 2026-01-11 05:51:06 -03:00
parent ff1f19bb08
commit 79a3617782
25 changed files with 2673 additions and 3 deletions

View file

@ -104,6 +104,34 @@ haiku-rag visualize <chunk_id>
This renders the source document pages with the chunk's location highlighted. Useful for verifying chunk boundaries and understanding document structure. This renders the source document pages with the chunk's location highlighted. Useful for verifying chunk boundaries and understanding document structure.
!!! note
Requires a terminal with image support (iTerm2, Kitty, WezTerm, etc.) and documents processed with docling that have page images stored.
## Search Images (Multimodal)
These commands require `multimodal.enabled: true` and an embedding server configured under `multimodal.model` (typically vLLM).
### Text → image
```bash
haiku-rag search-image-text "architecture diagram"
```
### Image → image
```bash
haiku-rag search-image /path/to/query.png
```
## Visualize Asset (Multimodal)
After `search-image*`, you can visualize an `asset_id`:
```bash
haiku-rag visualize-asset <asset_id> --mode crop
haiku-rag visualize-asset <asset_id> --mode page
```
!!! note !!! note
Requires a terminal with image support (iTerm2, Kitty, WezTerm, etc.) and documents processed with docling that have page images stored. Requires a terminal with image support (iTerm2, Kitty, WezTerm, etc.) and documents processed with docling that have page images stored.

View file

@ -190,6 +190,27 @@ embeddings:
**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.
## Multimodal Embeddings (vLLM)
Multimodal embeddings are configured separately under `multimodal` and are used for **image search** (not text chunk search).
```yaml
multimodal:
enabled: true
model:
provider: vllm
name: Qwen/Qwen3-VL-Embedding-2B
vector_dim: 2048
base_url: http://localhost:8000 # accepts with or without trailing /v1
timeout: 60
encoding_format: float
# dimensions: 2048 # omit unless your backend/model supports matryoshka output
```
!!! note
Unlike `embeddings.model.base_url` (OpenAI provider), `multimodal.model.base_url` can be either `http://host:port`
or `http://host:port/v1`. haiku.rag will normalize it internally.
## Question Answering Providers ## Question Answering Providers
Configure which LLM provider to use for question answering. Any provider and model supported by [Pydantic AI](https://ai.pydantic.dev/models/) can be used. Configure which LLM provider to use for question answering. Any provider and model supported by [Pydantic AI](https://ai.pydantic.dev/models/) can be used.

View file

@ -38,6 +38,8 @@ Three panels display your data:
- `/` - Open search modal - `/` - Open search modal
- `c` - Open context expansion modal (when viewing a chunk) - `c` - Open context expansion modal (when viewing a chunk)
- `v` - Open visual grounding modal (when viewing a chunk) - `v` - Open visual grounding modal (when viewing a chunk)
- `m` - Browse multimodal image assets (`mm_assets`) for the selected document
- `d` - Show whole-document index overlay (all chunk boxes + all image asset boxes)
- `q` - Quit - `q` - Quit
**Mouse:** Click to select, scroll to view content **Mouse:** Click to select, scroll to view content
@ -71,7 +73,31 @@ Press `v` while viewing a chunk to open the visual grounding modal:
- Shows page images from the source document with the chunk's location highlighted in yellow/orange - Shows page images from the source document with the chunk's location highlighted in yellow/orange
- Use `←` / `→` arrow keys to navigate between pages (when chunk spans multiple pages) - Use `←` / `→` arrow keys to navigate between pages (when chunk spans multiple pages)
- Press `o` to open the current page image in your OS image viewer (recommended for terminals that don't render images inline)
- Press `Esc` to close the modal - Press `Esc` to close the modal
!!! note !!! note
Visual grounding requires documents with a stored DoclingDocument that includes page images. Text-only documents or documents imported without DoclingDocument won't have visual grounding available. Visual and export images are written to a temporary directory and **deleted when you close the modal**.
If you want to keep an image, press `o` to open it and save/copy it from your OS viewer before closing.
## Multimodal Assets (mm_assets)
Press `m` to open the multimodal assets modal for the currently selected document:
- Lists all `mm_assets` rows (one per indexed picture crop)
- Press `p` to toggle between `crop` and `page` modes
- Press `o` to export + open the current selection in your OS image viewer
(exports are stored in a temporary directory and deleted when you close the modal)
## Whole-Document Index Overlay
Press `d` to open the whole-document overlay:
- Exports per-page PNGs with **all chunk bounding boxes** (yellow/orange) and **all mm_asset bounding boxes** (cyan/blue)
- Use `←` / `→` to navigate pages
- Press `o` to open the current page PNG in your OS viewer
(exports are stored in a temporary directory and deleted when you close the modal)
!!! note
Visual grounding requires documents with a stored DoclingDocument that includes page images.
Text-only documents or documents imported without DoclingDocument won't have visual grounding available.

175
docs/multimodal.md Normal file
View file

@ -0,0 +1,175 @@
# Multimodal Search (Images)
haiku.rag supports **multimodal embeddings** for:
- **text → image** search (`mm_assets` table)
- **image → image** search (`mm_assets` table)
This feature is **orthogonal** to the existing text chunk search (`chunks` table). Your existing `haiku-rag search` behavior does not change.
---
## What gets indexed (Phase 1)
When enabled, haiku.rag extracts Docling `PictureItem`s and indexes:
- An **image crop** from the Docling page image using the item bbox (plus padding).
- Optional text metadata if available (caption/description) stored alongside the vector.
Each indexed image becomes one row in `mm_assets`, linked to its `document_id` and bbox provenance (page + coordinates), so it can be visualized later.
---
## Requirements
- Your converter must produce **Docling page images** (needed for bbox crops and visualization).
- A multimodal embedding backend must be reachable.
**Current supported backend**
- `provider: vllm` via OpenAI-compatible `POST /v1/embeddings`
---
## Configuration
Add the following to your `haiku.rag.yaml`:
```yaml
multimodal:
enabled: true
index_pictures: true
# bbox → crop settings
image_crop_padding_px: 8
# guardrails / throughput knobs
image_max_side_px: 1024
embed_batch_size: 8
model:
provider: vllm
name: Qwen/Qwen3-VL-Embedding-2B
vector_dim: 2048
base_url: http://localhost:8000 # with or without trailing /v1
timeout: 60
# dimensions: 1024 # only if your backend/model supports it
encoding_format: float # float | base64
```
### Multimodal toggles
- **`multimodal.enabled`**: master switch. When `false`, haiku.rag will not create/use `mm_assets`.
- **`multimodal.index_pictures`**: when `true`, Docling `PictureItem`s are indexed during `add-src` / update.
### Crop / image settings
- **`image_crop_padding_px`**: extra pixels added around the bbox crop. Helps avoid tight crops clipping labels/axes.
- **`image_max_side_px`**: resize guardrail before upload (keeps requests smaller, avoids backend limits). Set `0` to disable resizing.
### Throughput settings
- **`embed_batch_size`**:
- Used as the **maximum in-flight requests** for image embedding to avoid overwhelming the backend.
- Text embeddings may be batched using `input=[...]`, but image embeddings currently run as **one request per image** (because the backend expects multimodal inputs via `messages`).
### Model / backend settings
- **`provider`**: currently only `vllm` is supported (OpenAI-compatible server).
- **`name`**: model identifier sent to the backend (e.g. `Qwen/Qwen3-VL-Embedding-2B`).
- **`vector_dim`**: expected embedding size. This must stay stable for an existing DB (same rule as text embeddings).
- **`base_url`**: server base URL; haiku.rag accepts either `http://host:port` or `http://host:port/v1`.
- **`timeout`**: request timeout in seconds.
- **`dimensions`**: optional request-time output dimension (only for models that support it).
- **`encoding_format`**: `float` is recommended; `base64` is supported by some servers.
!!! note
Some servers/models reject the presence of `dimensions` unless matryoshka output is supported.
For vLLM + Qwen3VLEmbedding2B, the safe default is to **omit** it and validate the returned vector length (2048).
---
## Using a different model (same backend)
To switch models (still using vLLM), change:
- `multimodal.model.name`
- `multimodal.model.vector_dim`
Best practice: **verify the output dimension** before ingesting a large corpus:
```bash
curl -s http://127.0.0.1:8000/v1/embeddings \
-H "Content-Type: application/json" \
-d '{"model":"<YOUR_MODEL>","input":["hello"],"encoding_format":"float"}' \
| python3 -c 'import sys,json; print(len(json.load(sys.stdin)["data"][0]["embedding"]))'
```
---
## Run vLLM (example)
Serve the embedding model:
```bash
vllm serve Qwen/Qwen3-VL-Embedding-2B --runner=pooling --host 127.0.0.1 --port 8000
```
Quick checks:
```bash
curl -s http://127.0.0.1:8000/v1/models
curl -s http://127.0.0.1:8000/v1/embeddings \
-H "Content-Type: application/json" \
-d '{"model":"Qwen/Qwen3-VL-Embedding-2B","input":["hello"],"encoding_format":"float"}'
```
---
## Index images (Phase 1)
Index documents as usual (multimodal indexing runs during document create/update if enabled):
```bash
haiku-rag add-src /path/to/doc.pdf
```
---
## Search images
### Text → image
```bash
haiku-rag search-image-text "architecture diagram"
```
### Image → image
```bash
haiku-rag search-image /path/to/query.png
```
---
## Visualize results
Search returns `asset_id` values. To see what was actually indexed:
```bash
haiku-rag visualize-asset <asset_id> --mode crop
haiku-rag visualize-asset <asset_id> --mode page
```
!!! note
Terminals often do not render images inline (you may see gray blocks).
Prefer opening the saved PNG in your OS viewer via:
- `haiku-rag visualize-asset ...` (prints the saved image path)
- `haiku-rag inspect` (`o` key inside the visual modals)
---
## Minimal accuracy validation (recommended)
Generate a quick “sanity” dataset from your DB and compute recall@k + MRR:
```bash
uv run evaluations mm-build --db /path/to/your.lancedb --config /path/to/haiku.rag.yaml --out ./mm_eval_out --n 50
uv run evaluations mm ./mm_eval_out/mm_eval.jsonl --db /path/to/your.lancedb --config /path/to/haiku.rag.yaml --k 1,5,10 --limit 10
```

View file

@ -22,6 +22,8 @@ from haiku.rag.config.models import ModelConfig
from haiku.rag.logging import configure_cli_logging from haiku.rag.logging import configure_cli_logging
from haiku.rag.qa import get_qa_agent from haiku.rag.qa import get_qa_agent
from haiku.rag.utils import get_model from haiku.rag.utils import get_model
from evaluations.mm_eval import run_mm_eval_sync
from evaluations.mm_dataset_builder import build_mm_dataset_sync
load_dotenv() load_dotenv()
@ -409,5 +411,124 @@ def run(
) )
@app.command()
def mm(
dataset: Path = typer.Argument(..., help="Path to multimodal eval dataset JSONL."),
config: Path | None = typer.Option(
None, "--config", help="Path to haiku.rag YAML config file."
),
db: Path | None = typer.Option(None, "--db", help="Override the database path."),
k: str = typer.Option(
"1,5,10", "--k", help="Comma-separated recall@k values to report."
),
limit: int = typer.Option(
10, "--limit", help="Top-N retrieved mm_assets per query."
),
) -> None:
"""Run multimodal retrieval benchmarks (mm_assets) and report recall@k + MRR.
Dataset format is JSONL; each line is a case:
- type: "text" | "image"
- instruction: optional string (recommended for Qwen3-VL)
- query_text or query_image_path
- relevant: [{document_uri, doc_item_ref, item_index?}, ...]
"""
if not dataset.exists():
raise typer.BadParameter(f"Dataset file not found: {dataset}")
ks = [int(x.strip()) for x in k.split(",") if x.strip()]
if not ks:
raise typer.BadParameter("--k must include at least one integer.")
# Load config from file or use defaults (mirror `run` command behavior).
if config:
if not config.exists():
raise typer.BadParameter(f"Config file not found: {config}")
console.print(f"Loading config from: {config}", style="dim")
yaml_data = load_yaml_config(config)
app_config = AppConfig.model_validate(yaml_data)
else:
config_path = find_config_file(None)
if config_path:
console.print(f"Loading config from: {config_path}", style="dim")
yaml_data = load_yaml_config(config_path)
app_config = AppConfig.model_validate(yaml_data)
else:
console.print("No config file found, using defaults", style="dim")
app_config = AppConfig()
raise SystemExit(
run_mm_eval_sync(
dataset_path=dataset,
config=app_config,
db_path=db,
ks=ks,
limit=int(limit),
)
)
@app.command()
def mm_build(
out: Path = typer.Option(
Path("./mm_eval_out"),
"--out",
help="Output directory (writes mm_eval.jsonl + query_images/ crops).",
),
config: Path | None = typer.Option(
None, "--config", help="Path to haiku.rag YAML config file."
),
db: Path | None = typer.Option(None, "--db", help="Database path (LanceDB dir)."),
n: int = typer.Option(50, "--n", help="How many mm_assets rows to sample."),
seed: int = typer.Option(0, "--seed", help="Random seed."),
include_text: bool = typer.Option(True, "--include-text/--no-include-text"),
include_image: bool = typer.Option(True, "--include-image/--no-include-image"),
instruction: str = typer.Option(
"Retrieve images matching this description.",
"--instruction",
help="Instruction prefix for text queries (folded into query_text).",
),
) -> None:
"""Build a *sanity* multimodal eval dataset from an existing DB.
This generates:
- textimage cases using mm_assets.caption/description when available
- imageimage cases by exporting bbox crops from Docling page images
"""
# Load config from file or use defaults (mirror `run` command behavior).
if config:
if not config.exists():
raise typer.BadParameter(f"Config file not found: {config}")
console.print(f"Loading config from: {config}", style="dim")
yaml_data = load_yaml_config(config)
app_config = AppConfig.model_validate(yaml_data)
else:
config_path = find_config_file(None)
if config_path:
console.print(f"Loading config from: {config_path}", style="dim")
yaml_data = load_yaml_config(config_path)
app_config = AppConfig.model_validate(yaml_data)
else:
console.print("No config file found, using defaults", style="dim")
app_config = AppConfig()
dataset_path = build_mm_dataset_sync(
config=app_config,
db_path=db,
out_dir=out,
n=int(n),
seed=int(seed),
include_text=bool(include_text),
include_image=bool(include_image),
instruction=str(instruction),
)
console.print("\nNow run:", style="bold")
console.print(
f"uv run evaluations mm {dataset_path} --config <config> --db <db>",
style="cyan",
)
if __name__ == "__main__": if __name__ == "__main__":
app() app()

View file

@ -0,0 +1,3 @@
{"id":"t1","type":"text","instruction":"Retrieve images matching this description.","query_text":"architecture diagram","relevant":[{"document_uri":"file:///path/to/your.pdf","doc_item_ref":"#/pictures/0","item_index":123}]}
{"id":"i1","type":"image","query_image_path":"/absolute/path/to/query_crop.png","relevant":[{"document_uri":"file:///path/to/your.pdf","doc_item_ref":"#/pictures/0","item_index":123}]}

View file

@ -0,0 +1,239 @@
import asyncio
import json
import random
from dataclasses import dataclass
from pathlib import Path
from rich.console import Console
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig
console = Console()
@dataclass(frozen=True)
class BuiltCase:
id: str
type: str # "text" | "image"
instruction: str | None
query_text: str | None
query_image_path: str | None
relevant: list[dict]
def to_jsonl(self) -> str:
obj = {
"id": self.id,
"type": self.type,
"relevant": self.relevant,
}
if self.instruction is not None:
obj["instruction"] = self.instruction
if self.query_text is not None:
obj["query_text"] = self.query_text
if self.query_image_path is not None:
obj["query_image_path"] = self.query_image_path
return json.dumps(obj, ensure_ascii=False)
def _crop_from_doc_bbox(
*,
pil_image,
page_width: float,
page_height: float,
bbox: dict,
padding_px: int,
):
# bbox is in Docling page coordinates with bottom-left origin.
left = float(bbox["left"])
top = float(bbox["top"])
right = float(bbox["right"])
bottom = float(bbox["bottom"])
scale_x = pil_image.width / page_width
scale_y = pil_image.height / page_height
x0 = left * scale_x
x1 = right * scale_x
# invert Y axis: doc bottom-left -> PIL top-left
y0 = (page_height - top) * scale_y
y1 = (page_height - bottom) * scale_y
if y0 > y1:
y0, y1 = y1, y0
x0 = max(0, int(x0) - padding_px)
y0 = max(0, int(y0) - padding_px)
x1 = min(pil_image.width, int(x1) + padding_px)
y1 = min(pil_image.height, int(y1) + padding_px)
if x1 <= x0 or y1 <= y0:
return None
return pil_image.crop((x0, y0, x1, y1))
async def build_mm_dataset(
*,
config: AppConfig,
db_path: Path | None,
out_dir: Path,
n: int = 50,
seed: int = 0,
include_text: bool = True,
include_image: bool = True,
instruction: str = "Retrieve images matching this description.",
) -> Path:
"""Auto-generate a JSONL dataset for `evaluations mm` from an existing DB.
This produces a *sanity* dataset:
- imageimage queries use the exact stored bbox crop (self-retrieval)
- textimage queries use caption/description when available (self-retrieval)
You can then hand-edit queries/relevants for a more realistic benchmark.
"""
out_dir.mkdir(parents=True, exist_ok=True)
crops_dir = out_dir / "query_images"
crops_dir.mkdir(parents=True, exist_ok=True)
# Optional dependency: Pillow is required to write crops.
try:
from PIL import Image # noqa: F401
except Exception as e:
raise ImportError(
"Building a multimodal eval dataset requires Pillow. "
"Install it (e.g. `uv pip install pillow`)."
) from e
async with HaikuRAG(db_path, config=config) as rag:
if not rag._config.multimodal.enabled:
raise ValueError(
"Multimodal is disabled in config. Set `multimodal.enabled: true` "
"and open the DB in writable mode at least once to create mm_assets."
)
if rag.store.mm_assets_table is None:
raise ValueError(
"mm_assets table is not available. Enable multimodal and open the DB "
"in writable mode at least once to create it."
)
# Recommended way to materialize the *entire table* is via Arrow/Pandas.
# We prefer Arrow here to avoid any implicit limits in query paths.
table = rag.store.mm_assets_table
try:
rows = table.to_arrow().to_pylist()
except Exception:
# Fallback: scan via query builder.
rows = table.search().limit(100_000).to_list()
if not rows:
raise ValueError("mm_assets is empty; index documents with pictures first.")
rnd = random.Random(seed)
rnd.shuffle(rows)
cases: list[BuiltCase] = []
used = 0
for row in rows:
if used >= n:
break
asset_id = row.get("id")
doc_id = row.get("document_id")
doc_item_ref = row.get("doc_item_ref")
item_index = row.get("item_index")
page_no = row.get("page_no")
bbox_raw = row.get("bbox")
caption = row.get("caption")
description = row.get("description")
if not asset_id or not doc_id or not doc_item_ref:
continue
doc = await rag.document_repository.get_by_id(str(doc_id))
if doc is None or not doc.uri:
continue
relevant = [
{
"document_uri": doc.uri,
"doc_item_ref": str(doc_item_ref),
"item_index": int(item_index) if item_index is not None else None,
}
]
# text→image case (when we have any text signal)
if include_text:
q = None
if caption:
q = str(caption).strip()
elif description:
q = str(description).strip()
if q:
cases.append(
BuiltCase(
id=f"auto-text-{asset_id}",
type="text",
instruction=instruction,
query_text=q,
query_image_path=None,
relevant=relevant,
)
)
# image→image case (requires bbox + page image)
if include_image and page_no is not None and bbox_raw:
try:
bbox = json.loads(bbox_raw) if isinstance(bbox_raw, str) else bbox_raw
except Exception:
bbox = None
if isinstance(bbox, dict) and all(
k in bbox for k in ("left", "top", "right", "bottom")
):
docling_doc = doc.get_docling_document()
if docling_doc and int(page_no) in docling_doc.pages:
page = docling_doc.pages[int(page_no)]
if (
page.image is not None
and page.image.pil_image is not None
and page.size is not None
):
crop = _crop_from_doc_bbox(
pil_image=page.image.pil_image,
page_width=float(page.size.width),
page_height=float(page.size.height),
bbox=bbox,
padding_px=int(rag._config.multimodal.image_crop_padding_px),
)
if crop is not None:
crop_path = crops_dir / f"{asset_id}.png"
crop.save(crop_path, format="PNG")
cases.append(
BuiltCase(
id=f"auto-image-{asset_id}",
type="image",
instruction=None,
query_text=None,
query_image_path=str(crop_path.resolve()),
relevant=relevant,
)
)
used += 1
dataset_path = out_dir / "mm_eval.jsonl"
dataset_path.write_text(
"\n".join(c.to_jsonl() for c in cases) + "\n", encoding="utf-8"
)
console.print(
f"Wrote {len(cases)} cases to {dataset_path} "
f"(from {min(n, len(rows))} sampled mm_assets rows).",
style="green",
)
return dataset_path
def build_mm_dataset_sync(**kwargs) -> Path:
return asyncio.run(build_mm_dataset(**kwargs))

View file

@ -0,0 +1,263 @@
import asyncio
import json
from dataclasses import dataclass
from pathlib import Path
from rich.console import Console
from haiku.rag.client import HaikuRAG
from haiku.rag.config import AppConfig
console = Console()
@dataclass(frozen=True)
class MMRelevantRef:
"""A stable reference to an mm_asset row (avoid UUIDs in datasets)."""
document_uri: str
doc_item_ref: str
item_index: int | None = None
@dataclass(frozen=True)
class MMEvalCase:
"""One multimodal retrieval evaluation case."""
case_id: str
query_type: str # "text" | "image"
instruction: str | None
query_text: str | None
query_image_path: str | None
relevant: tuple[MMRelevantRef, ...]
@dataclass(frozen=True)
class CaseMetrics:
recall_at_k: dict[int, float]
mrr: float
def _quote_lancedb(s: str) -> str:
# LanceDB uses SQL-ish string quoting with single quotes.
return s.replace("'", "''")
def load_mm_dataset(dataset_path: Path) -> list[MMEvalCase]:
cases: list[MMEvalCase] = []
for raw_line in dataset_path.read_text(encoding="utf-8").splitlines():
line = raw_line.strip()
if not line or line.startswith("#"):
continue
obj = json.loads(line)
case_id = str(obj.get("id") or obj.get("case_id") or "")
if not case_id:
raise ValueError("Each JSONL row must have a non-empty 'id'.")
query_type = obj.get("type")
if query_type not in ("text", "image"):
raise ValueError(f"Case {case_id}: 'type' must be 'text' or 'image'.")
instruction = obj.get("instruction")
query_text = obj.get("query_text")
query_image_path = obj.get("query_image_path")
if query_type == "text" and not query_text:
raise ValueError(f"Case {case_id}: missing 'query_text'.")
if query_type == "image" and not query_image_path:
raise ValueError(f"Case {case_id}: missing 'query_image_path'.")
rel = obj.get("relevant") or []
if not isinstance(rel, list) or not rel:
raise ValueError(f"Case {case_id}: 'relevant' must be a non-empty list.")
relevant: list[MMRelevantRef] = []
for r in rel:
if not isinstance(r, dict):
raise ValueError(f"Case {case_id}: relevant entries must be objects.")
document_uri = r.get("document_uri")
doc_item_ref = r.get("doc_item_ref")
if not document_uri or not doc_item_ref:
raise ValueError(
f"Case {case_id}: relevant entries must include "
"'document_uri' and 'doc_item_ref'."
)
item_index = r.get("item_index")
relevant.append(
MMRelevantRef(
document_uri=str(document_uri),
doc_item_ref=str(doc_item_ref),
item_index=int(item_index) if item_index is not None else None,
)
)
cases.append(
MMEvalCase(
case_id=case_id,
query_type=query_type,
instruction=str(instruction) if instruction is not None else None,
query_text=query_text,
query_image_path=query_image_path,
relevant=tuple(relevant),
)
)
return cases
async def _resolve_relevant_asset_ids(
rag: HaikuRAG, cases: list[MMEvalCase]
) -> dict[str, set[str]]:
"""Resolve stable refs (uri + doc_item_ref [+ item_index]) into asset UUIDs."""
if rag.store.mm_assets_table is None:
raise ValueError("mm_assets table is not available (multimodal not enabled?).")
# Resolve document_id per uri.
uris = sorted({r.document_uri for c in cases for r in c.relevant})
uri_to_doc_id: dict[str, str] = {}
for uri in uris:
doc = await rag.get_document_by_uri(uri)
if doc is None or not doc.id:
raise ValueError(f"Dataset references unknown document uri: {uri}")
uri_to_doc_id[uri] = doc.id
# Load mm_assets rows once (evaluation datasets are small; this is simplest).
# Recommended way to materialize the *entire table* is via Arrow/Pandas.
table = rag.store.mm_assets_table
try:
rows = table.to_arrow().to_pylist()
except Exception:
rows = table.search().limit(100_000).to_list()
# Build lookup maps.
by_doc_ref_idx: dict[tuple[str, str, int], str] = {}
by_doc_ref: dict[tuple[str, str], list[str]] = {}
for row in rows:
asset_id = row.get("id")
doc_id = row.get("document_id")
doc_item_ref = row.get("doc_item_ref")
item_index = row.get("item_index")
if not asset_id or not doc_id or not doc_item_ref:
continue
if item_index is not None:
try:
by_doc_ref_idx[(str(doc_id), str(doc_item_ref), int(item_index))] = str(
asset_id
)
except Exception:
pass
by_doc_ref.setdefault((str(doc_id), str(doc_item_ref)), []).append(str(asset_id))
case_to_relevant_ids: dict[str, set[str]] = {}
for c in cases:
ids: set[str] = set()
for r in c.relevant:
doc_id = uri_to_doc_id[r.document_uri]
if r.item_index is not None:
k = (doc_id, r.doc_item_ref, int(r.item_index))
asset_id = by_doc_ref_idx.get(k)
if asset_id:
ids.add(asset_id)
continue
# Fallback: any asset(s) matching doc+ref.
ids.update(by_doc_ref.get((doc_id, r.doc_item_ref), []))
if not ids:
raise ValueError(
f"Case {c.case_id}: could not resolve any relevant asset ids from refs. "
"Make sure the DB was built with multimodal indexing enabled and the "
"doc_item_ref/item_index match your Docling conversion."
)
case_to_relevant_ids[c.case_id] = ids
return case_to_relevant_ids
def _compute_case_metrics(
*, retrieved_ids: list[str], relevant_ids: set[str], ks: list[int]
) -> CaseMetrics:
recall_at_k: dict[int, float] = {}
for k in ks:
topk = retrieved_ids[:k]
hits = sum(1 for a in topk if a in relevant_ids)
recall_at_k[k] = hits / max(len(relevant_ids), 1)
rr = 0.0
for idx, a in enumerate(retrieved_ids, start=1):
if a in relevant_ids:
rr = 1.0 / idx
break
return CaseMetrics(recall_at_k=recall_at_k, mrr=rr)
async def run_mm_eval(
*,
dataset_path: Path,
config: AppConfig,
db_path: Path | None,
ks: list[int],
limit: int,
) -> int:
cases = load_mm_dataset(dataset_path)
if not cases:
raise ValueError("No cases found.")
async with HaikuRAG(db_path, config=config) as rag:
if not rag._config.multimodal.enabled:
raise ValueError(
"Multimodal is disabled in config. Set `multimodal.enabled: true` "
"and ensure mm_assets is created."
)
# Resolve ground-truth.
relevant_ids = await _resolve_relevant_asset_ids(rag, cases)
totals_recall = {k: 0.0 for k in ks}
totals_mrr = 0.0
for c in cases:
if c.query_type == "text":
# The Qwen3-VL embedding examples use an "instruction" field that can
# materially affect retrieval quality. vLLM /v1/embeddings typically
# doesn't expose a separate instruction channel, so we fold it into
# the query text in a stable way for evaluation.
q = c.query_text or ""
if c.instruction:
q = f"{c.instruction}\n{q}"
results = await rag.search_images_by_text(
q, limit=limit
)
else:
img_path = Path(c.query_image_path or "")
if not img_path.exists():
raise ValueError(
f"Case {c.case_id}: query_image_path not found: {img_path}"
)
results = await rag.search_images(img_path, limit=limit)
retrieved = [r.asset_id for r in results if r.asset_id]
m = _compute_case_metrics(
retrieved_ids=retrieved,
relevant_ids=relevant_ids[c.case_id],
ks=ks,
)
for k in ks:
totals_recall[k] += m.recall_at_k[k]
totals_mrr += m.mrr
n = float(len(cases))
console.print("\n=== Multimodal Retrieval Benchmark Results ===", style="bold cyan")
console.print(f"Dataset: {dataset_path}")
console.print(f"Total queries: {len(cases)}")
for k in ks:
console.print(f"recall@{k}: {totals_recall[k] / n:.4f}")
console.print(f"mrr: {totals_mrr / n:.4f}")
return 0
def run_mm_eval_sync(**kwargs) -> int:
return asyncio.run(run_mm_eval(**kwargs))

View file

@ -367,8 +367,50 @@ class HaikuRAGApp:
for result in results: for result in results:
self._rich_print_search_result(result) self._rich_print_search_result(result)
async def search_image_text(
self, query: str, limit: int | None = None, filter: str | None = None
):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
results = await self.client.search_images_by_text(
query, limit=limit, filter=filter
)
if not results:
self.console.print("[yellow]No image results found.[/yellow]")
return
for result in results:
self._rich_print_mm_search_result(result)
async def search_image(
self, image_path: str, limit: int | None = None, filter: str | None = None
):
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
results = await self.client.search_images(
Path(image_path), limit=limit, filter=filter
)
if not results:
self.console.print("[yellow]No image results found.[/yellow]")
return
for result in results:
self._rich_print_mm_search_result(result)
async def visualize_chunk(self, chunk_id: str): async def visualize_chunk(self, chunk_id: str):
"""Display visual grounding images for a chunk.""" """Display visual grounding images for a chunk."""
import os
import subprocess
import sys
import tempfile
from pathlib import Path
from textual_image.renderable import Image as RichImage from textual_image.renderable import Image as RichImage
async with HaikuRAG( async with HaikuRAG(
@ -404,6 +446,87 @@ class HaikuRAGApp:
) )
self.console.print(RichImage(img)) self.console.print(RichImage(img))
# Also save images to disk for reliable viewing (many terminals can't display images).
out_dir = Path(tempfile.mkdtemp(prefix="haiku-visualize-chunk-"))
saved: list[Path] = []
for i, img in enumerate(images):
p = out_dir / f"chunk_{chunk_id}_page_{i+1}.png"
try:
img.save(p, format="PNG")
saved.append(p)
except Exception:
continue
if saved:
self.console.print(
f"\n[green]Saved {len(saved)} image(s) to[/green] {out_dir}"
)
# Best-effort: open the first image on macOS.
# (DISPLAY isn't set on macOS; `open` is still fine.)
if sys.platform == "darwin":
try:
subprocess.run(["open", str(saved[0])], check=False)
except Exception:
pass
async def visualize_asset(self, asset_id: str, mode: str = "crop"):
"""Display a multimodal asset (mm_assets) as an image in the terminal.
Modes:
- crop: show the cropped image region (bbox + padding)
- page: show the full page with bbox overlay (if available)
"""
import os
import subprocess
import sys
import tempfile
from pathlib import Path
from textual_image.renderable import Image as RichImage
async with HaikuRAG(
db_path=self.db_path,
config=self.config,
read_only=self.read_only,
before=self.before,
) as self.client:
images = await self.client.visualize_mm_asset(asset_id=asset_id, mode=mode)
if not images:
self.console.print(
"[yellow]No image available for this asset.[/yellow]\n"
"This may be because the document was converted without page images "
"or the asset has no bbox provenance."
)
return
self.console.print(f"[bold]Visualize asset {asset_id}[/bold]")
self.console.print(f"[repr.attrib_name]mode[/repr.attrib_name]: {mode}")
for i, img in enumerate(images):
self.console.print(
f"\n[bold cyan]Image {i + 1}/{len(images)}[/bold cyan]"
)
self.console.print(RichImage(img))
# Also save images to disk for reliable viewing (many terminals can't display images).
out_dir = Path(tempfile.mkdtemp(prefix="haiku-visualize-asset-"))
saved: list[Path] = []
for i, img in enumerate(images):
p = out_dir / f"asset_{asset_id}_{mode}_{i+1}.png"
try:
img.save(p, format="PNG")
saved.append(p)
except Exception:
continue
if saved:
self.console.print(
f"\n[green]Saved {len(saved)} image(s) to[/green] {out_dir}"
)
# Best-effort: open the first image on macOS.
if sys.platform == "darwin":
try:
subprocess.run(["open", str(saved[0])], check=False)
except Exception:
pass
async def ask( async def ask(
self, self,
question: str, question: str,
@ -811,6 +934,22 @@ class HaikuRAGApp:
self.console.print(" > ".join(result.headings)) self.console.print(" > ".join(result.headings))
self.console.print("[repr.attrib_name]content[/repr.attrib_name]:") self.console.print("[repr.attrib_name]content[/repr.attrib_name]:")
self.console.print(content) self.console.print(content)
def _rich_print_mm_search_result(self, result):
"""Format a multimodal search result for display."""
from rich.text import Text
self.console.rule()
header = Text(f"Score: {result.score:.4f} ", style="bold")
header.append(f"asset={result.asset_id}", style="green")
header.append(f"doc={result.document_id}", style="cyan")
if result.page_no is not None:
header.append(f" page={result.page_no}", style="magenta")
self.console.print(header)
if result.caption:
self.console.print(f"[bold]caption[/bold]: {result.caption}")
if result.description:
self.console.print(f"[bold]description[/bold]: {result.description}")
self.console.rule() self.console.rule()
async def serve( async def serve(

View file

@ -297,6 +297,60 @@ def search(
asyncio.run(app.search(query=query, limit=limit, filter=filter)) asyncio.run(app.search(query=query, limit=limit, filter=filter))
@cli.command("search-image-text", help="Search indexed images using a text query")
def search_image_text(
query: str = typer.Argument(
help="Text query to embed and search against indexed image assets",
),
limit: int | None = typer.Option(
None,
"--limit",
"-l",
help="Maximum number of results to return (default: config search.default_limit)",
),
filter: str | None = typer.Option(
None,
"--filter",
"-f",
help="SQL WHERE clause to filter image assets (e.g., \"document_id = '...'\")",
),
db: Path | None = typer.Option(
None,
"--db",
help="Path to the LanceDB database file",
),
):
app = create_app(db)
asyncio.run(app.search_image_text(query=query, limit=limit, filter=filter))
@cli.command("search-image", help="Search indexed images using an image query (image-to-image)")
def search_image(
image_path: str = typer.Argument(
help="Path to an image file to embed and search against indexed image assets",
),
limit: int | None = typer.Option(
None,
"--limit",
"-l",
help="Maximum number of results to return (default: config search.default_limit)",
),
filter: str | None = typer.Option(
None,
"--filter",
"-f",
help="SQL WHERE clause to filter image assets (e.g., \"document_id = '...'\")",
),
db: Path | None = typer.Option(
None,
"--db",
help="Path to the LanceDB database file",
),
):
app = create_app(db)
asyncio.run(app.search_image(image_path=image_path, limit=limit, filter=filter))
@cli.command("visualize", help="Show visual grounding for a chunk") @cli.command("visualize", help="Show visual grounding for a chunk")
def visualize( def visualize(
chunk_id: str = typer.Argument( chunk_id: str = typer.Argument(
@ -312,6 +366,26 @@ def visualize(
asyncio.run(app.visualize_chunk(chunk_id=chunk_id)) asyncio.run(app.visualize_chunk(chunk_id=chunk_id))
@cli.command("visualize-asset", help="Show a multimodal asset image (mm_assets) by asset_id")
def visualize_asset(
asset_id: str = typer.Argument(
help="The asset_id returned by search-image/search-image-text",
),
mode: str = typer.Option(
"crop",
"--mode",
help="Visualization mode: crop (default) or page",
),
db: Path | None = typer.Option(
None,
"--db",
help="Path to the LanceDB database file",
),
):
app = create_app(db)
asyncio.run(app.visualize_asset(asset_id=asset_id, mode=mode))
@cli.command("ask", help="Ask a question using the QA agent") @cli.command("ask", help="Ask a question using the QA agent")
def ask( def ask(
question: str = typer.Argument( question: str = typer.Argument(

View file

@ -22,6 +22,7 @@ from haiku.rag.store.models.chunk import Chunk, SearchResult
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.repositories.mm_asset import MMAssetRepository
from haiku.rag.store.repositories.settings import SettingsRepository from haiku.rag.store.repositories.settings import SettingsRepository
if TYPE_CHECKING: if TYPE_CHECKING:
@ -88,6 +89,7 @@ class HaikuRAG:
) )
self.document_repository = DocumentRepository(self.store) self.document_repository = DocumentRepository(self.store)
self.chunk_repository = ChunkRepository(self.store) self.chunk_repository = ChunkRepository(self.store)
self.mm_asset_repository = MMAssetRepository(self.store)
@property @property
def is_read_only(self) -> bool: def is_read_only(self) -> bool:
@ -279,6 +281,9 @@ class HaikuRAG:
# Batch create all chunks in a single operation # Batch create all chunks in a single operation
await self.chunk_repository.create(chunks) await self.chunk_repository.create(chunks)
# Optional: index multimodal assets (images) into mm_assets table
await self._index_mm_assets_for_document(created_doc)
# Vacuum old versions in background (non-blocking) if auto_vacuum enabled # Vacuum old versions in background (non-blocking) if auto_vacuum enabled
if self._config.storage.auto_vacuum: if self._config.storage.auto_vacuum:
asyncio.create_task(self.store.vacuum()) asyncio.create_task(self.store.vacuum())
@ -317,6 +322,9 @@ class HaikuRAG:
# Delete existing chunks before writing new ones # Delete existing chunks before writing new ones
await self.chunk_repository.delete_by_document_id(document.id) await self.chunk_repository.delete_by_document_id(document.id)
# Delete existing multimodal assets before writing new ones (if table exists)
if self.store.mm_assets_table is not None:
await self.mm_asset_repository.delete_by_document_id(document.id)
try: try:
# Update the document # Update the document
@ -331,6 +339,9 @@ class HaikuRAG:
# Batch create all chunks in a single operation # Batch create all chunks in a single operation
await self.chunk_repository.create(chunks) await self.chunk_repository.create(chunks)
# Optional: re-index multimodal assets
await self._index_mm_assets_for_document(updated_doc)
# Vacuum old versions in background (non-blocking) if auto_vacuum enabled # Vacuum old versions in background (non-blocking) if auto_vacuum enabled
if self._config.storage.auto_vacuum: if self._config.storage.auto_vacuum:
asyncio.create_task(self.store.vacuum()) asyncio.create_task(self.store.vacuum())
@ -821,6 +832,150 @@ class HaikuRAG:
"""Delete a document by its ID.""" """Delete a document by its ID."""
return await self.document_repository.delete(document_id) return await self.document_repository.delete(document_id)
async def _index_mm_assets_for_document(self, document: Document) -> None:
"""Index Docling PictureItems into the mm_assets table (Phase 1).
This is best-effort and opt-in:
- No-op unless config.multimodal.enabled
- No-op in read-only mode
- No-op if mm_assets table is not available
Embedding is performed via the configured multimodal embedder (remote vLLM).
"""
if not self._config.multimodal.enabled:
return
if self.store.is_read_only:
return
if self.store.mm_assets_table is None:
return
if document.id is None:
return
docling_doc = document.get_docling_document()
if not docling_doc:
return
# Lazy imports (docling + pillow are optional in slim installs)
from haiku.rag.embeddings.multimodal import get_multimodal_embedder
embedder = get_multimodal_embedder(self._config)
if embedder is None:
return
try:
from docling_core.types.doc.document import PictureItem
except Exception:
# Docling not installed → nothing to index
return
padding_px = int(self._config.multimodal.image_crop_padding_px)
# Collect crops + asset metadata first, then embed in a batch.
crops = []
assets = []
all_items = list(docling_doc.iterate_items())
for idx, (item, _) in enumerate(all_items):
if not isinstance(item, PictureItem):
continue
self_ref = getattr(item, "self_ref", None)
if not self_ref:
continue
prov = getattr(item, "prov", None) or []
if not prov:
continue
# Prefer the first provenance entry for MVP; can be extended to multiple.
p = prov[0]
page_no = getattr(p, "page_no", None)
bbox = getattr(p, "bbox", None)
if page_no is None or bbox is None:
continue
if page_no not in docling_doc.pages:
continue
page = docling_doc.pages[page_no]
if page.image is None or page.image.pil_image is None:
continue
pil_image = page.image.pil_image
page_height = page.size.height
# Convert from Docling document coordinates to pixel coordinates (same as visualize_chunk)
scale_x = pil_image.width / page.size.width
scale_y = pil_image.height / page.size.height
x0 = bbox.l * scale_x
y0 = (page_height - bbox.t) * scale_y
x1 = bbox.r * scale_x
y1 = (page_height - bbox.b) * scale_y
if y0 > y1:
y0, y1 = y1, y0
# Apply pixel padding + clamp
x0 = max(0, int(x0) - padding_px)
y0 = max(0, int(y0) - padding_px)
x1 = min(pil_image.width, int(x1) + padding_px)
y1 = min(pil_image.height, int(y1) + padding_px)
if x1 <= x0 or y1 <= y0:
continue
crop = pil_image.crop((x0, y0, x1, y1))
crops.append(crop)
# Best-effort caption extraction
caption = None
cap = getattr(item, "caption", None)
if cap is not None and hasattr(cap, "text"):
caption = cap.text
assets.append(
{
"document_id": document.id,
"doc_item_ref": self_ref,
"item_index": idx,
"page_no": int(page_no),
"bbox": {
"left": float(bbox.l),
"top": float(bbox.t),
"right": float(bbox.r),
"bottom": float(bbox.b),
},
"caption": caption,
}
)
if not crops:
return
embeddings = await embedder.embed_images(crops)
if len(embeddings) != len(assets):
raise ValueError(
f"Multimodal embedder returned {len(embeddings)} embeddings for "
f"{len(assets)} image crops"
)
from haiku.rag.store.models.mm_asset import MMAsset
mm_entities = []
for a, e in zip(assets, embeddings):
mm_entities.append(
MMAsset(
document_id=a["document_id"],
doc_item_ref=a["doc_item_ref"],
item_index=a["item_index"],
page_no=a["page_no"],
bbox=a["bbox"],
caption=a["caption"],
metadata={},
embedding=e,
)
)
await self.mm_asset_repository.create(mm_entities)
async def list_documents( async def list_documents(
self, self,
limit: int | None = None, limit: int | None = None,
@ -1193,6 +1348,72 @@ class HaikuRAG:
return final_results + passthrough return final_results + passthrough
async def search_images_by_text(
self,
query: str,
*,
limit: int | None = None,
filter: str | None = None,
):
"""Search indexed multimodal assets using a text query (text→image).
Requires:
- config.multimodal.enabled
- mm_assets table present
- multimodal embedder configured (remote vLLM Mode A endpoint)
"""
if not self._config.multimodal.enabled or self.store.mm_assets_table is None:
return []
from haiku.rag.embeddings.multimodal import get_multimodal_embedder
embedder = get_multimodal_embedder(self._config)
if embedder is None:
return []
limit = limit or self._config.search.limit
vecs = await embedder.embed_texts([query])
if not vecs:
return []
return await self.mm_asset_repository.search_by_vector(
vecs[0], limit=int(limit), filter=filter
)
async def search_images(
self,
image_path: Path,
*,
limit: int | None = None,
filter: str | None = None,
):
"""Search indexed multimodal assets using an image query (image→image)."""
if not self._config.multimodal.enabled or self.store.mm_assets_table is None:
return []
from haiku.rag.embeddings.multimodal import get_multimodal_embedder
embedder = get_multimodal_embedder(self._config)
if embedder is None:
return []
try:
from PIL import Image
except Exception as e:
raise ImportError(
"Image search requires Pillow. Install an appropriate extra or dependency."
) from e
with Image.open(image_path) as img:
img = img.convert("RGB")
vecs = await embedder.embed_images([img])
if not vecs:
return []
limit = limit or self._config.search.limit
return await self.mm_asset_repository.search_by_vector(
vecs[0], limit=int(limit), filter=filter
)
async def _expand_with_chunks( async def _expand_with_chunks(
self, self,
doc_id: str, doc_id: str,
@ -1357,6 +1578,122 @@ class HaikuRAG:
return images return images
async def visualize_mm_asset(self, *, asset_id: str, mode: str = "crop") -> list:
"""Visualize a stored multimodal asset (mm_assets).
This is primarily a UX/debugging helper (CLI/inspector). It uses the stored
provenance (page_no + bbox) and the document's Docling page images.
Args:
asset_id: mm_assets row id (UUID).
mode: "crop" (return crop only) or "page" (full page with bbox overlay).
Returns:
List of PIL Image objects (usually length 1). Empty list if unavailable.
"""
# Optional dependency: only needed for visualization paths.
from copy import deepcopy
try:
from PIL import ImageDraw
except Exception as e:
raise ImportError(
"Visualizing multimodal assets requires Pillow. "
"Install an appropriate extra or dependency."
) from e
if self.store.mm_assets_table is None:
return []
if mode not in ("crop", "page"):
raise ValueError("mode must be 'crop' or 'page'")
# Fetch row by id. We keep this simple (scan) since this is a debug UX path.
# If this becomes hot, add an indexed lookup in MMAssetRepository.
rows = (
self.store.mm_assets_table.search()
.where(f"id = '{asset_id}'")
.limit(1)
.to_list()
)
row = rows[0] if rows else None
if row is None:
return []
doc_id = row.get("document_id")
page_no = row.get("page_no")
bbox_raw = row.get("bbox")
if not doc_id or page_no is None or not bbox_raw:
return []
# Parse bbox
bbox = None
try:
import json as _json
bbox = _json.loads(bbox_raw) if isinstance(bbox_raw, str) else bbox_raw
except Exception:
bbox = None
if not isinstance(bbox, dict) or not all(
k in bbox for k in ("left", "top", "right", "bottom")
):
return []
doc = await self.document_repository.get_by_id(str(doc_id))
if not doc:
return []
docling_doc = doc.get_docling_document()
if not docling_doc:
return []
if int(page_no) not in docling_doc.pages:
return []
page = docling_doc.pages[int(page_no)]
if page.image is None or page.image.pil_image is None or page.size is None:
return []
pil_image = page.image.pil_image
page_width = float(page.size.width)
page_height = float(page.size.height)
scale_x = pil_image.width / page_width
scale_y = pil_image.height / page_height
left = float(bbox["left"])
top = float(bbox["top"])
right = float(bbox["right"])
bottom = float(bbox["bottom"])
x0 = left * scale_x
x1 = right * scale_x
y0 = (page_height - top) * scale_y
y1 = (page_height - bottom) * scale_y
if y0 > y1:
y0, y1 = y1, y0
pad = int(self._config.multimodal.image_crop_padding_px)
x0i = max(0, int(x0) - pad)
y0i = max(0, int(y0) - pad)
x1i = min(pil_image.width, int(x1) + pad)
y1i = min(pil_image.height, int(y1) + pad)
if mode == "crop":
if x1i <= x0i or y1i <= y0i:
return []
return [pil_image.crop((x0i, y0i, x1i, y1i))]
# mode == "page": draw bbox overlay
image = deepcopy(pil_image)
draw = ImageDraw.Draw(image, "RGBA")
fill_color = (255, 255, 0, 80)
outline_color = (255, 165, 0, 255)
draw.rectangle([(x0, y0), (x1, y1)], fill=fill_color, outline=None)
draw.rectangle([(x0, y0), (x1, y1)], outline=outline_color, width=3)
return [image]
async def rebuild_database( async def rebuild_database(
self, mode: RebuildMode = RebuildMode.FULL self, mode: RebuildMode = RebuildMode.FULL
) -> AsyncGenerator[str, None]: ) -> AsyncGenerator[str, None]:

View file

@ -43,6 +43,64 @@ class EmbeddingModelConfig(BaseModel):
base_url: str | None = None base_url: str | None = None
class MultimodalEmbeddingModelConfig(BaseModel):
"""Configuration for a multimodal embedding model (text + image).
This is used for image-to-image and text-to-image retrieval.
Notes:
- This feature is opt-in and may require external services (e.g., vLLM).
- Keep vector_dim stable for a database once created, just like text embeddings.
"""
provider: Literal["vllm"] = "vllm"
name: str = "Qwen/Qwen3-VL-Embedding-2B"
vector_dim: int = 2048
# Base URL for the multimodal embedding service.
# vLLM default is typically http://localhost:8000
base_url: str = "http://localhost:8000"
# HTTP client settings
timeout: int = 60
# Optional: request server-side output dimension (if the backend supports it).
# For OpenAI-compatible /v1/embeddings this is usually called "dimensions".
#
# Note: Some servers/models (including vLLM + Qwen3-VL-Embedding-2B) reject the
# presence of this field unless the model explicitly supports matryoshka output.
# Keep this unset by default; haiku will validate the returned dimension via
# `vector_dim`.
dimensions: int | None = None
# OpenAI-compatible embeddings response format.
# vLLM/OpenAI supports "float"; some servers also support "base64".
encoding_format: Literal["float", "base64"] = "float"
class MultimodalConfig(BaseModel):
"""Opt-in multimodal embedding configuration."""
enabled: bool = False
index_pictures: bool = True
# Crop padding (in pixels) for bbox-based image crops.
# Kept here (rather than conversion options) because it affects embedding, not conversion.
image_crop_padding_px: int = 8
# Image resize guardrail before embedding (to keep payload sizes manageable and avoid
# backend limits). Set to 0 to disable resizing.
image_max_side_px: int = 1024
# Max number of images to send per /v1/embeddings call for mm_assets indexing.
# (Large PDFs can contain many figures; batching avoids oversized requests.)
embed_batch_size: int = 8
model: MultimodalEmbeddingModelConfig = Field(
default_factory=MultimodalEmbeddingModelConfig
)
class StorageConfig(BaseModel): class StorageConfig(BaseModel):
data_dir: Path = Field(default_factory=get_default_data_dir) data_dir: Path = Field(default_factory=get_default_data_dir)
auto_vacuum: bool = True auto_vacuum: bool = True
@ -198,6 +256,7 @@ class AppConfig(BaseModel):
monitor: MonitorConfig = Field(default_factory=MonitorConfig) monitor: MonitorConfig = Field(default_factory=MonitorConfig)
lancedb: LanceDBConfig = Field(default_factory=LanceDBConfig) lancedb: LanceDBConfig = Field(default_factory=LanceDBConfig)
embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig) embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig)
multimodal: MultimodalConfig = Field(default_factory=MultimodalConfig)
reranking: RerankingConfig = Field(default_factory=RerankingConfig) reranking: RerankingConfig = Field(default_factory=RerankingConfig)
qa: QAConfig = Field(default_factory=QAConfig) qa: QAConfig = Field(default_factory=QAConfig)
research: ResearchConfig = Field(default_factory=ResearchConfig) research: ResearchConfig = Field(default_factory=ResearchConfig)

View file

@ -0,0 +1,238 @@
import base64
import asyncio
from dataclasses import dataclass
import httpx
from haiku.rag.config import AppConfig, Config
@dataclass
class MultimodalEmbedInput:
"""Input for a multimodal embedding request."""
text: str | None = None
image_b64: str | None = None
class MultimodalEmbedderBase:
"""Base interface for multimodal embedders (text + image)."""
vector_dim: int
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
raise NotImplementedError
async def embed_images(self, images) -> list[list[float]]: # images: list[PIL.Image.Image]
raise NotImplementedError
class VLLMMultimodalEmbedder(MultimodalEmbedderBase):
"""vLLM-backed multimodal embedder via OpenAI-compatible `POST /v1/embeddings`."""
def __init__(
self,
*,
base_url: str,
model: str,
vector_dim: int,
timeout: int,
output_dim: int | None,
image_max_side_px: int,
embed_batch_size: int,
) -> None:
# Accept base_url with or without a trailing /v1 for convenience.
# (Docs for OpenAI-compatible servers often say to include /v1.)
cleaned = base_url.rstrip("/")
if cleaned.endswith("/v1"):
cleaned = cleaned[: -len("/v1")]
self.base_url = cleaned
self.model = model
self.vector_dim = vector_dim
self.timeout = timeout
self.output_dim = output_dim
self.image_max_side_px = int(image_max_side_px)
self.embed_batch_size = int(embed_batch_size)
self.encoding_format: str = "float"
async def _post_openai_embeddings_input(self, inputs: list[str]) -> list[list[float]]:
"""Batch text-only embeddings via OpenAI-compatible `input=[str, ...]`."""
payload: dict = {
"model": self.model,
"input": inputs,
"encoding_format": self.encoding_format,
}
# Only send "dimensions" when requesting a *different* output dimension.
# Some servers/models reject the presence of this field unless they support
# matryoshka output; for Qwen3-VL-Embedding-2B on vLLM, even sending the
# native dim can yield 400.
if self.output_dim is not None and int(self.output_dim) != int(self.vector_dim):
payload["dimensions"] = int(self.output_dim)
async with httpx.AsyncClient(timeout=self.timeout) as client:
resp = await client.post(f"{self.base_url}/v1/embeddings", json=payload)
try:
resp.raise_for_status()
except httpx.HTTPStatusError as e:
# Surface the backend error body; otherwise we lose crucial debugging context.
body = ""
try:
body = resp.text
except Exception:
body = "<unavailable>"
raise httpx.HTTPStatusError(
f"{e}. Response body: {body[:2000]}",
request=e.request,
response=e.response,
) from e
data = resp.json()
rows = data.get("data") or []
embeddings: list[list[float]] = []
for r in rows:
emb = r.get("embedding")
if emb is None:
continue
embeddings.append(list(emb))
if embeddings and len(embeddings[0]) != self.vector_dim:
raise ValueError(
f"Unexpected embedding dimension: got {len(embeddings[0])}, "
f"expected {self.vector_dim}"
)
return embeddings
async def _post_openai_embeddings_message(
self, inp: MultimodalEmbedInput
) -> list[float]:
"""Single multimodal embedding via OpenAI-compatible `messages=[...]`.
Note: vLLM treats `messages` as a *single* conversation and returns a single
embedding (i.e., it does not batch multiple independent inputs via `messages`).
"""
content: list[dict] = []
if inp.text is not None and inp.text != "":
content.append({"type": "text", "text": inp.text})
if inp.image_b64 is not None:
content.append(
{
"type": "image_url",
"image_url": {"url": f"data:image/png;base64,{inp.image_b64}"},
}
)
payload: dict = {
"model": self.model,
"messages": [{"role": "user", "content": content}],
"encoding_format": self.encoding_format,
}
if self.output_dim is not None and int(self.output_dim) != int(self.vector_dim):
payload["dimensions"] = int(self.output_dim)
async with httpx.AsyncClient(timeout=self.timeout) as client:
resp = await client.post(f"{self.base_url}/v1/embeddings", json=payload)
try:
resp.raise_for_status()
except httpx.HTTPStatusError as e:
body = ""
try:
body = resp.text
except Exception:
body = "<unavailable>"
raise httpx.HTTPStatusError(
f"{e}. Response body: {body[:2000]}",
request=e.request,
response=e.response,
) from e
data = resp.json()
rows = data.get("data") or []
if not rows or rows[0].get("embedding") is None:
raise ValueError("vLLM returned no embedding data")
emb = list(rows[0]["embedding"])
if emb and len(emb) != self.vector_dim:
raise ValueError(
f"Unexpected embedding dimension: got {len(emb)}, expected {self.vector_dim}"
)
return emb
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
if not texts:
return []
# Batch text embeddings via `input=[str,...]`.
bs = max(1, int(self.embed_batch_size))
out: list[list[float]] = []
for i in range(0, len(texts), bs):
out.extend(await self._post_openai_embeddings_input(texts[i : i + bs]))
return out
async def embed_images(self, images) -> list[list[float]]: # list[PIL.Image.Image]
if not images:
return []
# Lazy import: Pillow is optional in slim mode.
try:
from PIL import Image # noqa: F401
except Exception as e:
raise ImportError(
"Multimodal image embedding requires Pillow. "
"Install an appropriate extra or dependency (e.g. 'pillow')."
) from e
inputs: list[MultimodalEmbedInput] = []
for img in images:
# Encode to PNG for a stable transport format.
import io
# Resize guardrail: keep max side bounded to reduce payload sizes and avoid
# backend limits. (0 disables resizing.)
if self.image_max_side_px > 0:
max_side = max(int(img.size[0]), int(img.size[1]))
if max_side > self.image_max_side_px:
from PIL import Image as PILImage
scale = float(self.image_max_side_px) / float(max_side)
new_w = max(1, int(round(img.size[0] * scale)))
new_h = max(1, int(round(img.size[1] * scale)))
img = img.resize((new_w, new_h), resample=PILImage.Resampling.LANCZOS)
buf = io.BytesIO()
img.save(buf, format="PNG", optimize=True)
image_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
inputs.append(MultimodalEmbedInput(image_b64=image_b64))
# vLLM `messages=[...]` returns ONE embedding per call, so we run one request per image.
max_in_flight = max(1, int(self.embed_batch_size))
sem = asyncio.Semaphore(max_in_flight)
async def one(inp: MultimodalEmbedInput) -> list[float]:
async with sem:
return await self._post_openai_embeddings_message(inp)
return list(await asyncio.gather(*(one(i) for i in inputs)))
def get_multimodal_embedder(config: AppConfig = Config) -> MultimodalEmbedderBase | None:
"""Factory for multimodal embedders.
Returns None when multimodal indexing/search is disabled.
"""
if not config.multimodal.enabled:
return None
mm = config.multimodal.model
if mm.provider == "vllm":
emb = VLLMMultimodalEmbedder(
base_url=mm.base_url,
model=mm.name,
vector_dim=int(mm.vector_dim),
timeout=int(mm.timeout),
output_dim=mm.dimensions,
image_max_side_px=int(config.multimodal.image_max_side_px),
embed_batch_size=int(config.multimodal.embed_batch_size),
)
emb.encoding_format = mm.encoding_format
return emb
raise ValueError(f"Unsupported multimodal provider: {mm.provider}")

View file

@ -73,6 +73,8 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover
Binding("i", "show_info", "Info", show=True), Binding("i", "show_info", "Info", show=True),
Binding("v", "show_visual", "Visual", show=True), Binding("v", "show_visual", "Visual", show=True),
Binding("c", "show_context", "Context", show=True), Binding("c", "show_context", "Context", show=True),
Binding("m", "show_mm_assets", "MM Assets", show=True),
Binding("d", "show_doc_index", "Doc Index", show=True),
] ]
def __init__( def __init__(
@ -83,6 +85,9 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover
self.read_only = read_only self.read_only = read_only
self.before = before self.before = before
self.client: HaikuRAG | None = None self.client: HaikuRAG | None = None
self._selected_document_id: str | None = None
self._selected_document_uri: str | None = None
self._selected_document_title: str | None = None
def compose(self) -> "ComposeResult": def compose(self) -> "ComposeResult":
"""Compose the UI layout.""" """Compose the UI layout."""
@ -100,6 +105,14 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover
config=config, config=config,
read_only=self.read_only, read_only=self.read_only,
before=self.before, before=self.before,
# The inspector is a read-only browsing tool; it should be able to open a DB
# even when the current environment config doesn't match the DB settings.
#
# This commonly happens when you created a DB with one embedder (e.g. vLLM/OpenAI)
# and later run the inspector without passing the same `--config`.
#
# We skip config compatibility validation here to avoid hard-failing the TUI.
skip_validation=True,
) )
await self.client.__aenter__() await self.client.__aenter__()
@ -191,6 +204,9 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover
# Load chunks for this document # Load chunks for this document
if message.document.id: if message.document.id:
self._selected_document_id = str(message.document.id)
self._selected_document_uri = message.document.uri
self._selected_document_title = message.document.title
chunk_list = self.query_one(ChunkList) chunk_list = self.query_one(ChunkList)
await chunk_list.load_chunks_for_document(self.client, message.document.id) await chunk_list.load_chunks_for_document(self.client, message.document.id)
@ -238,6 +254,36 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover
await self._switch_modal(ContextModal(chunk=chunk, client=self.client)) await self._switch_modal(ContextModal(chunk=chunk, client=self.client))
async def action_show_mm_assets(self) -> None:
"""Show multimodal image assets (mm_assets) for the selected document."""
if not self.client or not self._selected_document_id:
return
from haiku.rag.inspector.widgets.mm_assets_modal import MMAssetsModal
await self._switch_modal(
MMAssetsModal(
client=self.client,
document_id=self._selected_document_id,
document_uri=self._selected_document_uri,
document_title=self._selected_document_title,
)
)
async def action_show_doc_index(self) -> None:
"""Show a whole-document overlay (all chunk boxes + mm_asset boxes)."""
if not self.client or not self._selected_document_id:
return
from haiku.rag.inspector.widgets.doc_index_modal import DocIndexModal
await self._switch_modal(
DocIndexModal(
client=self.client,
document_id=self._selected_document_id,
document_uri=self._selected_document_uri,
document_title=self._selected_document_title,
)
)
def run_inspector( def run_inspector(
db_path: Path | None = None, db_path: Path | None = None,

View file

@ -0,0 +1,269 @@
from __future__ import annotations
import subprocess
import sys
import tempfile
import shutil
from copy import deepcopy
from pathlib import Path
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical
from textual.screen import Screen
from textual.widget import Widget
from textual.widgets import Static
from textual_image.widget import Image as TextualImage
class DocIndexModal(Screen): # pragma: no cover
"""Whole-document overlay: all chunk boxes + all mm_asset boxes per page."""
BINDINGS = [
Binding("escape", "dismiss", "Close", show=True),
Binding("d", "dismiss", "Close", show=True),
Binding("left", "prev_page", "Previous Page"),
Binding("right", "next_page", "Next Page"),
Binding("o", "open_image", "Open image", show=True),
]
CSS = """
DocIndexModal {
background: $surface;
layout: vertical;
}
#doc-index-header {
dock: top;
height: auto;
padding: 1;
}
#doc-index-content {
height: 1fr;
width: 100%;
align: center middle;
}
#doc-index-content Image {
width: auto;
height: 100%;
}
#doc-index-footer {
dock: bottom;
height: auto;
padding: 1;
}
"""
def __init__(
self,
*,
client,
document_id: str,
document_uri: str | None = None,
document_title: str | None = None,
) -> None:
super().__init__()
self.client = client
self.document_id = document_id
self.document_uri = document_uri
self.document_title = document_title
self._image_widget: Widget = Static("Loading...", id="image-display")
self._page_info = Static("", id="page-info")
self._out_dir: Path = Path(tempfile.mkdtemp(prefix="haiku-inspector-doc-index-"))
self._pages: list[tuple[int, Path]] = [] # (page_no, png_path)
self._page_idx: int = 0
def compose(self) -> ComposeResult:
title = self.document_title or self.document_uri or self.document_id
with Vertical(id="doc-index-header"):
yield Static(f"[bold]Doc Index Overlay[/bold] - {title}")
yield Static(f"[dim]saved:[/dim] {self._out_dir}")
with Horizontal(id="doc-index-content"):
yield self._image_widget
with Horizontal(id="doc-index-footer"):
yield self._page_info
async def on_mount(self) -> None:
# Build overlay images once.
try:
from PIL import ImageDraw
except Exception as e:
if isinstance(self._image_widget, Static):
self._image_widget.update(f"[red]Pillow required: {e}[/red]")
return
doc = await self.client.document_repository.get_by_id(self.document_id)
if not doc:
if isinstance(self._image_widget, Static):
self._image_widget.update("[red]Document not found[/red]")
return
docling_doc = doc.get_docling_document()
if not docling_doc:
if isinstance(self._image_widget, Static):
self._image_widget.update("[yellow]No DoclingDocument stored.[/yellow]")
return
# Gather chunk bounding boxes grouped by page.
boxes_by_page_chunks: dict[int, list[tuple[float, float, float, float]]] = {}
# Load all chunks (paged).
chunks: list = []
offset = 0
batch = 200
while True:
part = await self.client.chunk_repository.get_by_document_id(
self.document_id, limit=batch, offset=offset
)
if not part:
break
chunks.extend(part)
offset += len(part)
if len(part) < batch:
break
for ch in chunks:
meta = ch.get_chunk_metadata()
for bb in meta.resolve_bounding_boxes(docling_doc):
boxes_by_page_chunks.setdefault(int(bb.page_no), []).append(
(float(bb.left), float(bb.top), float(bb.right), float(bb.bottom))
)
# Gather mm_asset bboxes grouped by page.
boxes_by_page_assets: dict[int, list[tuple[float, float, float, float]]] = {}
if self.client.store.mm_assets_table is not None:
rows = (
self.client.store.mm_assets_table.search()
.where(f"document_id = '{self.document_id}'")
.to_list()
)
for r in rows:
page_no = r.get("page_no")
bbox = r.get("bbox")
if page_no is None or not bbox:
continue
# bbox is stored as a dict or JSON string
if isinstance(bbox, str):
import json as _json
try:
bbox = _json.loads(bbox)
except Exception:
continue
if not isinstance(bbox, dict):
continue
if not all(k in bbox for k in ("left", "top", "right", "bottom")):
continue
boxes_by_page_assets.setdefault(int(page_no), []).append(
(
float(bbox["left"]),
float(bbox["top"]),
float(bbox["right"]),
float(bbox["bottom"]),
)
)
# Render all pages with available images.
for page_no, page in sorted(docling_doc.pages.items(), key=lambda kv: kv[0]):
if page.image is None or page.image.pil_image is None or page.size is None:
continue
pil = page.image.pil_image
page_h = float(page.size.height)
scale_x = pil.width / float(page.size.width)
scale_y = pil.height / float(page.size.height)
img = deepcopy(pil)
draw = ImageDraw.Draw(img, "RGBA")
# Chunks: yellow/orange
for left, top, right, bottom in boxes_by_page_chunks.get(int(page_no), []):
x0 = left * scale_x
y0 = (page_h - top) * scale_y
x1 = right * scale_x
y1 = (page_h - bottom) * scale_y
if y0 > y1:
y0, y1 = y1, y0
draw.rectangle([(x0, y0), (x1, y1)], fill=(255, 255, 0, 40), outline=None)
draw.rectangle([(x0, y0), (x1, y1)], outline=(255, 165, 0, 255), width=2)
# Assets: cyan/blue
for left, top, right, bottom in boxes_by_page_assets.get(int(page_no), []):
x0 = left * scale_x
y0 = (page_h - top) * scale_y
x1 = right * scale_x
y1 = (page_h - bottom) * scale_y
if y0 > y1:
y0, y1 = y1, y0
draw.rectangle([(x0, y0), (x1, y1)], fill=(0, 255, 255, 30), outline=None)
draw.rectangle([(x0, y0), (x1, y1)], outline=(0, 180, 255, 255), width=2)
out = self._out_dir / f"page_{int(page_no)}.png"
try:
img.save(out, format="PNG")
except Exception:
continue
self._pages.append((int(page_no), out))
await self._render_current_page()
async def _render_current_page(self) -> None:
if not self._pages:
if isinstance(self._image_widget, Static):
self._image_widget.update(
"[yellow]No page images available.[/yellow]\n"
"This document was converted without page images."
)
self._page_info.update("")
return
page_no, path = self._pages[self._page_idx]
self._page_info.update(
f"Page {page_no} ({self._page_idx+1}/{len(self._pages)}) | "
f"←/→ navigate | o open"
)
try:
from PIL import Image
pil = Image.open(path)
new_widget = TextualImage(pil, id="rendered-image")
await self._image_widget.remove()
content = self.query_one("#doc-index-content", Horizontal)
await content.mount(new_widget)
self._image_widget = new_widget
except Exception as e:
if isinstance(self._image_widget, Static):
self._image_widget.update(f"[red]Error: {e}[/red]")
async def action_prev_page(self) -> None:
if self._pages and self._page_idx > 0:
self._page_idx -= 1
await self._render_current_page()
async def action_next_page(self) -> None:
if self._pages and self._page_idx < len(self._pages) - 1:
self._page_idx += 1
await self._render_current_page()
async def action_open_image(self) -> None:
if not self._pages:
return
_, path = self._pages[self._page_idx]
try:
if sys.platform == "darwin":
subprocess.run(["open", str(path)], check=False)
else:
subprocess.run(["xdg-open", str(path)], check=False)
except Exception:
return
async def action_dismiss(self, result=None) -> None:
# Best-effort cleanup of temporary images.
try:
shutil.rmtree(self._out_dir, ignore_errors=True)
except Exception:
pass
self.app.pop_screen()

View file

@ -0,0 +1,196 @@
from __future__ import annotations
import subprocess
import sys
import tempfile
import shutil
from pathlib import Path
from textual.app import ComposeResult
from textual.binding import Binding
from textual.containers import Horizontal, Vertical, VerticalScroll
from textual.screen import Screen
from textual.widgets import ListItem, ListView, Markdown, Static
class MMAssetsModal(Screen): # pragma: no cover
"""Browse multimodal assets (mm_assets) for a document."""
BINDINGS = [
Binding("escape", "dismiss", "Close", show=True),
Binding("m", "dismiss", "Close", show=True),
Binding("o", "open", "Open image", show=True),
Binding("p", "toggle_mode", "Toggle page/crop", show=True),
]
CSS = """
MMAssetsModal {
background: $surface;
layout: vertical;
}
#mm-header {
dock: top;
height: auto;
padding: 1;
}
#mm-content {
height: 1fr;
width: 100%;
}
#mm-list {
width: 1fr;
border: solid $primary;
}
#mm-detail {
width: 2fr;
border: solid $accent;
}
"""
def __init__(
self,
*,
client,
document_id: str,
document_uri: str | None = None,
document_title: str | None = None,
) -> None:
super().__init__()
self.client = client
self.document_id = document_id
self.document_uri = document_uri
self.document_title = document_title
self.mode: str = "crop"
self.rows: list[dict] = []
self._out_dir: Path = Path(tempfile.mkdtemp(prefix="haiku-inspector-mm-assets-"))
def compose(self) -> ComposeResult:
title = self.document_title or self.document_uri or self.document_id
with Vertical(id="mm-header"):
yield Static(f"[bold]MM Assets[/bold] - {title}")
yield Static(
f"[dim]mode:[/dim] {self.mode} | "
f"press [bold]p[/bold] to toggle, [bold]o[/bold] to open image",
id="mm-status",
)
yield Static(f"[dim]saved:[/dim] {self._out_dir}")
with Horizontal(id="mm-content"):
with VerticalScroll(id="mm-list"):
yield ListView(id="mm-assets-list")
with VerticalScroll(id="mm-detail"):
yield Markdown("Select an asset…", id="mm-detail-md")
async def on_mount(self) -> None:
list_view = self.query_one("#mm-assets-list", ListView)
await list_view.clear()
if self.client.store.mm_assets_table is None:
await self.query_one("#mm-detail-md", Markdown).update(
"[yellow]mm_assets table is not available.[/yellow]\n\n"
"Enable multimodal indexing and ingest documents with pictures first."
)
return
# Load rows for this document (simple scan via query).
self.rows = (
self.client.store.mm_assets_table.search()
.where(f"document_id = '{self.document_id}'")
.to_list()
)
if not self.rows:
await self.query_one("#mm-detail-md", Markdown).update(
"[yellow]No mm_assets for this document.[/yellow]"
)
return
for r in self.rows:
asset_id = r.get("id")
page = r.get("page_no")
ref = r.get("doc_item_ref")
label = f"p.{page} {ref} {asset_id}"
await list_view.append(ListItem(Static(label)))
list_view.index = 0
await self._render_detail()
list_view.focus()
async def on_list_view_highlighted(self, event: ListView.Highlighted) -> None:
if event.list_view.id != "mm-assets-list":
return
await self._render_detail()
async def _render_detail(self) -> None:
list_view = self.query_one("#mm-assets-list", ListView)
idx = list_view.index
if idx is None or idx >= len(self.rows):
return
r = self.rows[idx]
asset_id = r.get("id")
page_no = r.get("page_no")
ref = r.get("doc_item_ref")
caption = r.get("caption")
description = r.get("description")
md = f"""\
**asset_id**: `{asset_id}`
**page**: {page_no}
**doc_item_ref**: `{ref}`
**caption**:
{caption or "*none*"}
**description**:
{description or "*none*"}
**saved_dir**:
`{self._out_dir}`
"""
await self.query_one("#mm-detail-md", Markdown).update(md)
async def action_toggle_mode(self) -> None:
self.mode = "page" if self.mode == "crop" else "crop"
self.query_one("#mm-status", Static).update(
f"[dim]mode:[/dim] {self.mode} | "
f"press [bold]p[/bold] to toggle, [bold]o[/bold] to open image"
)
async def action_open(self) -> None:
list_view = self.query_one("#mm-assets-list", ListView)
idx = list_view.index
if idx is None or idx >= len(self.rows):
return
asset_id = self.rows[idx].get("id")
if not asset_id:
return
images = await self.client.visualize_mm_asset(asset_id=str(asset_id), mode=self.mode)
if not images:
return
# Save first image and open via OS viewer.
path = self._out_dir / f"{asset_id}_{self.mode}.png"
try:
images[0].save(path, format="PNG")
except Exception:
return
try:
if sys.platform == "darwin":
subprocess.run(["open", str(path)], check=False)
else:
subprocess.run(["xdg-open", str(path)], check=False)
except Exception:
return
async def action_dismiss(self, result=None) -> None:
# Best-effort cleanup of temporary images.
try:
shutil.rmtree(self._out_dir, ignore_errors=True)
except Exception:
pass
self.app.pop_screen()

View file

@ -7,6 +7,11 @@ from textual.screen import Screen
from textual.widget import Widget from textual.widget import Widget
from textual.widgets import Static from textual.widgets import Static
from textual_image.widget import Image as TextualImage from textual_image.widget import Image as TextualImage
import shutil
import tempfile
from pathlib import Path
import subprocess
import sys
if TYPE_CHECKING: if TYPE_CHECKING:
from PIL.Image import Image as PILImage from PIL.Image import Image as PILImage
@ -23,6 +28,7 @@ class VisualGroundingModal(Screen): # pragma: no cover
Binding("v", "dismiss", "Close", show=True), Binding("v", "dismiss", "Close", show=True),
Binding("left", "prev_page", "Previous Page"), Binding("left", "prev_page", "Previous Page"),
Binding("right", "next_page", "Next Page"), Binding("right", "next_page", "Next Page"),
Binding("o", "open_image", "Open image", show=True),
] ]
CSS = """ CSS = """
@ -66,6 +72,8 @@ class VisualGroundingModal(Screen): # pragma: no cover
self.client = client self.client = client
self.document_uri = document_uri or chunk.document_uri self.document_uri = document_uri or chunk.document_uri
self.images: list[PILImage] = [] self.images: list[PILImage] = []
self._saved_paths: list[Path] = []
self._out_dir: Path | None = None
self.current_page_idx = 0 self.current_page_idx = 0
self._image_widget: Widget = Static("Loading...", id="image-display") self._image_widget: Widget = Static("Loading...", id="image-display")
self._page_info = Static("", id="page-info") self._page_info = Static("", id="page-info")
@ -82,6 +90,18 @@ class VisualGroundingModal(Screen): # pragma: no cover
async def on_mount(self) -> None: async def on_mount(self) -> None:
"""Load images and display the first page.""" """Load images and display the first page."""
self.images = await self.client.visualize_chunk(self.chunk) self.images = await self.client.visualize_chunk(self.chunk)
# Always save to disk as a reliable fallback: many terminals (incl. Cursor/VSCode)
# cannot render images inline via terminal graphics protocols.
if self.images:
self._out_dir = Path(tempfile.mkdtemp(prefix="haiku-inspector-visual-"))
self._saved_paths = []
for i, img in enumerate(self.images):
p = self._out_dir / f"page_{i+1}.png"
try:
img.save(p, format="PNG")
self._saved_paths.append(p)
except Exception:
continue
await self._render_current_page() await self._render_current_page()
async def _render_current_page(self) -> None: async def _render_current_page(self) -> None:
@ -95,8 +115,11 @@ class VisualGroundingModal(Screen): # pragma: no cover
self._page_info.update("") self._page_info.update("")
return return
extra = ""
if self._out_dir is not None:
extra = f" | Saved to: {self._out_dir} (press 'o' to open)"
self._page_info.update( self._page_info.update(
f"Page {self.current_page_idx + 1}/{len(self.images)} - Use ←/→ to navigate" f"Page {self.current_page_idx + 1}/{len(self.images)} - Use ←/→ to navigate{extra}"
) )
try: try:
@ -111,6 +134,12 @@ class VisualGroundingModal(Screen): # pragma: no cover
self._image_widget.update(f"[red]Error: {e}[/red]") self._image_widget.update(f"[red]Error: {e}[/red]")
async def action_dismiss(self, result=None) -> None: async def action_dismiss(self, result=None) -> None:
# Best-effort cleanup of temporary images.
if self._out_dir is not None:
try:
shutil.rmtree(self._out_dir, ignore_errors=True)
except Exception:
pass
self.app.pop_screen() self.app.pop_screen()
async def action_prev_page(self) -> None: async def action_prev_page(self) -> None:
@ -124,3 +153,16 @@ class VisualGroundingModal(Screen): # pragma: no cover
if self.current_page_idx < len(self.images) - 1: if self.current_page_idx < len(self.images) - 1:
self.current_page_idx += 1 self.current_page_idx += 1
await self._render_current_page() await self._render_current_page()
async def action_open_image(self) -> None:
"""Open the current page image in the OS viewer (best-effort)."""
if not self._saved_paths or self.current_page_idx >= len(self._saved_paths):
return
p = self._saved_paths[self.current_page_idx]
try:
if sys.platform == "darwin":
subprocess.run(["open", str(p)], check=False)
else:
subprocess.run(["xdg-open", str(p)], check=False)
except Exception:
return

View file

@ -53,6 +53,24 @@ class SettingsRecord(LanceModel):
settings: str = Field(default="{}") settings: str = Field(default="{}")
def create_mm_asset_model(vector_dim: int):
"""Create an MMAssetRecord model with the specified vector dimension."""
class MMAssetRecord(LanceModel):
id: str = Field(default_factory=lambda: str(uuid4()))
document_id: str
doc_item_ref: str
item_index: int | None = None
page_no: int | None = None
bbox: str | None = None # JSON string
caption: str | None = None
description: str | None = None
metadata: str = Field(default="{}")
vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim) # type: ignore
return MMAssetRecord
class Store: class Store:
def __init__( def __init__(
self, self,
@ -73,6 +91,9 @@ class Store:
# Create the ChunkRecord model with the correct vector dimension # Create the ChunkRecord model with the correct vector dimension
self.ChunkRecord = create_chunk_model(self.embedder._vector_dim) self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
self.MMAssetRecord = create_mm_asset_model(
int(self._config.multimodal.model.vector_dim)
)
# Check if database exists (for local filesystem only) # Check if database exists (for local filesystem only)
is_new_db = False is_new_db = False
@ -156,6 +177,12 @@ class Store:
self.documents_table, self.documents_table,
self.chunks_table, self.chunks_table,
self.settings_table, self.settings_table,
# Optional multimodal table
*( # type: ignore[misc]
[self.mm_assets_table]
if getattr(self, "mm_assets_table", None) is not None
else []
),
]: ]:
table.optimize(cleanup_older_than=retention) table.optimize(cleanup_older_than=retention)
except (RuntimeError, OSError) as e: except (RuntimeError, OSError) as e:
@ -195,6 +222,7 @@ class Store:
stats_dict: dict = { stats_dict: dict = {
"documents": {"exists": False}, "documents": {"exists": False},
"chunks": {"exists": False}, "chunks": {"exists": False},
"mm_assets": {"exists": False},
} }
# Documents table stats # Documents table stats
@ -307,6 +335,54 @@ class Store:
[SettingsRecord(id="settings", settings=json.dumps(settings_data))] [SettingsRecord(id="settings", settings=json.dumps(settings_data))]
) )
# Ensure `ChunkRecord` matches the DB schema for the `chunks.vector` column.
#
# This matters for read-only tools (e.g. inspector) where the runtime config
# may not match the DB's stored embedding config. If these differ, Pydantic
# validation will fail when materializing rows.
try:
stored_dim: int | None = None
# Prefer the DB's stored settings when available (works even if chunks is empty).
recs = list(
self.settings_table.search().where("id = 'settings'").limit(1).to_list()
)
if recs:
raw = recs[0].get("settings") or "{}"
settings_obj = json.loads(raw) if isinstance(raw, str) else raw
stored_dim_val = (
settings_obj.get("embeddings", {})
.get("model", {})
.get("vector_dim", None)
)
if stored_dim_val is not None:
stored_dim = int(stored_dim_val)
# If we have at least one chunk row, infer from the stored vector length.
inferred_dim: int | None = None
try:
one = self.chunks_table.search().limit(1).to_arrow().to_pylist()
if one and isinstance(one[0].get("vector"), list) and one[0]["vector"]:
inferred_dim = int(len(one[0]["vector"]))
except Exception:
inferred_dim = None
dim = inferred_dim or stored_dim or int(getattr(self.embedder, "_vector_dim", 1024))
self.ChunkRecord = create_chunk_model(int(dim))
except Exception:
# Best-effort: never fail DB opening due to inference.
pass
# Create or get multimodal assets table (optional)
self.mm_assets_table = None
if "mm_assets" in existing_tables:
self.mm_assets_table = self.db.open_table("mm_assets")
else:
# Only create this table if multimodal is enabled AND store is writable.
if self._config.multimodal.enabled and not self._read_only:
self.mm_assets_table = self.db.create_table(
"mm_assets", schema=self.MMAssetRecord
)
def _set_initial_version(self): def _set_initial_version(self):
"""Set the initial version for a new database.""" """Set the initial version for a new database."""
self.set_haiku_version(metadata.version("haiku.rag-slim")) self.set_haiku_version(metadata.version("haiku.rag-slim"))
@ -405,11 +481,14 @@ class Store:
def current_table_versions(self) -> dict[str, int]: def current_table_versions(self) -> dict[str, int]:
"""Capture current versions of key tables for rollback using LanceDB's API.""" """Capture current versions of key tables for rollback using LanceDB's API."""
return { versions = {
"documents": int(self.documents_table.version), "documents": int(self.documents_table.version),
"chunks": int(self.chunks_table.version), "chunks": int(self.chunks_table.version),
"settings": int(self.settings_table.version), "settings": int(self.settings_table.version),
} }
if self.mm_assets_table is not None:
versions["mm_assets"] = int(self.mm_assets_table.version)
return versions
def restore_table_versions(self, versions: dict[str, int]) -> bool: def restore_table_versions(self, versions: dict[str, int]) -> bool:
"""Restore tables to the provided versions using LanceDB's API. """Restore tables to the provided versions using LanceDB's API.
@ -421,6 +500,8 @@ class Store:
self.documents_table.restore(int(versions["documents"])) self.documents_table.restore(int(versions["documents"]))
self.chunks_table.restore(int(versions["chunks"])) self.chunks_table.restore(int(versions["chunks"]))
self.settings_table.restore(int(versions["settings"])) self.settings_table.restore(int(versions["settings"]))
if self.mm_assets_table is not None and "mm_assets" in versions:
self.mm_assets_table.restore(int(versions["mm_assets"]))
return True return True
@property @property
@ -450,6 +531,11 @@ class Store:
("documents", self.documents_table), ("documents", self.documents_table),
("chunks", self.chunks_table), ("chunks", self.chunks_table),
("settings", self.settings_table), ("settings", self.settings_table),
*(
[("mm_assets", self.mm_assets_table)]
if self.mm_assets_table is not None
else []
),
] ]
for table_name, table in tables: for table_name, table in tables:
@ -501,6 +587,7 @@ class Store:
"documents": self.documents_table, "documents": self.documents_table,
"chunks": self.chunks_table, "chunks": self.chunks_table,
"settings": self.settings_table, "settings": self.settings_table,
"mm_assets": self.mm_assets_table,
} }
table = table_map.get(table_name) table = table_map.get(table_name)
if table is None: if table is None:

View file

@ -1,10 +1,13 @@
from .chunk import BoundingBox, Chunk, ChunkMetadata, SearchResult from .chunk import BoundingBox, Chunk, ChunkMetadata, SearchResult
from .document import Document from .document import Document
from .mm_asset import MMAsset, MMSearchResult
__all__ = [ __all__ = [
"BoundingBox", "BoundingBox",
"Chunk", "Chunk",
"ChunkMetadata", "ChunkMetadata",
"Document", "Document",
"MMAsset",
"MMSearchResult",
"SearchResult", "SearchResult",
] ]

View file

@ -0,0 +1,41 @@
from pydantic import BaseModel
class MMAsset(BaseModel):
"""A multimodal asset (typically a Docling PictureItem crop) stored in LanceDB."""
id: str | None = None
document_id: str
# Docling pointer (e.g. "#/pictures/3") used to resolve provenance / ordering.
doc_item_ref: str
# Optional ordering anchor (index from DoclingDocument.iterate_items()).
item_index: int | None = None
# Provenance for layout-aware UX (visual grounding, grouping, etc.)
page_no: int | None = None
bbox: dict | None = None # {"left":..,"top":..,"right":..,"bottom":..}
# Text fields for display/debugging (NOT used as the embedding input for Phase 1).
caption: str | None = None
description: str | None = None
metadata: dict = {}
# Multimodal embedding vector (dimension is configured in Store via LanceDB schema)
embedding: list[float] | None = None
class MMSearchResult(BaseModel):
"""Search result for multimodal assets."""
asset_id: str
document_id: str
score: float
doc_item_ref: str
item_index: int | None = None
page_no: int | None = None
bbox: dict | None = None
caption: str | None = None
description: str | None = None

View file

@ -1,9 +1,11 @@
from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.document import DocumentRepository
from haiku.rag.store.repositories.mm_asset import MMAssetRepository
from haiku.rag.store.repositories.settings import SettingsRepository from haiku.rag.store.repositories.settings import SettingsRepository
__all__ = [ __all__ = [
"ChunkRepository", "ChunkRepository",
"DocumentRepository", "DocumentRepository",
"MMAssetRepository",
"SettingsRepository", "SettingsRepository",
] ]

View file

@ -17,6 +17,7 @@ class DocumentRepository:
def __init__(self, store: Store) -> None: def __init__(self, store: Store) -> None:
self.store = store self.store = store
self._chunk_repository = None self._chunk_repository = None
self._mm_asset_repository = None
@property @property
def chunk_repository(self): def chunk_repository(self):
@ -27,6 +28,15 @@ class DocumentRepository:
self._chunk_repository = ChunkRepository(self.store) self._chunk_repository = ChunkRepository(self.store)
return self._chunk_repository return self._chunk_repository
@property
def mm_asset_repository(self):
"""Lazy-load MMAssetRepository when needed."""
if self._mm_asset_repository is None:
from haiku.rag.store.repositories.mm_asset import MMAssetRepository
self._mm_asset_repository = MMAssetRepository(self.store)
return self._mm_asset_repository
def _record_to_document(self, record: DocumentRecord) -> Document: def _record_to_document(self, record: DocumentRecord) -> Document:
"""Convert a DocumentRecord to a Document model.""" """Convert a DocumentRecord to a Document model."""
return Document( return Document(
@ -135,6 +145,10 @@ class DocumentRepository:
# Delete associated chunks first # Delete associated chunks first
await self.chunk_repository.delete_by_document_id(entity_id) await self.chunk_repository.delete_by_document_id(entity_id)
# Delete associated multimodal assets (if table exists)
if getattr(self.store, "mm_assets_table", None) is not None:
await self.mm_asset_repository.delete_by_document_id(entity_id)
# Delete the document # Delete the document
self.store.documents_table.delete(f"id = '{entity_id}'") self.store.documents_table.delete(f"id = '{entity_id}'")
return True return True

View file

@ -0,0 +1,130 @@
import json
import logging
from uuid import uuid4
from haiku.rag.store.engine import Store
from haiku.rag.store.models.mm_asset import MMAsset, MMSearchResult
logger = logging.getLogger(__name__)
class MMAssetRepository:
"""Repository for multimodal asset operations."""
def __init__(self, store: Store) -> None:
self.store = store
def _require_table(self):
if self.store.mm_assets_table is None:
raise ValueError(
"mm_assets table is not available. Enable config.multimodal.enabled "
"and open the database in writable mode at least once to create it."
)
return self.store.mm_assets_table
async def create(self, entity: MMAsset | list[MMAsset]) -> MMAsset | list[MMAsset]:
"""Create one or more multimodal assets.
Assets must have embedding set before calling this method.
"""
self.store._assert_writable()
table = self._require_table()
if isinstance(entity, MMAsset):
assert entity.embedding is not None, "MMAsset must have an embedding"
asset_id = str(uuid4())
rec = self.store.MMAssetRecord(
id=asset_id,
document_id=entity.document_id,
doc_item_ref=entity.doc_item_ref,
item_index=entity.item_index,
page_no=entity.page_no,
bbox=json.dumps(entity.bbox) if entity.bbox else None,
caption=entity.caption,
description=entity.description,
metadata=json.dumps(entity.metadata or {}),
vector=entity.embedding,
)
table.add([rec])
entity.id = asset_id
return entity
assets = entity
if not assets:
return []
for a in assets:
assert a.embedding is not None, "All MMAssets must have embeddings"
records = []
for a in assets:
asset_id = str(uuid4())
records.append(
self.store.MMAssetRecord(
id=asset_id,
document_id=a.document_id,
doc_item_ref=a.doc_item_ref,
item_index=a.item_index,
page_no=a.page_no,
bbox=json.dumps(a.bbox) if a.bbox else None,
caption=a.caption,
description=a.description,
metadata=json.dumps(a.metadata or {}),
vector=a.embedding,
)
)
a.id = asset_id
table.add(records)
return assets
async def delete_by_document_id(self, document_id: str) -> bool:
"""Delete all multimodal assets for a document."""
self.store._assert_writable()
table = self._require_table()
# Fast path: delete without pre-check; table.delete is idempotent.
table.delete(f"document_id = '{document_id}'")
return True
async def search_by_vector(
self,
query_vector: list[float],
*,
limit: int = 5,
filter: str | None = None,
) -> list[MMSearchResult]:
"""Vector search over multimodal assets."""
table = self._require_table()
q = table.search(query_vector).limit(limit)
if filter:
q = q.where(filter)
# LanceDB returns _distance for vector search; smaller is better.
# Convert to a similarity-like score for UX consistency.
rows = q.to_list()
results: list[MMSearchResult] = []
for row in rows:
distance = float(row.get("_distance", 0.0))
score = 1.0 / (1.0 + distance)
bbox = None
try:
bbox = json.loads(row["bbox"]) if row.get("bbox") else None
except Exception:
bbox = None
results.append(
MMSearchResult(
asset_id=row.get("id"),
document_id=row.get("document_id"),
score=score,
doc_item_ref=row.get("doc_item_ref", ""),
item_index=row.get("item_index"),
page_no=row.get("page_no"),
bbox=bbox,
caption=row.get("caption"),
description=row.get("description"),
)
)
return results

View file

@ -73,6 +73,7 @@ nav:
- Agents: agents.md - Agents: agents.md
- Server: server.md - Server: server.md
- Remote processing: remote-processing.md - Remote processing: remote-processing.md
- Multimodal search: multimodal.md
- MCP: mcp.md - MCP: mcp.md
- Inspector: inspector.md - Inspector: inspector.md
- Benchmarks: benchmarks.md - Benchmarks: benchmarks.md

View file

@ -0,0 +1,116 @@
import pytest
from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.store.engine import Store
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.mm_asset import MMAsset
def _mm_enabled_config():
cfg = Config.model_copy(deep=True)
cfg.multimodal.enabled = True
# Ensure dim matches our dummy vectors in tests
cfg.multimodal.model.vector_dim = 2048
return cfg
def test_mm_assets_table_created_when_enabled(temp_db_path):
cfg = _mm_enabled_config()
store = Store(temp_db_path, create=True, config=cfg)
assert store.mm_assets_table is not None
assert "mm_assets" in store.db.table_names()
async def test_version_rollback_on_mm_index_failure(temp_db_path, monkeypatch):
cfg = _mm_enabled_config()
async with HaikuRAG(db_path=temp_db_path, create=True, config=cfg) as client:
dim = int(client.store.embedder._vector_dim)
async def fail_after_writing_one_asset(doc):
# write one asset, then fail -> should rollback all tables
assert doc.id is not None
await client.mm_asset_repository.create(
MMAsset(
document_id=doc.id,
doc_item_ref="#/pictures/0",
item_index=0,
page_no=1,
bbox={"left": 0.0, "top": 1.0, "right": 2.0, "bottom": 3.0},
caption="dummy",
embedding=[0.0] * 2048,
)
)
raise RuntimeError("boom")
monkeypatch.setattr(client, "_index_mm_assets_for_document", fail_after_writing_one_asset)
# Instead of calling create_document (which would hit external embedders),
# call the internal store path with pre-embedded chunks.
from haiku.rag.store.models.document import Document
with pytest.raises(RuntimeError, match="boom"):
await client._store_document_with_chunks(
Document(content="x"),
[Chunk(content="c", metadata={}, order=0, embedding=[0.0] * dim)],
)
# After failure, database should be rolled back to empty state
store = Store(temp_db_path, create=True, config=cfg)
assert store.documents_table.count_rows() == 0
assert store.chunks_table.count_rows() == 0
assert store.mm_assets_table is not None
assert store.mm_assets_table.count_rows() == 0
async def test_update_rollback_restores_mm_assets(temp_db_path, monkeypatch):
cfg = _mm_enabled_config()
async with HaikuRAG(db_path=temp_db_path, create=True, config=cfg) as client:
dim = int(client.store.embedder._vector_dim)
# Create baseline document without calling external embedders
from haiku.rag.store.models.document import Document
doc = await client._store_document_with_chunks(
Document(content="base"),
[Chunk(content="c", metadata={}, order=0, embedding=[0.0] * dim)],
)
assert doc.id is not None
# Seed mm_assets with one row
await client.mm_asset_repository.create(
MMAsset(
document_id=doc.id,
doc_item_ref="#/pictures/seed",
item_index=0,
page_no=1,
bbox={"left": 0.0, "top": 1.0, "right": 2.0, "bottom": 3.0},
caption="seed",
embedding=[1.0] * 2048,
)
)
async def fail_index(_doc):
raise RuntimeError("mm index fail")
monkeypatch.setattr(client, "_index_mm_assets_for_document", fail_index)
with pytest.raises(RuntimeError, match="mm index fail"):
# Call internal update path to avoid external embedder usage
await client._update_document_with_chunks(
Document(
id=doc.id,
content="changed",
docling_document_json=doc.docling_document_json,
docling_version=doc.docling_version,
),
[Chunk(content="c2", metadata={}, order=0, embedding=[0.0] * dim)],
)
# After rollback, original mm_assets row should still exist
store = Store(temp_db_path, create=True, config=cfg)
assert store.mm_assets_table is not None
assert store.mm_assets_table.count_rows() == 1