feat(multimodal): add mm_assets image retrieval via vLLM
Phase 1: orthogonal multimodal index (mm_assets) for text→image and image→image search. - Add multimodal config + vLLM embedder client - Index Docling PictureItems (bbox crops) into mm_assets with rollback support - Add CLI commands: search-image-text, search-image, visualize-asset - Extend inspector with mm_assets browser + doc index overlay - Add multimodal eval commands (mm-build/mm) and docs
This commit is contained in:
parent
ff1f19bb08
commit
79a3617782
25 changed files with 2673 additions and 3 deletions
28
docs/cli.md
28
docs/cli.md
|
|
@ -104,6 +104,34 @@ haiku-rag visualize <chunk_id>
|
|||
|
||||
This renders the source document pages with the chunk's location highlighted. Useful for verifying chunk boundaries and understanding document structure.
|
||||
|
||||
!!! note
|
||||
Requires a terminal with image support (iTerm2, Kitty, WezTerm, etc.) and documents processed with docling that have page images stored.
|
||||
|
||||
## Search Images (Multimodal)
|
||||
|
||||
These commands require `multimodal.enabled: true` and an embedding server configured under `multimodal.model` (typically vLLM).
|
||||
|
||||
### Text → image
|
||||
|
||||
```bash
|
||||
haiku-rag search-image-text "architecture diagram"
|
||||
```
|
||||
|
||||
### Image → image
|
||||
|
||||
```bash
|
||||
haiku-rag search-image /path/to/query.png
|
||||
```
|
||||
|
||||
## Visualize Asset (Multimodal)
|
||||
|
||||
After `search-image*`, you can visualize an `asset_id`:
|
||||
|
||||
```bash
|
||||
haiku-rag visualize-asset <asset_id> --mode crop
|
||||
haiku-rag visualize-asset <asset_id> --mode page
|
||||
```
|
||||
|
||||
!!! note
|
||||
Requires a terminal with image support (iTerm2, Kitty, WezTerm, etc.) and documents processed with docling that have page images stored.
|
||||
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
175
docs/multimodal.md
Normal file
175
docs/multimodal.md
Normal file
|
|
@ -0,0 +1,175 @@
|
|||
# Multimodal Search (Images)
|
||||
|
||||
haiku.rag supports **multimodal embeddings** for:
|
||||
- **text → image** search (`mm_assets` table)
|
||||
- **image → image** search (`mm_assets` table)
|
||||
|
||||
This feature is **orthogonal** to the existing text chunk search (`chunks` table). Your existing `haiku-rag search` behavior does not change.
|
||||
|
||||
---
|
||||
|
||||
## What gets indexed (Phase 1)
|
||||
|
||||
When enabled, haiku.rag extracts Docling `PictureItem`s and indexes:
|
||||
- An **image crop** from the Docling page image using the item bbox (plus padding).
|
||||
- Optional text metadata if available (caption/description) stored alongside the vector.
|
||||
|
||||
Each indexed image becomes one row in `mm_assets`, linked to its `document_id` and bbox provenance (page + coordinates), so it can be visualized later.
|
||||
|
||||
---
|
||||
|
||||
## Requirements
|
||||
|
||||
- Your converter must produce **Docling page images** (needed for bbox crops and visualization).
|
||||
- A multimodal embedding backend must be reachable.
|
||||
|
||||
**Current supported backend**
|
||||
- `provider: vllm` via OpenAI-compatible `POST /v1/embeddings`
|
||||
|
||||
---
|
||||
|
||||
## Configuration
|
||||
|
||||
Add the following to your `haiku.rag.yaml`:
|
||||
|
||||
```yaml
|
||||
multimodal:
|
||||
enabled: true
|
||||
index_pictures: true
|
||||
|
||||
# bbox → crop settings
|
||||
image_crop_padding_px: 8
|
||||
|
||||
# guardrails / throughput knobs
|
||||
image_max_side_px: 1024
|
||||
embed_batch_size: 8
|
||||
|
||||
model:
|
||||
provider: vllm
|
||||
name: Qwen/Qwen3-VL-Embedding-2B
|
||||
vector_dim: 2048
|
||||
base_url: http://localhost:8000 # with or without trailing /v1
|
||||
timeout: 60
|
||||
# dimensions: 1024 # only if your backend/model supports it
|
||||
encoding_format: float # float | base64
|
||||
```
|
||||
|
||||
### Multimodal toggles
|
||||
|
||||
- **`multimodal.enabled`**: master switch. When `false`, haiku.rag will not create/use `mm_assets`.
|
||||
- **`multimodal.index_pictures`**: when `true`, Docling `PictureItem`s are indexed during `add-src` / update.
|
||||
|
||||
### Crop / image settings
|
||||
|
||||
- **`image_crop_padding_px`**: extra pixels added around the bbox crop. Helps avoid tight crops clipping labels/axes.
|
||||
- **`image_max_side_px`**: resize guardrail before upload (keeps requests smaller, avoids backend limits). Set `0` to disable resizing.
|
||||
|
||||
### Throughput settings
|
||||
|
||||
- **`embed_batch_size`**:
|
||||
- Used as the **maximum in-flight requests** for image embedding to avoid overwhelming the backend.
|
||||
- Text embeddings may be batched using `input=[...]`, but image embeddings currently run as **one request per image** (because the backend expects multimodal inputs via `messages`).
|
||||
|
||||
### Model / backend settings
|
||||
|
||||
- **`provider`**: currently only `vllm` is supported (OpenAI-compatible server).
|
||||
- **`name`**: model identifier sent to the backend (e.g. `Qwen/Qwen3-VL-Embedding-2B`).
|
||||
- **`vector_dim`**: expected embedding size. This must stay stable for an existing DB (same rule as text embeddings).
|
||||
- **`base_url`**: server base URL; haiku.rag accepts either `http://host:port` or `http://host:port/v1`.
|
||||
- **`timeout`**: request timeout in seconds.
|
||||
- **`dimensions`**: optional request-time output dimension (only for models that support it).
|
||||
- **`encoding_format`**: `float` is recommended; `base64` is supported by some servers.
|
||||
|
||||
!!! note
|
||||
Some servers/models reject the presence of `dimensions` unless matryoshka output is supported.
|
||||
For vLLM + 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":"<YOUR_MODEL>","input":["hello"],"encoding_format":"float"}' \
|
||||
| python3 -c 'import sys,json; print(len(json.load(sys.stdin)["data"][0]["embedding"]))'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Run vLLM (example)
|
||||
|
||||
Serve the embedding model:
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-VL-Embedding-2B --runner=pooling --host 127.0.0.1 --port 8000
|
||||
```
|
||||
|
||||
Quick checks:
|
||||
|
||||
```bash
|
||||
curl -s http://127.0.0.1:8000/v1/models
|
||||
curl -s http://127.0.0.1:8000/v1/embeddings \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model":"Qwen/Qwen3-VL-Embedding-2B","input":["hello"],"encoding_format":"float"}'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Index images (Phase 1)
|
||||
|
||||
Index documents as usual (multimodal indexing runs during document create/update if enabled):
|
||||
|
||||
```bash
|
||||
haiku-rag add-src /path/to/doc.pdf
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Search images
|
||||
|
||||
### Text → image
|
||||
|
||||
```bash
|
||||
haiku-rag search-image-text "architecture diagram"
|
||||
```
|
||||
|
||||
### Image → image
|
||||
|
||||
```bash
|
||||
haiku-rag search-image /path/to/query.png
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Visualize results
|
||||
|
||||
Search returns `asset_id` values. To see what was actually indexed:
|
||||
|
||||
```bash
|
||||
haiku-rag visualize-asset <asset_id> --mode crop
|
||||
haiku-rag visualize-asset <asset_id> --mode page
|
||||
```
|
||||
|
||||
!!! note
|
||||
Terminals often do not render images inline (you may see gray blocks).
|
||||
Prefer opening the saved PNG in your OS viewer via:
|
||||
- `haiku-rag visualize-asset ...` (prints the saved image path)
|
||||
- `haiku-rag inspect` (`o` key inside the visual modals)
|
||||
|
||||
---
|
||||
|
||||
## Minimal accuracy validation (recommended)
|
||||
|
||||
Generate a quick “sanity” dataset from your DB and compute recall@k + MRR:
|
||||
|
||||
```bash
|
||||
uv run evaluations mm-build --db /path/to/your.lancedb --config /path/to/haiku.rag.yaml --out ./mm_eval_out --n 50
|
||||
uv run evaluations mm ./mm_eval_out/mm_eval.jsonl --db /path/to/your.lancedb --config /path/to/haiku.rag.yaml --k 1,5,10 --limit 10
|
||||
```
|
||||
|
|
@ -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 <config> --db <db>",
|
||||
style="cyan",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
|
|
|
|||
3
evaluations/evaluations/datasets/mm_eval_example.jsonl
Normal file
3
evaluations/evaluations/datasets/mm_eval_example.jsonl
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
{"id":"t1","type":"text","instruction":"Retrieve images matching this description.","query_text":"architecture diagram","relevant":[{"document_uri":"file:///path/to/your.pdf","doc_item_ref":"#/pictures/0","item_index":123}]}
|
||||
{"id":"i1","type":"image","query_image_path":"/absolute/path/to/query_crop.png","relevant":[{"document_uri":"file:///path/to/your.pdf","doc_item_ref":"#/pictures/0","item_index":123}]}
|
||||
|
||||
239
evaluations/evaluations/mm_dataset_builder.py
Normal file
239
evaluations/evaluations/mm_dataset_builder.py
Normal file
|
|
@ -0,0 +1,239 @@
|
|||
import asyncio
|
||||
import json
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import AppConfig
|
||||
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BuiltCase:
|
||||
id: str
|
||||
type: str # "text" | "image"
|
||||
instruction: str | None
|
||||
query_text: str | None
|
||||
query_image_path: str | None
|
||||
relevant: list[dict]
|
||||
|
||||
def to_jsonl(self) -> str:
|
||||
obj = {
|
||||
"id": self.id,
|
||||
"type": self.type,
|
||||
"relevant": self.relevant,
|
||||
}
|
||||
if self.instruction is not None:
|
||||
obj["instruction"] = self.instruction
|
||||
if self.query_text is not None:
|
||||
obj["query_text"] = self.query_text
|
||||
if self.query_image_path is not None:
|
||||
obj["query_image_path"] = self.query_image_path
|
||||
return json.dumps(obj, ensure_ascii=False)
|
||||
|
||||
|
||||
def _crop_from_doc_bbox(
|
||||
*,
|
||||
pil_image,
|
||||
page_width: float,
|
||||
page_height: float,
|
||||
bbox: dict,
|
||||
padding_px: int,
|
||||
):
|
||||
# bbox is in Docling page coordinates with bottom-left origin.
|
||||
left = float(bbox["left"])
|
||||
top = float(bbox["top"])
|
||||
right = float(bbox["right"])
|
||||
bottom = float(bbox["bottom"])
|
||||
|
||||
scale_x = pil_image.width / page_width
|
||||
scale_y = pil_image.height / page_height
|
||||
|
||||
x0 = left * scale_x
|
||||
x1 = right * scale_x
|
||||
# invert Y axis: doc bottom-left -> PIL top-left
|
||||
y0 = (page_height - top) * scale_y
|
||||
y1 = (page_height - bottom) * scale_y
|
||||
if y0 > y1:
|
||||
y0, y1 = y1, y0
|
||||
|
||||
x0 = max(0, int(x0) - padding_px)
|
||||
y0 = max(0, int(y0) - padding_px)
|
||||
x1 = min(pil_image.width, int(x1) + padding_px)
|
||||
y1 = min(pil_image.height, int(y1) + padding_px)
|
||||
|
||||
if x1 <= x0 or y1 <= y0:
|
||||
return None
|
||||
return pil_image.crop((x0, y0, x1, y1))
|
||||
|
||||
|
||||
async def build_mm_dataset(
|
||||
*,
|
||||
config: AppConfig,
|
||||
db_path: Path | None,
|
||||
out_dir: Path,
|
||||
n: int = 50,
|
||||
seed: int = 0,
|
||||
include_text: bool = True,
|
||||
include_image: bool = True,
|
||||
instruction: str = "Retrieve images matching this description.",
|
||||
) -> Path:
|
||||
"""Auto-generate a JSONL dataset for `evaluations mm` from an existing DB.
|
||||
|
||||
This produces a *sanity* dataset:
|
||||
- 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))
|
||||
|
||||
263
evaluations/evaluations/mm_eval.py
Normal file
263
evaluations/evaluations/mm_eval.py
Normal file
|
|
@ -0,0 +1,263 @@
|
|||
import asyncio
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from rich.console import Console
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import AppConfig
|
||||
|
||||
|
||||
console = Console()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MMRelevantRef:
|
||||
"""A stable reference to an mm_asset row (avoid UUIDs in datasets)."""
|
||||
|
||||
document_uri: str
|
||||
doc_item_ref: str
|
||||
item_index: int | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MMEvalCase:
|
||||
"""One multimodal retrieval evaluation case."""
|
||||
|
||||
case_id: str
|
||||
query_type: str # "text" | "image"
|
||||
instruction: str | None
|
||||
query_text: str | None
|
||||
query_image_path: str | None
|
||||
relevant: tuple[MMRelevantRef, ...]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CaseMetrics:
|
||||
recall_at_k: dict[int, float]
|
||||
mrr: float
|
||||
|
||||
|
||||
def _quote_lancedb(s: str) -> str:
|
||||
# LanceDB uses SQL-ish string quoting with single quotes.
|
||||
return s.replace("'", "''")
|
||||
|
||||
|
||||
def load_mm_dataset(dataset_path: Path) -> list[MMEvalCase]:
|
||||
cases: list[MMEvalCase] = []
|
||||
for raw_line in dataset_path.read_text(encoding="utf-8").splitlines():
|
||||
line = raw_line.strip()
|
||||
if not line or line.startswith("#"):
|
||||
continue
|
||||
obj = json.loads(line)
|
||||
|
||||
case_id = str(obj.get("id") or obj.get("case_id") or "")
|
||||
if not case_id:
|
||||
raise ValueError("Each JSONL row must have a non-empty 'id'.")
|
||||
|
||||
query_type = obj.get("type")
|
||||
if query_type not in ("text", "image"):
|
||||
raise ValueError(f"Case {case_id}: 'type' must be 'text' or 'image'.")
|
||||
|
||||
instruction = obj.get("instruction")
|
||||
query_text = obj.get("query_text")
|
||||
query_image_path = obj.get("query_image_path")
|
||||
|
||||
if query_type == "text" and not query_text:
|
||||
raise ValueError(f"Case {case_id}: missing 'query_text'.")
|
||||
if query_type == "image" and not query_image_path:
|
||||
raise ValueError(f"Case {case_id}: missing 'query_image_path'.")
|
||||
|
||||
rel = obj.get("relevant") or []
|
||||
if not isinstance(rel, list) or not rel:
|
||||
raise ValueError(f"Case {case_id}: 'relevant' must be a non-empty list.")
|
||||
|
||||
relevant: list[MMRelevantRef] = []
|
||||
for r in rel:
|
||||
if not isinstance(r, dict):
|
||||
raise ValueError(f"Case {case_id}: relevant entries must be objects.")
|
||||
document_uri = r.get("document_uri")
|
||||
doc_item_ref = r.get("doc_item_ref")
|
||||
if not document_uri or not doc_item_ref:
|
||||
raise ValueError(
|
||||
f"Case {case_id}: relevant entries must include "
|
||||
"'document_uri' and 'doc_item_ref'."
|
||||
)
|
||||
item_index = r.get("item_index")
|
||||
relevant.append(
|
||||
MMRelevantRef(
|
||||
document_uri=str(document_uri),
|
||||
doc_item_ref=str(doc_item_ref),
|
||||
item_index=int(item_index) if item_index is not None else None,
|
||||
)
|
||||
)
|
||||
|
||||
cases.append(
|
||||
MMEvalCase(
|
||||
case_id=case_id,
|
||||
query_type=query_type,
|
||||
instruction=str(instruction) if instruction is not None else None,
|
||||
query_text=query_text,
|
||||
query_image_path=query_image_path,
|
||||
relevant=tuple(relevant),
|
||||
)
|
||||
)
|
||||
return cases
|
||||
|
||||
|
||||
async def _resolve_relevant_asset_ids(
|
||||
rag: HaikuRAG, cases: list[MMEvalCase]
|
||||
) -> dict[str, set[str]]:
|
||||
"""Resolve stable refs (uri + doc_item_ref [+ item_index]) into asset UUIDs."""
|
||||
if rag.store.mm_assets_table is None:
|
||||
raise ValueError("mm_assets table is not available (multimodal not enabled?).")
|
||||
|
||||
# Resolve document_id per uri.
|
||||
uris = sorted({r.document_uri for c in cases for r in c.relevant})
|
||||
uri_to_doc_id: dict[str, str] = {}
|
||||
for uri in uris:
|
||||
doc = await rag.get_document_by_uri(uri)
|
||||
if doc is None or not doc.id:
|
||||
raise ValueError(f"Dataset references unknown document uri: {uri}")
|
||||
uri_to_doc_id[uri] = doc.id
|
||||
|
||||
# Load mm_assets rows once (evaluation datasets are small; this is simplest).
|
||||
# Recommended way to materialize the *entire table* is via Arrow/Pandas.
|
||||
table = rag.store.mm_assets_table
|
||||
try:
|
||||
rows = table.to_arrow().to_pylist()
|
||||
except Exception:
|
||||
rows = table.search().limit(100_000).to_list()
|
||||
|
||||
# Build lookup maps.
|
||||
by_doc_ref_idx: dict[tuple[str, str, int], str] = {}
|
||||
by_doc_ref: dict[tuple[str, str], list[str]] = {}
|
||||
for row in rows:
|
||||
asset_id = row.get("id")
|
||||
doc_id = row.get("document_id")
|
||||
doc_item_ref = row.get("doc_item_ref")
|
||||
item_index = row.get("item_index")
|
||||
if not asset_id or not doc_id or not doc_item_ref:
|
||||
continue
|
||||
if item_index is not None:
|
||||
try:
|
||||
by_doc_ref_idx[(str(doc_id), str(doc_item_ref), int(item_index))] = str(
|
||||
asset_id
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
by_doc_ref.setdefault((str(doc_id), str(doc_item_ref)), []).append(str(asset_id))
|
||||
|
||||
case_to_relevant_ids: dict[str, set[str]] = {}
|
||||
for c in cases:
|
||||
ids: set[str] = set()
|
||||
for r in c.relevant:
|
||||
doc_id = uri_to_doc_id[r.document_uri]
|
||||
if r.item_index is not None:
|
||||
k = (doc_id, r.doc_item_ref, int(r.item_index))
|
||||
asset_id = by_doc_ref_idx.get(k)
|
||||
if asset_id:
|
||||
ids.add(asset_id)
|
||||
continue
|
||||
# Fallback: any asset(s) matching doc+ref.
|
||||
ids.update(by_doc_ref.get((doc_id, r.doc_item_ref), []))
|
||||
if not ids:
|
||||
raise ValueError(
|
||||
f"Case {c.case_id}: could not resolve any relevant asset ids from refs. "
|
||||
"Make sure the DB was built with multimodal indexing enabled and the "
|
||||
"doc_item_ref/item_index match your Docling conversion."
|
||||
)
|
||||
case_to_relevant_ids[c.case_id] = ids
|
||||
|
||||
return case_to_relevant_ids
|
||||
|
||||
|
||||
def _compute_case_metrics(
|
||||
*, retrieved_ids: list[str], relevant_ids: set[str], ks: list[int]
|
||||
) -> CaseMetrics:
|
||||
recall_at_k: dict[int, float] = {}
|
||||
for k in ks:
|
||||
topk = retrieved_ids[:k]
|
||||
hits = sum(1 for a in topk if a in relevant_ids)
|
||||
recall_at_k[k] = hits / max(len(relevant_ids), 1)
|
||||
|
||||
rr = 0.0
|
||||
for idx, a in enumerate(retrieved_ids, start=1):
|
||||
if a in relevant_ids:
|
||||
rr = 1.0 / idx
|
||||
break
|
||||
|
||||
return CaseMetrics(recall_at_k=recall_at_k, mrr=rr)
|
||||
|
||||
|
||||
async def run_mm_eval(
|
||||
*,
|
||||
dataset_path: Path,
|
||||
config: AppConfig,
|
||||
db_path: Path | None,
|
||||
ks: list[int],
|
||||
limit: int,
|
||||
) -> int:
|
||||
cases = load_mm_dataset(dataset_path)
|
||||
if not cases:
|
||||
raise ValueError("No cases found.")
|
||||
|
||||
async with HaikuRAG(db_path, config=config) as rag:
|
||||
if not rag._config.multimodal.enabled:
|
||||
raise ValueError(
|
||||
"Multimodal is disabled in config. Set `multimodal.enabled: true` "
|
||||
"and ensure mm_assets is created."
|
||||
)
|
||||
|
||||
# Resolve ground-truth.
|
||||
relevant_ids = await _resolve_relevant_asset_ids(rag, cases)
|
||||
|
||||
totals_recall = {k: 0.0 for k in ks}
|
||||
totals_mrr = 0.0
|
||||
|
||||
for c in cases:
|
||||
if c.query_type == "text":
|
||||
# The Qwen3-VL embedding examples use an "instruction" field that can
|
||||
# materially affect retrieval quality. vLLM /v1/embeddings typically
|
||||
# doesn't expose a separate instruction channel, so we fold it into
|
||||
# the query text in a stable way for evaluation.
|
||||
q = c.query_text or ""
|
||||
if c.instruction:
|
||||
q = f"{c.instruction}\n{q}"
|
||||
results = await rag.search_images_by_text(
|
||||
q, limit=limit
|
||||
)
|
||||
else:
|
||||
img_path = Path(c.query_image_path or "")
|
||||
if not img_path.exists():
|
||||
raise ValueError(
|
||||
f"Case {c.case_id}: query_image_path not found: {img_path}"
|
||||
)
|
||||
results = await rag.search_images(img_path, limit=limit)
|
||||
|
||||
retrieved = [r.asset_id for r in results if r.asset_id]
|
||||
m = _compute_case_metrics(
|
||||
retrieved_ids=retrieved,
|
||||
relevant_ids=relevant_ids[c.case_id],
|
||||
ks=ks,
|
||||
)
|
||||
|
||||
for k in ks:
|
||||
totals_recall[k] += m.recall_at_k[k]
|
||||
totals_mrr += m.mrr
|
||||
|
||||
n = float(len(cases))
|
||||
console.print("\n=== Multimodal Retrieval Benchmark Results ===", style="bold cyan")
|
||||
console.print(f"Dataset: {dataset_path}")
|
||||
console.print(f"Total queries: {len(cases)}")
|
||||
for k in ks:
|
||||
console.print(f"recall@{k}: {totals_recall[k] / n:.4f}")
|
||||
console.print(f"mrr: {totals_mrr / n:.4f}")
|
||||
|
||||
return 0
|
||||
|
||||
|
||||
def run_mm_eval_sync(**kwargs) -> int:
|
||||
return asyncio.run(run_mm_eval(**kwargs))
|
||||
|
||||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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]:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
238
haiku_rag_slim/haiku/rag/embeddings/multimodal.py
Normal file
238
haiku_rag_slim/haiku/rag/embeddings/multimodal.py
Normal file
|
|
@ -0,0 +1,238 @@
|
|||
import base64
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
|
||||
import httpx
|
||||
|
||||
from haiku.rag.config import AppConfig, Config
|
||||
|
||||
|
||||
@dataclass
|
||||
class MultimodalEmbedInput:
|
||||
"""Input for a multimodal embedding request."""
|
||||
|
||||
text: str | None = None
|
||||
image_b64: str | None = None
|
||||
|
||||
|
||||
class MultimodalEmbedderBase:
|
||||
"""Base interface for multimodal embedders (text + image)."""
|
||||
|
||||
vector_dim: int
|
||||
|
||||
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
|
||||
raise NotImplementedError
|
||||
|
||||
async def embed_images(self, images) -> list[list[float]]: # images: list[PIL.Image.Image]
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class VLLMMultimodalEmbedder(MultimodalEmbedderBase):
|
||||
"""vLLM-backed multimodal embedder via OpenAI-compatible `POST /v1/embeddings`."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
base_url: str,
|
||||
model: str,
|
||||
vector_dim: int,
|
||||
timeout: int,
|
||||
output_dim: int | None,
|
||||
image_max_side_px: int,
|
||||
embed_batch_size: int,
|
||||
) -> None:
|
||||
# Accept base_url with or without a trailing /v1 for convenience.
|
||||
# (Docs for OpenAI-compatible servers often say to include /v1.)
|
||||
cleaned = base_url.rstrip("/")
|
||||
if cleaned.endswith("/v1"):
|
||||
cleaned = cleaned[: -len("/v1")]
|
||||
self.base_url = cleaned
|
||||
self.model = model
|
||||
self.vector_dim = vector_dim
|
||||
self.timeout = timeout
|
||||
self.output_dim = output_dim
|
||||
self.image_max_side_px = int(image_max_side_px)
|
||||
self.embed_batch_size = int(embed_batch_size)
|
||||
self.encoding_format: str = "float"
|
||||
|
||||
async def _post_openai_embeddings_input(self, inputs: list[str]) -> list[list[float]]:
|
||||
"""Batch text-only embeddings via OpenAI-compatible `input=[str, ...]`."""
|
||||
payload: dict = {
|
||||
"model": self.model,
|
||||
"input": inputs,
|
||||
"encoding_format": self.encoding_format,
|
||||
}
|
||||
# Only send "dimensions" when requesting a *different* output dimension.
|
||||
# Some servers/models reject the presence of this field unless they support
|
||||
# matryoshka output; for Qwen3-VL-Embedding-2B on vLLM, even sending the
|
||||
# native dim can yield 400.
|
||||
if self.output_dim is not None and int(self.output_dim) != int(self.vector_dim):
|
||||
payload["dimensions"] = int(self.output_dim)
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
resp = await client.post(f"{self.base_url}/v1/embeddings", json=payload)
|
||||
try:
|
||||
resp.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
# Surface the backend error body; otherwise we lose crucial debugging context.
|
||||
body = ""
|
||||
try:
|
||||
body = resp.text
|
||||
except Exception:
|
||||
body = "<unavailable>"
|
||||
raise httpx.HTTPStatusError(
|
||||
f"{e}. Response body: {body[:2000]}",
|
||||
request=e.request,
|
||||
response=e.response,
|
||||
) from e
|
||||
data = resp.json()
|
||||
|
||||
rows = data.get("data") or []
|
||||
embeddings: list[list[float]] = []
|
||||
for r in rows:
|
||||
emb = r.get("embedding")
|
||||
if emb is None:
|
||||
continue
|
||||
embeddings.append(list(emb))
|
||||
|
||||
if embeddings and len(embeddings[0]) != self.vector_dim:
|
||||
raise ValueError(
|
||||
f"Unexpected embedding dimension: got {len(embeddings[0])}, "
|
||||
f"expected {self.vector_dim}"
|
||||
)
|
||||
return embeddings
|
||||
|
||||
async def _post_openai_embeddings_message(
|
||||
self, inp: MultimodalEmbedInput
|
||||
) -> list[float]:
|
||||
"""Single multimodal embedding via OpenAI-compatible `messages=[...]`.
|
||||
|
||||
Note: vLLM treats `messages` as a *single* conversation and returns a single
|
||||
embedding (i.e., it does not batch multiple independent inputs via `messages`).
|
||||
"""
|
||||
content: list[dict] = []
|
||||
if inp.text is not None and inp.text != "":
|
||||
content.append({"type": "text", "text": inp.text})
|
||||
if inp.image_b64 is not None:
|
||||
content.append(
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {"url": f"data:image/png;base64,{inp.image_b64}"},
|
||||
}
|
||||
)
|
||||
|
||||
payload: dict = {
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"encoding_format": self.encoding_format,
|
||||
}
|
||||
if self.output_dim is not None and int(self.output_dim) != int(self.vector_dim):
|
||||
payload["dimensions"] = int(self.output_dim)
|
||||
|
||||
async with httpx.AsyncClient(timeout=self.timeout) as client:
|
||||
resp = await client.post(f"{self.base_url}/v1/embeddings", json=payload)
|
||||
try:
|
||||
resp.raise_for_status()
|
||||
except httpx.HTTPStatusError as e:
|
||||
body = ""
|
||||
try:
|
||||
body = resp.text
|
||||
except Exception:
|
||||
body = "<unavailable>"
|
||||
raise httpx.HTTPStatusError(
|
||||
f"{e}. Response body: {body[:2000]}",
|
||||
request=e.request,
|
||||
response=e.response,
|
||||
) from e
|
||||
data = resp.json()
|
||||
|
||||
rows = data.get("data") or []
|
||||
if not rows or rows[0].get("embedding") is None:
|
||||
raise ValueError("vLLM returned no embedding data")
|
||||
|
||||
emb = list(rows[0]["embedding"])
|
||||
if emb and len(emb) != self.vector_dim:
|
||||
raise ValueError(
|
||||
f"Unexpected embedding dimension: got {len(emb)}, expected {self.vector_dim}"
|
||||
)
|
||||
return emb
|
||||
|
||||
async def embed_texts(self, texts: list[str]) -> list[list[float]]:
|
||||
if not texts:
|
||||
return []
|
||||
# Batch text embeddings via `input=[str,...]`.
|
||||
bs = max(1, int(self.embed_batch_size))
|
||||
out: list[list[float]] = []
|
||||
for i in range(0, len(texts), bs):
|
||||
out.extend(await self._post_openai_embeddings_input(texts[i : i + bs]))
|
||||
return out
|
||||
|
||||
async def embed_images(self, images) -> list[list[float]]: # list[PIL.Image.Image]
|
||||
if not images:
|
||||
return []
|
||||
# Lazy import: Pillow is optional in slim mode.
|
||||
try:
|
||||
from PIL import Image # noqa: F401
|
||||
except Exception as e:
|
||||
raise ImportError(
|
||||
"Multimodal image embedding requires Pillow. "
|
||||
"Install an appropriate extra or dependency (e.g. 'pillow')."
|
||||
) from e
|
||||
|
||||
inputs: list[MultimodalEmbedInput] = []
|
||||
for img in images:
|
||||
# Encode to PNG for a stable transport format.
|
||||
import io
|
||||
|
||||
# Resize guardrail: keep max side bounded to reduce payload sizes and avoid
|
||||
# backend limits. (0 disables resizing.)
|
||||
if self.image_max_side_px > 0:
|
||||
max_side = max(int(img.size[0]), int(img.size[1]))
|
||||
if max_side > self.image_max_side_px:
|
||||
from PIL import Image as PILImage
|
||||
|
||||
scale = float(self.image_max_side_px) / float(max_side)
|
||||
new_w = max(1, int(round(img.size[0] * scale)))
|
||||
new_h = max(1, int(round(img.size[1] * scale)))
|
||||
img = img.resize((new_w, new_h), resample=PILImage.Resampling.LANCZOS)
|
||||
|
||||
buf = io.BytesIO()
|
||||
img.save(buf, format="PNG", optimize=True)
|
||||
image_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
|
||||
inputs.append(MultimodalEmbedInput(image_b64=image_b64))
|
||||
|
||||
# vLLM `messages=[...]` returns ONE embedding per call, so we run one request per image.
|
||||
max_in_flight = max(1, int(self.embed_batch_size))
|
||||
sem = asyncio.Semaphore(max_in_flight)
|
||||
|
||||
async def one(inp: MultimodalEmbedInput) -> list[float]:
|
||||
async with sem:
|
||||
return await self._post_openai_embeddings_message(inp)
|
||||
|
||||
return list(await asyncio.gather(*(one(i) for i in inputs)))
|
||||
|
||||
|
||||
def get_multimodal_embedder(config: AppConfig = Config) -> MultimodalEmbedderBase | None:
|
||||
"""Factory for multimodal embedders.
|
||||
|
||||
Returns None when multimodal indexing/search is disabled.
|
||||
"""
|
||||
if not config.multimodal.enabled:
|
||||
return None
|
||||
|
||||
mm = config.multimodal.model
|
||||
if mm.provider == "vllm":
|
||||
emb = VLLMMultimodalEmbedder(
|
||||
base_url=mm.base_url,
|
||||
model=mm.name,
|
||||
vector_dim=int(mm.vector_dim),
|
||||
timeout=int(mm.timeout),
|
||||
output_dim=mm.dimensions,
|
||||
image_max_side_px=int(config.multimodal.image_max_side_px),
|
||||
embed_batch_size=int(config.multimodal.embed_batch_size),
|
||||
)
|
||||
emb.encoding_format = mm.encoding_format
|
||||
return emb
|
||||
|
||||
raise ValueError(f"Unsupported multimodal provider: {mm.provider}")
|
||||
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
269
haiku_rag_slim/haiku/rag/inspector/widgets/doc_index_modal.py
Normal file
269
haiku_rag_slim/haiku/rag/inspector/widgets/doc_index_modal.py
Normal file
|
|
@ -0,0 +1,269 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import shutil
|
||||
from copy import deepcopy
|
||||
from pathlib import Path
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.containers import Horizontal, Vertical
|
||||
from textual.screen import Screen
|
||||
from textual.widget import Widget
|
||||
from textual.widgets import Static
|
||||
from textual_image.widget import Image as TextualImage
|
||||
|
||||
|
||||
class DocIndexModal(Screen): # pragma: no cover
|
||||
"""Whole-document overlay: all chunk boxes + all mm_asset boxes per page."""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "dismiss", "Close", show=True),
|
||||
Binding("d", "dismiss", "Close", show=True),
|
||||
Binding("left", "prev_page", "Previous Page"),
|
||||
Binding("right", "next_page", "Next Page"),
|
||||
Binding("o", "open_image", "Open image", show=True),
|
||||
]
|
||||
|
||||
CSS = """
|
||||
DocIndexModal {
|
||||
background: $surface;
|
||||
layout: vertical;
|
||||
}
|
||||
|
||||
#doc-index-header {
|
||||
dock: top;
|
||||
height: auto;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
#doc-index-content {
|
||||
height: 1fr;
|
||||
width: 100%;
|
||||
align: center middle;
|
||||
}
|
||||
|
||||
#doc-index-content Image {
|
||||
width: auto;
|
||||
height: 100%;
|
||||
}
|
||||
|
||||
#doc-index-footer {
|
||||
dock: bottom;
|
||||
height: auto;
|
||||
padding: 1;
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client,
|
||||
document_id: str,
|
||||
document_uri: str | None = None,
|
||||
document_title: str | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.client = client
|
||||
self.document_id = document_id
|
||||
self.document_uri = document_uri
|
||||
self.document_title = document_title
|
||||
self._image_widget: Widget = Static("Loading...", id="image-display")
|
||||
self._page_info = Static("", id="page-info")
|
||||
self._out_dir: Path = Path(tempfile.mkdtemp(prefix="haiku-inspector-doc-index-"))
|
||||
self._pages: list[tuple[int, Path]] = [] # (page_no, png_path)
|
||||
self._page_idx: int = 0
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
title = self.document_title or self.document_uri or self.document_id
|
||||
with Vertical(id="doc-index-header"):
|
||||
yield Static(f"[bold]Doc Index Overlay[/bold] - {title}")
|
||||
yield Static(f"[dim]saved:[/dim] {self._out_dir}")
|
||||
with Horizontal(id="doc-index-content"):
|
||||
yield self._image_widget
|
||||
with Horizontal(id="doc-index-footer"):
|
||||
yield self._page_info
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
# Build overlay images once.
|
||||
try:
|
||||
from PIL import ImageDraw
|
||||
except Exception as e:
|
||||
if isinstance(self._image_widget, Static):
|
||||
self._image_widget.update(f"[red]Pillow required: {e}[/red]")
|
||||
return
|
||||
|
||||
doc = await self.client.document_repository.get_by_id(self.document_id)
|
||||
if not doc:
|
||||
if isinstance(self._image_widget, Static):
|
||||
self._image_widget.update("[red]Document not found[/red]")
|
||||
return
|
||||
|
||||
docling_doc = doc.get_docling_document()
|
||||
if not docling_doc:
|
||||
if isinstance(self._image_widget, Static):
|
||||
self._image_widget.update("[yellow]No DoclingDocument stored.[/yellow]")
|
||||
return
|
||||
|
||||
# Gather chunk bounding boxes grouped by page.
|
||||
boxes_by_page_chunks: dict[int, list[tuple[float, float, float, float]]] = {}
|
||||
# Load all chunks (paged).
|
||||
chunks: list = []
|
||||
offset = 0
|
||||
batch = 200
|
||||
while True:
|
||||
part = await self.client.chunk_repository.get_by_document_id(
|
||||
self.document_id, limit=batch, offset=offset
|
||||
)
|
||||
if not part:
|
||||
break
|
||||
chunks.extend(part)
|
||||
offset += len(part)
|
||||
if len(part) < batch:
|
||||
break
|
||||
|
||||
for ch in chunks:
|
||||
meta = ch.get_chunk_metadata()
|
||||
for bb in meta.resolve_bounding_boxes(docling_doc):
|
||||
boxes_by_page_chunks.setdefault(int(bb.page_no), []).append(
|
||||
(float(bb.left), float(bb.top), float(bb.right), float(bb.bottom))
|
||||
)
|
||||
|
||||
# Gather mm_asset bboxes grouped by page.
|
||||
boxes_by_page_assets: dict[int, list[tuple[float, float, float, float]]] = {}
|
||||
if self.client.store.mm_assets_table is not None:
|
||||
rows = (
|
||||
self.client.store.mm_assets_table.search()
|
||||
.where(f"document_id = '{self.document_id}'")
|
||||
.to_list()
|
||||
)
|
||||
for r in rows:
|
||||
page_no = r.get("page_no")
|
||||
bbox = r.get("bbox")
|
||||
if page_no is None or not bbox:
|
||||
continue
|
||||
# bbox is stored as a dict or JSON string
|
||||
if isinstance(bbox, str):
|
||||
import json as _json
|
||||
|
||||
try:
|
||||
bbox = _json.loads(bbox)
|
||||
except Exception:
|
||||
continue
|
||||
if not isinstance(bbox, dict):
|
||||
continue
|
||||
if not all(k in bbox for k in ("left", "top", "right", "bottom")):
|
||||
continue
|
||||
boxes_by_page_assets.setdefault(int(page_no), []).append(
|
||||
(
|
||||
float(bbox["left"]),
|
||||
float(bbox["top"]),
|
||||
float(bbox["right"]),
|
||||
float(bbox["bottom"]),
|
||||
)
|
||||
)
|
||||
|
||||
# Render all pages with available images.
|
||||
for page_no, page in sorted(docling_doc.pages.items(), key=lambda kv: kv[0]):
|
||||
if page.image is None or page.image.pil_image is None or page.size is None:
|
||||
continue
|
||||
pil = page.image.pil_image
|
||||
page_h = float(page.size.height)
|
||||
scale_x = pil.width / float(page.size.width)
|
||||
scale_y = pil.height / float(page.size.height)
|
||||
|
||||
img = deepcopy(pil)
|
||||
draw = ImageDraw.Draw(img, "RGBA")
|
||||
|
||||
# Chunks: yellow/orange
|
||||
for left, top, right, bottom in boxes_by_page_chunks.get(int(page_no), []):
|
||||
x0 = left * scale_x
|
||||
y0 = (page_h - top) * scale_y
|
||||
x1 = right * scale_x
|
||||
y1 = (page_h - bottom) * scale_y
|
||||
if y0 > y1:
|
||||
y0, y1 = y1, y0
|
||||
draw.rectangle([(x0, y0), (x1, y1)], fill=(255, 255, 0, 40), outline=None)
|
||||
draw.rectangle([(x0, y0), (x1, y1)], outline=(255, 165, 0, 255), width=2)
|
||||
|
||||
# Assets: cyan/blue
|
||||
for left, top, right, bottom in boxes_by_page_assets.get(int(page_no), []):
|
||||
x0 = left * scale_x
|
||||
y0 = (page_h - top) * scale_y
|
||||
x1 = right * scale_x
|
||||
y1 = (page_h - bottom) * scale_y
|
||||
if y0 > y1:
|
||||
y0, y1 = y1, y0
|
||||
draw.rectangle([(x0, y0), (x1, y1)], fill=(0, 255, 255, 30), outline=None)
|
||||
draw.rectangle([(x0, y0), (x1, y1)], outline=(0, 180, 255, 255), width=2)
|
||||
|
||||
out = self._out_dir / f"page_{int(page_no)}.png"
|
||||
try:
|
||||
img.save(out, format="PNG")
|
||||
except Exception:
|
||||
continue
|
||||
self._pages.append((int(page_no), out))
|
||||
|
||||
await self._render_current_page()
|
||||
|
||||
async def _render_current_page(self) -> None:
|
||||
if not self._pages:
|
||||
if isinstance(self._image_widget, Static):
|
||||
self._image_widget.update(
|
||||
"[yellow]No page images available.[/yellow]\n"
|
||||
"This document was converted without page images."
|
||||
)
|
||||
self._page_info.update("")
|
||||
return
|
||||
|
||||
page_no, path = self._pages[self._page_idx]
|
||||
self._page_info.update(
|
||||
f"Page {page_no} ({self._page_idx+1}/{len(self._pages)}) | "
|
||||
f"←/→ navigate | o open"
|
||||
)
|
||||
|
||||
try:
|
||||
from PIL import Image
|
||||
|
||||
pil = Image.open(path)
|
||||
new_widget = TextualImage(pil, id="rendered-image")
|
||||
await self._image_widget.remove()
|
||||
content = self.query_one("#doc-index-content", Horizontal)
|
||||
await content.mount(new_widget)
|
||||
self._image_widget = new_widget
|
||||
except Exception as e:
|
||||
if isinstance(self._image_widget, Static):
|
||||
self._image_widget.update(f"[red]Error: {e}[/red]")
|
||||
|
||||
async def action_prev_page(self) -> None:
|
||||
if self._pages and self._page_idx > 0:
|
||||
self._page_idx -= 1
|
||||
await self._render_current_page()
|
||||
|
||||
async def action_next_page(self) -> None:
|
||||
if self._pages and self._page_idx < len(self._pages) - 1:
|
||||
self._page_idx += 1
|
||||
await self._render_current_page()
|
||||
|
||||
async def action_open_image(self) -> None:
|
||||
if not self._pages:
|
||||
return
|
||||
_, path = self._pages[self._page_idx]
|
||||
try:
|
||||
if sys.platform == "darwin":
|
||||
subprocess.run(["open", str(path)], check=False)
|
||||
else:
|
||||
subprocess.run(["xdg-open", str(path)], check=False)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
async def action_dismiss(self, result=None) -> None:
|
||||
# Best-effort cleanup of temporary images.
|
||||
try:
|
||||
shutil.rmtree(self._out_dir, ignore_errors=True)
|
||||
except Exception:
|
||||
pass
|
||||
self.app.pop_screen()
|
||||
|
||||
196
haiku_rag_slim/haiku/rag/inspector/widgets/mm_assets_modal.py
Normal file
196
haiku_rag_slim/haiku/rag/inspector/widgets/mm_assets_modal.py
Normal file
|
|
@ -0,0 +1,196 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
|
||||
from textual.app import ComposeResult
|
||||
from textual.binding import Binding
|
||||
from textual.containers import Horizontal, Vertical, VerticalScroll
|
||||
from textual.screen import Screen
|
||||
from textual.widgets import ListItem, ListView, Markdown, Static
|
||||
|
||||
|
||||
class MMAssetsModal(Screen): # pragma: no cover
|
||||
"""Browse multimodal assets (mm_assets) for a document."""
|
||||
|
||||
BINDINGS = [
|
||||
Binding("escape", "dismiss", "Close", show=True),
|
||||
Binding("m", "dismiss", "Close", show=True),
|
||||
Binding("o", "open", "Open image", show=True),
|
||||
Binding("p", "toggle_mode", "Toggle page/crop", show=True),
|
||||
]
|
||||
|
||||
CSS = """
|
||||
MMAssetsModal {
|
||||
background: $surface;
|
||||
layout: vertical;
|
||||
}
|
||||
|
||||
#mm-header {
|
||||
dock: top;
|
||||
height: auto;
|
||||
padding: 1;
|
||||
}
|
||||
|
||||
#mm-content {
|
||||
height: 1fr;
|
||||
width: 100%;
|
||||
}
|
||||
|
||||
#mm-list {
|
||||
width: 1fr;
|
||||
border: solid $primary;
|
||||
}
|
||||
|
||||
#mm-detail {
|
||||
width: 2fr;
|
||||
border: solid $accent;
|
||||
}
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client,
|
||||
document_id: str,
|
||||
document_uri: str | None = None,
|
||||
document_title: str | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.client = client
|
||||
self.document_id = document_id
|
||||
self.document_uri = document_uri
|
||||
self.document_title = document_title
|
||||
self.mode: str = "crop"
|
||||
self.rows: list[dict] = []
|
||||
self._out_dir: Path = Path(tempfile.mkdtemp(prefix="haiku-inspector-mm-assets-"))
|
||||
|
||||
def compose(self) -> ComposeResult:
|
||||
title = self.document_title or self.document_uri or self.document_id
|
||||
with Vertical(id="mm-header"):
|
||||
yield Static(f"[bold]MM Assets[/bold] - {title}")
|
||||
yield Static(
|
||||
f"[dim]mode:[/dim] {self.mode} | "
|
||||
f"press [bold]p[/bold] to toggle, [bold]o[/bold] to open image",
|
||||
id="mm-status",
|
||||
)
|
||||
yield Static(f"[dim]saved:[/dim] {self._out_dir}")
|
||||
with Horizontal(id="mm-content"):
|
||||
with VerticalScroll(id="mm-list"):
|
||||
yield ListView(id="mm-assets-list")
|
||||
with VerticalScroll(id="mm-detail"):
|
||||
yield Markdown("Select an asset…", id="mm-detail-md")
|
||||
|
||||
async def on_mount(self) -> None:
|
||||
list_view = self.query_one("#mm-assets-list", ListView)
|
||||
await list_view.clear()
|
||||
|
||||
if self.client.store.mm_assets_table is None:
|
||||
await self.query_one("#mm-detail-md", Markdown).update(
|
||||
"[yellow]mm_assets table is not available.[/yellow]\n\n"
|
||||
"Enable multimodal indexing and ingest documents with pictures first."
|
||||
)
|
||||
return
|
||||
|
||||
# Load rows for this document (simple scan via query).
|
||||
self.rows = (
|
||||
self.client.store.mm_assets_table.search()
|
||||
.where(f"document_id = '{self.document_id}'")
|
||||
.to_list()
|
||||
)
|
||||
|
||||
if not self.rows:
|
||||
await self.query_one("#mm-detail-md", Markdown).update(
|
||||
"[yellow]No mm_assets for this document.[/yellow]"
|
||||
)
|
||||
return
|
||||
|
||||
for r in self.rows:
|
||||
asset_id = r.get("id")
|
||||
page = r.get("page_no")
|
||||
ref = r.get("doc_item_ref")
|
||||
label = f"p.{page} {ref} {asset_id}"
|
||||
await list_view.append(ListItem(Static(label)))
|
||||
|
||||
list_view.index = 0
|
||||
await self._render_detail()
|
||||
list_view.focus()
|
||||
|
||||
async def on_list_view_highlighted(self, event: ListView.Highlighted) -> None:
|
||||
if event.list_view.id != "mm-assets-list":
|
||||
return
|
||||
await self._render_detail()
|
||||
|
||||
async def _render_detail(self) -> None:
|
||||
list_view = self.query_one("#mm-assets-list", ListView)
|
||||
idx = list_view.index
|
||||
if idx is None or idx >= len(self.rows):
|
||||
return
|
||||
r = self.rows[idx]
|
||||
asset_id = r.get("id")
|
||||
page_no = r.get("page_no")
|
||||
ref = r.get("doc_item_ref")
|
||||
caption = r.get("caption")
|
||||
description = r.get("description")
|
||||
md = f"""\
|
||||
**asset_id**: `{asset_id}`
|
||||
**page**: {page_no}
|
||||
**doc_item_ref**: `{ref}`
|
||||
|
||||
**caption**:
|
||||
{caption or "*none*"}
|
||||
|
||||
**description**:
|
||||
{description or "*none*"}
|
||||
|
||||
**saved_dir**:
|
||||
`{self._out_dir}`
|
||||
"""
|
||||
await self.query_one("#mm-detail-md", Markdown).update(md)
|
||||
|
||||
async def action_toggle_mode(self) -> None:
|
||||
self.mode = "page" if self.mode == "crop" else "crop"
|
||||
self.query_one("#mm-status", Static).update(
|
||||
f"[dim]mode:[/dim] {self.mode} | "
|
||||
f"press [bold]p[/bold] to toggle, [bold]o[/bold] to open image"
|
||||
)
|
||||
|
||||
async def action_open(self) -> None:
|
||||
list_view = self.query_one("#mm-assets-list", ListView)
|
||||
idx = list_view.index
|
||||
if idx is None or idx >= len(self.rows):
|
||||
return
|
||||
asset_id = self.rows[idx].get("id")
|
||||
if not asset_id:
|
||||
return
|
||||
|
||||
images = await self.client.visualize_mm_asset(asset_id=str(asset_id), mode=self.mode)
|
||||
if not images:
|
||||
return
|
||||
|
||||
# Save first image and open via OS viewer.
|
||||
path = self._out_dir / f"{asset_id}_{self.mode}.png"
|
||||
try:
|
||||
images[0].save(path, format="PNG")
|
||||
except Exception:
|
||||
return
|
||||
|
||||
try:
|
||||
if sys.platform == "darwin":
|
||||
subprocess.run(["open", str(path)], check=False)
|
||||
else:
|
||||
subprocess.run(["xdg-open", str(path)], check=False)
|
||||
except Exception:
|
||||
return
|
||||
|
||||
async def action_dismiss(self, result=None) -> None:
|
||||
# Best-effort cleanup of temporary images.
|
||||
try:
|
||||
shutil.rmtree(self._out_dir, ignore_errors=True)
|
||||
except Exception:
|
||||
pass
|
||||
self.app.pop_screen()
|
||||
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
41
haiku_rag_slim/haiku/rag/store/models/mm_asset.py
Normal file
41
haiku_rag_slim/haiku/rag/store/models/mm_asset.py
Normal file
|
|
@ -0,0 +1,41 @@
|
|||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class MMAsset(BaseModel):
|
||||
"""A multimodal asset (typically a Docling PictureItem crop) stored in LanceDB."""
|
||||
|
||||
id: str | None = None
|
||||
document_id: str
|
||||
|
||||
# Docling pointer (e.g. "#/pictures/3") used to resolve provenance / ordering.
|
||||
doc_item_ref: str
|
||||
|
||||
# Optional ordering anchor (index from DoclingDocument.iterate_items()).
|
||||
item_index: int | None = None
|
||||
|
||||
# Provenance for layout-aware UX (visual grounding, grouping, etc.)
|
||||
page_no: int | None = None
|
||||
bbox: dict | None = None # {"left":..,"top":..,"right":..,"bottom":..}
|
||||
|
||||
# Text fields for display/debugging (NOT used as the embedding input for Phase 1).
|
||||
caption: str | None = None
|
||||
description: str | None = None
|
||||
metadata: dict = {}
|
||||
|
||||
# Multimodal embedding vector (dimension is configured in Store via LanceDB schema)
|
||||
embedding: list[float] | None = None
|
||||
|
||||
|
||||
class MMSearchResult(BaseModel):
|
||||
"""Search result for multimodal assets."""
|
||||
|
||||
asset_id: str
|
||||
document_id: str
|
||||
score: float
|
||||
|
||||
doc_item_ref: str
|
||||
item_index: int | None = None
|
||||
page_no: int | None = None
|
||||
bbox: dict | None = None
|
||||
caption: str | None = None
|
||||
description: str | None = None
|
||||
|
|
@ -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",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
130
haiku_rag_slim/haiku/rag/store/repositories/mm_asset.py
Normal file
130
haiku_rag_slim/haiku/rag/store/repositories/mm_asset.py
Normal file
|
|
@ -0,0 +1,130 @@
|
|||
import json
|
||||
import logging
|
||||
from uuid import uuid4
|
||||
|
||||
from haiku.rag.store.engine import Store
|
||||
from haiku.rag.store.models.mm_asset import MMAsset, MMSearchResult
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MMAssetRepository:
|
||||
"""Repository for multimodal asset operations."""
|
||||
|
||||
def __init__(self, store: Store) -> None:
|
||||
self.store = store
|
||||
|
||||
def _require_table(self):
|
||||
if self.store.mm_assets_table is None:
|
||||
raise ValueError(
|
||||
"mm_assets table is not available. Enable config.multimodal.enabled "
|
||||
"and open the database in writable mode at least once to create it."
|
||||
)
|
||||
return self.store.mm_assets_table
|
||||
|
||||
async def create(self, entity: MMAsset | list[MMAsset]) -> MMAsset | list[MMAsset]:
|
||||
"""Create one or more multimodal assets.
|
||||
|
||||
Assets must have embedding set before calling this method.
|
||||
"""
|
||||
self.store._assert_writable()
|
||||
table = self._require_table()
|
||||
|
||||
if isinstance(entity, MMAsset):
|
||||
assert entity.embedding is not None, "MMAsset must have an embedding"
|
||||
asset_id = str(uuid4())
|
||||
rec = self.store.MMAssetRecord(
|
||||
id=asset_id,
|
||||
document_id=entity.document_id,
|
||||
doc_item_ref=entity.doc_item_ref,
|
||||
item_index=entity.item_index,
|
||||
page_no=entity.page_no,
|
||||
bbox=json.dumps(entity.bbox) if entity.bbox else None,
|
||||
caption=entity.caption,
|
||||
description=entity.description,
|
||||
metadata=json.dumps(entity.metadata or {}),
|
||||
vector=entity.embedding,
|
||||
)
|
||||
table.add([rec])
|
||||
entity.id = asset_id
|
||||
return entity
|
||||
|
||||
assets = entity
|
||||
if not assets:
|
||||
return []
|
||||
|
||||
for a in assets:
|
||||
assert a.embedding is not None, "All MMAssets must have embeddings"
|
||||
|
||||
records = []
|
||||
for a in assets:
|
||||
asset_id = str(uuid4())
|
||||
records.append(
|
||||
self.store.MMAssetRecord(
|
||||
id=asset_id,
|
||||
document_id=a.document_id,
|
||||
doc_item_ref=a.doc_item_ref,
|
||||
item_index=a.item_index,
|
||||
page_no=a.page_no,
|
||||
bbox=json.dumps(a.bbox) if a.bbox else None,
|
||||
caption=a.caption,
|
||||
description=a.description,
|
||||
metadata=json.dumps(a.metadata or {}),
|
||||
vector=a.embedding,
|
||||
)
|
||||
)
|
||||
a.id = asset_id
|
||||
|
||||
table.add(records)
|
||||
return assets
|
||||
|
||||
async def delete_by_document_id(self, document_id: str) -> bool:
|
||||
"""Delete all multimodal assets for a document."""
|
||||
self.store._assert_writable()
|
||||
table = self._require_table()
|
||||
# Fast path: delete without pre-check; table.delete is idempotent.
|
||||
table.delete(f"document_id = '{document_id}'")
|
||||
return True
|
||||
|
||||
async def search_by_vector(
|
||||
self,
|
||||
query_vector: list[float],
|
||||
*,
|
||||
limit: int = 5,
|
||||
filter: str | None = None,
|
||||
) -> list[MMSearchResult]:
|
||||
"""Vector search over multimodal assets."""
|
||||
table = self._require_table()
|
||||
|
||||
q = table.search(query_vector).limit(limit)
|
||||
if filter:
|
||||
q = q.where(filter)
|
||||
|
||||
# LanceDB returns _distance for vector search; smaller is better.
|
||||
# Convert to a similarity-like score for UX consistency.
|
||||
rows = q.to_list()
|
||||
results: list[MMSearchResult] = []
|
||||
for row in rows:
|
||||
distance = float(row.get("_distance", 0.0))
|
||||
score = 1.0 / (1.0 + distance)
|
||||
bbox = None
|
||||
try:
|
||||
bbox = json.loads(row["bbox"]) if row.get("bbox") else None
|
||||
except Exception:
|
||||
bbox = None
|
||||
results.append(
|
||||
MMSearchResult(
|
||||
asset_id=row.get("id"),
|
||||
document_id=row.get("document_id"),
|
||||
score=score,
|
||||
doc_item_ref=row.get("doc_item_ref", ""),
|
||||
item_index=row.get("item_index"),
|
||||
page_no=row.get("page_no"),
|
||||
bbox=bbox,
|
||||
caption=row.get("caption"),
|
||||
description=row.get("description"),
|
||||
)
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
|
@ -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
|
||||
|
|
|
|||
116
tests/test_multimodal_phase1.py
Normal file
116
tests/test_multimodal_phase1.py
Normal file
|
|
@ -0,0 +1,116 @@
|
|||
import pytest
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.store.engine import Store
|
||||
from haiku.rag.store.models.chunk import Chunk
|
||||
from haiku.rag.store.models.mm_asset import MMAsset
|
||||
|
||||
|
||||
def _mm_enabled_config():
|
||||
cfg = Config.model_copy(deep=True)
|
||||
cfg.multimodal.enabled = True
|
||||
# Ensure dim matches our dummy vectors in tests
|
||||
cfg.multimodal.model.vector_dim = 2048
|
||||
return cfg
|
||||
|
||||
|
||||
def test_mm_assets_table_created_when_enabled(temp_db_path):
|
||||
cfg = _mm_enabled_config()
|
||||
store = Store(temp_db_path, create=True, config=cfg)
|
||||
assert store.mm_assets_table is not None
|
||||
assert "mm_assets" in store.db.table_names()
|
||||
|
||||
|
||||
async def test_version_rollback_on_mm_index_failure(temp_db_path, monkeypatch):
|
||||
cfg = _mm_enabled_config()
|
||||
|
||||
async with HaikuRAG(db_path=temp_db_path, create=True, config=cfg) as client:
|
||||
dim = int(client.store.embedder._vector_dim)
|
||||
|
||||
async def fail_after_writing_one_asset(doc):
|
||||
# write one asset, then fail -> should rollback all tables
|
||||
assert doc.id is not None
|
||||
await client.mm_asset_repository.create(
|
||||
MMAsset(
|
||||
document_id=doc.id,
|
||||
doc_item_ref="#/pictures/0",
|
||||
item_index=0,
|
||||
page_no=1,
|
||||
bbox={"left": 0.0, "top": 1.0, "right": 2.0, "bottom": 3.0},
|
||||
caption="dummy",
|
||||
embedding=[0.0] * 2048,
|
||||
)
|
||||
)
|
||||
raise RuntimeError("boom")
|
||||
|
||||
monkeypatch.setattr(client, "_index_mm_assets_for_document", fail_after_writing_one_asset)
|
||||
|
||||
# Instead of calling create_document (which would hit external embedders),
|
||||
# call the internal store path with pre-embedded chunks.
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
||||
with pytest.raises(RuntimeError, match="boom"):
|
||||
await client._store_document_with_chunks(
|
||||
Document(content="x"),
|
||||
[Chunk(content="c", metadata={}, order=0, embedding=[0.0] * dim)],
|
||||
)
|
||||
|
||||
# After failure, database should be rolled back to empty state
|
||||
store = Store(temp_db_path, create=True, config=cfg)
|
||||
assert store.documents_table.count_rows() == 0
|
||||
assert store.chunks_table.count_rows() == 0
|
||||
assert store.mm_assets_table is not None
|
||||
assert store.mm_assets_table.count_rows() == 0
|
||||
|
||||
|
||||
async def test_update_rollback_restores_mm_assets(temp_db_path, monkeypatch):
|
||||
cfg = _mm_enabled_config()
|
||||
|
||||
async with HaikuRAG(db_path=temp_db_path, create=True, config=cfg) as client:
|
||||
dim = int(client.store.embedder._vector_dim)
|
||||
|
||||
# Create baseline document without calling external embedders
|
||||
from haiku.rag.store.models.document import Document
|
||||
|
||||
doc = await client._store_document_with_chunks(
|
||||
Document(content="base"),
|
||||
[Chunk(content="c", metadata={}, order=0, embedding=[0.0] * dim)],
|
||||
)
|
||||
assert doc.id is not None
|
||||
|
||||
# Seed mm_assets with one row
|
||||
await client.mm_asset_repository.create(
|
||||
MMAsset(
|
||||
document_id=doc.id,
|
||||
doc_item_ref="#/pictures/seed",
|
||||
item_index=0,
|
||||
page_no=1,
|
||||
bbox={"left": 0.0, "top": 1.0, "right": 2.0, "bottom": 3.0},
|
||||
caption="seed",
|
||||
embedding=[1.0] * 2048,
|
||||
)
|
||||
)
|
||||
|
||||
async def fail_index(_doc):
|
||||
raise RuntimeError("mm index fail")
|
||||
|
||||
monkeypatch.setattr(client, "_index_mm_assets_for_document", fail_index)
|
||||
|
||||
with pytest.raises(RuntimeError, match="mm index fail"):
|
||||
# Call internal update path to avoid external embedder usage
|
||||
await client._update_document_with_chunks(
|
||||
Document(
|
||||
id=doc.id,
|
||||
content="changed",
|
||||
docling_document_json=doc.docling_document_json,
|
||||
docling_version=doc.docling_version,
|
||||
),
|
||||
[Chunk(content="c2", metadata={}, order=0, embedding=[0.0] * dim)],
|
||||
)
|
||||
|
||||
# After rollback, original mm_assets row should still exist
|
||||
store = Store(temp_db_path, create=True, config=cfg)
|
||||
assert store.mm_assets_table is not None
|
||||
assert store.mm_assets_table.count_rows() == 1
|
||||
|
||||
Loading…
Reference in a new issue