diff --git a/CHANGELOG.md b/CHANGELOG.md index 255f8a71..a06db78b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ - `doctor`'s docling-serve probe sends `X-Api-Key`, so an instance requiring a key is reported reachable rather than unreachable. - The picture-description request to the public OpenAI endpoint sends `OPENAI_API_KEY`; it carried no authorization header. - `haiku-rag` prints the message and exits when the configured embedder does not match the database, instead of raising a traceback. +- A capability built without a client opens the databases the configuration places — `lancedb.uri` or the whole `lancedb.databases` set — instead of the default under `storage.data_dir`. - A `lancedb.uri` with no scheme is a local path, as it already is in `lancedb.databases`: `haiku-rag init` creates it and every command that opens an existing database requires it to exist, where a missing path was opened as object storage and became an empty database. `--db PATH` overrides `lancedb.uri`. ## [0.77.0] - 2026-08-21 diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index 4988e84d..7d30915b 100644 --- a/docs/capabilities/index.md +++ b/docs/capabilities/index.md @@ -143,10 +143,18 @@ Capabilities use a plain `state: dict[str, Any]` attribute on agent dependencies Applications serving AG-UI should adapt the agent with Pydantic AI's `AGUIAdapter`. Native model and tool events require no haiku.rag-specific bridge. -## Database path +## Which databases a capability covers -Both factories resolve their database in this order: +Both factories resolve this once, in order: -1. The `db_path` argument. -2. `HAIKU_RAG_DB`. -3. `config.storage.data_dir / "haiku.rag.lancedb"`. +1. The `db_path` argument, which covers that one database. +2. `HAIKU_RAG_DB`, the same way. +3. [`lancedb.databases`](../configuration/storage.md#several-databases), covering the + whole configured set. A capability covering several says so in its instructions, so + the model can attribute evidence to one while it answers. +4. [`lancedb.uri`](../configuration/storage.md#changing-the-default-database-path), + covering the one database it places. +5. `config.storage.data_dir / "haiku.rag.lancedb"`. + +Passing a live client through `rag=` overrides all of it: the capability reads what that +client covers, and never closes it. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 76ba1bf8..ce333881 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -34,6 +34,7 @@ from haiku.rag.capabilities._tools import ( ) from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord, EvidenceRef from haiku.rag.client import HaikuRAG, all_found +from haiku.rag.client.scope import DatabaseScope from haiku.rag.config.models import AppConfig from haiku.rag.store.exceptions import AmbiguousCitationError from haiku.rag.store.models.chunk import SearchResult @@ -87,20 +88,16 @@ def _nearest_known_id(chunk_id: str, known_ids: list[str]) -> str: return match[0] if match else chunk_id -def resolve_db_path(db_path: Path | str | None, config: AppConfig) -> Path | None: - """The database a capability opens for itself, or None to let the client decide. +def resolve_scope(db_path: Path | str | None, config: AppConfig) -> DatabaseScope: + """The databases a capability covers, resolved once at its entry point. - None where `lancedb.databases` names the databases: a path would name one of - them instead, and a capability nobody handed a client would search a single - database where the configuration says several. + ``HAIKU_RAG_DB`` is read here and nowhere else. ``DatabaseScope`` is + environment-agnostic on purpose, so honouring the variable inside it would + silently extend it to every other caller. """ - if db_path is not None: - return Path(db_path) - if env_db := os.environ.get("HAIKU_RAG_DB"): - return Path(env_db).expanduser() - if config.lancedb.databases: - return None - return config.storage.data_dir / "haiku.rag.lancedb" + if db_path is None and (env_db := os.environ.get("HAIKU_RAG_DB")): + db_path = Path(env_db).expanduser() + return DatabaseScope.resolve(config, database_path=db_path) class EvidenceState(BaseModel): @@ -134,18 +131,15 @@ class EvidenceState(BaseModel): self.searches.clear() -def covers_several_databases( - db_path: Path | None, config: AppConfig, rag: "HaikuRAG | None" -) -> bool: +def covers_several_databases(scope: DatabaseScope, rag: "HaikuRAG | None") -> bool: """Whether the capability will read from more than one database. - What the configuration names is not what a capability opens: an explicit - `db_path` opens that one database, and a lent client already knows what it - covers. Instructions follow coverage, not configuration. + A lent client already covers what it covers; otherwise the scope says. + Instructions follow coverage, not configuration. """ if rag is not None: return rag.covers_multiple - return db_path is None and len(config.lancedb.databases) > 1 + return scope.covers_multiple def _awaits_the_model(messages: list[ModelMessage]) -> bool: @@ -178,7 +172,7 @@ def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) - @dataclass class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): - db_path: Path | None + scope: DatabaseScope config: AppConfig state_type: type[StateT] state_namespace: str @@ -415,7 +409,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]): if self.rag is None: async with self.resource_lock: if self.rag is None: - rag = HaikuRAG(self.db_path, config=self.config, read_only=True) + rag = HaikuRAG._covering(self.scope, self.config, read_only=True) await rag.__aenter__() self.rag = rag return self.rag @@ -644,5 +638,5 @@ __all__ = [ "CodeExecutionEntry", "RAGCapabilityBase", "covers_several_databases", - "resolve_db_path", + "resolve_scope", ] diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index fdc3e9a1..ec5f44e2 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -16,7 +16,7 @@ from haiku.rag.capabilities._base import ( EvidenceState, RAGCapabilityBase, covers_several_databases, - resolve_db_path, + resolve_scope, ) from haiku.rag.capabilities._tools import merge_results from haiku.rag.config.models import AppConfig @@ -84,8 +84,8 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): if self.sandbox is None: rag = await self._ensure_rag() assert self.state is not None - self.sandbox = Sandbox( - db_path=self.db_path, + self.sandbox = Sandbox._covering( + scope=self.scope, config=self.config, context=AnalysisContext( filter=self.state.document_filter, @@ -213,12 +213,12 @@ def create_capability( config = get_config() analysis_model = config.analysis.model or config.qa.model - resolved_db_path = resolve_db_path(db_path, config) + scope = resolve_scope(db_path, config) instruction_text = instructions() - if covers_several_databases(resolved_db_path, config, rag): + if covers_several_databases(scope, rag): instruction_text += several_databases_instructions() return AnalysisCapability( - db_path=resolved_db_path, + scope=scope, config=config, borrowed_rag=rag, state_type=AnalysisState, diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index db3e4515..f452cec2 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -14,7 +14,7 @@ from haiku.rag.capabilities._base import ( EvidenceState, RAGCapabilityBase, covers_several_databases, - resolve_db_path, + resolve_scope, ) from haiku.rag.config.models import AppConfig @@ -116,12 +116,12 @@ def create_capability( from haiku.rag.config import get_config config = get_config() - resolved_db_path = resolve_db_path(db_path, config) + scope = resolve_scope(db_path, config) instruction_text = instructions() - if covers_several_databases(resolved_db_path, config, rag): + if covers_several_databases(scope, rag): instruction_text += several_databases_instructions() return RAGCapability( - db_path=resolved_db_path, + scope=scope, config=config, borrowed_rag=rag, state_type=RAGState, diff --git a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py index c424e891..e84f0806 100644 --- a/haiku_rag_slim/haiku/rag/sandbox/sandbox.py +++ b/haiku_rag_slim/haiku/rag/sandbox/sandbox.py @@ -5,7 +5,6 @@ from collections.abc import AsyncIterator, Callable, Coroutine from contextlib import asynccontextmanager, suppress from dataclasses import dataclass from itertools import zip_longest -from pathlib import Path from typing import TYPE_CHECKING, Any, Literal import pydantic_monty @@ -23,9 +22,10 @@ from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.document_item import PICTURE_REF_PREFIX, DocumentItem if TYPE_CHECKING: - from pathlib import PurePosixPath + from pathlib import Path, PurePosixPath from haiku.rag.client import HaikuRAG + from haiku.rag.client.scope import DatabaseScope @dataclass @@ -138,7 +138,7 @@ class Sandbox: each read opens an ephemeral read-only connection. """ - _db_path: Path | None + _scope: "DatabaseScope" _config: AppConfig _context: AnalysisContext _rag: "HaikuRAG | None" @@ -157,13 +157,52 @@ class Sandbox: def __init__( self, - db_path: Path | None, + db_path: "Path | str | None", config: AppConfig, context: AnalysisContext, rag: "HaikuRAG | None" = None, lock: "asyncio.Lock | None" = None, ): - self._db_path = db_path + from haiku.rag.client.scope import DatabaseScope + + self._configure( + DatabaseScope.resolve(config, database_path=db_path), + config, + context, + rag, + lock, + ) + + @classmethod + def _covering( + cls, + scope: "DatabaseScope", + config: AppConfig, + context: AnalysisContext, + rag: "HaikuRAG | None" = None, + lock: "asyncio.Lock | None" = None, + ) -> "Sandbox": + """A sandbox over databases someone already resolved. + + Internal: the public constructor takes a path and resolves it, which is + its own job. This is for callers that did the resolving, as + ``HaikuRAG._covering`` is. It sets the sandbox up directly rather than + through ``__init__``, so the scope it is handed is the only one resolved. + """ + sandbox = cls.__new__(cls) + sandbox._configure(scope, config, context, rag, lock) + return sandbox + + def _configure( + self, + scope: "DatabaseScope", + config: AppConfig, + context: AnalysisContext, + rag: "HaikuRAG | None", + lock: "asyncio.Lock | None", + ) -> None: + """The state every sandbox starts with, however its scope was reached.""" + self._scope = scope self._config = config self._context = context self._rag = rag @@ -202,7 +241,7 @@ class Sandbox: return from haiku.rag.client import HaikuRAG - async with HaikuRAG(self._db_path, config=self._config, read_only=True) as rag: + async with HaikuRAG._covering(self._scope, self._config, read_only=True) as rag: yield rag async def _documents(self) -> "tuple[list[Any], dict[str, HaikuRAG]]": diff --git a/tests/capabilities/test_agent_spec.py b/tests/capabilities/test_agent_spec.py index c8c46423..ce7eee57 100644 --- a/tests/capabilities/test_agent_spec.py +++ b/tests/capabilities/test_agent_spec.py @@ -14,6 +14,7 @@ from haiku.rag.capabilities.compaction import EvidenceCompactionCapability from haiku.rag.capabilities.policy import CAPABILITY_ID as POLICY_ID from haiku.rag.capabilities.policy import CitationPolicyCapability from haiku.rag.capabilities.rag import RAGCapability, RAGState +from haiku.rag.client.scope import DatabaseRef ALL_CAPABILITIES = [ RAGCapability, @@ -46,7 +47,7 @@ def test_rag_capability_is_built_from_a_spec(temp_db_path): ) assert isinstance(capability, RAGCapability) - assert capability.db_path == temp_db_path + assert capability.scope.databases == (DatabaseRef.at(temp_db_path),) assert capability.id == "haiku-rag" assert capability.state_type is RAGState assert capability.tool_names == {"rag_search", "rag_cite"} @@ -60,7 +61,7 @@ def test_analysis_capability_is_built_from_a_spec(temp_db_path): ) assert isinstance(capability, AnalysisCapability) - assert capability.db_path == temp_db_path + assert capability.scope.databases == (DatabaseRef.at(temp_db_path),) assert capability.id == "haiku-rag-analysis" assert capability.state_type is AnalysisState assert capability.request_limit == 30 diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index e11c51f8..f62810a3 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -39,6 +39,7 @@ from haiku.rag.capabilities.ledger import ( ) from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGCapability, RAGState from haiku.rag.capabilities.rag import create_capability as create_rag +from haiku.rag.client.scope import DatabaseRef from haiku.rag.config.models import AppConfig, PromptsConfig from haiku.rag.sandbox import Sandbox, SandboxResult from haiku.rag.store.models.chunk import Chunk, SearchResult @@ -91,20 +92,26 @@ def test_analysis_capability_api(temp_db_path): assert capability.request_limit == 30 +def _placed(capability) -> "Path | None": + """Where a capability covering one local database will open it.""" + [ref] = capability.scope.databases + return ref.db_path + + def test_capability_factories_resolve_environment_and_defaults( temp_db_path, monkeypatch ): config = AppConfig() monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path)) - assert create_rag(config=config).db_path == temp_db_path + assert _placed(create_rag(config=config)) == temp_db_path monkeypatch.delenv("HAIKU_RAG_DB") - assert create_rag(config=config).db_path == ( + assert _placed(create_rag(config=config)) == ( config.storage.data_dir / "haiku.rag.lancedb" ) for factory in (create_rag, create_analysis): - db_path = factory(db_path=str(temp_db_path), config=config).db_path + db_path = _placed(factory(db_path=str(temp_db_path), config=config)) assert db_path == temp_db_path assert isinstance(db_path, Path) @@ -113,6 +120,63 @@ def test_capability_factories_resolve_environment_and_defaults( assert create_analysis().config is config +class TestACapabilityFollowsTheConfiguredLocation: + """A capability nobody handed a client opens one for itself, and has to open + the database the configuration places rather than the default directory.""" + + def _config(self, tmp_path, uri: str) -> AppConfig: + from haiku.rag.config.models import LanceDBConfig, StorageConfig + + return AppConfig( + lancedb=LanceDBConfig(uri=uri), + storage=StorageConfig(data_dir=tmp_path / "elsewhere"), + ) + + def test_a_configured_uri_is_left_to_the_client(self, tmp_path): + """A path overrides a configured location, so manufacturing one would + send the capability to the default directory instead of the bucket.""" + located = tmp_path / "notes.lancedb" + for factory in (create_rag, create_analysis): + [local] = factory( + config=self._config(tmp_path, str(located)) + ).scope.databases + assert local == DatabaseRef.configured(None, str(located)) + + remote = self._config(tmp_path, "s3://bucket/one.lancedb") + [ref] = factory(config=remote).scope.databases + assert ref == DatabaseRef(None, "s3://bucket/one.lancedb", None) + + @pytest.mark.asyncio + async def test_it_opens_the_database_the_uri_places(self, tmp_path): + from haiku.rag.client import HaikuRAG + + located = tmp_path / "notes.lancedb" + config = self._config(tmp_path, str(located)) + async with HaikuRAG(config=config, create=True): + pass + + capability = create_rag(config=config) + rag = await capability._ensure_rag() + try: + assert rag.store.db_path == located + finally: + await capability._close() + + def test_an_explicit_path_still_overrides_the_configured_uri(self, tmp_path): + config = self._config(tmp_path, str(tmp_path / "notes.lancedb")) + chosen = tmp_path / "chosen.lancedb" + + assert _placed(create_rag(db_path=chosen, config=config)) == chosen + + def test_the_environment_still_overrides_the_configured_uri( + self, tmp_path, monkeypatch + ): + config = self._config(tmp_path, "s3://bucket/one.lancedb") + monkeypatch.setenv("HAIKU_RAG_DB", str(tmp_path / "from-env.lancedb")) + + assert _placed(create_rag(config=config)) == tmp_path / "from-env.lancedb" + + @pytest.mark.asyncio async def test_a_string_db_path_opens_a_store(temp_db_path): """Store calls `absolute()` and `exists()` on db_path, which a str lacks.""" diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 1991ccbb..cf6b730f 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -90,9 +90,12 @@ def test_chat_capabilities_read_the_named_database(tmp_path, monkeypatch): run_chat(scope=DatabaseScope.resolve(config)) [covering] = chat_app.call_args.kwargs["capabilities"] - assert named.db_path == tmp_path / "b.lancedb" + # The chat lends its own client, so this scope is the fallback: what matters + # is that it places the named database rather than the whole set. + [placed] = named.scope.databases + assert placed.db_path == tmp_path / "b.lancedb" assert named.config.lancedb.databases == {} - assert covering.db_path is None + assert covering.scope.names == ("a", "b") assert set(covering.config.lancedb.databases) == {"a", "b"} diff --git a/tests/sandbox/test_sandbox_multi_db.py b/tests/sandbox/test_sandbox_multi_db.py index d842ceaa..92c7c8cf 100644 --- a/tests/sandbox/test_sandbox_multi_db.py +++ b/tests/sandbox/test_sandbox_multi_db.py @@ -3,6 +3,7 @@ import shutil import pytest from haiku.rag.client import HaikuRAG +from haiku.rag.client.scope import DatabaseRef, DatabaseScope from haiku.rag.sandbox import AnalysisContext, Sandbox from tests.test_multi_db import _config, _seed @@ -10,7 +11,7 @@ from tests.test_multi_db import _config, _seed async def _mounted(rag, sources=None): """The sandbox's view of the corpus, and the sandbox itself.""" sandbox = Sandbox( - db_path=rag._db_path, + db_path=None, config=rag._config, context=AnalysisContext(sources=sources), rag=rag, @@ -80,6 +81,74 @@ class TestDocumentsAcrossDatabases: assert owners[doc.id].source in content +class TestTheSandboxConstructors: + """`Sandbox` is public and takes a path; `_covering` is for callers that + already resolved a scope, as `HaikuRAG._covering` is.""" + + @pytest.mark.asyncio + async def test_the_public_constructor_resolves_the_path_it_is_given(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + sandbox = Sandbox( + db_path=tmp_path / "alpha.lancedb", + config=config, + context=AnalysisContext(), + ) + + assert sandbox._scope.databases == (DatabaseRef.at(tmp_path / "alpha.lancedb"),) + + @pytest.mark.asyncio + async def test_no_path_covers_what_the_configuration_places(self, tmp_path): + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + sandbox = Sandbox(db_path=None, config=config, context=AnalysisContext()) + + assert sandbox._scope.names == ("alpha", "beta") + + @pytest.mark.asyncio + async def test_covering_resolves_nothing_of_its_own(self, tmp_path, monkeypatch): + """Handed a scope, it must not reach resolution again: resolving twice + is what let a capability's databases and its sandbox's disagree.""" + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + scope = DatabaseScope.resolve(config, database_name="alpha") + + def _refuse(*args, **kwargs): + raise AssertionError("resolved a scope it was already given") + + monkeypatch.setattr(DatabaseScope, "resolve", _refuse) + sandbox = Sandbox._covering(scope, config, AnalysisContext()) + + assert sandbox._scope is scope + assert sandbox._config is config + + +class TestTheSandboxCoversWhatTheCapabilityCovers: + @pytest.mark.asyncio + async def test_the_capability_hands_over_the_scope_it_resolved(self, tmp_path): + """The capability resolved its databases once. Letting the sandbox + resolve them again from the same configuration reaches a different + answer wherever a path or the environment named one of a set.""" + from haiku.rag.capabilities.analysis import AnalysisState, create_capability + + config = _config(tmp_path, ["alpha", "beta"]) + await _seed(config, "alpha", ["alpha document about cats"]) + + capability = create_capability( + db_path=tmp_path / "alpha.lancedb", config=config, defer_loading=False + ) + capability.state = AnalysisState() + + sandbox = await capability._ensure_sandbox() + try: + assert sandbox._scope is capability.scope + assert capability.scope.names == () + finally: + await capability._close() + + class TestExecutingAcrossDatabases: @pytest.mark.asyncio async def test_code_reads_documents_from_every_database(self, tmp_path): @@ -91,7 +160,7 @@ class TestExecutingAcrossDatabases: async with HaikuRAG(config=config) as rag: sandbox = Sandbox( - db_path=rag._db_path, + db_path=None, config=rag._config, context=AnalysisContext(), rag=rag, @@ -121,7 +190,7 @@ class TestExecutingAcrossDatabases: beta = (await rag.clients_for(["beta"]))[0] [outside] = await beta.document_repository.list_all(limit=1) sandbox = Sandbox( - db_path=rag._db_path, + db_path=None, config=rag._config, context=AnalysisContext(sources=["alpha"]), rag=rag, diff --git a/tests/test_multi_db_ask.py b/tests/test_multi_db_ask.py index fec7aee3..c0aa321d 100644 --- a/tests/test_multi_db_ask.py +++ b/tests/test_multi_db_ask.py @@ -512,7 +512,7 @@ class TestStandaloneCapabilities: await _seed(config, "beta", ["beta document about cats"]) capability = create_capability(config=config, defer_loading=False) - assert capability.db_path is None + assert capability.scope.names == ("alpha", "beta") run = await capability.for_run(make_context(Deps())) try: formatted = await run._search("cats", limit=10)