Merge branch 'main' into chunk-metadata
This commit is contained in:
Lawrence Akka
GitHub
commit
cdd1b99e6b
41 changed files with 1311 additions and 377 deletions
13
CHANGELOG.md
13
CHANGELOG.md
|
|
@ -6,16 +6,27 @@
|
|||
- `evaluations run --filter/-f CLAUSE`: SQL `WHERE` clause over document columns, applied to the retrieval benchmark's searches and to every capability search during QA. Recorded as `document_filter` in experiment metadata.
|
||||
- `mtrag_clapnq` / `mtrag_clapnq_rewrite` / `mtrag_clapnq_live` / `mtrag_clapnq_live_uncompacted` evaluation datasets: multi-turn QA with gold-prefix and live-session conversation replay, live arms with and without `EvidenceCompactionCapability`, Recall@k/nDCG@k retrieval metrics, eligibility-aware citation scoring, refusal precision/recall, per-turn tool-traffic attributes, and `citation_status` / `turn_citation_status` eval attributes.
|
||||
- Raw chunk metadata is now exposed to search and citation results, through `SearchResult.chunk_meta` and `Citation.chunk_meta`. For context-expanded results, the metadata is that of the anchor chunk.
|
||||
|
||||
- BTree indexes on `chunks.id`, `chunks.document_id` and `documents.id`, and a Bitmap index on `document_items.label`. Existing databases need `haiku-rag migrate`.
|
||||
- `lancedb.read_consistency_interval_seconds` (default 30), `lancedb.index_cache_size_bytes` and `lancedb.metadata_cache_size_bytes`. The LanceDB session is shared across connections in a process, so its index and metadata caches survive a connection being closed.
|
||||
|
||||
### Changed
|
||||
|
||||
- Eval judge pinned to `qwen3.8`: `DEFAULT_JUDGE_MODEL` is `ollama:qwen3.8`, and the reference configs use `Inferact/Qwen3.8-27B-NVFP4` with `extra_body.chat_template_kwargs.reasoning_effort: low`. Results in `docs/benchmarks.md` were judged by `Qwen3.6-35B-A3B-NVFP4` and are not re-judged.
|
||||
- `create_capability(rag=...)` lends a capability an open client rather than having it open its own; `client.ask`/`client.analyze` now pass theirs.
|
||||
- The MCP server opens one database client for its lifetime instead of one per tool call, and fails at startup if the database cannot be opened. `delete_document` no longer opens its own connection with `skip_validation=True`, so the server no longer opts out of embedding-config validation: drift that validation rejects (any `vector_dim` mismatch, or identity drift on a writable server) now fails MCP startup instead of serving a delete-only server. Use the CLI to delete under drift.
|
||||
- `import_documents` embeds chunks across the whole batch in one pass instead of per document.
|
||||
- `haiku.rag` and `haiku.rag-slim` summaries and keywords; both packages now publish `[project.urls]`.
|
||||
- `server.json` declares `title` and `websiteUrl`, and drops the `keywords` and `license` keys, which are not in the server schema.
|
||||
|
||||
### Removed
|
||||
|
||||
- `wix` evaluation dataset and its reference config `evaluations/configs/wix.yaml`.
|
||||
|
||||
### Fixed
|
||||
|
||||
- `DocumentRepository.delete_all` recreated `document_items` from `DocumentItemRecord` instead of `get_document_items_arrow_schema()`, returning `picture_data` as `binary` rather than `large_binary`.
|
||||
- `server.json` runtime arguments are `mcp --stdio`, was `serve --mcp`.
|
||||
|
||||
## [0.74.0] - 2026-08-13
|
||||
|
||||
### Added
|
||||
|
|
|
|||
10
README.md
10
README.md
|
|
@ -1,9 +1,15 @@
|
|||
# Haiku RAG
|
||||
# haiku.rag
|
||||
|
||||
[](https://pypi.org/project/haiku.rag/)
|
||||
[](https://pypi.org/project/haiku.rag/)
|
||||
[](https://pepy.tech/projects/haiku-rag-slim)
|
||||
[](https://ggozad.github.io/haiku.rag/)
|
||||
[](https://github.com/ggozad/haiku.rag/actions/workflows/test.yml)
|
||||
[](https://codecov.io/gh/ggozad/haiku.rag)
|
||||
|
||||
Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.pydantic.dev/), and [Docling](https://docling-project.github.io/docling/).
|
||||
Agentic RAG that answers questions about your own documents with citations to page numbers and section headings. Runs locally on an embedded database, no server required.
|
||||
|
||||
Built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.pydantic.dev/), and [Docling](https://docling-project.github.io/docling/). Full documentation at [ggozad.github.io/haiku.rag](https://ggozad.github.io/haiku.rag/).
|
||||
|
||||
> **New: vision and multimodal search.** Picture-aware ingestion captures embedded figure bytes; vision-capable QA models receive them alongside text. Multimodal embedders put picture vectors in the same space as text, enabling text-as-query → figure hits and image-as-query retrieval.
|
||||
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ dependencies = [
|
|||
"uvicorn[standard]>=0.40.0",
|
||||
"pydantic-ai-slim[ag-ui,anthropic,openai]>=2.18.0,<3.0.0",
|
||||
"python-dotenv>=1.2.1",
|
||||
"haiku.rag-slim>=0.74.0",
|
||||
"haiku.rag-slim>=0.75.0",
|
||||
"logfire[pydantic-ai]>=3.17.0",
|
||||
]
|
||||
|
||||
|
|
|
|||
|
|
@ -2,140 +2,9 @@
|
|||
|
||||
We evaluate `haiku.rag` on a small set of datasets that exercise different parts of the pipeline. OpenRAG Bench (ORB), T²-RAGBench, HotpotQA, and MTRAG are the datasets we currently track. Retrieval, QA accuracy, and citation retrieval are scored end-to-end through the RAG and analysis capabilities.
|
||||
|
||||
## Running Evaluations
|
||||
|
||||
You can run evaluations with the `evaluations` CLI:
|
||||
|
||||
```bash
|
||||
evaluations run hotpotqa
|
||||
evaluations run orb_text
|
||||
```
|
||||
|
||||
The evaluation flow is orchestrated with [`pydantic-evals`](https://github.com/pydantic/pydantic-ai/tree/main/libs/pydantic-evals), which we leverage for dataset management, scoring, and report generation.
|
||||
|
||||
### Pre-built Databases
|
||||
|
||||
Building evaluation databases from scratch can take a long time, especially for large datasets like OpenRAG Bench. Pre-built databases are available on HuggingFace:
|
||||
|
||||
```bash
|
||||
# Download a specific dataset
|
||||
evaluations download hotpotqa
|
||||
|
||||
# Download all datasets
|
||||
evaluations download all
|
||||
|
||||
# Force re-download (overwrite existing)
|
||||
evaluations download hotpotqa --force
|
||||
```
|
||||
|
||||
Active datasets:
|
||||
|
||||
| Dataset | Size |
|
||||
|---------|------|
|
||||
| `orb_text` — OpenRAG Bench, text embedder (`qwen3-embedding:4b`) with VLM picture descriptions baked into chunk content | ~18 GB |
|
||||
| `orb_multimodal` — OpenRAG Bench, multimodal embedder (`qwen3-vl-embedding-8b`); picture vectors live in the same space as text for cross-modal retrieval | ~16 GB |
|
||||
| `orb_multimodal_nemotron` — OpenRAG Bench, multimodal embedder (`nvidia/llama-nemotron-embed-vl-1b-v2`), the embedder behind the published headline results | ~16 GB |
|
||||
| `t2_finqa` — T²-RAGBench (FinQA) financial QA, text embedder (`qwen3-embedding:4b`); scored by exact numeric match, run with `--target analysis-capability` | ~2 GB |
|
||||
| `hotpotqa` — HotpotQA multi-hop QA over Wikipedia paragraphs, text embedder (`qwen3-embedding:4b`) | ~1.5 GB |
|
||||
| `mtrag_clapnq` — MTRAG multi-turn RAG, ClapNQ (Wikipedia) passages, text embedder (`qwen3-embedding:4b`); also serves the `mtrag_clapnq_rewrite`, `mtrag_clapnq_live` and `mtrag_clapnq_live_uncompacted` keys | ~2.8 GB |
|
||||
|
||||
After downloading, run benchmarks with `--skip-db`. Each database is built with a specific embedder, so pass its reference config from `evaluations/configs/` (a database only opens against a config whose embedder matches):
|
||||
|
||||
```bash
|
||||
evaluations run orb_multimodal_nemotron --skip-db --config configs/orb_multimodal_nemotron.yaml
|
||||
```
|
||||
|
||||
The configs use `vllm` as the model host. Point `base_url` at your own OpenAI-compatible endpoints to reproduce the numbers.
|
||||
|
||||
### Configuration
|
||||
|
||||
The benchmark script accepts several options:
|
||||
|
||||
```bash
|
||||
evaluations run hotpotqa --config /path/to/haiku.rag.yaml --db /path/to/custom.lancedb
|
||||
```
|
||||
|
||||
**Options:**
|
||||
|
||||
- `--config PATH` - Specify a custom `haiku.rag.yaml` configuration file
|
||||
- `--db PATH` - Override the database path (default: platform-specific user data directory)
|
||||
- `--skip-db` - Skip updating the evaluation database
|
||||
- `--skip-retrieval` - Skip retrieval benchmark
|
||||
- `--skip-qa` - Skip QA benchmark
|
||||
- `--limit N` - Limit number of test cases
|
||||
- `--name NAME` - Override the evaluation name
|
||||
- `--target {rag-capability,analysis-capability}` - Choose which [capability](capabilities/index.md) to benchmark end-to-end (default: `rag-capability`). The target names remain stable dataset identifiers.
|
||||
- `--capability-model PROVIDER:NAME` - Override the capability model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-capability`).
|
||||
- `--filter CLAUSE` / `-f CLAUSE` - Restrict every benchmark search to a subset of the database (see [Restricting the corpus](#restricting-the-corpus)).
|
||||
|
||||
If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults.
|
||||
|
||||
To pin the LLM judge in YAML (rather than the default `ollama:qwen3.6`). These are the recommended settings:
|
||||
|
||||
```yaml
|
||||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: RedHatAI/Qwen3.6-35B-A3B-NVFP4
|
||||
base_url: http://localhost:8000/v1 # optional, for OpenAI-compatible servers (vLLM, LM Studio, etc.)
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
top_p: 0.95
|
||||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
enable_thinking: true
|
||||
```
|
||||
|
||||
### Restricting the corpus
|
||||
|
||||
When a database holds documents from several corpora — only some of which a dataset's questions are drawn from — `--filter` restricts every benchmark search to a subset. It takes the same SQL `WHERE` clause as `haiku-rag search --filter`, over document columns (`id`, `uri`, `title`, `created_at`, `updated_at`, `metadata`). Each dataset writes its own URIs: `orb_text` uses bare arXiv ids such as `2407.01528v3`, `hotpotqa` uses page titles.
|
||||
|
||||
```bash
|
||||
evaluations run orb_text --skip-db --config haiku.rag.s3.yaml \
|
||||
--filter "uri LIKE '2407%'"
|
||||
```
|
||||
|
||||
If the corpora are distinguished by a tag rather than by URI, attach it at ingest time as document metadata and match it with `LIKE`. `metadata` is stored as a `json.dumps` string, so there is no JSON subfield access — match the serialized key/value, including the space after the colon:
|
||||
|
||||
```bash
|
||||
evaluations run orb_text --skip-db --filter "metadata LIKE '%\"corpus\": \"orb_text\"%'"
|
||||
```
|
||||
|
||||
The clause applies to both benchmark phases — the retrieval benchmark's searches and every search the capability runs during QA — so the two score the same subset. It is recorded as `document_filter` in the run's experiment metadata, so a filtered run is never mistaken for an unfiltered one when comparing results.
|
||||
|
||||
Filtering affects searches only — a run without `--skip-db` still populates the database with the dataset's full corpus.
|
||||
|
||||
## Methodology
|
||||
|
||||
### Retrieval Metrics
|
||||
|
||||
**Mean Average Precision (MAP)** scores ranked retrieval results against the gold `expected_uris`.
|
||||
|
||||
- For each relevant document at position k, calculate precision@k = (relevant docs in top k) / k
|
||||
- Average Precision (AP) = sum of these precision values / total relevant documents
|
||||
- MAP is the mean of AP scores across all queries
|
||||
- Range: 0 to 1. Rewards ranking relevant documents higher
|
||||
- For single-doc queries this collapses to `1/rank` (i.e. reciprocal rank)
|
||||
|
||||
### QA Accuracy
|
||||
|
||||
`pydantic-evals` coordinates an LLM judge to determine whether the capability's answer is correct. The default judge is `ollama:qwen3.6`, pinned so changes to the capability model don't change the judge underneath. Set `evaluations.judge` in `haiku.rag.yaml` to override (including a custom `base_url` for any OpenAI-compatible endpoint). Accuracy is the fraction of correctly answered questions.
|
||||
|
||||
A dataset that brings its own deterministic evaluator is scored by that evaluator instead, and no judge runs. T²-RAGBench is the only such dataset today, scored by `NumberMatchEvaluator`.
|
||||
|
||||
We picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibration (gpt-oss / qwen3.6 as both answerer and judge, with Claude Opus 4.7 as a reference). `qwen3.6` had κ ≥ 0.66 vs the reference on both same-family and cross-family answerers (vs ~0.39–0.55 for `gpt-oss`) and showed no measurable self-preference bias, while `gpt-oss` was ~10 pp more lenient on its own outputs.
|
||||
|
||||
### Citation Retrieval
|
||||
|
||||
Alongside QA accuracy, a second metric scores the URIs the capability registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MAP math as raw retrieval. The score key is `cited_map`. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case.
|
||||
|
||||
This is computed alongside QA accuracy from the same capability run, no extra invocations. The signal complements raw retrieval: where raw retrieval measures whether the retriever surfaced the gold document at any rank, citation retrieval measures whether the capability grounded its answer on it.
|
||||
|
||||
## Current results
|
||||
|
||||
Numbers measured under the current pinned judge (`ollama:qwen3.6`) on a recent `haiku.rag` version.
|
||||
Numbers below were measured under `Qwen3.6-35B-A3B-NVFP4` as judge, on a recent `haiku.rag` version. The pinned judge is now `qwen3.8`; rows are not re-judged, so compare rows to each other rather than to runs judged by `qwen3.8`.
|
||||
|
||||
### OpenRAG Bench (ORB)
|
||||
|
||||
|
|
@ -260,3 +129,136 @@ The two live arms replay the same 29 conversations (224 turns) and differ only i
|
|||
- Answer pass rate: 185/224 vs 175/224 turns. Of the 18 turns where the arms disagree, 14 pass only compacted and 4 only uncompacted. McNemar exact two-sided p = 0.031. The paired difference is +4.5pp with a Wald 95% CI of +0.8 to +8.1pp, so the honest claim is an improvement of roughly 1 to 8 points, not the point estimate.
|
||||
- Citation MAP, macro-averaged over conversations with 208 of 224 turns eligible (turns with gold passages) in each arm: 0.4174 compacted vs 0.4230 uncompacted. The gold-prefix 0.35 is over 208 of 224 eligible cases.
|
||||
- Refusal precision and recall against the answerability labels (16 UNANSWERABLE turns per arm): compacted 0.33 precision and 0.44 recall (21 refusals), uncompacted 0.23 and 0.31 (22 refusals). Gold-prefix: 0.24 and 0.44 (29 refusals).
|
||||
|
||||
## Methodology
|
||||
|
||||
### Retrieval Metrics
|
||||
|
||||
**Mean Average Precision (MAP)** scores ranked retrieval results against the gold `expected_uris`.
|
||||
|
||||
- For each relevant document at position k, calculate precision@k = (relevant docs in top k) / k
|
||||
- Average Precision (AP) = sum of these precision values / total relevant documents
|
||||
- MAP is the mean of AP scores across all queries
|
||||
- Range: 0 to 1. Rewards ranking relevant documents higher
|
||||
- For single-doc queries this collapses to `1/rank` (i.e. reciprocal rank)
|
||||
|
||||
### QA Accuracy
|
||||
|
||||
`pydantic-evals` coordinates an LLM judge to determine whether the capability's answer is correct. The default judge is `ollama:qwen3.8`, pinned so changes to the capability model don't change the judge underneath. Set `evaluations.judge` in `haiku.rag.yaml` to override (including a custom `base_url` for any OpenAI-compatible endpoint). Accuracy is the fraction of correctly answered questions.
|
||||
|
||||
A dataset that brings its own deterministic evaluator is scored by that evaluator instead, and no judge runs. T²-RAGBench is the only such dataset today, scored by `NumberMatchEvaluator`.
|
||||
|
||||
`qwen3.8` replaced `qwen3.6` after a 120-case calibration on ORB, stratified 60 pass / 60 fail: agreement 0.950, Cohen's κ 0.900, and in all 6 disagreements it matched or beat `qwen3.6` (4 were `qwen3.6` failing answers that were equivalent in different notation). It emits no reasoning content, so it avoids the thinking spirals that made `qwen3.6` exceed its output budget and drop verdicts. `reasoning_effort` changes its verdicts in 1 case per 120, so the cheaper `low` is pinned.
|
||||
|
||||
Before that, we picked `qwen3.6` over the previously-pinned `gpt-oss` after a 4-cell calibration (gpt-oss / qwen3.6 as both answerer and judge, with Claude Opus 4.7 as a reference). `qwen3.6` had κ ≥ 0.66 vs the reference on both same-family and cross-family answerers (vs ~0.39–0.55 for `gpt-oss`) and showed no measurable self-preference bias, while `gpt-oss` was ~10 pp more lenient on its own outputs.
|
||||
|
||||
### Citation Retrieval
|
||||
|
||||
Alongside QA accuracy, a second metric scores the URIs the capability registered via the `cite` tool against each dataset's gold `expected_uris`, using the same MAP math as raw retrieval. The score key is `cited_map`. Console output also includes the cite rate (% of cases with at least one citation) and the mean number of citations per case.
|
||||
|
||||
This is computed alongside QA accuracy from the same capability run, no extra invocations. The signal complements raw retrieval: where raw retrieval measures whether the retriever surfaced the gold document at any rank, citation retrieval measures whether the capability grounded its answer on it.
|
||||
|
||||
## Running Evaluations
|
||||
|
||||
You can run evaluations with the `evaluations` CLI:
|
||||
|
||||
```bash
|
||||
evaluations run hotpotqa
|
||||
evaluations run orb_text
|
||||
```
|
||||
|
||||
The evaluation flow is orchestrated with [`pydantic-evals`](https://github.com/pydantic/pydantic-ai/tree/main/libs/pydantic-evals), which we leverage for dataset management, scoring, and report generation.
|
||||
|
||||
### Pre-built Databases
|
||||
|
||||
Building evaluation databases from scratch can take a long time, especially for large datasets like OpenRAG Bench. Pre-built databases are available on HuggingFace:
|
||||
|
||||
```bash
|
||||
# Download a specific dataset
|
||||
evaluations download hotpotqa
|
||||
|
||||
# Download all datasets
|
||||
evaluations download all
|
||||
|
||||
# Force re-download (overwrite existing)
|
||||
evaluations download hotpotqa --force
|
||||
```
|
||||
|
||||
Active datasets:
|
||||
|
||||
| Dataset | Size |
|
||||
|---------|------|
|
||||
| `orb_text` — OpenRAG Bench, text embedder (`qwen3-embedding:4b`) with VLM picture descriptions baked into chunk content | ~18 GB |
|
||||
| `orb_multimodal` — OpenRAG Bench, multimodal embedder (`qwen3-vl-embedding-8b`); picture vectors live in the same space as text for cross-modal retrieval | ~16 GB |
|
||||
| `orb_multimodal_nemotron` — OpenRAG Bench, multimodal embedder (`nvidia/llama-nemotron-embed-vl-1b-v2`), the embedder behind the published headline results | ~16 GB |
|
||||
| `t2_finqa` — T²-RAGBench (FinQA) financial QA, text embedder (`qwen3-embedding:4b`); scored by exact numeric match, run with `--target analysis-capability` | ~2 GB |
|
||||
| `hotpotqa` — HotpotQA multi-hop QA over Wikipedia paragraphs, text embedder (`qwen3-embedding:4b`) | ~1.5 GB |
|
||||
| `mtrag_clapnq` — MTRAG multi-turn RAG, ClapNQ (Wikipedia) passages, text embedder (`qwen3-embedding:4b`); also serves the `mtrag_clapnq_rewrite`, `mtrag_clapnq_live` and `mtrag_clapnq_live_uncompacted` keys | ~2.8 GB |
|
||||
|
||||
After downloading, run benchmarks with `--skip-db`. Each database is built with a specific embedder, so pass its reference config from `evaluations/configs/` (a database only opens against a config whose embedder matches):
|
||||
|
||||
```bash
|
||||
evaluations run orb_multimodal_nemotron --skip-db --config configs/orb_multimodal_nemotron.yaml
|
||||
```
|
||||
|
||||
The configs use `vllm` as the model host. Point `base_url` at your own OpenAI-compatible endpoints to reproduce the numbers.
|
||||
|
||||
### Configuration
|
||||
|
||||
The benchmark script accepts several options:
|
||||
|
||||
```bash
|
||||
evaluations run hotpotqa --config /path/to/haiku.rag.yaml --db /path/to/custom.lancedb
|
||||
```
|
||||
|
||||
**Options:**
|
||||
|
||||
- `--config PATH` - Specify a custom `haiku.rag.yaml` configuration file
|
||||
- `--db PATH` - Override the database path (default: platform-specific user data directory)
|
||||
- `--skip-db` - Skip updating the evaluation database
|
||||
- `--skip-retrieval` - Skip retrieval benchmark
|
||||
- `--skip-qa` - Skip QA benchmark
|
||||
- `--limit N` - Limit number of test cases
|
||||
- `--name NAME` - Override the evaluation name
|
||||
- `--target {rag-capability,analysis-capability}` - Choose which [capability](capabilities/index.md) to benchmark end-to-end (default: `rag-capability`). The target names remain stable dataset identifiers.
|
||||
- `--capability-model PROVIDER:NAME` - Override the capability model independently from the judge (default: `config.qa.model`, or `config.analysis.model` when set for `--target analysis-capability`).
|
||||
- `--filter CLAUSE` / `-f CLAUSE` - Restrict every benchmark search to a subset of the database (see [Restricting the corpus](#restricting-the-corpus)).
|
||||
|
||||
If no config file is specified, the script searches standard locations: `./haiku.rag.yaml`, user config directory, then falls back to defaults.
|
||||
|
||||
To pin the LLM judge in YAML (rather than the default `ollama:qwen3.8`). These are the recommended settings:
|
||||
|
||||
```yaml
|
||||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://localhost:8000/v1 # optional, for OpenAI-compatible servers (vLLM, LM Studio, etc.)
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
top_p: 0.95
|
||||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
reasoning_effort: low # qwen3.8: low | medium | xhigh (default)
|
||||
```
|
||||
|
||||
### Restricting the corpus
|
||||
|
||||
When a database holds documents from several corpora — only some of which a dataset's questions are drawn from — `--filter` restricts every benchmark search to a subset. It takes the same SQL `WHERE` clause as `haiku-rag search --filter`, over document columns (`id`, `uri`, `title`, `created_at`, `updated_at`, `metadata`). Each dataset writes its own URIs: `orb_text` uses bare arXiv ids such as `2407.01528v3`, `hotpotqa` uses page titles.
|
||||
|
||||
```bash
|
||||
evaluations run orb_text --skip-db --config haiku.rag.s3.yaml \
|
||||
--filter "uri LIKE '2407%'"
|
||||
```
|
||||
|
||||
If the corpora are distinguished by a tag rather than by URI, attach it at ingest time as document metadata and match it with `LIKE`. `metadata` is stored as a `json.dumps` string, so there is no JSON subfield access — match the serialized key/value, including the space after the colon:
|
||||
|
||||
```bash
|
||||
evaluations run orb_text --skip-db --filter "metadata LIKE '%\"corpus\": \"orb_text\"%'"
|
||||
```
|
||||
|
||||
The clause applies to both benchmark phases — the retrieval benchmark's searches and every search the capability runs during QA — so the two score the same subset. It is recorded as `document_filter` in the run's experiment metadata, so a filtered run is never mistaken for an unfiltered one when comparing results.
|
||||
|
||||
Filtering affects searches only — a run without `--skip-db` still populates the database with the dataset's full corpus.
|
||||
|
|
|
|||
|
|
@ -122,6 +122,18 @@ The `storage_options` keys are case-insensitive and passed directly to the under
|
|||
|
||||
**Note:** Table optimization is automatically handled by LanceDB Cloud (`db://` URIs) and is disabled for better performance. For object storage backends (S3, Azure, GCS), optimization and vector indexing are still performed normally.
|
||||
|
||||
### Caching and Read Consistency
|
||||
|
||||
```yaml
|
||||
lancedb:
|
||||
read_consistency_interval_seconds: 30 # null to never re-check
|
||||
index_cache_size_bytes: 536870912 # null for the LanceDB default
|
||||
metadata_cache_size_bytes: 268435456
|
||||
```
|
||||
|
||||
- **read_consistency_interval_seconds**: how often a connection checks for writes from another process. `null` never checks, so a long-lived reader never sees the ingester's writes. `0` checks on every read.
|
||||
- **index_cache_size_bytes** / **metadata_cache_size_bytes**: sizes for the caches held by the LanceDB session, which is shared across every connection in the process. The first vector query loads the index into it, so on object storage the cache is what stops the next connection refetching it. Size it for the total set of indexes a process keeps warm, against the memory available to it.
|
||||
|
||||
### Deployment Pattern: One Writer, Many Readers
|
||||
|
||||
LanceDB on S3 supports **exactly one writer + N readers** per database URI. Multiple writers against the same URI can race on the manifest commit and corrupt state. This is a LanceDB property, not something `haiku.rag` enforces.
|
||||
|
|
|
|||
|
|
@ -1,3 +1,37 @@
|
|||
---
|
||||
title: haiku.rag
|
||||
description: Local-first agentic RAG. Index PDFs, web pages, and whole directories, then ask questions and get answers cited to page numbers and section headings. Hybrid search, reranking, and multimodal retrieval on embedded LanceDB.
|
||||
---
|
||||
|
||||
haiku.rag indexes PDFs, web pages, and whole directories, retrieves with hybrid search, and answers with citations down to the page number and section heading. It runs on an embedded database with open models, so your documents stay on your machine and there is no server to operate.
|
||||
|
||||
```bash
|
||||
uv pip install haiku.rag
|
||||
|
||||
haiku-rag init
|
||||
haiku-rag add-src ~/Documents/some-paper.pdf
|
||||
haiku-rag ask "what does it conclude?"
|
||||
```
|
||||
|
||||
[Quickstart](tutorial.md) covers provider setup and the first ingestion.
|
||||
|
||||
## Why haiku.rag
|
||||
|
||||
**Answers you can check.** Every answer carries citations with page numbers and section headings. Visual grounding shows the cited chunk highlighted on the original page image. Optional capabilities require an answer to declare what grounds it, including declaring that nothing does.
|
||||
|
||||
**Local-first, no server.** Embedded [LanceDB](https://lancedb.com/) and open models through [Ollama](https://ollama.com/) by default. No database to run and no API keys required. The same code runs against S3, GCS, Azure, LanceDB Cloud, or any provider Pydantic AI supports.
|
||||
|
||||
**Built for agents.** Native [Pydantic AI](https://ai.pydantic.dev/) capabilities compose into your own agents. An [MCP server](mcp.md) exposes the same database to Claude Desktop and other assistants. The analysis capability runs sandboxed Python across documents for questions that need computation rather than retrieval.
|
||||
|
||||
**Measured, not asserted.** Retrieval and answer quality are tracked against public benchmarks with runnable configs. See [Benchmarks](benchmarks.md).
|
||||
|
||||
## Start here
|
||||
|
||||
- [Quickstart](tutorial.md): install, index, chat.
|
||||
- [Overview](overview.md): what haiku.rag does, end to end.
|
||||
- [Capabilities](capabilities/index.md): native RAG and analysis capabilities for Pydantic AI agents.
|
||||
- [Python API](python.md): use haiku.rag from code.
|
||||
- [MCP server](mcp.md): expose haiku.rag to Claude Desktop or other AI assistants.
|
||||
- [Configuration](configuration/index.md): every setting.
|
||||
|
||||
MIT licensed. Source on [GitHub](https://github.com/ggozad/haiku.rag).
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
# Haiku RAG - Evaluations
|
||||
# haiku.rag - Evaluations
|
||||
|
||||
Internal benchmarking and evaluation scripts for haiku.rag.
|
||||
|
||||
|
|
|
|||
|
|
@ -25,8 +25,8 @@ qa:
|
|||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: RedHatAI/Qwen3.6-35B-A3B-NVFP4
|
||||
base_url: http://vllm:11430/v1
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
|
|
@ -34,4 +34,4 @@ evaluations:
|
|||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
enable_thinking: true
|
||||
reasoning_effort: low
|
||||
|
|
|
|||
|
|
@ -44,8 +44,8 @@ qa:
|
|||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: RedHatAI/Qwen3.6-35B-A3B-NVFP4
|
||||
base_url: http://vllm:11430/v1
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
|
|
@ -53,4 +53,4 @@ evaluations:
|
|||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
enable_thinking: true
|
||||
reasoning_effort: low
|
||||
|
|
|
|||
|
|
@ -29,8 +29,8 @@ qa:
|
|||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: RedHatAI/Qwen3.6-35B-A3B-NVFP4
|
||||
base_url: http://vllm:11430/v1
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
|
|
@ -38,4 +38,4 @@ evaluations:
|
|||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
enable_thinking: true
|
||||
reasoning_effort: low
|
||||
|
|
|
|||
|
|
@ -30,8 +30,8 @@ qa:
|
|||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: RedHatAI/Qwen3.6-35B-A3B-NVFP4
|
||||
base_url: http://vllm:11430/v1
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
|
|
@ -39,4 +39,4 @@ evaluations:
|
|||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
enable_thinking: true
|
||||
reasoning_effort: low
|
||||
|
|
|
|||
|
|
@ -31,8 +31,8 @@ qa:
|
|||
evaluations:
|
||||
judge:
|
||||
provider: openai
|
||||
name: RedHatAI/Qwen3.6-35B-A3B-NVFP4
|
||||
base_url: http://vllm:11430/v1
|
||||
name: Inferact/Qwen3.8-27B-NVFP4
|
||||
base_url: http://vllm:11439/v1
|
||||
temperature: 0.6
|
||||
max_tokens: 16384
|
||||
extra_body:
|
||||
|
|
@ -40,4 +40,4 @@ evaluations:
|
|||
top_k: 20
|
||||
min_p: 0
|
||||
chat_template_kwargs:
|
||||
enable_thinking: true
|
||||
reasoning_effort: low
|
||||
|
|
|
|||
|
|
@ -47,10 +47,11 @@ TARGETS: tuple[Target, ...] = ("rag-capability", "analysis-capability")
|
|||
# Sampling follows Qwen's recommendation for thinking mode; its model cards
|
||||
# forbid greedy decoding. Only the keys ollama honours are set: it silently
|
||||
# ignores `top_k`, `min_p` and `chat_template_kwargs`. The vLLM reference
|
||||
# configs under `evaluations/configs/` carry those too.
|
||||
# configs under `evaluations/configs/` carry those too, plus
|
||||
# `reasoning_effort`, which qwen3.8 reads from `chat_template_kwargs`.
|
||||
DEFAULT_JUDGE_MODEL = ModelConfig(
|
||||
provider="ollama",
|
||||
name="qwen3.6",
|
||||
name="qwen3.8",
|
||||
temperature=0.6,
|
||||
max_tokens=16384,
|
||||
extra_body={"top_p": 0.95},
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@
|
|||
|
||||
name = "haiku.rag-evals"
|
||||
description = "Benchmarking and evaluation scripts for haiku.rag"
|
||||
version = "0.74.0"
|
||||
version = "0.75.0"
|
||||
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
|
||||
license = { text = "MIT" }
|
||||
requires-python = ">=3.12"
|
||||
|
|
|
|||
|
|
@ -677,6 +677,7 @@ class TestRunQaBenchmarkJudgeModel:
|
|||
from evaluations.benchmark import DEFAULT_JUDGE_MODEL
|
||||
|
||||
assert DEFAULT_JUDGE_MODEL.temperature == 0.6
|
||||
assert DEFAULT_JUDGE_MODEL.name == "qwen3.8"
|
||||
assert DEFAULT_JUDGE_MODEL.max_tokens == 16384
|
||||
assert DEFAULT_JUDGE_MODEL.extra_body == {"top_p": 0.95}
|
||||
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ PINNED_JUDGE_SAMPLING = {
|
|||
"top_p": 0.95,
|
||||
"top_k": 20,
|
||||
"min_p": 0,
|
||||
"chat_template_kwargs": {"enable_thinking": True},
|
||||
"chat_template_kwargs": {"reasoning_effort": "low"},
|
||||
},
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -129,6 +129,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
state: StateT | None = field(default=None, repr=False)
|
||||
outer_state: dict[str, Any] | None = field(default=None, repr=False)
|
||||
rag: HaikuRAG | None = field(default=None, repr=False)
|
||||
"""A connection this capability opened, and must close."""
|
||||
borrowed_rag: HaikuRAG | None = field(default=None, repr=False)
|
||||
"""A caller's connection, reused and never closed here."""
|
||||
rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
|
||||
resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
|
||||
search_count: int = field(default=0, repr=False)
|
||||
|
|
@ -347,6 +350,8 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
|
|||
raise error
|
||||
|
||||
async def _ensure_rag(self) -> HaikuRAG:
|
||||
if self.borrowed_rag is not None:
|
||||
return self.borrowed_rag
|
||||
if self.rag is None:
|
||||
async with self.resource_lock:
|
||||
if self.rag is None:
|
||||
|
|
|
|||
|
|
@ -1,13 +1,16 @@
|
|||
from dataclasses import dataclass, field
|
||||
from functools import cache
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic_ai import RunContext, ToolFailed
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
from pydantic_ai.toolsets import FunctionToolset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
from haiku.rag.capabilities._base import (
|
||||
CodeExecutionEntry,
|
||||
RAGCapabilityBase,
|
||||
|
|
@ -162,6 +165,7 @@ def create_capability(
|
|||
config: AppConfig | None = None,
|
||||
*,
|
||||
defer_loading: bool = True,
|
||||
rag: "HaikuRAG | None" = None,
|
||||
request_limit: int | None = 30,
|
||||
vision: bool | None = None,
|
||||
) -> AnalysisCapability:
|
||||
|
|
@ -180,6 +184,7 @@ def create_capability(
|
|||
return AnalysisCapability(
|
||||
db_path=resolve_db_path(db_path, config),
|
||||
config=config,
|
||||
borrowed_rag=rag,
|
||||
state_type=AnalysisState,
|
||||
state_namespace=STATE_NAMESPACE,
|
||||
instruction_text=instructions(),
|
||||
|
|
|
|||
|
|
@ -1,13 +1,16 @@
|
|||
from dataclasses import dataclass
|
||||
from functools import cache
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic_ai import RunContext
|
||||
from pydantic_ai.messages import ToolReturn
|
||||
from pydantic_ai.toolsets import FunctionToolset
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
from haiku.rag.capabilities._base import (
|
||||
RAGCapabilityBase,
|
||||
resolve_db_path,
|
||||
|
|
@ -72,6 +75,7 @@ def create_capability(
|
|||
config: AppConfig | None = None,
|
||||
*,
|
||||
defer_loading: bool = True,
|
||||
rag: "HaikuRAG | None" = None,
|
||||
request_limit: int | None = 20,
|
||||
vision: bool | None = None,
|
||||
) -> RAGCapability:
|
||||
|
|
@ -88,6 +92,7 @@ def create_capability(
|
|||
return RAGCapability(
|
||||
db_path=resolve_db_path(db_path, config),
|
||||
config=config,
|
||||
borrowed_rag=rag,
|
||||
state_type=RAGState,
|
||||
state_namespace=STATE_NAMESPACE,
|
||||
instruction_text=instructions(),
|
||||
|
|
|
|||
|
|
@ -63,6 +63,7 @@ async def ask(
|
|||
capability = create_capability(
|
||||
db_path=client.store.db_path,
|
||||
config=client._config,
|
||||
rag=client,
|
||||
defer_loading=False,
|
||||
)
|
||||
deps = _AgentDeps(
|
||||
|
|
@ -115,6 +116,7 @@ async def analyze(
|
|||
capability = create_capability(
|
||||
db_path=client.store.db_path,
|
||||
config=client._config,
|
||||
rag=client,
|
||||
defer_loading=False,
|
||||
)
|
||||
deps = _AgentDeps(
|
||||
|
|
|
|||
|
|
@ -62,10 +62,21 @@ class StorageConfig(BaseModel):
|
|||
|
||||
|
||||
class LanceDBConfig(BaseModel):
|
||||
"""LanceDB connection settings.
|
||||
|
||||
read_consistency_interval_seconds bounds how stale a reader may be. None
|
||||
never re-checks, so a long-lived reader never sees another process's writes.
|
||||
The cache sizes are per process, since the session is shared across
|
||||
connections.
|
||||
"""
|
||||
|
||||
uri: str = ""
|
||||
api_key: str = ""
|
||||
region: str = ""
|
||||
storage_options: dict[str, str] = Field(default_factory=dict)
|
||||
read_consistency_interval_seconds: float | None = Field(default=30, ge=0)
|
||||
index_cache_size_bytes: int | None = Field(default=None, ge=0)
|
||||
metadata_cache_size_bytes: int | None = Field(default=None, ge=0)
|
||||
|
||||
|
||||
class EmbeddingsConfig(BaseModel):
|
||||
|
|
|
|||
|
|
@ -1,3 +1,6 @@
|
|||
import asyncio
|
||||
from collections.abc import AsyncIterator
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
|
@ -28,7 +31,42 @@ def create_mcp_server(
|
|||
config: Configuration to use.
|
||||
read_only: If True, write tools (add_document_*, delete_document) are not registered.
|
||||
"""
|
||||
mcp = FastMCP("haiku-rag")
|
||||
client: HaikuRAG | None = None
|
||||
stack = AsyncExitStack()
|
||||
client_lock = asyncio.Lock()
|
||||
|
||||
async def _client() -> HaikuRAG:
|
||||
"""The server's client, opened once.
|
||||
|
||||
Opening cost is per connection, and on object storage the first vector
|
||||
query loads the index into the session cache, so a client per tool call
|
||||
pays that repeatedly.
|
||||
"""
|
||||
nonlocal client
|
||||
async with client_lock:
|
||||
if client is None:
|
||||
client = await stack.enter_async_context(
|
||||
HaikuRAG(db_path, config=config, read_only=read_only)
|
||||
)
|
||||
return client
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_server: FastMCP) -> AsyncIterator[None]:
|
||||
# Open eagerly so an unopenable database fails startup rather than
|
||||
# every tool call.
|
||||
nonlocal client
|
||||
await _client()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
# The lifespan can be re-entered; without the reset the next cycle
|
||||
# hands out the closed client, including when aclose itself fails.
|
||||
try:
|
||||
await stack.aclose()
|
||||
finally:
|
||||
client = None
|
||||
|
||||
mcp = FastMCP("haiku-rag", lifespan=lifespan)
|
||||
|
||||
# Write tools - only registered when not in read-only mode
|
||||
if not read_only:
|
||||
|
|
@ -41,14 +79,14 @@ def create_mcp_server(
|
|||
) -> str | None:
|
||||
"""Add a document to the RAG system from a file path."""
|
||||
try:
|
||||
async with HaikuRAG(db_path, config=config) as rag:
|
||||
result = await rag.create_document_from_source(
|
||||
Path(file_path), title=title, metadata=metadata or {}
|
||||
)
|
||||
# Handle both single document and list of documents (directories)
|
||||
if isinstance(result, list):
|
||||
return result[0].id if result else None
|
||||
return result.id
|
||||
rag = await _client()
|
||||
result = await rag.create_document_from_source(
|
||||
Path(file_path), title=title, metadata=metadata or {}
|
||||
)
|
||||
# Handle both single document and list of documents (directories)
|
||||
if isinstance(result, list):
|
||||
return result[0].id if result else None
|
||||
return result.id
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
|
@ -58,14 +96,14 @@ def create_mcp_server(
|
|||
) -> str | None:
|
||||
"""Add a document to the RAG system from a URL."""
|
||||
try:
|
||||
async with HaikuRAG(db_path, config=config) as rag:
|
||||
result = await rag.create_document_from_source(
|
||||
url, title=title, metadata=metadata or {}
|
||||
)
|
||||
# Handle both single document and list of documents
|
||||
if isinstance(result, list):
|
||||
return result[0].id if result else None
|
||||
return result.id
|
||||
rag = await _client()
|
||||
result = await rag.create_document_from_source(
|
||||
url, title=title, metadata=metadata or {}
|
||||
)
|
||||
# Handle both single document and list of documents
|
||||
if isinstance(result, list):
|
||||
return result[0].id if result else None
|
||||
return result.id
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
|
@ -78,11 +116,11 @@ def create_mcp_server(
|
|||
) -> str | None:
|
||||
"""Add a document to the RAG system from text content."""
|
||||
try:
|
||||
async with HaikuRAG(db_path, config=config) as rag:
|
||||
document = await rag.create_document(
|
||||
content, uri, title=title, metadata=metadata or {}
|
||||
)
|
||||
return document.id
|
||||
rag = await _client()
|
||||
document = await rag.create_document(
|
||||
content, uri, title=title, metadata=metadata or {}
|
||||
)
|
||||
return document.id
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
|
@ -90,10 +128,8 @@ def create_mcp_server(
|
|||
async def delete_document(document_id: str) -> bool:
|
||||
"""Delete a document by its ID."""
|
||||
try:
|
||||
async with HaikuRAG(
|
||||
db_path, config=config, skip_validation=True
|
||||
) as rag:
|
||||
return await rag.delete_document(document_id)
|
||||
rag = await _client()
|
||||
return await rag.delete_document(document_id)
|
||||
except Exception:
|
||||
return False
|
||||
|
||||
|
|
@ -110,10 +146,8 @@ def create_mcp_server(
|
|||
response (smaller JSON payload for plain-text consumers).
|
||||
"""
|
||||
try:
|
||||
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
||||
return await rag.search(
|
||||
query, limit=limit, include_images=include_images
|
||||
)
|
||||
rag = await _client()
|
||||
return await rag.search(query, limit=limit, include_images=include_images)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
|
@ -145,10 +179,8 @@ def create_mcp_server(
|
|||
except Exception:
|
||||
return []
|
||||
try:
|
||||
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
||||
return await rag.search(
|
||||
raw, limit=limit, include_images=include_images
|
||||
)
|
||||
rag = await _client()
|
||||
return await rag.search(raw, limit=limit, include_images=include_images)
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
|
@ -156,8 +188,8 @@ def create_mcp_server(
|
|||
async def get_document(document_id: str) -> Document | None:
|
||||
"""Get a document by its ID."""
|
||||
try:
|
||||
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
||||
return await rag.get_document_by_id(document_id)
|
||||
rag = await _client()
|
||||
return await rag.get_document_by_id(document_id)
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
|
@ -175,18 +207,18 @@ def create_mcp_server(
|
|||
filter: Optional SQL WHERE clause to filter documents.
|
||||
"""
|
||||
try:
|
||||
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
||||
documents = await rag.list_documents(limit, offset, filter)
|
||||
rag = await _client()
|
||||
documents = await rag.list_documents(limit, offset, filter)
|
||||
|
||||
return [
|
||||
DocumentInfo(
|
||||
id=doc.id,
|
||||
title=doc.title or "Untitled",
|
||||
uri=doc.uri or "",
|
||||
created=doc.created_at.strftime("%Y-%m-%d"),
|
||||
)
|
||||
for doc in documents
|
||||
]
|
||||
return [
|
||||
DocumentInfo(
|
||||
id=doc.id,
|
||||
title=doc.title or "Untitled",
|
||||
uri=doc.uri or "",
|
||||
created=doc.created_at.strftime("%Y-%m-%d"),
|
||||
)
|
||||
for doc in documents
|
||||
]
|
||||
except Exception:
|
||||
return []
|
||||
|
||||
|
|
@ -209,11 +241,11 @@ def create_mcp_server(
|
|||
"""
|
||||
try:
|
||||
images = _decode_images(images_base64)
|
||||
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
||||
answer, citations = await rag.ask(question, images=images)
|
||||
if cite and citations:
|
||||
answer += "\n\n" + format_citations(citations)
|
||||
return answer
|
||||
rag = await _client()
|
||||
answer, citations = await rag.ask(question, images=images)
|
||||
if cite and citations:
|
||||
answer += "\n\n" + format_citations(citations)
|
||||
return answer
|
||||
except Exception as e:
|
||||
return f"Error answering question: {e!s}"
|
||||
|
||||
|
|
@ -240,9 +272,9 @@ def create_mcp_server(
|
|||
"""
|
||||
try:
|
||||
images = _decode_images(images_base64)
|
||||
async with HaikuRAG(db_path, config=config, read_only=read_only) as rag:
|
||||
result = await rag.analyze(question, filter=filter, images=images)
|
||||
return result.answer
|
||||
rag = await _client()
|
||||
result = await rag.analyze(question, filter=filter, images=images)
|
||||
return result.answer
|
||||
except Exception as e:
|
||||
return f"Error running analysis capability: {e!s}"
|
||||
|
||||
|
|
|
|||
|
|
@ -12,7 +12,7 @@ from uuid import uuid4
|
|||
|
||||
import lancedb
|
||||
import pyarrow as pa
|
||||
from lancedb.index import FTS, BTree, IvfPq
|
||||
from lancedb.index import FTS, Bitmap, BTree, IvfPq
|
||||
from lancedb.pydantic import LanceModel, Vector
|
||||
from packaging.version import parse
|
||||
from pydantic import BaseModel, Field
|
||||
|
|
@ -55,25 +55,56 @@ class ConnectionMode(Enum):
|
|||
return ConnectionMode.OBJECT_STORAGE
|
||||
|
||||
|
||||
_sessions: dict[tuple[int | None, int | None], lancedb.Session] = {}
|
||||
|
||||
|
||||
def _session(config: AppConfig) -> lancedb.Session:
|
||||
"""The process's session for these cache sizes.
|
||||
|
||||
Sessions hold the index and metadata caches. Sharing one across connections
|
||||
is what keeps a cached index from being refetched per connection, which on
|
||||
object storage is the dominant cost of the first query.
|
||||
"""
|
||||
key = (
|
||||
config.lancedb.index_cache_size_bytes,
|
||||
config.lancedb.metadata_cache_size_bytes,
|
||||
)
|
||||
if key not in _sessions:
|
||||
kwargs = {}
|
||||
if key[0] is not None:
|
||||
kwargs["index_cache_size_bytes"] = key[0]
|
||||
if key[1] is not None:
|
||||
kwargs["metadata_cache_size_bytes"] = key[1]
|
||||
_sessions[key] = lancedb.Session(**kwargs)
|
||||
return _sessions[key]
|
||||
|
||||
|
||||
async def connect_lancedb(
|
||||
config: AppConfig, db_path: Path | None = None
|
||||
) -> lancedb.AsyncConnection:
|
||||
interval = config.lancedb.read_consistency_interval_seconds
|
||||
kwargs: dict[str, Any] = {
|
||||
"session": _session(config),
|
||||
"read_consistency_interval": (
|
||||
timedelta(seconds=interval) if interval is not None else None
|
||||
),
|
||||
}
|
||||
mode = ConnectionMode.from_config(config)
|
||||
if mode == ConnectionMode.CLOUD:
|
||||
return await lancedb.connect_async(
|
||||
uri=config.lancedb.uri,
|
||||
api_key=config.lancedb.api_key,
|
||||
region=config.lancedb.region,
|
||||
**kwargs,
|
||||
)
|
||||
elif mode == ConnectionMode.OBJECT_STORAGE:
|
||||
kwargs: dict[str, Any] = {"uri": config.lancedb.uri}
|
||||
if config.lancedb.storage_options:
|
||||
kwargs["storage_options"] = config.lancedb.storage_options
|
||||
return await lancedb.connect_async(**kwargs)
|
||||
return await lancedb.connect_async(uri=config.lancedb.uri, **kwargs)
|
||||
else:
|
||||
if db_path is None:
|
||||
raise ValueError("No lancedb.uri configured and no db_path provided")
|
||||
return await lancedb.connect_async(db_path.absolute())
|
||||
return await lancedb.connect_async(db_path.absolute(), **kwargs)
|
||||
|
||||
|
||||
class DocumentRecord(LanceModel):
|
||||
|
|
@ -176,6 +207,64 @@ def get_document_items_arrow_schema() -> pa.Schema:
|
|||
return pa.schema(fields)
|
||||
|
||||
|
||||
def _stored_vector_dim(settings: dict) -> int | None:
|
||||
"""The vector dimension a database's chunks were written at."""
|
||||
return settings.get("embeddings", {}).get("model", {}).get("vector_dim")
|
||||
|
||||
|
||||
def index_specs(table_name: str) -> list[tuple[str, Bitmap | BTree | FTS]]:
|
||||
"""The index set each table carries."""
|
||||
match table_name:
|
||||
case "documents":
|
||||
return [("id", BTree())]
|
||||
case "document_meta":
|
||||
return [("id", BTree()), ("uri", BTree())]
|
||||
case "chunks":
|
||||
return [
|
||||
# Positions and stop words are required for phrase queries.
|
||||
("content_fts", FTS(with_position=True, remove_stop_words=False)),
|
||||
("id", BTree()),
|
||||
("document_id", BTree()),
|
||||
]
|
||||
case "document_items":
|
||||
return [
|
||||
("document_id", BTree()),
|
||||
("position", BTree()),
|
||||
("self_ref", BTree()),
|
||||
("label", Bitmap()),
|
||||
]
|
||||
case _:
|
||||
return []
|
||||
|
||||
|
||||
async def ensure_indexes(table: lancedb.AsyncTable, table_name: str) -> list[str]:
|
||||
"""Create any declared index missing from a column. Returns the columns indexed.
|
||||
|
||||
Matches on index type, not column coverage, so a BTree does not satisfy a
|
||||
declared Bitmap. Never drops or converts an index it did not declare.
|
||||
Re-creating is not free: `create_index(replace=True)` rebuilds.
|
||||
"""
|
||||
covering: dict[str, set[str]] = {}
|
||||
for index in await table.list_indices():
|
||||
for column in index.columns:
|
||||
covering.setdefault(column, set()).add(index.index_type)
|
||||
|
||||
applied: list[str] = []
|
||||
for column, config in index_specs(table_name):
|
||||
declared = type(config).__name__
|
||||
present = covering.get(column, set())
|
||||
if declared in present:
|
||||
continue
|
||||
if present:
|
||||
logger.info(
|
||||
f"Adding {declared} index on {table_name}.{column}, which carries "
|
||||
f"{', '.join(sorted(present))}"
|
||||
)
|
||||
await table.create_index(column, config=config, replace=True)
|
||||
applied.append(column)
|
||||
return applied
|
||||
|
||||
|
||||
class SettingsRecord(LanceModel):
|
||||
id: str = Field(default="settings")
|
||||
settings: str = Field(default="{}")
|
||||
|
|
@ -449,29 +538,27 @@ class Store:
|
|||
self._config, self.db_path
|
||||
)
|
||||
|
||||
# For remote stores (and as a safety net for local paths that exist but
|
||||
# have no tables — e.g. a previously failed init), detect new DB by
|
||||
# checking whether any tables exist.
|
||||
is_new_db = self._is_new_db
|
||||
if not is_new_db:
|
||||
existing_tables = (await self.db.list_tables()).tables
|
||||
if not existing_tables:
|
||||
is_new_db = True
|
||||
# Read once and thread onward: on object storage each of these is a
|
||||
# round trip. A local path that exists with no tables is a failed init,
|
||||
# so treat it as new.
|
||||
existing_tables = (await self.db.list_tables()).tables
|
||||
is_new_db = self._is_new_db or not existing_tables
|
||||
|
||||
# For existing databases, read stored vector dimension to create ChunkRecord
|
||||
# that can read existing chunks. For new databases, use config's dimension.
|
||||
stored_vector_dim = None
|
||||
if not is_new_db:
|
||||
stored_vector_dim = await self._get_stored_vector_dim()
|
||||
stored_settings: dict = {}
|
||||
if not is_new_db and "settings" in existing_tables:
|
||||
self.settings_table = await self.db.open_table("settings")
|
||||
stored_settings = await self._read_stored_settings()
|
||||
|
||||
# Create ChunkRecord with stored dimension (for reading) or config dimension (for new DB)
|
||||
# An existing database's chunks can only be read with the dimension they
|
||||
# were written at.
|
||||
stored_vector_dim = _stored_vector_dim(stored_settings)
|
||||
chunk_vector_dim = stored_vector_dim or self.embedder._vector_dim
|
||||
self.ChunkRecord: type[ChunkRecordBase] = create_chunk_model(chunk_vector_dim)
|
||||
|
||||
# Initialize tables (creates them if they don't exist). For an existing
|
||||
# DB this raises MigrationRequiredError up front when migrations are
|
||||
# pending, before creating any newly-introduced table.
|
||||
await self._init_tables(is_new_db)
|
||||
await self._init_tables(is_new_db, existing_tables, stored_settings)
|
||||
|
||||
# Set version for new databases.
|
||||
if is_new_db and not self._read_only:
|
||||
|
|
@ -479,7 +566,7 @@ class Store:
|
|||
|
||||
# Validate config compatibility after connection is established
|
||||
if not self._skip_validation:
|
||||
await self._validate_configuration()
|
||||
await self._validate_configuration(stored_settings)
|
||||
|
||||
async def __aenter__(self):
|
||||
# If _initialize connects to LanceDB but then fails (e.g. migration
|
||||
|
|
@ -501,33 +588,26 @@ class Store:
|
|||
"""Whether the store is in read-only mode."""
|
||||
return self._read_only
|
||||
|
||||
async def _get_stored_vector_dim(self) -> int | None:
|
||||
"""Read the stored vector dimension from the settings table.
|
||||
async def _read_stored_settings(self) -> dict:
|
||||
"""The stored settings blob, or {} if it is absent or not a JSON object.
|
||||
|
||||
Returns:
|
||||
The stored vector dimension, or None if not found.
|
||||
Only decoding failures are tolerated. A storage failure must propagate:
|
||||
read as empty settings it would look like version 0.0.0, and the
|
||||
migration check would declare every migration pending.
|
||||
"""
|
||||
rows = (
|
||||
await self.settings_table.query()
|
||||
.where("id = 'settings'")
|
||||
.limit(1)
|
||||
.to_arrow()
|
||||
).to_pylist()
|
||||
if not rows or not rows[0].get("settings"):
|
||||
return {}
|
||||
try:
|
||||
existing_tables = (await self.db.list_tables()).tables
|
||||
if "settings" not in existing_tables:
|
||||
return None
|
||||
|
||||
settings_table = await self.db.open_table("settings")
|
||||
rows = (
|
||||
await settings_table.query()
|
||||
.where("id = 'settings'")
|
||||
.limit(1)
|
||||
.to_arrow()
|
||||
).to_pylist()
|
||||
if not rows or not rows[0].get("settings"):
|
||||
return None
|
||||
|
||||
settings = json.loads(rows[0]["settings"])
|
||||
embeddings = settings.get("embeddings", {})
|
||||
model = embeddings.get("model", {})
|
||||
return model.get("vector_dim")
|
||||
except Exception:
|
||||
return None
|
||||
decoded = json.loads(rows[0]["settings"])
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return {}
|
||||
return decoded if isinstance(decoded, dict) else {}
|
||||
|
||||
def _assert_writable(self) -> None:
|
||||
"""Raise ReadOnlyError if the store is in read-only mode."""
|
||||
|
|
@ -658,16 +738,19 @@ class Store:
|
|||
except Exception as e:
|
||||
logger.warning(f"Could not create vector index: {e}")
|
||||
|
||||
async def _validate_configuration(self) -> None:
|
||||
async def _validate_configuration(
|
||||
self, stored_settings: dict | None = None
|
||||
) -> None:
|
||||
"""Validate that the configuration is compatible with the database."""
|
||||
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||
|
||||
settings_repo = SettingsRepository(self)
|
||||
await settings_repo.validate_config_compatibility()
|
||||
await settings_repo.validate_config_compatibility(stored_settings)
|
||||
|
||||
async def _init_tables(self, is_new_db: bool):
|
||||
async def _init_tables(
|
||||
self, is_new_db: bool, existing_tables: list[str], stored_settings: dict
|
||||
):
|
||||
"""Initialize database tables (create if they don't exist)."""
|
||||
existing_tables = (await self.db.list_tables()).tables
|
||||
|
||||
# Surface pending migrations BEFORE creating any newly-introduced table.
|
||||
# Otherwise opening a legacy DB would either mutate it (creating an empty
|
||||
|
|
@ -679,8 +762,7 @@ class Store:
|
|||
and not self._skip_migration_check
|
||||
and "settings" in existing_tables
|
||||
):
|
||||
self.settings_table = await self.db.open_table("settings")
|
||||
await self._check_migrations()
|
||||
await self._check_migrations(stored_settings.get("version", "0.0.0"))
|
||||
|
||||
missing_tables = set(REQUIRED_TABLES) - set(existing_tables)
|
||||
|
||||
|
|
@ -697,22 +779,17 @@ class Store:
|
|||
self.documents_table = await self.db.create_table(
|
||||
"documents", schema=get_documents_arrow_schema()
|
||||
)
|
||||
await ensure_indexes(self.documents_table, "documents")
|
||||
|
||||
# Create or open document_meta table (mutable attributes kept out of the
|
||||
# blob-bearing documents row). Indexed by document_id and uri — both are
|
||||
# hot look-up keys (get_by_id, get_by_uri).
|
||||
# blob-bearing documents row).
|
||||
if "document_meta" in existing_tables:
|
||||
self.document_meta_table = await self.db.open_table("document_meta")
|
||||
else:
|
||||
self.document_meta_table = await self.db.create_table(
|
||||
"document_meta", schema=DocumentMetaRecord
|
||||
)
|
||||
await self.document_meta_table.create_index(
|
||||
"id", config=BTree(), replace=True
|
||||
)
|
||||
await self.document_meta_table.create_index(
|
||||
"uri", config=BTree(), replace=True
|
||||
)
|
||||
await ensure_indexes(self.document_meta_table, "document_meta")
|
||||
|
||||
# Create or open chunks table
|
||||
if "chunks" in existing_tables:
|
||||
|
|
@ -721,12 +798,7 @@ class Store:
|
|||
self.chunks_table = await self.db.create_table(
|
||||
"chunks", schema=self.ChunkRecord
|
||||
)
|
||||
# Create FTS index on content_fts (contextualized content) for better search
|
||||
await self.chunks_table.create_index(
|
||||
"content_fts",
|
||||
config=FTS(with_position=True, remove_stop_words=False),
|
||||
replace=True,
|
||||
)
|
||||
await ensure_indexes(self.chunks_table, "chunks")
|
||||
|
||||
# Create or open document_items table
|
||||
if "document_items" in existing_tables:
|
||||
|
|
@ -735,20 +807,10 @@ class Store:
|
|||
self.document_items_table = await self.db.create_table(
|
||||
"document_items", schema=get_document_items_arrow_schema()
|
||||
)
|
||||
await self.document_items_table.create_index(
|
||||
"document_id", config=BTree(), replace=True
|
||||
)
|
||||
await self.document_items_table.create_index(
|
||||
"position", config=BTree(), replace=True
|
||||
)
|
||||
await self.document_items_table.create_index(
|
||||
"self_ref", config=BTree(), replace=True
|
||||
)
|
||||
await ensure_indexes(self.document_items_table, "document_items")
|
||||
|
||||
# Create or open settings table
|
||||
if "settings" in existing_tables:
|
||||
self.settings_table = await self.db.open_table("settings")
|
||||
else:
|
||||
# _initialize opened the settings table when the database had one.
|
||||
if "settings" not in existing_tables:
|
||||
self.settings_table = await self.db.create_table(
|
||||
"settings", schema=SettingsRecord
|
||||
)
|
||||
|
|
@ -762,7 +824,7 @@ class Store:
|
|||
"""Set the initial version for a new database."""
|
||||
await self.set_haiku_version(metadata.version("haiku.rag-slim"))
|
||||
|
||||
async def _check_migrations(self) -> None:
|
||||
async def _check_migrations(self, db_version: str) -> None:
|
||||
"""Raise if migrations are pending. Opening never writes the version.
|
||||
|
||||
Raises:
|
||||
|
|
@ -771,7 +833,6 @@ class Store:
|
|||
from haiku.rag.store.upgrades import get_pending_upgrades
|
||||
|
||||
current_version = metadata.version("haiku.rag-slim")
|
||||
db_version = await self.get_haiku_version()
|
||||
|
||||
pending = get_pending_upgrades(db_version)
|
||||
|
||||
|
|
@ -872,13 +933,7 @@ class Store:
|
|||
self.chunks_table = await self.db.create_table(
|
||||
"chunks", schema=self.ChunkRecord
|
||||
)
|
||||
|
||||
# Create FTS index on content_fts (contextualized content) for better search
|
||||
await self.chunks_table.create_index(
|
||||
"content_fts",
|
||||
config=FTS(with_position=True, remove_stop_words=False),
|
||||
replace=True,
|
||||
)
|
||||
await ensure_indexes(self.chunks_table, "chunks")
|
||||
|
||||
def close(self):
|
||||
"""Close the database connection."""
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ if TYPE_CHECKING:
|
|||
from lancedb.index import FTS
|
||||
from lancedb.rerankers import RRFReranker
|
||||
|
||||
from haiku.rag.store.engine import Store, query_to_pydantic
|
||||
from haiku.rag.store.engine import Store, ensure_indexes, query_to_pydantic
|
||||
from haiku.rag.store.models.chunk import Chunk, SearchType
|
||||
from haiku.rag.utils import escape_sql_string
|
||||
|
||||
|
|
@ -187,12 +187,7 @@ class ChunkRepository:
|
|||
self.store.chunks_table = await self.store.db.create_table(
|
||||
"chunks", schema=self.store.ChunkRecord
|
||||
)
|
||||
# Create FTS index on content_fts (contextualized content) for better search
|
||||
await self.store.chunks_table.create_index(
|
||||
"content_fts",
|
||||
config=FTS(with_position=True, remove_stop_words=False),
|
||||
replace=True,
|
||||
)
|
||||
await ensure_indexes(self.store.chunks_table, "chunks")
|
||||
|
||||
async def delete_by_document_id(self, document_id: str) -> bool:
|
||||
"""Delete all chunks for a document."""
|
||||
|
|
|
|||
|
|
@ -3,12 +3,12 @@ from datetime import datetime
|
|||
from typing import overload
|
||||
from uuid import uuid4
|
||||
|
||||
from lancedb.index import BTree
|
||||
|
||||
from haiku.rag.store.engine import (
|
||||
DocumentMetaRecord,
|
||||
DocumentRecord,
|
||||
Store,
|
||||
ensure_indexes,
|
||||
get_document_items_arrow_schema,
|
||||
get_documents_arrow_schema,
|
||||
query_to_pydantic,
|
||||
)
|
||||
|
|
@ -407,23 +407,14 @@ class DocumentRepository:
|
|||
async def delete_all(self) -> None:
|
||||
"""Delete all documents from the database."""
|
||||
self.store._assert_writable()
|
||||
from haiku.rag.store.engine import DocumentItemRecord
|
||||
|
||||
# Delete all chunks and items first
|
||||
await self.chunk_repository.delete_all()
|
||||
await self.store.db.drop_table("document_items")
|
||||
self.store.document_items_table = await self.store.db.create_table(
|
||||
"document_items", schema=DocumentItemRecord
|
||||
)
|
||||
await self.store.document_items_table.create_index(
|
||||
"document_id", config=BTree(), replace=True
|
||||
)
|
||||
await self.store.document_items_table.create_index(
|
||||
"position", config=BTree(), replace=True
|
||||
)
|
||||
await self.store.document_items_table.create_index(
|
||||
"self_ref", config=BTree(), replace=True
|
||||
"document_items", schema=get_document_items_arrow_schema()
|
||||
)
|
||||
await ensure_indexes(self.store.document_items_table, "document_items")
|
||||
|
||||
# Get count before deletion
|
||||
count = len(
|
||||
|
|
@ -437,13 +428,9 @@ class DocumentRepository:
|
|||
self.store.documents_table = await self.store.db.create_table(
|
||||
"documents", schema=get_documents_arrow_schema()
|
||||
)
|
||||
await ensure_indexes(self.store.documents_table, "documents")
|
||||
await self.store.db.drop_table("document_meta")
|
||||
self.store.document_meta_table = await self.store.db.create_table(
|
||||
"document_meta", schema=DocumentMetaRecord
|
||||
)
|
||||
await self.store.document_meta_table.create_index(
|
||||
"id", config=BTree(), replace=True
|
||||
)
|
||||
await self.store.document_meta_table.create_index(
|
||||
"uri", config=BTree(), replace=True
|
||||
)
|
||||
await ensure_indexes(self.store.document_meta_table, "document_meta")
|
||||
|
|
|
|||
|
|
@ -60,7 +60,9 @@ class SettingsRepository:
|
|||
)
|
||||
await self.store.settings_table.add([settings_record])
|
||||
|
||||
async def validate_config_compatibility(self) -> None:
|
||||
async def validate_config_compatibility(
|
||||
self, stored_settings: dict | None = None
|
||||
) -> None:
|
||||
"""Validate the current configuration against stored settings without writing.
|
||||
|
||||
Opening a database never modifies it. ``vector_dim`` mismatches raise —
|
||||
|
|
@ -72,7 +74,8 @@ class SettingsRepository:
|
|||
while a read-only open continues. Stored settings are reconciled
|
||||
explicitly via ``haiku-rag rebuild --set-embedder``, never on open.
|
||||
"""
|
||||
stored_settings = await self.get_current_settings()
|
||||
if stored_settings is None:
|
||||
stored_settings = await self.get_current_settings()
|
||||
|
||||
# Nothing stored to validate against — never write on open.
|
||||
if not stored_settings:
|
||||
|
|
|
|||
|
|
@ -99,6 +99,9 @@ from haiku.rag.store.upgrades.v0_58_0 import (
|
|||
from haiku.rag.store.upgrades.v0_64_0 import (
|
||||
upgrade_rename_document_meta_id as upgrade_0_64_0_rename_document_meta_id,
|
||||
)
|
||||
from haiku.rag.store.upgrades.v0_75_0 import (
|
||||
upgrade_index_hot_lookup_keys as upgrade_0_75_0_index_hot_lookup_keys,
|
||||
)
|
||||
|
||||
upgrades.append(upgrade_0_20_0_docling)
|
||||
upgrades.append(upgrade_0_23_1_contextualize)
|
||||
|
|
@ -110,3 +113,4 @@ upgrades.append(upgrade_0_48_0_heading_hierarchy)
|
|||
upgrades.append(upgrade_0_50_0_canonical_metadata_keys)
|
||||
upgrades.append(upgrade_0_58_0_split_document_meta)
|
||||
upgrades.append(upgrade_0_64_0_rename_document_meta_id)
|
||||
upgrades.append(upgrade_0_75_0_index_hot_lookup_keys)
|
||||
|
|
|
|||
24
haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py
Normal file
24
haiku_rag_slim/haiku/rag/store/upgrades/v0_75_0.py
Normal file
|
|
@ -0,0 +1,24 @@
|
|||
import logging
|
||||
|
||||
from haiku.rag.store.engine import Store, ensure_indexes
|
||||
from haiku.rag.store.upgrades import Upgrade
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
async def _apply_index_hot_lookup_keys(store: Store) -> None:
|
||||
"""Add the declared indexes to a database created before 0.75.0.
|
||||
|
||||
Rewrites no rows. Each index build reads the column it indexes.
|
||||
"""
|
||||
for table_name, table in store._tables().items():
|
||||
applied = await ensure_indexes(table, table_name)
|
||||
if applied:
|
||||
logger.info(f"Indexed {table_name}: {', '.join(sorted(applied))}")
|
||||
|
||||
|
||||
upgrade_index_hot_lookup_keys = Upgrade(
|
||||
version="0.75.0",
|
||||
apply=_apply_index_hot_lookup_keys,
|
||||
description="Index documents.id, chunks.id, chunks.document_id and document_items.label",
|
||||
)
|
||||
|
|
@ -1,13 +1,28 @@
|
|||
[project]
|
||||
|
||||
name = "haiku.rag-slim"
|
||||
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling - Minimal dependencies"
|
||||
version = "0.74.0"
|
||||
description = "Local-first agentic RAG with citations - hybrid search, reranking and multimodal retrieval over your own documents, no database server required - Minimal dependencies"
|
||||
version = "0.75.0"
|
||||
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
|
||||
license = { text = "MIT" }
|
||||
readme = { file = "README.md", content-type = "text/markdown" }
|
||||
requires-python = ">=3.12"
|
||||
keywords = ["RAG", "lancedb", "vector-database", "ml", "mcp"]
|
||||
keywords = [
|
||||
"RAG",
|
||||
"agentic-rag",
|
||||
"lancedb",
|
||||
"vector-database",
|
||||
"hybrid-search",
|
||||
"reranking",
|
||||
"multimodal-rag",
|
||||
"embeddings",
|
||||
"citations",
|
||||
"document-ingestion",
|
||||
"mcp",
|
||||
"mcp-server",
|
||||
"pydantic-ai",
|
||||
"docling",
|
||||
]
|
||||
classifiers = [
|
||||
"Development Status :: 4 - Beta",
|
||||
"Environment :: Console",
|
||||
|
|
@ -41,6 +56,13 @@ dependencies = [
|
|||
"zstandard>=0.23.0; python_version<'3.14'",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://ggozad.github.io/haiku.rag/"
|
||||
Documentation = "https://ggozad.github.io/haiku.rag/"
|
||||
Repository = "https://github.com/ggozad/haiku.rag"
|
||||
Issues = "https://github.com/ggozad/haiku.rag/issues"
|
||||
Changelog = "https://ggozad.github.io/haiku.rag/changelog/"
|
||||
|
||||
[project.optional-dependencies]
|
||||
# Document processing
|
||||
docling = ["docling>=2.102.2,<3.0.0", "opencv-python-headless>=4.6.0.66,<5.0.0.0"]
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@
|
|||
{% block extrahead %}
|
||||
<meta property="og:type" content="website">
|
||||
<meta property="og:title" content="{% if page.title and page.title != config.site_name %}{{ page.title }} – {{ config.site_name }}{% else %}{{ config.site_name }}{% endif %}">
|
||||
<meta property="og:description" content="{{ config.site_description }}">
|
||||
<meta property="og:description" content="{% if page.meta and page.meta.description %}{{ page.meta.description }}{% else %}{{ config.site_description }}{% endif %}">
|
||||
<meta property="og:url" content="{{ page.canonical_url }}">
|
||||
<meta property="og:image" content="{{ config.site_url }}img/chat-qa.png">
|
||||
<meta name="twitter:card" content="summary_large_image">
|
||||
|
|
@ -17,7 +17,7 @@
|
|||
<div class="haiku-rag-hero__inner">
|
||||
<div class="haiku-rag-hero__text">
|
||||
<h1 class="haiku-rag-hero__title">haiku.rag</h1>
|
||||
<p class="haiku-rag-hero__tagline">Opinionated agentic RAG powered by <a href="https://lancedb.com/">LanceDB</a>, <a href="https://ai.pydantic.dev/">Pydantic AI</a>, and <a href="https://docling-project.github.io/docling/">Docling</a>. Runs locally, scales to production.</p>
|
||||
<p class="haiku-rag-hero__tagline">Ask questions about your own documents and get answers that cite their sources. Agentic RAG on <a href="https://lancedb.com/">LanceDB</a>, <a href="https://ai.pydantic.dev/">Pydantic AI</a>, and <a href="https://docling-project.github.io/docling/">Docling</a>. Runs locally, scales to production.</p>
|
||||
<div class="haiku-rag-hero__actions">
|
||||
<a class="md-button md-button--primary" href="tutorial/">Get started</a>
|
||||
<a class="md-button" href="overview/">Learn more</a>
|
||||
|
|
|
|||
|
|
@ -1,18 +1,25 @@
|
|||
[project]
|
||||
|
||||
name = "haiku.rag"
|
||||
description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling"
|
||||
version = "0.74.0"
|
||||
description = "Local-first agentic RAG with citations - hybrid search, reranking and multimodal retrieval over your own documents, no database server required"
|
||||
version = "0.75.0"
|
||||
authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
|
||||
license = { text = "MIT" }
|
||||
readme = { file = "README.md", content-type = "text/markdown" }
|
||||
requires-python = ">=3.12"
|
||||
keywords = [
|
||||
"RAG",
|
||||
"agentic-rag",
|
||||
"lancedb",
|
||||
"vector-database",
|
||||
"ml",
|
||||
"hybrid-search",
|
||||
"reranking",
|
||||
"multimodal-rag",
|
||||
"embeddings",
|
||||
"citations",
|
||||
"document-ingestion",
|
||||
"mcp",
|
||||
"mcp-server",
|
||||
"pydantic-ai",
|
||||
"docling",
|
||||
]
|
||||
|
|
@ -30,17 +37,24 @@ classifiers = [
|
|||
]
|
||||
|
||||
dependencies = [
|
||||
"haiku.rag-slim[docling,voyageai,cohere,zeroentropy,tui,cross-encoder]==0.74.0",
|
||||
"haiku.rag-slim[docling,voyageai,cohere,zeroentropy,tui,cross-encoder]==0.75.0",
|
||||
]
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://ggozad.github.io/haiku.rag/"
|
||||
Documentation = "https://ggozad.github.io/haiku.rag/"
|
||||
Repository = "https://github.com/ggozad/haiku.rag"
|
||||
Issues = "https://github.com/ggozad/haiku.rag/issues"
|
||||
Changelog = "https://ggozad.github.io/haiku.rag/changelog/"
|
||||
|
||||
[project.scripts]
|
||||
haiku-rag = "haiku.rag.cli:cli"
|
||||
|
||||
[project.optional-dependencies]
|
||||
tui = ["textual>=8.2.4"]
|
||||
s3 = ["haiku.rag-slim[s3]==0.74.0"]
|
||||
cross-encoder = ["haiku.rag-slim[cross-encoder]==0.74.0"]
|
||||
ingester = ["haiku.rag-slim[ingester]==0.74.0"]
|
||||
s3 = ["haiku.rag-slim[s3]==0.75.0"]
|
||||
cross-encoder = ["haiku.rag-slim[cross-encoder]==0.75.0"]
|
||||
ingester = ["haiku.rag-slim[ingester]==0.75.0"]
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
|
|
|
|||
20
server.json
20
server.json
|
|
@ -1,24 +1,14 @@
|
|||
{
|
||||
"$schema": "https://static.modelcontextprotocol.io/schemas/2025-12-11/server.schema.json",
|
||||
"name": "io.github.ggozad/haiku-rag",
|
||||
"title": "haiku.rag",
|
||||
"description": "Local-first agentic RAG with citations - hybrid search, reranking, multimodal document retrieval",
|
||||
"version": "{{VERSION}}",
|
||||
"description": "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling",
|
||||
"websiteUrl": "https://ggozad.github.io/haiku.rag/",
|
||||
"repository": {
|
||||
"url": "https://github.com/ggozad/haiku.rag",
|
||||
"source": "github"
|
||||
},
|
||||
"license": "MIT",
|
||||
"keywords": [
|
||||
"rag",
|
||||
"lancedb",
|
||||
"vector-database",
|
||||
"embeddings",
|
||||
"search",
|
||||
"qa",
|
||||
"research",
|
||||
"docling",
|
||||
"pydantic-ai"
|
||||
],
|
||||
"packages": [
|
||||
{
|
||||
"registryType": "pypi",
|
||||
|
|
@ -29,11 +19,11 @@
|
|||
"runtimeArguments": [
|
||||
{
|
||||
"type": "positional",
|
||||
"value": "serve"
|
||||
"value": "mcp"
|
||||
},
|
||||
{
|
||||
"type": "named",
|
||||
"name": "--mcp"
|
||||
"name": "--stdio"
|
||||
}
|
||||
],
|
||||
"transport": {
|
||||
|
|
|
|||
132
tests/capabilities/test_borrowed_client.py
Normal file
132
tests/capabilities/test_borrowed_client.py
Normal file
|
|
@ -0,0 +1,132 @@
|
|||
import pytest
|
||||
|
||||
from haiku.rag.capabilities.rag import create_capability
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_borrowed_client_is_reused_not_reopened(temp_db_path, monkeypatch):
|
||||
"""A capability handed a client must not open a second connection to the
|
||||
same database."""
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
opens = 0
|
||||
initialize = Store._initialize
|
||||
|
||||
async def counted(self):
|
||||
nonlocal opens
|
||||
opens += 1
|
||||
return await initialize(self)
|
||||
|
||||
monkeypatch.setattr(Store, "_initialize", counted)
|
||||
|
||||
capability = create_capability(
|
||||
db_path=client.store.db_path, config=client._config, rag=client
|
||||
)
|
||||
|
||||
assert await capability._ensure_rag() is client
|
||||
assert opens == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_closing_never_closes_a_borrowed_client(temp_db_path):
|
||||
"""`_close` owns only what it opened. Closing the caller's client would be a
|
||||
use-after-close for the caller."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
capability = create_capability(
|
||||
db_path=client.store.db_path, config=client._config, rag=client
|
||||
)
|
||||
await capability._ensure_rag()
|
||||
|
||||
await capability._close()
|
||||
|
||||
# Still usable by its owner.
|
||||
assert await client.list_documents() == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_borrowed_client_survives_for_run(temp_db_path):
|
||||
"""for_run clears the owned connection per run; a borrowed one is the
|
||||
caller's and carries into the run copy."""
|
||||
from tests.capabilities.test_capabilities import Deps, make_context
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
capability = create_capability(
|
||||
db_path=client.store.db_path, config=client._config, rag=client
|
||||
)
|
||||
|
||||
run_capability = await capability.for_run(make_context(Deps()))
|
||||
|
||||
assert run_capability is not capability
|
||||
assert run_capability.rag is None
|
||||
assert run_capability.borrowed_rag is client
|
||||
assert await run_capability._ensure_rag() is client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ask_hands_its_client_to_the_capability(temp_db_path, monkeypatch):
|
||||
"""`ask` built the capability from a db_path alone, so the capability opened
|
||||
its own connection to a database the client already had open."""
|
||||
from haiku.rag.capabilities import rag as rag_capability
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
real = rag_capability.create_capability
|
||||
built = {}
|
||||
|
||||
def spy(**kwargs):
|
||||
built["capability"] = real(**kwargs)
|
||||
raise RuntimeError("stop before running the agent")
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
monkeypatch.setattr(rag_capability, "create_capability", spy)
|
||||
|
||||
with pytest.raises(RuntimeError, match="stop before running the agent"):
|
||||
await client.ask("anything")
|
||||
|
||||
capability = built["capability"]
|
||||
opens = 0
|
||||
initialize = Store._initialize
|
||||
|
||||
async def counted(self):
|
||||
nonlocal opens
|
||||
opens += 1
|
||||
return await initialize(self)
|
||||
|
||||
monkeypatch.setattr(Store, "_initialize", counted)
|
||||
|
||||
assert await capability._ensure_rag() is client
|
||||
assert opens == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_analyze_hands_its_client_to_the_capability(temp_db_path, monkeypatch):
|
||||
from haiku.rag.capabilities import analysis as analysis_capability
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
real = analysis_capability.create_capability
|
||||
built = {}
|
||||
|
||||
def spy(**kwargs):
|
||||
built["capability"] = real(**kwargs)
|
||||
raise RuntimeError("stop before running the agent")
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
monkeypatch.setattr(analysis_capability, "create_capability", spy)
|
||||
|
||||
with pytest.raises(RuntimeError, match="stop before running the agent"):
|
||||
await client.analyze("anything")
|
||||
|
||||
capability = built["capability"]
|
||||
opens = 0
|
||||
initialize = Store._initialize
|
||||
|
||||
async def counted(self):
|
||||
nonlocal opens
|
||||
opens += 1
|
||||
return await initialize(self)
|
||||
|
||||
monkeypatch.setattr(Store, "_initialize", counted)
|
||||
|
||||
assert await capability._ensure_rag() is client
|
||||
assert opens == 0
|
||||
131
tests/store/test_indexes.py
Normal file
131
tests/store/test_indexes.py
Normal file
|
|
@ -0,0 +1,131 @@
|
|||
import pyarrow as pa
|
||||
import pytest
|
||||
from lancedb.index import BTree
|
||||
|
||||
from haiku.rag.store.engine import Store, ensure_indexes
|
||||
from haiku.rag.store.models import Document
|
||||
from haiku.rag.store.repositories.document import DocumentRepository
|
||||
|
||||
EXPECTED_INDEXED_COLUMNS = {
|
||||
"documents": {"id"},
|
||||
"document_meta": {"id", "uri"},
|
||||
"chunks": {"content_fts", "id", "document_id"},
|
||||
"document_items": {"document_id", "position", "self_ref", "label"},
|
||||
}
|
||||
|
||||
|
||||
async def _indexed_columns(table) -> set[str]:
|
||||
return {column for index in await table.list_indices() for column in index.columns}
|
||||
|
||||
|
||||
async def _index_type(table, column: str) -> str | None:
|
||||
for index in await table.list_indices():
|
||||
if column in index.columns:
|
||||
return index.index_type
|
||||
return None
|
||||
|
||||
|
||||
async def _covering(table, column: str) -> list[tuple[str, str]]:
|
||||
"""Every index over `column`, as (name, index_type)."""
|
||||
return [
|
||||
(index.name, index.index_type)
|
||||
for index in await table.list_indices()
|
||||
if column in index.columns
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fresh_database_indexes_every_hot_lookup_key(temp_db_path):
|
||||
"""A new database carries the full index set."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
for name, table in store._tables().items():
|
||||
expected = EXPECTED_INDEXED_COLUMNS.get(name, set())
|
||||
assert await _indexed_columns(table) == expected, name
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_indexes_skips_existing_instead_of_rebuilding(temp_db_path):
|
||||
"""A second pass must not rebuild: replace=True writes a new version."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
table = store.chunks_table
|
||||
version_before = await table.version()
|
||||
|
||||
await ensure_indexes(table, "chunks")
|
||||
|
||||
assert await table.version() == version_before
|
||||
assert await _indexed_columns(table) == EXPECTED_INDEXED_COLUMNS["chunks"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_indexes_corrects_an_index_of_the_wrong_type(temp_db_path):
|
||||
"""A wrong-typed index does not satisfy the declared one."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
table = store.document_items_table
|
||||
await table.create_index("label", config=BTree(), replace=True)
|
||||
assert await _index_type(table, "label") == "BTree"
|
||||
|
||||
await ensure_indexes(table, "document_items")
|
||||
|
||||
assert await _index_type(table, "label") == "Bitmap"
|
||||
assert (
|
||||
await _indexed_columns(table) == EXPECTED_INDEXED_COLUMNS["document_items"]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_indexes_adds_the_declared_type_beside_a_custom_index(
|
||||
temp_db_path,
|
||||
):
|
||||
"""A custom-named index neither satisfies the check nor is destroyed."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
table = store.document_items_table
|
||||
await table.drop_index("label_idx")
|
||||
await table.create_index("label", config=BTree(), name="operator_label")
|
||||
|
||||
await ensure_indexes(table, "document_items")
|
||||
|
||||
covering = dict(await _covering(table, "label"))
|
||||
assert covering["operator_label"] == "BTree"
|
||||
assert "Bitmap" in covering.values()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ensure_indexes_keeps_an_operator_index_on_a_declared_column(
|
||||
temp_db_path,
|
||||
):
|
||||
"""An index we did not declare survives, even on a declared column."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
table = store.document_items_table
|
||||
await table.create_index("label", config=BTree(), name="operator_label")
|
||||
|
||||
await ensure_indexes(table, "document_items")
|
||||
|
||||
covering = dict(await _covering(table, "label"))
|
||||
assert covering == {"label_idx": "Bitmap", "operator_label": "BTree"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_all_restores_the_full_index_set(temp_db_path):
|
||||
"""Recreated tables come back with the full index set."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
repo = DocumentRepository(store)
|
||||
await repo.create(Document(content="A document"))
|
||||
|
||||
await repo.delete_all()
|
||||
|
||||
for name, table in store._tables().items():
|
||||
expected = EXPECTED_INDEXED_COLUMNS.get(name, set())
|
||||
assert await _indexed_columns(table) == expected, name
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_delete_all_keeps_picture_data_as_large_binary(temp_db_path):
|
||||
"""picture_data must survive delete_all as large_binary, not binary."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
repo = DocumentRepository(store)
|
||||
await repo.create(Document(content="A document"))
|
||||
|
||||
await repo.delete_all()
|
||||
|
||||
schema = await store.document_items_table.schema()
|
||||
assert schema.field("picture_data").type == pa.large_binary()
|
||||
71
tests/store/test_open_path.py
Normal file
71
tests/store/test_open_path.py
Normal file
|
|
@ -0,0 +1,71 @@
|
|||
import lancedb
|
||||
import pytest
|
||||
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def counts(monkeypatch):
|
||||
"""Count the connection-level calls an open makes."""
|
||||
tally: dict[str, int] = {"list_tables": 0, "open_settings": 0, "settings_query": 0}
|
||||
|
||||
list_tables = lancedb.AsyncConnection.list_tables
|
||||
open_table = lancedb.AsyncConnection.open_table
|
||||
query = lancedb.AsyncTable.query
|
||||
|
||||
async def counted_list_tables(self, *args, **kwargs):
|
||||
tally["list_tables"] += 1
|
||||
return await list_tables(self, *args, **kwargs)
|
||||
|
||||
async def counted_open_table(self, name, *args, **kwargs):
|
||||
if name == "settings":
|
||||
tally["open_settings"] += 1
|
||||
return await open_table(self, name, *args, **kwargs)
|
||||
|
||||
def counted_query(self):
|
||||
if self.name == "settings":
|
||||
tally["settings_query"] += 1
|
||||
return query(self)
|
||||
|
||||
monkeypatch.setattr(lancedb.AsyncConnection, "list_tables", counted_list_tables)
|
||||
monkeypatch.setattr(lancedb.AsyncConnection, "open_table", counted_open_table)
|
||||
monkeypatch.setattr(lancedb.AsyncTable, "query", counted_query)
|
||||
return tally
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reopening_reads_the_table_list_and_settings_once(temp_db_path, counts):
|
||||
async with Store(temp_db_path, create=True):
|
||||
pass
|
||||
for key in counts:
|
||||
counts[key] = 0
|
||||
|
||||
async with Store(temp_db_path):
|
||||
pass
|
||||
|
||||
assert counts["list_tables"] == 1
|
||||
assert counts["open_settings"] == 1
|
||||
assert counts["settings_query"] == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_storage_failures_propagate(temp_db_path):
|
||||
"""A read failure must not read as empty settings: the migration check would
|
||||
then see version 0.0.0 and declare every migration pending."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
|
||||
def boom():
|
||||
raise RuntimeError("s3 is having a day")
|
||||
|
||||
store.settings_table.query = boom
|
||||
|
||||
with pytest.raises(RuntimeError, match="s3 is having a day"):
|
||||
await store._read_stored_settings()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_non_dict_settings_read_as_empty(temp_db_path):
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
await store.settings_table.update({"settings": "[]"}, where="id = 'settings'")
|
||||
|
||||
assert await store._read_stored_settings() == {}
|
||||
105
tests/store/test_v0_75_0_migration.py
Normal file
105
tests/store/test_v0_75_0_migration.py
Normal file
|
|
@ -0,0 +1,105 @@
|
|||
import pytest
|
||||
from lancedb.index import BTree
|
||||
|
||||
from haiku.rag.store.engine import Store
|
||||
from haiku.rag.store.models import Document
|
||||
from haiku.rag.store.repositories.document import DocumentRepository
|
||||
from haiku.rag.store.upgrades.v0_75_0 import _apply_index_hot_lookup_keys
|
||||
|
||||
# The indexes a pre-0.75.0 database lacked.
|
||||
LEGACY_DROPPED = {
|
||||
"documents": ["id_idx"],
|
||||
"chunks": ["id_idx", "document_id_idx"],
|
||||
"document_items": ["label_idx"],
|
||||
}
|
||||
|
||||
|
||||
async def _indexed(table) -> dict[str, str]:
|
||||
return {
|
||||
column: index.index_type
|
||||
for index in await table.list_indices()
|
||||
for column in index.columns
|
||||
}
|
||||
|
||||
|
||||
async def _make_legacy(store: Store) -> None:
|
||||
for table_name, indexes in LEGACY_DROPPED.items():
|
||||
table = store._tables()[table_name]
|
||||
for index in indexes:
|
||||
await table.drop_index(index)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_adds_the_missing_indexes(temp_db_path):
|
||||
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||
await _make_legacy(store)
|
||||
|
||||
await _apply_index_hot_lookup_keys(store)
|
||||
|
||||
assert await _indexed(store.documents_table) == {"id": "BTree"}
|
||||
assert await _indexed(store.chunks_table) == {
|
||||
"content_fts": "FTS",
|
||||
"id": "BTree",
|
||||
"document_id": "BTree",
|
||||
}
|
||||
assert await _indexed(store.document_items_table) == {
|
||||
"document_id": "BTree",
|
||||
"position": "BTree",
|
||||
"self_ref": "BTree",
|
||||
"label": "Bitmap",
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_keeps_every_row(temp_db_path):
|
||||
"""Indexing must not touch data."""
|
||||
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||
repo = DocumentRepository(store)
|
||||
await repo.create(Document(content="Kept", uri="test://kept"))
|
||||
await _make_legacy(store)
|
||||
|
||||
await _apply_index_hot_lookup_keys(store)
|
||||
|
||||
docs = await repo.list_all(include_content=True)
|
||||
assert [d.content for d in docs] == ["Kept"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_is_a_no_op_on_an_already_indexed_database(temp_db_path):
|
||||
"""A second run must not re-index: replace=True rebuilds."""
|
||||
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||
await _apply_index_hot_lookup_keys(store)
|
||||
versions = {name: await t.version() for name, t in store._tables().items()}
|
||||
|
||||
await _apply_index_hot_lookup_keys(store)
|
||||
|
||||
assert {name: await t.version() for name, t in store._tables().items()} == (
|
||||
versions
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_replaces_a_wrong_typed_legacy_index(temp_db_path):
|
||||
"""A BTree on `label` does not satisfy the declared Bitmap."""
|
||||
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||
await store.document_items_table.create_index(
|
||||
"label", config=BTree(), replace=True
|
||||
)
|
||||
|
||||
await _apply_index_hot_lookup_keys(store)
|
||||
|
||||
indexed = await _indexed(store.document_items_table)
|
||||
assert indexed["label"] == "Bitmap"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_leaves_undeclared_indexes_alone(temp_db_path):
|
||||
"""Indexes haiku.rag never declared are not dropped."""
|
||||
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||
await store.document_meta_table.create_index(
|
||||
"title", config=BTree(), replace=True
|
||||
)
|
||||
|
||||
await _apply_index_hot_lookup_keys(store)
|
||||
|
||||
assert "title" in await _indexed(store.document_meta_table)
|
||||
|
|
@ -993,10 +993,15 @@ async def test_client_import_documents_mixed_embeddings(temp_db_path):
|
|||
async def test_client_update_document_replaces_rows_with_bounded_versions(
|
||||
temp_db_path,
|
||||
):
|
||||
"""Updating one document should replace stale rows with bounded versions."""
|
||||
dim = Config.embeddings.model.vector_dim
|
||||
"""Updating one document should replace stale rows with bounded versions.
|
||||
|
||||
async with HaikuRAG(temp_db_path, create=True) as client:
|
||||
auto_vacuum is off: its writes would land inside the measured window.
|
||||
"""
|
||||
dim = Config.embeddings.model.vector_dim
|
||||
config = Config.model_copy(deep=True)
|
||||
config.storage.auto_vacuum = False
|
||||
|
||||
async with HaikuRAG(temp_db_path, config=config, create=True) as client:
|
||||
created = await client.import_document(
|
||||
_docling_doc("original", "Original body"),
|
||||
[Chunk(content="Original body", embedding=[0.1] * dim, order=0)],
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
from datetime import timedelta
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
|
|
@ -49,7 +51,8 @@ class TestConnectLancedb:
|
|||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config, db_path=temp_db_path)
|
||||
mock_connect.assert_called_once_with(temp_db_path.absolute())
|
||||
mock_connect.assert_awaited_once()
|
||||
assert mock_connect.call_args.args == (temp_db_path.absolute(),)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_resolves_relative_db_path(self, tmp_path, monkeypatch):
|
||||
|
|
@ -62,7 +65,8 @@ class TestConnectLancedb:
|
|||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config, db_path=relative)
|
||||
mock_connect.assert_called_once_with(relative.absolute())
|
||||
mock_connect.assert_awaited_once()
|
||||
assert mock_connect.call_args.args == (relative.absolute(),)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_passes_uri_api_key_region(self):
|
||||
|
|
@ -75,9 +79,11 @@ class TestConnectLancedb:
|
|||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
mock_connect.assert_called_once_with(
|
||||
uri="db://my-database", api_key="test-key", region="us-west-2"
|
||||
)
|
||||
mock_connect.assert_awaited_once()
|
||||
kwargs = mock_connect.call_args.kwargs
|
||||
assert kwargs["uri"] == "db://my-database"
|
||||
assert kwargs["api_key"] == "test-key"
|
||||
assert kwargs["region"] == "us-west-2"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_object_storage_passes_uri_and_storage_options(self):
|
||||
|
|
@ -94,13 +100,13 @@ class TestConnectLancedb:
|
|||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
mock_connect.assert_called_once_with(
|
||||
uri="s3://bucket/path",
|
||||
storage_options={
|
||||
"endpoint": "http://minio:9000",
|
||||
"region": "us-east-1",
|
||||
},
|
||||
)
|
||||
mock_connect.assert_awaited_once()
|
||||
kwargs = mock_connect.call_args.kwargs
|
||||
assert kwargs["uri"] == "s3://bucket/path"
|
||||
assert kwargs["storage_options"] == {
|
||||
"endpoint": "http://minio:9000",
|
||||
"region": "us-east-1",
|
||||
}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_object_storage_without_storage_options(self):
|
||||
|
|
@ -109,7 +115,10 @@ class TestConnectLancedb:
|
|||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
mock_connect.assert_called_once_with(uri="s3://bucket/path")
|
||||
mock_connect.assert_awaited_once()
|
||||
kwargs = mock_connect.call_args.kwargs
|
||||
assert kwargs["uri"] == "s3://bucket/path"
|
||||
assert "storage_options" not in kwargs
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_without_db_path_raises(self):
|
||||
|
|
@ -267,7 +276,7 @@ class TestInitFailureCleanup:
|
|||
async def fake_connect(*args, **kwargs):
|
||||
return mock_conn
|
||||
|
||||
async def failing_init_tables(self, is_new_db):
|
||||
async def failing_init_tables(self, *args):
|
||||
raise RuntimeError("simulated table init failure")
|
||||
|
||||
monkeypatch.setattr("haiku.rag.store.engine.connect_lancedb", fake_connect)
|
||||
|
|
@ -381,7 +390,7 @@ class TestStoreMiscellany:
|
|||
{"settings": "not json at all"}, where="id = 'settings'"
|
||||
)
|
||||
|
||||
assert await store._get_stored_vector_dim() is None
|
||||
assert await store._read_stored_settings() == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_vacuum_skips_when_already_running(self, temp_db_path):
|
||||
|
|
@ -398,3 +407,119 @@ class TestStoreMiscellany:
|
|||
async with Store(temp_db_path, create=True) as store:
|
||||
with pytest.raises(ValueError, match="Unknown table"):
|
||||
await store.list_table_versions("not_a_table")
|
||||
|
||||
|
||||
class TestSessionAndConsistency:
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_is_shared_across_connections(self):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
await connect_lancedb(config)
|
||||
|
||||
sessions = [c.kwargs["session"] for c in mock_connect.call_args_list]
|
||||
assert sessions[0] is sessions[1]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_sizes_select_distinct_sessions(self):
|
||||
small = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
uri="s3://bucket/path", index_cache_size_bytes=1 << 20
|
||||
)
|
||||
)
|
||||
large = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
uri="s3://bucket/path", index_cache_size_bytes=1 << 30
|
||||
)
|
||||
)
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(small)
|
||||
await connect_lancedb(large)
|
||||
|
||||
sessions = [c.kwargs["session"] for c in mock_connect.call_args_list]
|
||||
assert sessions[0] is not sessions[1]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_both_cache_sizes_are_applied(self):
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
uri="s3://bucket/path",
|
||||
index_cache_size_bytes=2 << 20,
|
||||
metadata_cache_size_bytes=4 << 20,
|
||||
)
|
||||
)
|
||||
with (
|
||||
patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
),
|
||||
patch("haiku.rag.store.engine.lancedb.Session") as mock_session,
|
||||
):
|
||||
await connect_lancedb(config)
|
||||
|
||||
mock_session.assert_called_once_with(
|
||||
index_cache_size_bytes=2 << 20, metadata_cache_size_bytes=4 << 20
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_consistency_interval_is_forwarded(self):
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
uri="s3://bucket/path", read_consistency_interval_seconds=5
|
||||
)
|
||||
)
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
|
||||
assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta(
|
||||
seconds=5
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_read_consistency_interval_omitted_when_disabled(self):
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
uri="s3://bucket/path", read_consistency_interval_seconds=None
|
||||
)
|
||||
)
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
|
||||
assert mock_connect.call_args.kwargs["read_consistency_interval"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_connection_also_gets_session_and_consistency(self, tmp_path):
|
||||
config = AppConfig()
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config, tmp_path / "db.lancedb")
|
||||
|
||||
assert mock_connect.call_args.kwargs["session"] is not None
|
||||
assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta(
|
||||
seconds=30
|
||||
)
|
||||
|
||||
|
||||
class TestLanceDBConfigValidation:
|
||||
def test_negative_values_are_rejected(self):
|
||||
"""Negatives overflow or panic inside Lance, so reject them here."""
|
||||
with pytest.raises(ValidationError):
|
||||
LanceDBConfig(read_consistency_interval_seconds=-1)
|
||||
with pytest.raises(ValidationError):
|
||||
LanceDBConfig(index_cache_size_bytes=-1)
|
||||
with pytest.raises(ValidationError):
|
||||
LanceDBConfig(metadata_cache_size_bytes=-1)
|
||||
|
||||
def test_zero_is_allowed(self):
|
||||
config = LanceDBConfig(
|
||||
read_consistency_interval_seconds=0, index_cache_size_bytes=0
|
||||
)
|
||||
assert config.read_consistency_interval_seconds == 0
|
||||
|
|
|
|||
|
|
@ -516,3 +516,147 @@ class TestMCPToolsDegradeOnError:
|
|||
assert "AI Overview" in with_cite
|
||||
|
||||
assert await ask(question="q", cite=False) == "the answer"
|
||||
|
||||
|
||||
class TestMCPClientLifetime:
|
||||
@pytest.mark.asyncio
|
||||
async def test_tool_calls_share_one_database_open(self, mcp_db, monkeypatch):
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
opens = 0
|
||||
initialize = Store._initialize
|
||||
|
||||
async def counted(self):
|
||||
nonlocal opens
|
||||
opens += 1
|
||||
return await initialize(self)
|
||||
|
||||
monkeypatch.setattr(Store, "_initialize", counted)
|
||||
|
||||
mcp = create_mcp_server(mcp_db, read_only=True)
|
||||
search = await _get_tool(mcp, "search_documents")
|
||||
list_docs = await _get_tool(mcp, "list_documents")
|
||||
await search(query="artificial intelligence")
|
||||
await list_docs()
|
||||
await search(query="machine learning")
|
||||
|
||||
assert opens == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_reads_share_one_open(self, mcp_db, monkeypatch):
|
||||
import asyncio
|
||||
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
opens = 0
|
||||
initialize = Store._initialize
|
||||
|
||||
async def counted(self):
|
||||
nonlocal opens
|
||||
opens += 1
|
||||
return await initialize(self)
|
||||
|
||||
monkeypatch.setattr(Store, "_initialize", counted)
|
||||
|
||||
mcp = create_mcp_server(mcp_db, read_only=True)
|
||||
list_docs = await _get_tool(mcp, "list_documents")
|
||||
|
||||
results = await asyncio.gather(*(list_docs() for _ in range(5)))
|
||||
|
||||
assert opens == 1
|
||||
assert all(len(r) == 2 for r in results)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_write_is_visible_to_the_next_read(self, mcp_db):
|
||||
"""One connection sees its own writes, whatever the consistency interval."""
|
||||
mcp = create_mcp_server(mcp_db, read_only=False)
|
||||
list_docs = await _get_tool(mcp, "list_documents")
|
||||
delete_doc = await _get_tool(mcp, "delete_document")
|
||||
|
||||
docs = await list_docs()
|
||||
assert await delete_doc(document_id=docs[0].id) is True
|
||||
|
||||
assert len(await list_docs()) == len(docs) - 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_lifespan_opens_and_closes_once(self, mcp_db, monkeypatch):
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
opens = 0
|
||||
initialize = Store._initialize
|
||||
|
||||
async def counted(self):
|
||||
nonlocal opens
|
||||
opens += 1
|
||||
return await initialize(self)
|
||||
|
||||
monkeypatch.setattr(Store, "_initialize", counted)
|
||||
|
||||
mcp = create_mcp_server(mcp_db, read_only=True)
|
||||
# _lifespan_manager is what every transport enters; the public
|
||||
# lifespan() combines provider lifespans only.
|
||||
async with mcp._lifespan_manager():
|
||||
assert opens == 1, "startup should open the database, not the first call"
|
||||
search = await _get_tool(mcp, "search_documents")
|
||||
await search(query="artificial intelligence")
|
||||
assert opens == 1
|
||||
|
||||
assert opens == 1
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_startup_fails_when_the_database_cannot_open(self, tmp_path):
|
||||
mcp = create_mcp_server(tmp_path / "does-not-exist.lancedb", read_only=True)
|
||||
|
||||
with pytest.raises(FileNotFoundError):
|
||||
async with mcp._lifespan_manager():
|
||||
pass
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_second_lifespan_cycle_opens_a_fresh_client(
|
||||
self, mcp_db, monkeypatch
|
||||
):
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
opens = 0
|
||||
initialize = Store._initialize
|
||||
|
||||
async def counted(self):
|
||||
nonlocal opens
|
||||
opens += 1
|
||||
return await initialize(self)
|
||||
|
||||
monkeypatch.setattr(Store, "_initialize", counted)
|
||||
|
||||
mcp = create_mcp_server(mcp_db, read_only=True)
|
||||
search = await _get_tool(mcp, "search_documents")
|
||||
|
||||
async with mcp._lifespan_manager():
|
||||
await search(query="artificial intelligence")
|
||||
assert opens == 1
|
||||
|
||||
async with mcp._lifespan_manager():
|
||||
results = await search(query="artificial intelligence")
|
||||
assert opens == 2
|
||||
assert len(results) > 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_same_dim_drift_starts_read_only_but_not_writable(self, mcp_db):
|
||||
"""Validation is unchanged: same-dimension identity drift warns in
|
||||
read-only mode and raises in writable mode. The MCP server no longer
|
||||
opts out of it for deletion."""
|
||||
from haiku.rag.config import Config
|
||||
from haiku.rag.store.repositories.settings import ConfigMismatchError
|
||||
|
||||
drifted = Config.model_copy(deep=True)
|
||||
drifted.embeddings.model.name = "a-different-model"
|
||||
|
||||
async with create_mcp_server(
|
||||
mcp_db, config=drifted, read_only=True
|
||||
)._lifespan_manager():
|
||||
pass
|
||||
|
||||
with pytest.raises(ConfigMismatchError):
|
||||
async with create_mcp_server(
|
||||
mcp_db, config=drifted, read_only=False
|
||||
)._lifespan_manager():
|
||||
pass
|
||||
|
|
|
|||
6
uv.lock
6
uv.lock
|
|
@ -1577,7 +1577,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "haiku-rag"
|
||||
version = "0.74.0"
|
||||
version = "0.75.0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "haiku-rag-slim", extra = ["cohere", "cross-encoder", "docling", "tui", "voyageai", "zeroentropy"] },
|
||||
|
|
@ -1644,7 +1644,7 @@ dev = [
|
|||
|
||||
[[package]]
|
||||
name = "haiku-rag-evals"
|
||||
version = "0.74.0"
|
||||
version = "0.75.0"
|
||||
source = { editable = "evaluations" }
|
||||
dependencies = [
|
||||
{ name = "datasets" },
|
||||
|
|
@ -1667,7 +1667,7 @@ requires-dist = [
|
|||
|
||||
[[package]]
|
||||
name = "haiku-rag-slim"
|
||||
version = "0.74.0"
|
||||
version = "0.75.0"
|
||||
source = { editable = "haiku_rag_slim" }
|
||||
dependencies = [
|
||||
{ name = "docling-core" },
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
site_name = "haiku.rag"
|
||||
site_description = "Opinionated agentic RAG powered by LanceDB, Pydantic AI, and Docling."
|
||||
site_description = "Local-first agentic RAG. Index your documents, then ask questions and get answers cited to page numbers and section headings. Hybrid search, reranking, and multimodal retrieval on embedded LanceDB."
|
||||
site_url = "https://ggozad.github.io/haiku.rag/"
|
||||
repo_url = "https://github.com/ggozad/haiku.rag"
|
||||
repo_name = "ggozad/haiku.rag"
|
||||
|
|
@ -46,8 +46,8 @@ nav = [
|
|||
{ Toolsets = "tools.md" },
|
||||
{ "Web app" = "apps.md" },
|
||||
] },
|
||||
{ Benchmarks = "benchmarks.md" },
|
||||
{ Reference = [
|
||||
{ Benchmarks = "benchmarks.md" },
|
||||
{ Development = "development.md" },
|
||||
{ Changelog = "changelog.md" },
|
||||
] },
|
||||
|
|
|
|||
Loading…
Reference in a new issue