Ask questions about your own documents and get answers that cite their sources. Agentic RAG on LanceDB, Pydantic AI, and Docling. Runs locally, scales to production.
Get started
Learn more
diff --git a/pyproject.toml b/pyproject.toml
index 194fc7d1..e6836e02 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -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"]
diff --git a/server.json b/server.json
index 5c33f50a..23d187e9 100644
--- a/server.json
+++ b/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": {
diff --git a/tests/capabilities/test_borrowed_client.py b/tests/capabilities/test_borrowed_client.py
new file mode 100644
index 00000000..af1dcec6
--- /dev/null
+++ b/tests/capabilities/test_borrowed_client.py
@@ -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
diff --git a/tests/store/test_indexes.py b/tests/store/test_indexes.py
new file mode 100644
index 00000000..3fc2c542
--- /dev/null
+++ b/tests/store/test_indexes.py
@@ -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()
diff --git a/tests/store/test_open_path.py b/tests/store/test_open_path.py
new file mode 100644
index 00000000..887db575
--- /dev/null
+++ b/tests/store/test_open_path.py
@@ -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() == {}
diff --git a/tests/store/test_v0_75_0_migration.py b/tests/store/test_v0_75_0_migration.py
new file mode 100644
index 00000000..3b996d51
--- /dev/null
+++ b/tests/store/test_v0_75_0_migration.py
@@ -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)
diff --git a/tests/test_client.py b/tests/test_client.py
index 939d1f07..e6e379a2 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -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)],
diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py
index 8e624250..ed65a5f9 100644
--- a/tests/test_lancedb_connection.py
+++ b/tests/test_lancedb_connection.py
@@ -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
diff --git a/tests/test_mcp.py b/tests/test_mcp.py
index 9c3b4553..044c5779 100644
--- a/tests/test_mcp.py
+++ b/tests/test_mcp.py
@@ -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
diff --git a/uv.lock b/uv.lock
index c7c19729..99dac8fa 100644
--- a/uv.lock
+++ b/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" },
diff --git a/zensical.toml b/zensical.toml
index fbf59e7f..033af248 100644
--- a/zensical.toml
+++ b/zensical.toml
@@ -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" },
] },