Hand storage the database location, not the configuration that placed it
Store, connect_lancedb, gather_database_info and run_doctor take a location, a path or a URI, and classify it with ConnectionMode.of. SingleDatabaseSession owns the resolved DatabaseRef and passes its location down. This removes DatabaseRef.connection(), default_db_path, the placeholder path for URI-backed databases and the per-database config copies, so the configuration a client holds is the one the caller gave it. The chat hands its capabilities the scope it opened along with the client it lends, and the v0.58.0 migration no longer checks local free disk for a database behind a URI.
This commit is contained in:
Yiorgis Gozadinos
parent
4f7f69aaf0
commit
9baa213b34
23 changed files with 424 additions and 386 deletions
|
|
@ -10,6 +10,11 @@
|
|||
- Searches in one model response deduplicate their results: evidence a sibling
|
||||
search already showed collapses to a reference line, and a picture attaches
|
||||
once per response.
|
||||
- `Store(location, config)`, `connect_lancedb(location, config)`,
|
||||
`gather_database_info(location, config)` and `run_doctor(config, location, ...)`
|
||||
take the database location, a path or a URI. `ConnectionMode.of(location)`
|
||||
replaces `ConnectionMode.from_config`. `DatabaseRef.connection()` and
|
||||
`default_db_path` removed.
|
||||
|
||||
## [0.81.0] - 2026-09-01
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import logging
|
||||
from functools import cached_property
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
|
|
@ -64,21 +63,10 @@ class HaikuRAGApp:
|
|||
[ref] = self.scope.databases
|
||||
return ref
|
||||
|
||||
@cached_property
|
||||
def _connection(self) -> "tuple[AppConfig, Path]":
|
||||
"""How to open the one database this command works on, directly.
|
||||
|
||||
Derived per database from its configured location.
|
||||
"""
|
||||
from haiku.rag.client.session import default_db_path
|
||||
|
||||
config, db_path = self._one.connection(self.config)
|
||||
return config, db_path or default_db_path(config)
|
||||
|
||||
@property
|
||||
def _store_config(self) -> AppConfig:
|
||||
"""The configuration for opening the one database directly."""
|
||||
return self._connection[0]
|
||||
def _location(self) -> "Path | str":
|
||||
"""Where the one database this command works on is."""
|
||||
return self._one.location
|
||||
|
||||
@property
|
||||
def _is_local(self) -> bool:
|
||||
|
|
@ -91,12 +79,9 @@ class HaikuRAGApp:
|
|||
|
||||
@property
|
||||
def _path(self) -> Path:
|
||||
"""The path of the one database this command works on.
|
||||
|
||||
A database behind a URI has none of its own, and the default stands in:
|
||||
the URI in `_store_config` is what decides where it connects.
|
||||
"""
|
||||
return self._connection[1]
|
||||
"""The path of the one local database this command works on."""
|
||||
assert self._one.db_path is not None
|
||||
return self._one.db_path
|
||||
|
||||
@property
|
||||
def display_path(self) -> "Path | str":
|
||||
|
|
@ -140,7 +125,7 @@ class HaikuRAGApp:
|
|||
self.console.print("[red]Database path does not exist.[/red]")
|
||||
return
|
||||
|
||||
info = await gather_database_info(self._store_config, self._path)
|
||||
info = await gather_database_info(self._location, self.config)
|
||||
|
||||
if not info.exists:
|
||||
self.console.print(
|
||||
|
|
@ -282,8 +267,8 @@ class HaikuRAGApp:
|
|||
cm = status if status is not None else nullcontext()
|
||||
with cm:
|
||||
report = await run_doctor(
|
||||
self._store_config,
|
||||
self._path,
|
||||
self.config,
|
||||
self._location,
|
||||
dict(os.environ),
|
||||
duplicates_out=duplicates_out,
|
||||
on_progress=on_progress,
|
||||
|
|
@ -340,8 +325,8 @@ class HaikuRAGApp:
|
|||
return
|
||||
|
||||
async with Store(
|
||||
self._path,
|
||||
config=self._store_config,
|
||||
self._location,
|
||||
config=self.config,
|
||||
skip_validation=True,
|
||||
read_only=True,
|
||||
skip_migration_check=True,
|
||||
|
|
@ -415,15 +400,15 @@ class HaikuRAGApp:
|
|||
"""
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
return Store(self._path, config=self._store_config, read_only=self.read_only)
|
||||
return Store(self._location, config=self.config, read_only=self.read_only)
|
||||
|
||||
def _tag_read_store(self) -> "Store":
|
||||
"""Read-only store for tag inspection; works on old or drifted DBs."""
|
||||
from haiku.rag.store.engine import Store
|
||||
|
||||
return Store(
|
||||
self._path,
|
||||
config=self._store_config,
|
||||
self._location,
|
||||
config=self.config,
|
||||
skip_validation=True,
|
||||
skip_migration_check=True,
|
||||
read_only=True,
|
||||
|
|
@ -760,8 +745,8 @@ class HaikuRAGApp:
|
|||
from haiku.rag.store.engine import Store
|
||||
|
||||
async with Store(
|
||||
self._path,
|
||||
config=self._store_config,
|
||||
self._location,
|
||||
config=self.config,
|
||||
skip_validation=True,
|
||||
skip_migration_check=True,
|
||||
read_only=self.read_only,
|
||||
|
|
|
|||
|
|
@ -42,14 +42,8 @@ def run_chat(
|
|||
config.qa.model = model_config
|
||||
config.analysis.model = model_config
|
||||
|
||||
# The capabilities read the databases the scope covers, not what the
|
||||
# configuration names: a `--db PATH` selection is outside the
|
||||
# configuration, and a `--db-name NAME` selection is narrower than it.
|
||||
if scope.covers_multiple:
|
||||
capability_config, capability_db_path = config, None
|
||||
else:
|
||||
capability_config, capability_db_path = scope.databases[0].connection(config)
|
||||
|
||||
# The app opens the scope and lends that client to the capabilities, which
|
||||
# read what `--db PATH` or `--db-name NAME` selected.
|
||||
enabled = capabilities or ["rag"]
|
||||
capability_list = []
|
||||
defer_loading = len(enabled) > 1
|
||||
|
|
@ -68,8 +62,7 @@ def run_chat(
|
|||
|
||||
capability_list.append(
|
||||
create_capability(
|
||||
db_path=capability_db_path,
|
||||
config=capability_config,
|
||||
config=config,
|
||||
defer_loading=defer_loading,
|
||||
vision=driving_model.vision,
|
||||
)
|
||||
|
|
@ -80,8 +73,7 @@ def run_chat(
|
|||
|
||||
capability_list.append(
|
||||
create_capability(
|
||||
db_path=capability_db_path,
|
||||
config=capability_config,
|
||||
config=config,
|
||||
defer_loading=defer_loading,
|
||||
vision=driving_model.vision,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -148,10 +148,12 @@ class ChatApp(App):
|
|||
# a client whose __aenter__ failed.
|
||||
await client.__aenter__()
|
||||
self.client = client
|
||||
# Lent to the capabilities: already the databases they were built for,
|
||||
# and one connection per database however many capabilities read it.
|
||||
# Lent to the capabilities, with the scope it covers: one connection
|
||||
# per database however many capabilities read it, and the analysis
|
||||
# sandbox is built over the same selection.
|
||||
for capability in self._capabilities:
|
||||
capability.borrowed_rag = client
|
||||
capability.scope = self.scope
|
||||
|
||||
self._agent = Agent(
|
||||
self._model,
|
||||
|
|
|
|||
|
|
@ -21,7 +21,6 @@ from haiku.rag.client.session import (
|
|||
FederatedSession,
|
||||
SingleDatabaseSession,
|
||||
aclose_quietly,
|
||||
default_db_path,
|
||||
)
|
||||
from haiku.rag.config import AppConfig, get_config
|
||||
from haiku.rag.converters import get_converter
|
||||
|
|
@ -145,9 +144,6 @@ class HaikuRAG:
|
|||
of nothing to search.
|
||||
"""
|
||||
self._configured = config if config is not None else get_config()
|
||||
# What the caller configured, kept intact: entering derives a
|
||||
# single-database configuration from it, and every re-entry derives
|
||||
# from the configured set.
|
||||
self._config = self._configured
|
||||
self._requested_db_path = Path(db_path) if db_path is not None else None
|
||||
if self._requested_db_path is not None and sources is not None:
|
||||
|
|
@ -339,15 +335,12 @@ class HaikuRAG:
|
|||
return self
|
||||
|
||||
[ref] = scope.databases
|
||||
self._config, db_path = ref.connection(self._configured)
|
||||
|
||||
self._session = await SingleDatabaseSession(
|
||||
db_path if db_path is not None else default_db_path(self._config),
|
||||
ref,
|
||||
self._config,
|
||||
skip_validation=self._skip_validation,
|
||||
create=self._create,
|
||||
read_only=self._read_only,
|
||||
source=ref.name,
|
||||
).open()
|
||||
return self
|
||||
|
||||
|
|
|
|||
|
|
@ -43,15 +43,10 @@ class DatabaseRef:
|
|||
uri, db_path = locate_database(location)
|
||||
return cls(name=name, uri=uri, db_path=db_path)
|
||||
|
||||
def connection(self, config: AppConfig) -> tuple[AppConfig, Path | None]:
|
||||
"""The configuration and path to open this one database with.
|
||||
|
||||
A copy: the caller's configuration still names whatever set it named.
|
||||
"""
|
||||
one = config.model_copy(deep=True)
|
||||
one.lancedb.databases = {}
|
||||
one.lancedb.uri = self.uri
|
||||
return one, self.db_path
|
||||
@property
|
||||
def location(self) -> Path | str:
|
||||
"""Where the database is: its path, or its URI."""
|
||||
return self.db_path if self.db_path is not None else self.uri
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
|
|
|||
|
|
@ -43,36 +43,30 @@ async def aclose_quietly(closeable: Any, what: str) -> None:
|
|||
logger.debug("Closing the %s failed on teardown", what, exc_info=True)
|
||||
|
||||
|
||||
def default_db_path(config: AppConfig) -> Path:
|
||||
"""Where a database lives when its location names no path."""
|
||||
return config.storage.data_dir / "haiku.rag.lancedb"
|
||||
|
||||
|
||||
class SingleDatabaseSession:
|
||||
"""One database: its store, its repositories, and their lifecycle.
|
||||
|
||||
Everything that needs a store lives here, so nothing above has to ask whether
|
||||
it has one. ``source`` is the configured name this database answers to, or
|
||||
None where nothing names it.
|
||||
it has one. Built from the resolved reference: ``source`` is the configured
|
||||
name it answers to, or None where nothing names it, and the store receives
|
||||
its location.
|
||||
|
||||
``db_path``, ``config``, ``read_only`` and ``source`` are readable: a client
|
||||
``ref``, ``config``, ``read_only`` and ``source`` are readable: a client
|
||||
borrowing this session reports them as its own.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
db_path: Path | str,
|
||||
ref: DatabaseRef,
|
||||
config: AppConfig,
|
||||
*,
|
||||
skip_validation: bool = False,
|
||||
create: bool = False,
|
||||
read_only: bool = False,
|
||||
source: str | None = None,
|
||||
) -> None:
|
||||
self.db_path = db_path
|
||||
self.ref = ref
|
||||
self.config = config
|
||||
self.read_only = read_only
|
||||
self.source = source
|
||||
self._skip_validation = skip_validation
|
||||
self._create = create
|
||||
self._vacuum_tasks: set[asyncio.Task] = set()
|
||||
|
|
@ -80,19 +74,25 @@ class SingleDatabaseSession:
|
|||
self._vacuum_dirty = False
|
||||
|
||||
@property
|
||||
def location(self) -> Path | str:
|
||||
"""Configured URI or local path for this database.
|
||||
def source(self) -> str | None:
|
||||
return self.ref.name
|
||||
|
||||
Not `db_path`, which is a placeholder where a URI holds the database.
|
||||
"""
|
||||
return self.config.lancedb.uri or self.db_path
|
||||
@property
|
||||
def location(self) -> Path | str:
|
||||
"""Where this database is: its path, or its URI."""
|
||||
return self.ref.location
|
||||
|
||||
@property
|
||||
def db_path(self) -> Path | None:
|
||||
"""The local path, or None for a database behind a URI."""
|
||||
return self.ref.db_path
|
||||
|
||||
async def open(self) -> "SingleDatabaseSession":
|
||||
"""Connect, validate, and build the repositories."""
|
||||
failure: str | None = None
|
||||
try:
|
||||
self.store = Store(
|
||||
self.db_path,
|
||||
self.location,
|
||||
config=self.config,
|
||||
skip_validation=self._skip_validation,
|
||||
create=self._create,
|
||||
|
|
@ -309,14 +309,11 @@ class FederatedSession:
|
|||
|
||||
Registered here because a cancelled `gather` discards its results.
|
||||
"""
|
||||
ref = self._refs[name]
|
||||
one, db_path = ref.connection(self._config)
|
||||
self._sessions[name] = await SingleDatabaseSession(
|
||||
db_path if db_path is not None else default_db_path(one),
|
||||
one,
|
||||
self._refs[name],
|
||||
self._config,
|
||||
skip_validation=self._skip_validation,
|
||||
read_only=self._read_only,
|
||||
source=ref.name,
|
||||
).open()
|
||||
|
||||
async def aclose(self) -> None:
|
||||
|
|
|
|||
|
|
@ -1080,7 +1080,7 @@ async def run_provider_checks(
|
|||
|
||||
async def run_doctor(
|
||||
config: AppConfig,
|
||||
db_path: Path,
|
||||
location: Path | str,
|
||||
environ: dict[str, str],
|
||||
duplicates_out: Path | None = None,
|
||||
on_progress: Callable[[str], None] | None = None,
|
||||
|
|
@ -1092,7 +1092,7 @@ async def run_doctor(
|
|||
"""
|
||||
notify = on_progress or (lambda _label: None)
|
||||
notify("Inspecting tables")
|
||||
db = await connect_lancedb(config, db_path)
|
||||
db = await connect_lancedb(location, config)
|
||||
stats = await get_database_stats(db)
|
||||
|
||||
results: list[CheckResult] = []
|
||||
|
|
@ -1110,7 +1110,7 @@ async def run_doctor(
|
|||
missing = [name for name in REQUIRED_TABLES if not stats[name]["exists"]]
|
||||
if not missing:
|
||||
async with Store(
|
||||
db_path,
|
||||
location,
|
||||
config=config,
|
||||
skip_validation=True,
|
||||
read_only=True,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,5 @@
|
|||
from fastapi import APIRouter, Depends, HTTPException, status
|
||||
|
||||
from haiku.rag.client.session import default_db_path
|
||||
from haiku.rag.ingester.api.server import APIState, get_state
|
||||
from haiku.rag.store.info import DatabaseInfo, gather_database_info
|
||||
|
||||
|
|
@ -21,5 +20,4 @@ async def database(state: APIState = Depends(get_state)) -> DatabaseInfo:
|
|||
detail="database not configured",
|
||||
)
|
||||
[ref] = state.scope.databases
|
||||
one, db_path = ref.connection(state.config)
|
||||
return await gather_database_info(one, db_path or default_db_path(one))
|
||||
return await gather_database_info(ref.location, state.config)
|
||||
|
|
|
|||
|
|
@ -20,13 +20,12 @@ async def database_lines(client: "HaikuRAG") -> list[str]:
|
|||
Reported through the connection the client already holds. A failure becomes
|
||||
a line of the report, and the other databases still report.
|
||||
"""
|
||||
from haiku.rag.store.engine import ConnectionMode
|
||||
from haiku.rag.store.info import get_database_stats
|
||||
|
||||
lines: list[str] = []
|
||||
db_path = client.store.db_path
|
||||
|
||||
if client.store._connection_mode == ConnectionMode.LOCAL and not db_path.exists():
|
||||
if db_path is not None and not db_path.exists():
|
||||
return ["[red]Database path does not exist.[/red]"]
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -38,11 +38,12 @@ class ConnectionMode(Enum):
|
|||
OBJECT_STORAGE = "object_storage"
|
||||
|
||||
@staticmethod
|
||||
def from_config(config: AppConfig) -> "ConnectionMode":
|
||||
uri = config.lancedb.uri
|
||||
if not uri:
|
||||
def of(location: Path | str) -> "ConnectionMode":
|
||||
"""How a location is connected to: a path is local, `db://` is LanceDB
|
||||
Cloud, any other scheme is object storage."""
|
||||
if isinstance(location, Path) or "://" not in location:
|
||||
return ConnectionMode.LOCAL
|
||||
if uri.startswith("db://"):
|
||||
if location.startswith("db://"):
|
||||
return ConnectionMode.CLOUD
|
||||
return ConnectionMode.OBJECT_STORAGE
|
||||
|
||||
|
|
@ -72,8 +73,10 @@ def _session(config: AppConfig) -> lancedb.Session:
|
|||
|
||||
|
||||
async def connect_lancedb(
|
||||
config: AppConfig, db_path: Path | None = None
|
||||
location: Path | str, config: AppConfig
|
||||
) -> lancedb.AsyncConnection:
|
||||
"""Connect to the database at `location`, with the connection settings
|
||||
(credentials, storage options, caches, consistency) from `config`."""
|
||||
interval = config.lancedb.read_consistency_interval_seconds
|
||||
kwargs: dict[str, Any] = {
|
||||
"session": _session(config),
|
||||
|
|
@ -81,22 +84,19 @@ async def connect_lancedb(
|
|||
timedelta(seconds=interval) if interval is not None else None
|
||||
),
|
||||
}
|
||||
mode = ConnectionMode.from_config(config)
|
||||
mode = ConnectionMode.of(location)
|
||||
if mode == ConnectionMode.CLOUD:
|
||||
return await lancedb.connect_async(
|
||||
uri=config.lancedb.uri,
|
||||
uri=str(location),
|
||||
api_key=config.lancedb.api_key,
|
||||
region=config.lancedb.region,
|
||||
**kwargs,
|
||||
)
|
||||
elif mode == ConnectionMode.OBJECT_STORAGE:
|
||||
if mode == ConnectionMode.OBJECT_STORAGE:
|
||||
if config.lancedb.storage_options:
|
||||
kwargs["storage_options"] = config.lancedb.storage_options
|
||||
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(), **kwargs)
|
||||
return await lancedb.connect_async(uri=str(location), **kwargs)
|
||||
return await lancedb.connect_async(Path(location).absolute(), **kwargs)
|
||||
|
||||
|
||||
def _stored_vector_dim(settings: dict) -> int | None:
|
||||
|
|
@ -180,14 +180,24 @@ class TagInfo:
|
|||
class Store:
|
||||
def __init__(
|
||||
self,
|
||||
db_path: Path | str,
|
||||
location: Path | str,
|
||||
config: AppConfig | None = None,
|
||||
skip_validation: bool = False,
|
||||
create: bool = False,
|
||||
read_only: bool = False,
|
||||
skip_migration_check: bool = False,
|
||||
):
|
||||
self.db_path: Path = Path(db_path)
|
||||
"""A store over the database at `location`, a local path or a URI.
|
||||
|
||||
`config` supplies connection settings; where the database is comes
|
||||
from `location` alone.
|
||||
"""
|
||||
self._location: Path | str = location
|
||||
self.db_path: Path | None = (
|
||||
Path(location)
|
||||
if ConnectionMode.of(location) == ConnectionMode.LOCAL
|
||||
else None
|
||||
)
|
||||
self._config = config if config is not None else get_config()
|
||||
self._read_only = read_only
|
||||
self._create = create
|
||||
|
|
@ -200,7 +210,7 @@ class Store:
|
|||
self._rebuild_lock = asyncio.Lock()
|
||||
self._is_new_db = False
|
||||
|
||||
if self._connection_mode == ConnectionMode.LOCAL:
|
||||
if self.db_path is not None:
|
||||
if not self.db_path.exists():
|
||||
if not create:
|
||||
raise FileNotFoundError(
|
||||
|
|
@ -231,7 +241,7 @@ class Store:
|
|||
async def _initialize(self):
|
||||
"""Perform async initialization: connect to LanceDB, init tables, validate."""
|
||||
self.db: lancedb.AsyncConnection = await connect_lancedb(
|
||||
self._config, self.db_path
|
||||
self.location, self._config
|
||||
)
|
||||
|
||||
# Read once and thread onward: on object storage each of these is a
|
||||
|
|
@ -392,9 +402,14 @@ class Store:
|
|||
needed = datetime.now() - oldest + TAG_RETENTION_MARGIN
|
||||
return max(retention, needed)
|
||||
|
||||
@property
|
||||
def location(self) -> Path | str:
|
||||
"""Where this store connected: a local path, or a URI."""
|
||||
return self._location
|
||||
|
||||
@property
|
||||
def _connection_mode(self) -> ConnectionMode:
|
||||
return ConnectionMode.from_config(self._config)
|
||||
return ConnectionMode.of(self._location)
|
||||
|
||||
async def _ensure_vector_index(self) -> None:
|
||||
"""Create or rebuild vector index on chunks table.
|
||||
|
|
|
|||
|
|
@ -96,15 +96,15 @@ class DatabaseInfo(BaseModel):
|
|||
packages: dict[str, str] = Field(default_factory=dict)
|
||||
|
||||
|
||||
async def gather_database_info(config: AppConfig, db_path: Path) -> DatabaseInfo:
|
||||
async def gather_database_info(location: Path | str, config: AppConfig) -> DatabaseInfo:
|
||||
"""Collect read-only database state without going through Store, so a
|
||||
database missing tables (e.g. pre-migration) still reports what it can."""
|
||||
from haiku.rag.store.upgrades import get_pending_upgrades
|
||||
from haiku.rag.utils import get_package_versions
|
||||
|
||||
display_path = config.lancedb.uri or str(db_path)
|
||||
display_path = str(location)
|
||||
|
||||
db = await connect_lancedb(config, db_path)
|
||||
db = await connect_lancedb(location, config)
|
||||
stats = await get_database_stats(db)
|
||||
|
||||
if not any(entry["exists"] for entry in stats.values()):
|
||||
|
|
|
|||
|
|
@ -77,8 +77,11 @@ async def _apply_split_document_meta(store: Store) -> None:
|
|||
Exception
|
||||
): # pragma: no cover - defensive; stats() failure shouldn't block the split
|
||||
live_bytes = 0
|
||||
free_bytes = shutil.disk_usage(store.db_path).free
|
||||
if live_bytes and free_bytes < live_bytes:
|
||||
# A database behind a URI has no local disk to run out of.
|
||||
free_bytes = (
|
||||
shutil.disk_usage(store.db_path).free if store.db_path is not None else None
|
||||
)
|
||||
if live_bytes and free_bytes is not None and free_bytes < live_bytes:
|
||||
logger.warning(
|
||||
"Skipping post-migration vacuum: need ~%.2f GB free to compact the "
|
||||
"documents table, have %.2f GB. Run `haiku-rag vacuum` once you have "
|
||||
|
|
|
|||
|
|
@ -84,17 +84,18 @@ def test_chat_capabilities_read_the_named_database(tmp_path, monkeypatch):
|
|||
from haiku.rag.chat import run_chat
|
||||
|
||||
run_chat(scope=DatabaseScope.resolve(config, database_name="b"))
|
||||
named_scope = chat_app.call_args.kwargs["scope"]
|
||||
[named] = chat_app.call_args.kwargs["capabilities"]
|
||||
|
||||
run_chat(scope=DatabaseScope.resolve(config))
|
||||
covering_scope = chat_app.call_args.kwargs["scope"]
|
||||
[covering] = chat_app.call_args.kwargs["capabilities"]
|
||||
|
||||
# The chat lends its own client, so this scope is the fallback: it places
|
||||
# the named database alone.
|
||||
[placed] = named.scope.databases
|
||||
assert placed.db_path == tmp_path / "b.lancedb"
|
||||
assert named.config.lancedb.databases == {}
|
||||
assert covering.scope.names == ("a", "b")
|
||||
# The app opens the scope it is handed and lends that client to the
|
||||
# capabilities, which keep the configuration as the caller named it.
|
||||
assert named_scope.names == ("b",)
|
||||
assert covering_scope.names == ("a", "b")
|
||||
assert set(named.config.lancedb.databases) == {"a", "b"}
|
||||
assert set(covering.config.lancedb.databases) == {"a", "b"}
|
||||
|
||||
|
||||
|
|
@ -684,6 +685,37 @@ class TestLendingTheClient:
|
|||
assert borrowed == [client] * len(app._capabilities)
|
||||
assert borrowed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mounting_gives_every_capability_the_apps_scope(self, tmp_path):
|
||||
"""A capability built over the configured set covers what the chat
|
||||
selected once mounted: the analysis sandbox is built over that scope."""
|
||||
from haiku.rag.chat.app import ChatApp
|
||||
from haiku.rag.client.scope import DatabaseScope
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
databases={
|
||||
"a": str(tmp_path / "a.lancedb"),
|
||||
"b": str(tmp_path / "b.lancedb"),
|
||||
}
|
||||
)
|
||||
)
|
||||
selected = DatabaseScope.resolve(config, database_name="b")
|
||||
capability = create_capability(config=config)
|
||||
assert capability.scope.covers_multiple
|
||||
|
||||
client = _make_mock_client()
|
||||
app = ChatApp(scope=selected, capabilities=[capability], read_only=True)
|
||||
with (
|
||||
patch("haiku.rag.chat.app.HaikuRAG") as stub_rag,
|
||||
_covering_returns(stub_rag, client),
|
||||
):
|
||||
async with app.run_test():
|
||||
pass
|
||||
|
||||
assert capability.scope == selected
|
||||
|
||||
|
||||
class TestDocumentSelectionIdentity:
|
||||
"""Two documents can share a title, within a corpus and across databases, so
|
||||
|
|
|
|||
|
|
@ -7,7 +7,6 @@ import pytest
|
|||
from haiku.rag.capabilities._tools import search_corpus
|
||||
from haiku.rag.capabilities.rag import RAGState, create_capability
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.client.scope import DatabaseScope
|
||||
from haiku.rag.client.session import FederatedSession
|
||||
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
||||
from haiku.rag.store.exceptions import UnknownDatabaseError
|
||||
|
|
@ -256,12 +255,9 @@ class TestLendingANamedClient:
|
|||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
# `run_chat` derives these for a single-database scope.
|
||||
scope = DatabaseScope.resolve(config, database_name="alpha")
|
||||
one_config, one_path = scope.databases[0].connection(config)
|
||||
capability = create_capability(
|
||||
db_path=one_path, config=one_config, defer_loading=False
|
||||
)
|
||||
# What `run_chat` builds: the capability's own scope is the set, and
|
||||
# the lent client is what narrows it.
|
||||
capability = create_capability(config=config, defer_loading=False)
|
||||
|
||||
async with HaikuRAG(config=config, sources=["alpha"]) as client:
|
||||
# What `ChatApp.on_mount` does.
|
||||
|
|
|
|||
|
|
@ -125,10 +125,44 @@ class TestOneConfiguredLocation:
|
|||
config = self._config("s3://bucket/one.lancedb")
|
||||
|
||||
[ref] = DatabaseScope.resolve(config).databases
|
||||
one, db_path = ref.connection(config)
|
||||
|
||||
assert db_path is None
|
||||
assert ConnectionMode.from_config(one) == ConnectionMode.OBJECT_STORAGE
|
||||
assert ref.location == "s3://bucket/one.lancedb"
|
||||
assert ConnectionMode.of(ref.location) == ConnectionMode.OBJECT_STORAGE
|
||||
|
||||
|
||||
class TestSessionsOwnTheRef:
|
||||
"""A session is built from the resolved reference and hands storage only
|
||||
its location; the configuration it keeps is the one the caller named."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_session_opens_the_location_with_the_undivided_config(
|
||||
self, tmp_path
|
||||
):
|
||||
from haiku.rag.client.session import SingleDatabaseSession
|
||||
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
[ref] = DatabaseScope.resolve(config, database_name="alpha").databases
|
||||
|
||||
session = await SingleDatabaseSession(ref, config, read_only=True).open()
|
||||
try:
|
||||
assert session.source == "alpha"
|
||||
assert session.location == ref.location
|
||||
assert session.db_path == ref.location
|
||||
assert session.store.location == ref.location
|
||||
assert session.store._config is config
|
||||
finally:
|
||||
await session.aclose()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_client_keeps_the_configuration_it_was_given(self, tmp_path):
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config, sources=["alpha"]) as rag:
|
||||
assert rag._config is config
|
||||
assert set(rag._config.lancedb.databases) == {"alpha", "beta"}
|
||||
assert rag.store.location == tmp_path / "alpha.lancedb"
|
||||
|
||||
|
||||
class TestLocate:
|
||||
|
|
|
|||
|
|
@ -69,7 +69,7 @@ async def _seed(temp_db_path, *, version: str, with_items: bool = True):
|
|||
async def test_gather_database_info_reports_tables_and_settings(temp_db_path):
|
||||
await _seed(temp_db_path, version="1.2.3")
|
||||
|
||||
info = await gather_database_info(AppConfig(), temp_db_path)
|
||||
info = await gather_database_info(temp_db_path, AppConfig())
|
||||
|
||||
assert info.exists is True
|
||||
assert info.path == str(temp_db_path)
|
||||
|
|
@ -98,7 +98,7 @@ async def test_gather_database_info_flags_missing_table_and_pending_migrations(
|
|||
):
|
||||
await _seed(temp_db_path, version="0.39.0", with_items=False)
|
||||
|
||||
info = await gather_database_info(AppConfig(), temp_db_path)
|
||||
info = await gather_database_info(temp_db_path, AppConfig())
|
||||
|
||||
tables = {t.name: t for t in info.tables}
|
||||
assert tables["document_items"].exists is False
|
||||
|
|
@ -112,7 +112,30 @@ async def test_gather_database_info_empty_database(temp_db_path):
|
|||
|
||||
await lancedb.connect_async(temp_db_path) # creates the dir, no tables
|
||||
|
||||
info = await gather_database_info(AppConfig(), temp_db_path)
|
||||
info = await gather_database_info(temp_db_path, AppConfig())
|
||||
|
||||
assert info.exists is False
|
||||
assert info.path == str(temp_db_path)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_gather_database_info_connects_to_the_location_it_is_given():
|
||||
"""A remote location is passed to the connection as is and reported back
|
||||
as the path; the configuration's own `uri` plays no part."""
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from haiku.rag.config.models import LanceDBConfig
|
||||
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri="s3://elsewhere/other.lancedb"))
|
||||
with patch(
|
||||
"haiku.rag.store.info.connect_lancedb", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
listing = MagicMock()
|
||||
listing.tables = []
|
||||
mock_connect.return_value.list_tables = AsyncMock(return_value=listing)
|
||||
|
||||
info = await gather_database_info("s3://bucket/papers.lancedb", config)
|
||||
|
||||
assert mock_connect.call_args.args[0] == "s3://bucket/papers.lancedb"
|
||||
assert info.path == "s3://bucket/papers.lancedb"
|
||||
assert info.exists is False
|
||||
|
|
|
|||
|
|
@ -126,6 +126,43 @@ class TestV0_58_0MigrationEdgeCases:
|
|||
assert set(by_id) == {"a", "b"} # exactly one row each, no duplicates
|
||||
assert len(rows) == 2
|
||||
|
||||
async def test_a_remote_store_has_no_disk_to_check(self, temp_db_path, monkeypatch):
|
||||
"""A store behind a URI has no local path: the reclaim vacuum runs
|
||||
without a free-disk check."""
|
||||
from haiku.rag.store.upgrades import v0_58_0
|
||||
|
||||
async with Store(temp_db_path, create=True, skip_migration_check=True) as store:
|
||||
await seed_legacy_documents(
|
||||
store,
|
||||
[LegacyDocumentRecord(id="a", content="x", uri="u", metadata="{}")],
|
||||
)
|
||||
await store.set_haiku_version("0.57.0")
|
||||
|
||||
def no_disk(_path):
|
||||
raise AssertionError("disk_usage consulted for a remote store")
|
||||
|
||||
monkeypatch.setattr(v0_58_0.shutil, "disk_usage", no_disk)
|
||||
vacuum_calls: list[int] = []
|
||||
|
||||
async with Store(temp_db_path, skip_migration_check=True) as store:
|
||||
store.db_path = None
|
||||
|
||||
async def fake_stats():
|
||||
return {"total_bytes": 10_000_000}
|
||||
|
||||
monkeypatch.setattr(store.documents_table, "stats", fake_stats)
|
||||
orig_vacuum = store.vacuum
|
||||
|
||||
async def tracking_vacuum(*args, **kwargs):
|
||||
vacuum_calls.append(1)
|
||||
return await orig_vacuum(*args, **kwargs)
|
||||
|
||||
monkeypatch.setattr(store, "vacuum", tracking_vacuum)
|
||||
|
||||
await store.migrate()
|
||||
|
||||
assert vacuum_calls == [1]
|
||||
|
||||
async def test_skips_vacuum_when_disk_is_tight(self, temp_db_path, monkeypatch):
|
||||
"""When free disk can't cover one compacted copy, the split still
|
||||
completes but the reclaim vacuum is skipped."""
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ from pathlib import Path
|
|||
import pytest
|
||||
|
||||
from haiku.rag.client.scope import DatabaseRef, DatabaseScope
|
||||
from haiku.rag.client.session import default_db_path
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig, StorageConfig
|
||||
from haiku.rag.store.exceptions import (
|
||||
AmbiguousDatabaseError,
|
||||
|
|
@ -89,16 +88,13 @@ class TestResolution:
|
|||
assert ref.uri == ""
|
||||
|
||||
def test_a_path_selects_the_database_over_a_configured_uri(self):
|
||||
"""`--db` exists to override what is configured, and the configuration
|
||||
derived from the ref is what makes the connection follow it."""
|
||||
"""`--db` exists to override what is configured."""
|
||||
config = _config(uri="s3://bucket/one.lancedb")
|
||||
|
||||
scope = DatabaseScope.resolve(config, database_path=Path("/data/local"))
|
||||
|
||||
[ref] = scope.databases
|
||||
assert ref.db_path == Path("/data/local")
|
||||
one, _ = ref.connection(config)
|
||||
assert one.lancedb.uri == ""
|
||||
assert ref.location == Path("/data/local")
|
||||
|
||||
def test_nothing_configured_falls_back_to_the_data_directory(self, tmp_path):
|
||||
config = AppConfig(storage=StorageConfig(data_dir=tmp_path))
|
||||
|
|
@ -158,61 +154,23 @@ class TestResolution:
|
|||
DatabaseScope(())
|
||||
|
||||
|
||||
class TestConnectionDerivation:
|
||||
"""Opening one of a set must not disturb the configuration it came from."""
|
||||
class TestLocation:
|
||||
"""One value says where a database is: a path for a local one, a URI string
|
||||
for a remote one. Storage connects to it as given."""
|
||||
|
||||
def test_a_local_location_becomes_a_path(self):
|
||||
def test_a_local_location_is_a_path(self):
|
||||
config = _config(databases={"alpha": "/data/alpha.lancedb"})
|
||||
[ref] = DatabaseScope.resolve(config).databases
|
||||
|
||||
one, db_path = ref.connection(config)
|
||||
assert ref.location == Path("/data/alpha.lancedb")
|
||||
|
||||
assert db_path == Path("/data/alpha.lancedb")
|
||||
assert one.lancedb.uri == ""
|
||||
assert one.lancedb.databases == {}
|
||||
|
||||
def test_a_uri_location_stays_a_uri(self):
|
||||
def test_a_uri_location_is_the_uri(self):
|
||||
config = _config(databases={"alpha": "s3://bucket/alpha.lancedb"})
|
||||
[ref] = DatabaseScope.resolve(config).databases
|
||||
|
||||
one, db_path = ref.connection(config)
|
||||
assert ref.location == "s3://bucket/alpha.lancedb"
|
||||
|
||||
assert db_path is None
|
||||
assert one.lancedb.uri == "s3://bucket/alpha.lancedb"
|
||||
|
||||
def test_the_original_configuration_is_untouched(self):
|
||||
"""Rewriting it in place is what left downstream code unable to tell a set
|
||||
had been named."""
|
||||
config = _config(databases={"alpha": "/a.lancedb", "beta": "/b.lancedb"})
|
||||
|
||||
for ref in DatabaseScope.resolve(config).databases:
|
||||
ref.connection(config)
|
||||
|
||||
assert config.lancedb.databases == {"alpha": "/a.lancedb", "beta": "/b.lancedb"}
|
||||
assert config.lancedb.uri == ""
|
||||
|
||||
def test_each_derived_configuration_is_its_own_copy(self):
|
||||
config = _config(databases={"alpha": "/a.lancedb", "beta": "s3://b/b.lancedb"})
|
||||
alpha, beta = DatabaseScope.resolve(config).databases
|
||||
|
||||
one, _ = alpha.connection(config)
|
||||
other, _ = beta.connection(config)
|
||||
|
||||
assert one is not other
|
||||
assert one.lancedb.uri == ""
|
||||
assert other.lancedb.uri == "s3://b/b.lancedb"
|
||||
|
||||
|
||||
def test_a_database_behind_a_uri_has_no_path_of_its_own(tmp_path):
|
||||
"""`connection` hands back no path for a URI, and the store still needs one:
|
||||
the default stands in, and the URI is what decides where it connects."""
|
||||
config = AppConfig(
|
||||
storage=StorageConfig(data_dir=tmp_path),
|
||||
lancedb=LanceDBConfig(databases={"alpha": "s3://bucket/alpha.lancedb"}),
|
||||
)
|
||||
[ref] = DatabaseScope.resolve(config).databases
|
||||
|
||||
one, db_path = ref.connection(config)
|
||||
|
||||
assert db_path is None
|
||||
assert default_db_path(one) == tmp_path / "haiku.rag.lancedb"
|
||||
def test_a_path_the_caller_gave_is_its_location(self):
|
||||
assert DatabaseRef.at("/data/other.lancedb").location == Path(
|
||||
"/data/other.lancedb"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -177,8 +177,7 @@ async def test_app_info_opens_a_named_remote_database(tmp_path):
|
|||
app = HaikuRAGApp(scope=scope, config=config)
|
||||
|
||||
assert app._is_local is False
|
||||
assert app._store_config.lancedb.uri == "s3://bucket/papers.lancedb"
|
||||
assert app._store_config.lancedb.databases == {}
|
||||
assert app._location == "s3://bucket/papers.lancedb"
|
||||
|
||||
with patch(
|
||||
"haiku.rag.store.info.connect_lancedb", new_callable=AsyncMock
|
||||
|
|
@ -189,13 +188,12 @@ async def test_app_info_opens_a_named_remote_database(tmp_path):
|
|||
mock_db.list_tables = AsyncMock(return_value=mock_list_result)
|
||||
await app.info()
|
||||
|
||||
opened = mock_connect.call_args.args[0]
|
||||
assert opened.lancedb.uri == "s3://bucket/papers.lancedb"
|
||||
assert mock_connect.call_args.args[0] == "s3://bucket/papers.lancedb"
|
||||
|
||||
|
||||
async def test_app_doctor_opens_a_named_remote_database():
|
||||
"""`run_doctor` connects with the configuration it is handed: the one
|
||||
derived for the database, not the one naming the set."""
|
||||
"""`run_doctor` is handed the database's location, not the configuration
|
||||
naming the set."""
|
||||
from haiku.rag.client.scope import DatabaseScope
|
||||
|
||||
config = AppConfig(
|
||||
|
|
@ -207,7 +205,7 @@ async def test_app_doctor_opens_a_named_remote_database():
|
|||
run.return_value = MagicMock(checks=[], ok=True, duplicates=None)
|
||||
await app.doctor()
|
||||
|
||||
assert run.call_args.args[0].lancedb.uri == "s3://bucket/papers.lancedb"
|
||||
assert run.call_args.args[1] == "s3://bucket/papers.lancedb"
|
||||
|
||||
|
||||
async def test_app_info_uses_connect_lancedb_for_remote(tmp_path):
|
||||
|
|
@ -230,9 +228,8 @@ async def test_app_info_uses_connect_lancedb_for_remote(tmp_path):
|
|||
mock_db.list_tables = AsyncMock(return_value=mock_list_result)
|
||||
await app.info()
|
||||
|
||||
# The uri decides where it connects; the path argument is not read.
|
||||
mock_connect.assert_called_once()
|
||||
assert mock_connect.call_args.args[0].lancedb.uri == "s3://bucket/path"
|
||||
assert mock_connect.call_args.args[0] == "s3://bucket/path"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -340,23 +340,19 @@ class TestReportedLocation:
|
|||
@staticmethod
|
||||
def _session(location: str):
|
||||
from haiku.rag.client.scope import DatabaseScope
|
||||
from haiku.rag.client.session import SingleDatabaseSession, default_db_path
|
||||
from haiku.rag.client.session import SingleDatabaseSession
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
|
||||
config = AppConfig(lancedb=LanceDBConfig(databases={"alpha": location}))
|
||||
[ref] = DatabaseScope.resolve(config, database_name="alpha").databases
|
||||
one, db_path = ref.connection(config)
|
||||
return SingleDatabaseSession(
|
||||
db_path if db_path is not None else default_db_path(one),
|
||||
one,
|
||||
source="alpha",
|
||||
)
|
||||
return SingleDatabaseSession(ref, config)
|
||||
|
||||
def test_a_named_remote_database_reports_its_uri(self):
|
||||
session = self._session("s3://bucket/alpha.lancedb")
|
||||
|
||||
assert isinstance(session.db_path, Path)
|
||||
assert session.db_path is None
|
||||
assert session.location == "s3://bucket/alpha.lancedb"
|
||||
assert session.source == "alpha"
|
||||
|
||||
def test_a_named_local_database_reports_its_path(self):
|
||||
session = self._session("/data/alpha.lancedb")
|
||||
|
|
|
|||
|
|
@ -4,53 +4,44 @@ from unittest.mock import AsyncMock, patch
|
|||
import pytest
|
||||
from pydantic import ValidationError
|
||||
|
||||
from haiku.rag.config import get_config
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
from haiku.rag.store.engine import ConnectionMode, Store, connect_lancedb
|
||||
|
||||
|
||||
class TestConnectionMode:
|
||||
def test_local_when_uri_empty(self):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri=""))
|
||||
assert ConnectionMode.from_config(config) == ConnectionMode.LOCAL
|
||||
"""A location is classified by itself: a path is local, `db://` is LanceDB
|
||||
Cloud, any other scheme is object storage."""
|
||||
|
||||
def test_a_path_is_local(self, tmp_path):
|
||||
assert ConnectionMode.of(tmp_path / "db.lancedb") == ConnectionMode.LOCAL
|
||||
|
||||
def test_a_schemeless_string_is_local(self):
|
||||
assert ConnectionMode.of("/data/db.lancedb") == ConnectionMode.LOCAL
|
||||
|
||||
def test_cloud_when_db_uri(self):
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
uri="db://my-database", api_key="key", region="us-east-1"
|
||||
)
|
||||
)
|
||||
assert ConnectionMode.from_config(config) == ConnectionMode.CLOUD
|
||||
assert ConnectionMode.of("db://my-database") == ConnectionMode.CLOUD
|
||||
|
||||
def test_object_storage_s3(self):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
|
||||
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
|
||||
|
||||
def test_object_storage_gs(self):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri="gs://bucket/path"))
|
||||
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
|
||||
|
||||
def test_object_storage_az(self):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri="az://container/path"))
|
||||
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
|
||||
|
||||
def test_object_storage_hdfs(self):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri="hdfs://namenode/path"))
|
||||
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
|
||||
|
||||
def test_unknown_uri_treated_as_object_storage(self):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri="custom://something"))
|
||||
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
|
||||
@pytest.mark.parametrize(
|
||||
"uri",
|
||||
[
|
||||
"s3://bucket/path",
|
||||
"gs://bucket/path",
|
||||
"az://container/path",
|
||||
"hdfs://namenode/path",
|
||||
"custom://something",
|
||||
],
|
||||
)
|
||||
def test_any_other_scheme_is_object_storage(self, uri):
|
||||
assert ConnectionMode.of(uri) == ConnectionMode.OBJECT_STORAGE
|
||||
|
||||
|
||||
class TestConnectLancedb:
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_passes_absolute_db_path(self, temp_db_path):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri=""))
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config, db_path=temp_db_path)
|
||||
await connect_lancedb(temp_db_path, AppConfig())
|
||||
mock_connect.assert_awaited_once()
|
||||
assert mock_connect.call_args.args == (temp_db_path.absolute(),)
|
||||
|
||||
|
|
@ -60,25 +51,34 @@ class TestConnectLancedb:
|
|||
|
||||
monkeypatch.chdir(tmp_path)
|
||||
relative = Path("db/rag.lancedb")
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri=""))
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config, db_path=relative)
|
||||
await connect_lancedb(relative, AppConfig())
|
||||
mock_connect.assert_awaited_once()
|
||||
assert mock_connect.call_args.args == (relative.absolute(),)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_configured_uri_is_not_consulted(self, temp_db_path):
|
||||
"""Storage connects to the location it is handed; placement is the
|
||||
caller's, and the configuration's own `uri` never redirects it."""
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri="s3://elsewhere/db.lancedb"))
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(temp_db_path, config)
|
||||
assert mock_connect.call_args.args == (temp_db_path.absolute(),)
|
||||
assert "uri" not in mock_connect.call_args.kwargs
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_passes_uri_api_key_region(self):
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
uri="db://my-database", api_key="test-key", region="us-west-2"
|
||||
)
|
||||
lancedb=LanceDBConfig(api_key="test-key", region="us-west-2")
|
||||
)
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
await connect_lancedb("db://my-database", config)
|
||||
mock_connect.assert_awaited_once()
|
||||
kwargs = mock_connect.call_args.kwargs
|
||||
assert kwargs["uri"] == "db://my-database"
|
||||
|
|
@ -89,7 +89,6 @@ class TestConnectLancedb:
|
|||
async def test_object_storage_passes_uri_and_storage_options(self):
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
uri="s3://bucket/path",
|
||||
storage_options={
|
||||
"endpoint": "http://minio:9000",
|
||||
"region": "us-east-1",
|
||||
|
|
@ -99,7 +98,7 @@ class TestConnectLancedb:
|
|||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
await connect_lancedb("s3://bucket/path", config)
|
||||
mock_connect.assert_awaited_once()
|
||||
kwargs = mock_connect.call_args.kwargs
|
||||
assert kwargs["uri"] == "s3://bucket/path"
|
||||
|
|
@ -110,23 +109,25 @@ class TestConnectLancedb:
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_object_storage_without_storage_options(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("s3://bucket/path", AppConfig())
|
||||
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):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri=""))
|
||||
with pytest.raises(
|
||||
ValueError, match="No lancedb.uri configured and no db_path provided"
|
||||
):
|
||||
await connect_lancedb(config)
|
||||
|
||||
def _remote_store(location: str, config: AppConfig | None = None) -> Store:
|
||||
"""A store over a remote location, opened against a mocked connection."""
|
||||
return Store(
|
||||
location,
|
||||
config=config,
|
||||
create=True,
|
||||
skip_validation=True,
|
||||
skip_migration_check=True,
|
||||
)
|
||||
|
||||
|
||||
class TestStoreConnectionMode:
|
||||
|
|
@ -134,132 +135,127 @@ class TestStoreConnectionMode:
|
|||
async def test_store_connection_mode_local(self, temp_db_path):
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
assert store._connection_mode == ConnectionMode.LOCAL
|
||||
assert store.location == temp_db_path
|
||||
assert store.db_path == temp_db_path
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_connection_mode_cloud(self, temp_db_path):
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
with (
|
||||
patch.object(get_config().lancedb, "uri", "db://test-database"),
|
||||
patch.object(get_config().lancedb, "api_key", "test-api-key"),
|
||||
patch.object(get_config().lancedb, "region", "us-east-1"),
|
||||
):
|
||||
async def test_a_local_store_ignores_the_configured_uri(self, temp_db_path):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri="s3://elsewhere/db.lancedb"))
|
||||
async with Store(temp_db_path, config=config, create=True) as store:
|
||||
assert store._connection_mode == ConnectionMode.LOCAL
|
||||
assert store.db_path == temp_db_path
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_connection_mode_cloud(self):
|
||||
config = AppConfig(lancedb=LanceDBConfig(api_key="key", region="us-east-1"))
|
||||
with (
|
||||
patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
),
|
||||
patch.object(Store, "_init_tables", new_callable=AsyncMock),
|
||||
):
|
||||
async with _remote_store("db://test-database", config) as store:
|
||||
assert store._connection_mode == ConnectionMode.CLOUD
|
||||
assert store.location == "db://test-database"
|
||||
assert store.db_path is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_connection_mode_object_storage(self, temp_db_path):
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
with patch.object(get_config().lancedb, "uri", "s3://bucket/path"):
|
||||
async def test_store_connection_mode_object_storage(self):
|
||||
with (
|
||||
patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
),
|
||||
patch.object(Store, "_init_tables", new_callable=AsyncMock),
|
||||
):
|
||||
async with _remote_store("s3://bucket/path") as store:
|
||||
assert store._connection_mode == ConnectionMode.OBJECT_STORAGE
|
||||
assert store.db_path is None
|
||||
|
||||
|
||||
def _remote_store_with_mock_tables(location: str) -> Store:
|
||||
"""A remote store whose tables are mocks: the mode decision is under test,
|
||||
not the tables."""
|
||||
store = _remote_store(location)
|
||||
store.chunks_table = AsyncMock()
|
||||
return store
|
||||
|
||||
|
||||
class TestVacuumByConnectionMode:
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_skips_vacuum(self, temp_db_path):
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
with (
|
||||
patch.object(get_config().lancedb, "uri", "db://test-database"),
|
||||
patch.object(get_config().lancedb, "api_key", "test-api-key"),
|
||||
patch.object(get_config().lancedb, "region", "us-east-1"),
|
||||
):
|
||||
with patch.object(
|
||||
store.chunks_table, "optimize", new_callable=AsyncMock
|
||||
) as mock_optimize:
|
||||
await store.vacuum()
|
||||
mock_optimize.assert_not_called()
|
||||
async def test_cloud_skips_vacuum(self):
|
||||
with (
|
||||
patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
),
|
||||
patch.object(Store, "_init_tables", new_callable=AsyncMock),
|
||||
):
|
||||
async with _remote_store_with_mock_tables("db://test-database") as store:
|
||||
await store.vacuum()
|
||||
store.chunks_table.optimize.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_object_storage_runs_vacuum(self, temp_db_path):
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
with patch.object(get_config().lancedb, "uri", "s3://bucket/path"):
|
||||
async def test_object_storage_runs_vacuum(self):
|
||||
with (
|
||||
patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
),
|
||||
patch.object(Store, "_init_tables", new_callable=AsyncMock),
|
||||
):
|
||||
async with _remote_store_with_mock_tables("s3://bucket/path") as store:
|
||||
store.chunks_table.tags.list = AsyncMock(return_value={})
|
||||
with patch.object(
|
||||
store.chunks_table, "optimize", new_callable=AsyncMock
|
||||
) as mock_optimize:
|
||||
store, "_tables", return_value={"chunks": store.chunks_table}
|
||||
):
|
||||
await store.vacuum()
|
||||
mock_optimize.assert_called()
|
||||
store.chunks_table.optimize.assert_awaited_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_local_runs_vacuum(self, temp_db_path):
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
with patch.object(get_config().lancedb, "uri", ""):
|
||||
with patch.object(
|
||||
store.chunks_table, "optimize", new_callable=AsyncMock
|
||||
) as mock_optimize:
|
||||
await store.vacuum()
|
||||
mock_optimize.assert_called()
|
||||
with patch.object(
|
||||
store.chunks_table, "optimize", new_callable=AsyncMock
|
||||
) as mock_optimize:
|
||||
await store.vacuum()
|
||||
mock_optimize.assert_called()
|
||||
|
||||
|
||||
class TestVectorIndexByConnectionMode:
|
||||
@pytest.mark.asyncio
|
||||
async def test_cloud_skips_index_creation(self, temp_db_path):
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
with (
|
||||
patch.object(get_config().lancedb, "uri", "db://test-database"),
|
||||
patch.object(get_config().lancedb, "api_key", "test-api-key"),
|
||||
patch.object(get_config().lancedb, "region", "us-east-1"),
|
||||
):
|
||||
with patch.object(
|
||||
store.chunks_table, "count_rows", new_callable=AsyncMock
|
||||
) as mock_count:
|
||||
await store._ensure_vector_index()
|
||||
mock_count.assert_not_called()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_object_storage_runs_index_creation(self, temp_db_path):
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
with patch.object(get_config().lancedb, "uri", "s3://bucket/path"):
|
||||
with patch.object(
|
||||
store.chunks_table,
|
||||
"count_rows",
|
||||
new_callable=AsyncMock,
|
||||
return_value=0,
|
||||
) as mock_count:
|
||||
await store._ensure_vector_index()
|
||||
mock_count.assert_called()
|
||||
|
||||
|
||||
class TestStoreSkipsPathValidationForRemote:
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_path_check_for_cloud(self, tmp_path):
|
||||
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
uri="db://test-database", api_key="key", region="us-east-1"
|
||||
)
|
||||
)
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
async def test_cloud_skips_index_creation(self):
|
||||
with (
|
||||
patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
),
|
||||
patch.object(Store, "_init_tables", new_callable=AsyncMock),
|
||||
):
|
||||
with patch.object(Store, "_init_tables", new_callable=AsyncMock):
|
||||
async with Store(
|
||||
nonexistent,
|
||||
config=config,
|
||||
create=True,
|
||||
skip_validation=True,
|
||||
skip_migration_check=True,
|
||||
) as store:
|
||||
assert store is not None
|
||||
async with _remote_store_with_mock_tables("db://test-database") as store:
|
||||
await store._ensure_vector_index()
|
||||
store.chunks_table.count_rows.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_skips_path_check_for_object_storage(self, tmp_path):
|
||||
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
|
||||
config = AppConfig(
|
||||
lancedb=LanceDBConfig(
|
||||
uri="s3://bucket/path",
|
||||
storage_options={"endpoint": "http://localhost:9000"},
|
||||
)
|
||||
)
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
async def test_object_storage_runs_index_creation(self):
|
||||
with (
|
||||
patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
),
|
||||
patch.object(Store, "_init_tables", new_callable=AsyncMock),
|
||||
):
|
||||
with patch.object(Store, "_init_tables", new_callable=AsyncMock):
|
||||
async with Store(
|
||||
nonexistent,
|
||||
config=config,
|
||||
create=True,
|
||||
skip_validation=True,
|
||||
skip_migration_check=True,
|
||||
) as store:
|
||||
assert store is not None
|
||||
async with _remote_store_with_mock_tables("s3://bucket/path") as store:
|
||||
store.chunks_table.count_rows = AsyncMock(return_value=0)
|
||||
await store._ensure_vector_index()
|
||||
store.chunks_table.count_rows.assert_awaited_once()
|
||||
|
||||
|
||||
class TestLocationIsFixed:
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_store_keeps_the_location_it_opened(self, temp_db_path):
|
||||
"""`db_path` and the connection mode derive from the location once; a
|
||||
store cannot be pointed elsewhere after it is built."""
|
||||
async with Store(temp_db_path, create=True) as store:
|
||||
with pytest.raises(AttributeError):
|
||||
store.location = "s3://bucket/path" # type: ignore[misc]
|
||||
assert store.location == temp_db_path
|
||||
assert store._connection_mode == ConnectionMode.LOCAL
|
||||
|
||||
|
||||
class TestInitFailureCleanup:
|
||||
|
|
@ -412,33 +408,25 @@ class TestStoreMiscellany:
|
|||
class TestSessionAndConsistency:
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_is_shared_across_connections(self):
|
||||
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
|
||||
config = AppConfig()
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
await connect_lancedb(config)
|
||||
await connect_lancedb("s3://bucket/path", config)
|
||||
await connect_lancedb("s3://bucket/path", 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
|
||||
)
|
||||
)
|
||||
small = AppConfig(lancedb=LanceDBConfig(index_cache_size_bytes=1 << 20))
|
||||
large = AppConfig(lancedb=LanceDBConfig(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)
|
||||
await connect_lancedb("s3://bucket/path", small)
|
||||
await connect_lancedb("s3://bucket/path", large)
|
||||
|
||||
sessions = [c.kwargs["session"] for c in mock_connect.call_args_list]
|
||||
assert sessions[0] is not sessions[1]
|
||||
|
|
@ -447,7 +435,6 @@ class TestSessionAndConsistency:
|
|||
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,
|
||||
)
|
||||
|
|
@ -458,7 +445,7 @@ class TestSessionAndConsistency:
|
|||
),
|
||||
patch("haiku.rag.store.engine.lancedb.Session") as mock_session,
|
||||
):
|
||||
await connect_lancedb(config)
|
||||
await connect_lancedb("s3://bucket/path", config)
|
||||
|
||||
mock_session.assert_called_once_with(
|
||||
index_cache_size_bytes=2 << 20, metadata_cache_size_bytes=4 << 20
|
||||
|
|
@ -466,15 +453,11 @@ class TestSessionAndConsistency:
|
|||
|
||||
@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
|
||||
)
|
||||
)
|
||||
config = AppConfig(lancedb=LanceDBConfig(read_consistency_interval_seconds=5))
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
await connect_lancedb("s3://bucket/path", config)
|
||||
|
||||
assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta(
|
||||
seconds=5
|
||||
|
|
@ -483,14 +466,12 @@ class TestSessionAndConsistency:
|
|||
@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
|
||||
)
|
||||
lancedb=LanceDBConfig(read_consistency_interval_seconds=None)
|
||||
)
|
||||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config)
|
||||
await connect_lancedb("s3://bucket/path", config)
|
||||
|
||||
assert mock_connect.call_args.kwargs["read_consistency_interval"] is None
|
||||
|
||||
|
|
@ -500,7 +481,7 @@ class TestSessionAndConsistency:
|
|||
with patch(
|
||||
"haiku.rag.store.engine.lancedb.connect_async", new_callable=AsyncMock
|
||||
) as mock_connect:
|
||||
await connect_lancedb(config, tmp_path / "db.lancedb")
|
||||
await connect_lancedb(tmp_path / "db.lancedb", config)
|
||||
|
||||
assert mock_connect.call_args.kwargs["session"] is not None
|
||||
assert mock_connect.call_args.kwargs["read_consistency_interval"] == timedelta(
|
||||
|
|
|
|||
|
|
@ -89,7 +89,7 @@ async def _remote_client(config: AppConfig):
|
|||
async def test_store_connect_and_create(tmp_path, config):
|
||||
from haiku.rag.store.info import get_database_stats
|
||||
|
||||
async with Store(tmp_path / "unused", config=config, create=True) as store:
|
||||
async with Store(config.lancedb.uri, config=config, create=True) as store:
|
||||
stats = await get_database_stats(store.db)
|
||||
assert stats["documents"]["exists"]
|
||||
assert stats["chunks"]["exists"]
|
||||
|
|
@ -97,7 +97,7 @@ async def test_store_connect_and_create(tmp_path, config):
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_store_vacuum(tmp_path, config):
|
||||
async with Store(tmp_path / "unused", config=config, create=True) as store:
|
||||
async with Store(config.lancedb.uri, config=config, create=True) as store:
|
||||
await store.vacuum()
|
||||
|
||||
|
||||
|
|
@ -106,7 +106,7 @@ async def test_store_add_document(tmp_path, config):
|
|||
from haiku.rag.store.info import get_database_stats
|
||||
from haiku.rag.store.schema import DocumentRecord
|
||||
|
||||
async with Store(tmp_path / "unused", config=config, create=True) as store:
|
||||
async with Store(config.lancedb.uri, config=config, create=True) as store:
|
||||
doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.")
|
||||
await store.documents_table.add([doc])
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue