Resolve a capability's databases once, into a scope

`resolve_db_path` manufactured the default path whenever `lancedb.databases`
was empty, and `covers_several_databases` read coverage back out of the
configuration, so a capability built without a client opened
`storage.data_dir/haiku.rag.lancedb` instead of what `lancedb.uri` placed.
The entry point resolves a `DatabaseScope` instead: instructions ask it what
it covers and `_ensure_rag` opens it through `HaikuRAG._covering`, so
coverage is decided once rather than encoded in a path and re-derived.
`Sandbox._covering` takes the scope the capability already resolved, beside
the public constructor that takes a path. The factory signatures are
unchanged.
This commit is contained in:
Yiorgis Gozadinos 2026-08-26 12:20:21 +03:00
parent 5b420a9b23
commit 71e4e4a40a
No known key found for this signature in database
11 changed files with 233 additions and 54 deletions

View file

@ -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

View file

@ -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.

View file

@ -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",
]

View file

@ -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,

View file

@ -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,

View file

@ -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]]":

View file

@ -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

View file

@ -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."""

View file

@ -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"}

View file

@ -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,

View file

@ -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)