diff --git a/docs/cli.md b/docs/cli.md index db1ced53..9e831644 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -104,6 +104,34 @@ haiku-rag visualize 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 --mode crop +haiku-rag visualize-asset --mode page +``` + !!! note Requires a terminal with image support (iTerm2, Kitty, WezTerm, etc.) and documents processed with docling that have page images stored. diff --git a/docs/configuration/providers.md b/docs/configuration/providers.md index 12bb9603..39fe667a 100644 --- a/docs/configuration/providers.md +++ b/docs/configuration/providers.md @@ -190,6 +190,27 @@ embeddings: **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 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. diff --git a/docs/inspector.md b/docs/inspector.md index a5250821..a5b8c365 100644 --- a/docs/inspector.md +++ b/docs/inspector.md @@ -38,6 +38,8 @@ Three panels display your data: - `/` - Open search modal - `c` - Open context expansion 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 **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 - 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 !!! 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. diff --git a/docs/multimodal.md b/docs/multimodal.md new file mode 100644 index 00000000..0ec568de --- /dev/null +++ b/docs/multimodal.md @@ -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 + Qwen3‑VL‑Embedding‑2B, 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":"","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 --mode crop +haiku-rag visualize-asset --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 +``` diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index c2e94cd4..1cbfe00e 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -22,6 +22,8 @@ from haiku.rag.config.models import ModelConfig from haiku.rag.logging import configure_cli_logging from haiku.rag.qa import get_qa_agent 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() @@ -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: + - text→image cases using mm_assets.caption/description when available + - image→image 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 --db ", + style="cyan", + ) + + if __name__ == "__main__": app() diff --git a/evaluations/evaluations/datasets/mm_eval_example.jsonl b/evaluations/evaluations/datasets/mm_eval_example.jsonl new file mode 100644 index 00000000..e0824363 --- /dev/null +++ b/evaluations/evaluations/datasets/mm_eval_example.jsonl @@ -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}]} + diff --git a/evaluations/evaluations/mm_dataset_builder.py b/evaluations/evaluations/mm_dataset_builder.py new file mode 100644 index 00000000..747fa54b --- /dev/null +++ b/evaluations/evaluations/mm_dataset_builder.py @@ -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: + - image→image queries use the exact stored bbox crop (self-retrieval) + - text→image 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)) + diff --git a/evaluations/evaluations/mm_eval.py b/evaluations/evaluations/mm_eval.py new file mode 100644 index 00000000..ea86b5aa --- /dev/null +++ b/evaluations/evaluations/mm_eval.py @@ -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)) + diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 1c76b4ff..cae1fff7 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -367,8 +367,50 @@ class HaikuRAGApp: for result in results: 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): """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 async with HaikuRAG( @@ -404,6 +446,87 @@ class HaikuRAGApp: ) 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( self, question: str, @@ -811,6 +934,22 @@ class HaikuRAGApp: self.console.print(" > ".join(result.headings)) self.console.print("[repr.attrib_name]content[/repr.attrib_name]:") 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() async def serve( diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index 65c0479a..79fc0061 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -297,6 +297,60 @@ def search( 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") def visualize( chunk_id: str = typer.Argument( @@ -312,6 +366,26 @@ def visualize( 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") def ask( question: str = typer.Argument( diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 48592048..b9c96546 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -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.repositories.chunk import ChunkRepository 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 if TYPE_CHECKING: @@ -88,6 +89,7 @@ class HaikuRAG: ) self.document_repository = DocumentRepository(self.store) self.chunk_repository = ChunkRepository(self.store) + self.mm_asset_repository = MMAssetRepository(self.store) @property def is_read_only(self) -> bool: @@ -279,6 +281,9 @@ class HaikuRAG: # Batch create all chunks in a single operation 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 if self._config.storage.auto_vacuum: asyncio.create_task(self.store.vacuum()) @@ -317,6 +322,9 @@ class HaikuRAG: # Delete existing chunks before writing new ones 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: # Update the document @@ -331,6 +339,9 @@ class HaikuRAG: # Batch create all chunks in a single operation 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 if self._config.storage.auto_vacuum: asyncio.create_task(self.store.vacuum()) @@ -821,6 +832,150 @@ class HaikuRAG: """Delete a document by its 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( self, limit: int | None = None, @@ -1193,6 +1348,72 @@ class HaikuRAG: 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( self, doc_id: str, @@ -1357,6 +1578,122 @@ class HaikuRAG: 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( self, mode: RebuildMode = RebuildMode.FULL ) -> AsyncGenerator[str, None]: diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 5462056f..37f16401 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -43,6 +43,64 @@ class EmbeddingModelConfig(BaseModel): 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): data_dir: Path = Field(default_factory=get_default_data_dir) auto_vacuum: bool = True @@ -198,6 +256,7 @@ class AppConfig(BaseModel): monitor: MonitorConfig = Field(default_factory=MonitorConfig) lancedb: LanceDBConfig = Field(default_factory=LanceDBConfig) embeddings: EmbeddingsConfig = Field(default_factory=EmbeddingsConfig) + multimodal: MultimodalConfig = Field(default_factory=MultimodalConfig) reranking: RerankingConfig = Field(default_factory=RerankingConfig) qa: QAConfig = Field(default_factory=QAConfig) research: ResearchConfig = Field(default_factory=ResearchConfig) diff --git a/haiku_rag_slim/haiku/rag/embeddings/multimodal.py b/haiku_rag_slim/haiku/rag/embeddings/multimodal.py new file mode 100644 index 00000000..d113fd5f --- /dev/null +++ b/haiku_rag_slim/haiku/rag/embeddings/multimodal.py @@ -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 = "" + 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 = "" + 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}") + diff --git a/haiku_rag_slim/haiku/rag/inspector/app.py b/haiku_rag_slim/haiku/rag/inspector/app.py index fd1b761d..95e8c913 100644 --- a/haiku_rag_slim/haiku/rag/inspector/app.py +++ b/haiku_rag_slim/haiku/rag/inspector/app.py @@ -73,6 +73,8 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover Binding("i", "show_info", "Info", show=True), Binding("v", "show_visual", "Visual", 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__( @@ -83,6 +85,9 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover self.read_only = read_only self.before = before 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": """Compose the UI layout.""" @@ -100,6 +105,14 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover config=config, read_only=self.read_only, 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__() @@ -191,6 +204,9 @@ class InspectorApp(App): # type: ignore[misc] # pragma: no cover # Load chunks for this document 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) 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)) + 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( db_path: Path | None = None, diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/doc_index_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/doc_index_modal.py new file mode 100644 index 00000000..c1c987bb --- /dev/null +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/doc_index_modal.py @@ -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() + diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/mm_assets_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/mm_assets_modal.py new file mode 100644 index 00000000..b7f5ead8 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/mm_assets_modal.py @@ -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() + diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py index eb09ba76..377fbd83 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/visual_modal.py @@ -7,6 +7,11 @@ from textual.screen import Screen from textual.widget import Widget from textual.widgets import Static from textual_image.widget import Image as TextualImage +import shutil +import tempfile +from pathlib import Path +import subprocess +import sys if TYPE_CHECKING: from PIL.Image import Image as PILImage @@ -23,6 +28,7 @@ class VisualGroundingModal(Screen): # pragma: no cover Binding("v", "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 = """ @@ -66,6 +72,8 @@ class VisualGroundingModal(Screen): # pragma: no cover self.client = client self.document_uri = document_uri or chunk.document_uri self.images: list[PILImage] = [] + self._saved_paths: list[Path] = [] + self._out_dir: Path | None = None self.current_page_idx = 0 self._image_widget: Widget = Static("Loading...", id="image-display") self._page_info = Static("", id="page-info") @@ -82,6 +90,18 @@ class VisualGroundingModal(Screen): # pragma: no cover async def on_mount(self) -> None: """Load images and display the first page.""" 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() async def _render_current_page(self) -> None: @@ -95,8 +115,11 @@ class VisualGroundingModal(Screen): # pragma: no cover self._page_info.update("") return + extra = "" + if self._out_dir is not None: + extra = f" | Saved to: {self._out_dir} (press 'o' to open)" 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: @@ -111,6 +134,12 @@ class VisualGroundingModal(Screen): # pragma: no cover self._image_widget.update(f"[red]Error: {e}[/red]") 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() 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: self.current_page_idx += 1 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 diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 6e2b9ace..b3b3713a 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -53,6 +53,24 @@ class SettingsRecord(LanceModel): 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: def __init__( self, @@ -73,6 +91,9 @@ class Store: # Create the ChunkRecord model with the correct vector dimension 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) is_new_db = False @@ -156,6 +177,12 @@ class Store: self.documents_table, self.chunks_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) except (RuntimeError, OSError) as e: @@ -195,6 +222,7 @@ class Store: stats_dict: dict = { "documents": {"exists": False}, "chunks": {"exists": False}, + "mm_assets": {"exists": False}, } # Documents table stats @@ -307,6 +335,54 @@ class Store: [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): """Set the initial version for a new database.""" self.set_haiku_version(metadata.version("haiku.rag-slim")) @@ -405,11 +481,14 @@ class Store: def current_table_versions(self) -> dict[str, int]: """Capture current versions of key tables for rollback using LanceDB's API.""" - return { + versions = { "documents": int(self.documents_table.version), "chunks": int(self.chunks_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: """Restore tables to the provided versions using LanceDB's API. @@ -421,6 +500,8 @@ class Store: self.documents_table.restore(int(versions["documents"])) self.chunks_table.restore(int(versions["chunks"])) 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 @property @@ -450,6 +531,11 @@ class Store: ("documents", self.documents_table), ("chunks", self.chunks_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: @@ -501,6 +587,7 @@ class Store: "documents": self.documents_table, "chunks": self.chunks_table, "settings": self.settings_table, + "mm_assets": self.mm_assets_table, } table = table_map.get(table_name) if table is None: diff --git a/haiku_rag_slim/haiku/rag/store/models/__init__.py b/haiku_rag_slim/haiku/rag/store/models/__init__.py index d1cc810b..fc9533ea 100644 --- a/haiku_rag_slim/haiku/rag/store/models/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/models/__init__.py @@ -1,10 +1,13 @@ from .chunk import BoundingBox, Chunk, ChunkMetadata, SearchResult from .document import Document +from .mm_asset import MMAsset, MMSearchResult __all__ = [ "BoundingBox", "Chunk", "ChunkMetadata", "Document", + "MMAsset", + "MMSearchResult", "SearchResult", ] diff --git a/haiku_rag_slim/haiku/rag/store/models/mm_asset.py b/haiku_rag_slim/haiku/rag/store/models/mm_asset.py new file mode 100644 index 00000000..38962e40 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/models/mm_asset.py @@ -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 diff --git a/haiku_rag_slim/haiku/rag/store/repositories/__init__.py b/haiku_rag_slim/haiku/rag/store/repositories/__init__.py index be72d771..3282693a 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/__init__.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/__init__.py @@ -1,9 +1,11 @@ from haiku.rag.store.repositories.chunk import ChunkRepository 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 __all__ = [ "ChunkRepository", "DocumentRepository", + "MMAssetRepository", "SettingsRepository", ] diff --git a/haiku_rag_slim/haiku/rag/store/repositories/document.py b/haiku_rag_slim/haiku/rag/store/repositories/document.py index 3e97042f..ec3b4f44 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/document.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/document.py @@ -17,6 +17,7 @@ class DocumentRepository: def __init__(self, store: Store) -> None: self.store = store self._chunk_repository = None + self._mm_asset_repository = None @property def chunk_repository(self): @@ -27,6 +28,15 @@ class DocumentRepository: self._chunk_repository = ChunkRepository(self.store) 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: """Convert a DocumentRecord to a Document model.""" return Document( @@ -135,6 +145,10 @@ class DocumentRepository: # Delete associated chunks first 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 self.store.documents_table.delete(f"id = '{entity_id}'") return True diff --git a/haiku_rag_slim/haiku/rag/store/repositories/mm_asset.py b/haiku_rag_slim/haiku/rag/store/repositories/mm_asset.py new file mode 100644 index 00000000..2a176b11 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/store/repositories/mm_asset.py @@ -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 + diff --git a/mkdocs.yml b/mkdocs.yml index 86a50ed6..a6ccc952 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -73,6 +73,7 @@ nav: - Agents: agents.md - Server: server.md - Remote processing: remote-processing.md + - Multimodal search: multimodal.md - MCP: mcp.md - Inspector: inspector.md - Benchmarks: benchmarks.md diff --git a/tests/test_multimodal_phase1.py b/tests/test_multimodal_phase1.py new file mode 100644 index 00000000..40f352a3 --- /dev/null +++ b/tests/test_multimodal_phase1.py @@ -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 +