diff --git a/CHANGELOG.md b/CHANGELOG.md index bbfe8308..bce01898 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,12 +5,14 @@ - `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. - - 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. diff --git a/docs/configuration/storage.md b/docs/configuration/storage.md index 4d86b215..57ab3aa4 100644 --- a/docs/configuration/storage.md +++ b/docs/configuration/storage.md @@ -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. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index b9f0ff60..2f0bea64 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -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: diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index e2de8868..b75fa06c 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -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(), diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index 79216a87..f8ea44e9 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -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(), diff --git a/haiku_rag_slim/haiku/rag/client/agents.py b/haiku_rag_slim/haiku/rag/client/agents.py index 9a739f74..c85fa1ab 100644 --- a/haiku_rag_slim/haiku/rag/client/agents.py +++ b/haiku_rag_slim/haiku/rag/client/agents.py @@ -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( diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index 0bfbfec7..2fa58a8e 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -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): diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 30c2ed11..a8a44b50 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -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}" diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 002283a1..08867fa6 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -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,11 @@ 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: @@ -502,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: @@ -532,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 @@ -554,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.""" @@ -711,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 @@ -732,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) @@ -780,10 +809,8 @@ class Store: ) 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 ) @@ -797,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: @@ -806,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) diff --git a/haiku_rag_slim/haiku/rag/store/repositories/settings.py b/haiku_rag_slim/haiku/rag/store/repositories/settings.py index d7b2986c..e38c17cf 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/settings.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/settings.py @@ -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: 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_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/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 f0c36474..9077586e 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -481,3 +481,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