Ask and analyze across several databases
Chunk 2 gave search a configured set to fan out over. ask and analyze
covered one database still: the RAG capability had no way to be told which
databases a question spanned, and the analysis sandbox mounted one
document tree.
The selection travels as sources on EvidenceState, beside the filter it
scopes with, so both capabilities read it the same way. clients_covering
is the one rule that turns a selection into clients, used by search, the
sandbox mount and the cite fallback, so a question scoped to some
databases cannot search, mount or cite another. Citations carry the
database they came from, and format_for_agent names it, so the model can
attribute evidence while it answers rather than only afterwards.
The sandbox keeps one flat /documents/{id}/ namespace and resolves each id
to the client holding it, which rests on ids being UUID4. A database
copied from another breaks that, so an id held twice is refused rather
than resolved to whichever arrived last.
On the CLI, search, ask and analyze cover the configured set and label
each result with its database. Every other command works on one, named
with --database NAME (a name reaches a database behind a URI, which --db
cannot) or --db PATH, and refuses a set it cannot choose from instead of
silently reading the default database. Cold databases open together, so a
first query costs the slowest open rather than their sum.
This commit is contained in:
Yiorgis Gozadinos
parent
f33b789a31
commit
fdb5710491
33 changed files with 1943 additions and 158 deletions
|
|
@ -5,6 +5,9 @@
|
|||
### Added
|
||||
|
||||
- `api_key` on model and embedding-model config, overriding the provider's environment variable. Honored on the `openai` and `ollama` providers, `vllm` embedders and rerankers, the picture-description VLM endpoint, and `doctor`'s endpoint probes; other providers raise.
|
||||
- `lancedb.databases`: a name-to-location mapping for searching several databases at once, mutually exclusive with `lancedb.uri`. `client.search(..., sources=[...])` selects which to search, `sources=None` searches all of them, and `SearchResult.source` carries the configured name a result came from. Candidates are fused by the configured reranker over the union, or by reciprocal rank fusion when none is configured. `SearchResult.format_for_agent` names the database, so the model can attribute evidence to one while it answers. `haiku-rag search`, `ask` and `analyze` cover the configured set and label each result and citation with its database; every other command works on one, named with `--database NAME` or `--db PATH`.
|
||||
- `client.ask(..., sources=[...])` asks across the selected databases, and `Citation.source` names the one a cited chunk came from. The cite fallback for an id absent from the run's results looks only in the selected databases, so a question scoped to some cannot cite another.
|
||||
- `client.analyze(..., sources=[...])` analyzes across the selected databases: the sandbox mounts their documents under one flat `/documents/{id}/` namespace, resolving each id to the database holding it, and in-code `search()` covers the same selection.
|
||||
|
||||
### Fixed
|
||||
|
||||
|
|
@ -38,8 +41,6 @@
|
|||
- `haiku.rag.capabilities.EvidenceState`: the state base `RAGState` and `AnalysisState` derive from, with `begin_invocation()` for the per-question reset. `RAGCapabilityBase.evidence_record()` and `citation_index()` expose what a capability recorded, so a host reads it without reaching into `capability.state`.
|
||||
- The `haiku.rag` package declares the `jina` extra, so `provider: jina-local` is supported by declaration rather than through `cross-encoder`'s transitive `transformers` and `torch`. Raises the full package's torch floor to 2.0.
|
||||
- `providers.docling_serve.timeout` (default 300 seconds), forwarded to the docling-serve client's per-request timeout.
|
||||
- `lancedb.databases`: a name-to-location mapping for searching several databases at once, mutually exclusive with `lancedb.uri`. `client.search(..., sources=[...])` selects which to search, `sources=None` searches all of them, and `SearchResult.source` carries the configured name a result came from. Candidates are fused by the configured reranker over the union, or by reciprocal rank fusion when none is configured.
|
||||
- `client.ask(..., sources=[...])` asks across the selected databases, and `Citation.source` names the one a cited chunk came from. The cite fallback for an id absent from the run's results looks only in the selected databases, so a question scoped to some cannot cite another.
|
||||
|
||||
### Changed
|
||||
|
||||
|
|
|
|||
|
|
@ -36,10 +36,12 @@ class HaikuRAGApp:
|
|||
db_path: Path,
|
||||
config: AppConfig | None = None,
|
||||
read_only: bool = False,
|
||||
federated: bool = False,
|
||||
):
|
||||
self.db_path = db_path
|
||||
self.config = config if config is not None else get_config()
|
||||
self.read_only = read_only
|
||||
self.federated = federated
|
||||
self.console = Console()
|
||||
|
||||
from haiku.rag.store.engine import ConnectionMode
|
||||
|
|
@ -47,6 +49,12 @@ class HaikuRAGApp:
|
|||
self._is_local = ConnectionMode.from_config(self.config) == ConnectionMode.LOCAL
|
||||
self._display_path = self.db_path if self._is_local else self.config.lancedb.uri
|
||||
|
||||
@property
|
||||
def _client_db_path(self) -> "Path | None":
|
||||
"""None where the command covers `lancedb.databases`, so the client
|
||||
opens the configured set rather than one path."""
|
||||
return None if self.federated else self.db_path
|
||||
|
||||
async def init(self):
|
||||
"""Initialize a new database."""
|
||||
if self._is_local and self.db_path.exists():
|
||||
|
|
@ -534,7 +542,7 @@ class HaikuRAGApp:
|
|||
search_input = query
|
||||
|
||||
async with HaikuRAG(
|
||||
db_path=self.db_path,
|
||||
db_path=self._client_db_path,
|
||||
config=self.config,
|
||||
read_only=True,
|
||||
) as self.client:
|
||||
|
|
@ -601,7 +609,7 @@ class HaikuRAGApp:
|
|||
images: Paths of images to attach to the question
|
||||
"""
|
||||
async with HaikuRAG(
|
||||
db_path=self.db_path,
|
||||
db_path=self._client_db_path,
|
||||
config=self.config,
|
||||
read_only=True,
|
||||
) as self.client:
|
||||
|
|
@ -634,7 +642,7 @@ class HaikuRAGApp:
|
|||
images: Paths of images to attach to the question
|
||||
"""
|
||||
async with HaikuRAG(
|
||||
db_path=self.db_path,
|
||||
db_path=self._client_db_path,
|
||||
config=self.config,
|
||||
read_only=True,
|
||||
) as self.client:
|
||||
|
|
@ -863,6 +871,10 @@ class HaikuRAGApp:
|
|||
f"[repr.attrib_name]chunk_id[/repr.attrib_name]: {result.chunk_id} "
|
||||
f"[repr.attrib_name]score[/repr.attrib_name]: {result.score:.4f}"
|
||||
)
|
||||
if result.source and len(self.config.lancedb.databases) > 1:
|
||||
self.console.print(
|
||||
f"[repr.attrib_name]database[/repr.attrib_name]: {result.source}"
|
||||
)
|
||||
if result.document_uri:
|
||||
self.console.print(
|
||||
f"[repr.attrib_name]document uri[/repr.attrib_name]: {result.document_uri}"
|
||||
|
|
|
|||
|
|
@ -70,11 +70,19 @@ 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:
|
||||
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.
|
||||
|
||||
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.
|
||||
"""
|
||||
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"
|
||||
|
||||
|
||||
|
|
@ -154,7 +162,7 @@ async def _first_holding(
|
|||
|
||||
@dataclass
|
||||
class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
|
||||
db_path: Path
|
||||
db_path: Path | None
|
||||
config: AppConfig
|
||||
state_type: type[StateT]
|
||||
state_namespace: str
|
||||
|
|
@ -396,7 +404,9 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
|
|||
self.rag = rag
|
||||
return self.rag
|
||||
|
||||
async def get_picture_bytes(self, document_id: str, self_ref: str) -> bytes | None:
|
||||
async def get_picture_bytes(
|
||||
self, document_id: str, self_ref: str, source: str | None = None
|
||||
) -> bytes | None:
|
||||
"""Fetch a picture of this capability's evidence, for whoever re-attaches it.
|
||||
|
||||
Public because compaction rehydrates cited pictures and this capability
|
||||
|
|
@ -404,9 +414,7 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
|
|||
"""
|
||||
async with self.rag_lock:
|
||||
rag = await self._ensure_rag()
|
||||
return await rag.document_item_repository.get_picture_bytes(
|
||||
document_id, self_ref
|
||||
)
|
||||
return await rag.get_picture_bytes(document_id, self_ref, source)
|
||||
|
||||
async def _close(self) -> None:
|
||||
if self.rag is not None:
|
||||
|
|
@ -521,14 +529,10 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
|
|||
if missing:
|
||||
async with self.rag_lock:
|
||||
rag = await self._ensure_rag()
|
||||
# A chunk id says nothing about which database holds it, so a
|
||||
# federating client looks through the ones it covers.
|
||||
if rag._federated:
|
||||
lookups = await rag.clients_for(
|
||||
getattr(self.state, "sources", None) or list(rag._federated)
|
||||
)
|
||||
else:
|
||||
lookups = [rag]
|
||||
# A chunk id says nothing about which database holds it, so the
|
||||
# fallback looks through everything the question covers — and
|
||||
# nothing it does not.
|
||||
lookups = await rag.clients_covering(self.state.sources)
|
||||
synthetic: list[SearchResult] = []
|
||||
documents: dict[tuple[str | None, str], Any] = {}
|
||||
for chunk_id in missing:
|
||||
|
|
|
|||
|
|
@ -76,7 +76,10 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
|
|||
self.sandbox = Sandbox(
|
||||
db_path=self.db_path,
|
||||
config=self.config,
|
||||
context=AnalysisContext(filter=self.state.document_filter),
|
||||
context=AnalysisContext(
|
||||
filter=self.state.document_filter,
|
||||
sources=self.state.sources,
|
||||
),
|
||||
rag=rag,
|
||||
lock=self.rag_lock,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -70,6 +70,7 @@ class RetainedPicture:
|
|||
document_id: str
|
||||
self_ref: str
|
||||
label: str
|
||||
source: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
|
@ -176,6 +177,7 @@ def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule:
|
|||
document_id=entry.citation.document_id,
|
||||
self_ref=self_ref,
|
||||
label=picture_label(entry.chunk_id, self_ref),
|
||||
source=entry.citation.source,
|
||||
)
|
||||
)
|
||||
return Capsule(text=ENTRY_SEPARATOR.join(lines), pictures=tuple(pictures))
|
||||
|
|
@ -406,7 +408,7 @@ class EvidenceCompactionCapability(AbstractCapability[Any]):
|
|||
owner = owners[retained.capability]
|
||||
try:
|
||||
data = await owner.get_picture_bytes(
|
||||
retained.document_id, retained.self_ref
|
||||
retained.document_id, retained.self_ref, retained.source
|
||||
)
|
||||
except Exception:
|
||||
# A read that fails costs this picture, not the answer.
|
||||
|
|
|
|||
|
|
@ -300,8 +300,8 @@ class ChatApp(App):
|
|||
continue
|
||||
blobs: list[bytes] = []
|
||||
for ref in refs:
|
||||
data = await self.client.document_item_repository.get_picture_bytes(
|
||||
citation.document_id, ref
|
||||
data = await self.client.get_picture_bytes(
|
||||
citation.document_id, ref, citation.source
|
||||
)
|
||||
if data:
|
||||
blobs.append(data)
|
||||
|
|
|
|||
|
|
@ -22,14 +22,17 @@ from haiku.rag.config import ( # noqa: E402
|
|||
)
|
||||
from haiku.rag.logging import configure_cli_logging # noqa: E402
|
||||
from haiku.rag.store.exceptions import ( # noqa: E402
|
||||
AmbiguousDatabaseError,
|
||||
MigrationRequiredError,
|
||||
ReadOnlyError,
|
||||
SourceUnavailableError,
|
||||
)
|
||||
from haiku.rag.store.models.chunk import SearchType # noqa: E402
|
||||
from haiku.rag.utils import is_up_to_date # noqa: E402
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from haiku.rag.app import HaikuRAGApp
|
||||
from haiku.rag.config.models import AppConfig
|
||||
|
||||
_cli = typer.Typer(
|
||||
context_settings={"help_option_names": ["-h", "--help"]},
|
||||
|
|
@ -41,29 +44,97 @@ _cli = typer.Typer(
|
|||
def cli():
|
||||
try:
|
||||
_cli()
|
||||
except (MigrationRequiredError, ReadOnlyError) as e:
|
||||
except (
|
||||
AmbiguousDatabaseError,
|
||||
MigrationRequiredError,
|
||||
ReadOnlyError,
|
||||
SourceUnavailableError,
|
||||
) as e:
|
||||
typer.echo(f"Error: {e}", err=True)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
# Module-level flags set by callback
|
||||
_read_only: bool = False
|
||||
_database: str | None = None
|
||||
_database_path: Path | None = None
|
||||
|
||||
|
||||
def create_app(db: Path | None = None) -> "HaikuRAGApp":
|
||||
def create_app(db: Path | None = None, *, federated: bool = False) -> "HaikuRAGApp":
|
||||
"""Create HaikuRAGApp with loaded config and resolved database path.
|
||||
|
||||
Args:
|
||||
db: Optional database path. If None, uses path from config.
|
||||
db: Optional database path. If None, uses `--database`, then the path
|
||||
from config.
|
||||
federated: Whether this command works across `lancedb.databases`.
|
||||
|
||||
Returns:
|
||||
HaikuRAGApp instance with proper config and db path.
|
||||
|
||||
Raises:
|
||||
AmbiguousDatabaseError: Several databases are configured and this
|
||||
command works on one, without `--db` or `--database` naming which.
|
||||
"""
|
||||
from haiku.rag.app import HaikuRAGApp
|
||||
|
||||
db_path = resolve_db_path(db, federated=federated)
|
||||
return HaikuRAGApp(
|
||||
db_path=db_path,
|
||||
config=get_config(),
|
||||
read_only=_read_only,
|
||||
federated=federated and db is None and _database is None,
|
||||
)
|
||||
|
||||
|
||||
def resolve_db_path(db: Path | None = None, *, federated: bool = False) -> Path:
|
||||
"""The database a command works on, from `--db`, `--database`, or config."""
|
||||
if db is not None and _database is not None:
|
||||
raise AmbiguousDatabaseError(
|
||||
"pass --db or --database, not both: they name the same thing"
|
||||
)
|
||||
require_one_database(get_config(), db, federated=federated)
|
||||
if db is not None:
|
||||
return db
|
||||
if _database_path is not None:
|
||||
return _database_path
|
||||
return get_config().storage.data_dir / "haiku.rag.lancedb"
|
||||
|
||||
|
||||
def require_one_database(
|
||||
config: "AppConfig", db: Path | None, *, federated: bool
|
||||
) -> None:
|
||||
"""Refuse a one-database command that cannot tell which one to use."""
|
||||
databases = config.lancedb.databases
|
||||
if federated or db is not None or not databases:
|
||||
return
|
||||
raise AmbiguousDatabaseError(
|
||||
f"lancedb.databases names {', '.join(sorted(databases))}; this command "
|
||||
"works on a single database: pass --database NAME, or --db PATH."
|
||||
)
|
||||
|
||||
|
||||
def select_database(name: str) -> Path | None:
|
||||
"""Point the configuration at one database from `lancedb.databases`.
|
||||
|
||||
Returns its local path, or None where it lives behind a URI. Rewriting the
|
||||
configuration is what lets every command, the TUIs included, work on the
|
||||
selected database without knowing the set exists.
|
||||
"""
|
||||
from haiku.rag.utils import locate_database
|
||||
|
||||
config = get_config()
|
||||
db_path = db if db else config.storage.data_dir / "haiku.rag.lancedb"
|
||||
return HaikuRAGApp(db_path=db_path, config=config, read_only=_read_only)
|
||||
databases = config.lancedb.databases
|
||||
if name not in databases:
|
||||
raise AmbiguousDatabaseError(
|
||||
f"unknown database {name!r}; lancedb.databases names "
|
||||
f"{', '.join(sorted(databases)) or 'nothing'}"
|
||||
)
|
||||
uri, db_path = locate_database(databases[name])
|
||||
selected = config.model_copy(deep=True)
|
||||
selected.lancedb.databases = {}
|
||||
selected.lancedb.uri = uri
|
||||
set_config(selected)
|
||||
return db_path
|
||||
|
||||
|
||||
async def check_version():
|
||||
|
|
@ -102,16 +173,27 @@ def main(
|
|||
"--read-only",
|
||||
help="Open database in read-only mode",
|
||||
),
|
||||
database: str | None = typer.Option(
|
||||
None,
|
||||
"--database",
|
||||
help="Name of a database from lancedb.databases to work on",
|
||||
),
|
||||
):
|
||||
"""haiku.rag CLI - Vector database RAG system"""
|
||||
global _read_only
|
||||
global _read_only, _database, _database_path
|
||||
_read_only = read_only
|
||||
_database = database
|
||||
_database_path = None
|
||||
# Load config from --config, local folder, or default directory
|
||||
config_path = find_config_file(cli_path=config)
|
||||
if config_path:
|
||||
yaml_data = load_yaml_config(config_path)
|
||||
loaded_config = AppConfig.model_validate(yaml_data)
|
||||
set_config(loaded_config)
|
||||
set_config(AppConfig.model_validate(yaml_data))
|
||||
else:
|
||||
set_config(AppConfig())
|
||||
|
||||
if database is not None:
|
||||
_database_path = select_database(database)
|
||||
|
||||
configure_cli_logging()
|
||||
|
||||
|
|
@ -307,7 +389,7 @@ def search(
|
|||
help="Path to the LanceDB database file",
|
||||
),
|
||||
):
|
||||
app = create_app(db)
|
||||
app = create_app(db, federated=True)
|
||||
asyncio.run(
|
||||
app.search(
|
||||
query=query,
|
||||
|
|
@ -361,7 +443,7 @@ def ask(
|
|||
help="Path to an image to attach to the question (repeatable; requires a vision-capable model)",
|
||||
),
|
||||
):
|
||||
app = create_app(db)
|
||||
app = create_app(db, federated=True)
|
||||
asyncio.run(
|
||||
app.ask(
|
||||
question=question,
|
||||
|
|
@ -393,7 +475,7 @@ def analyze(
|
|||
help="Path to an image to attach to the question (repeatable; requires a vision-capable model)",
|
||||
),
|
||||
):
|
||||
app = create_app(db)
|
||||
app = create_app(db, federated=True)
|
||||
asyncio.run(
|
||||
app.analyze(
|
||||
question=question,
|
||||
|
|
@ -749,8 +831,7 @@ def inspect(
|
|||
typer.echo(f"Error: {e}", err=True)
|
||||
raise typer.Exit(1) from e
|
||||
|
||||
db_path = db if db else get_config().storage.data_dir / "haiku.rag.lancedb"
|
||||
run_inspector(db_path, read_only=True)
|
||||
run_inspector(resolve_db_path(db), read_only=True)
|
||||
|
||||
|
||||
@_cli.command("chat", help="Launch interactive chat TUI for conversational RAG")
|
||||
|
|
@ -775,7 +856,7 @@ def chat(
|
|||
"""Launch the chat TUI for conversational RAG."""
|
||||
from haiku.rag.chat import run_chat
|
||||
|
||||
db_path = db if db else get_config().storage.data_dir / "haiku.rag.lancedb"
|
||||
db_path = resolve_db_path(db)
|
||||
capabilities = capability if capability else ["rag"]
|
||||
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -19,15 +19,22 @@ from haiku.rag.config import AppConfig, get_config
|
|||
from haiku.rag.converters import get_converter
|
||||
from haiku.rag.reranking import get_reranker
|
||||
from haiku.rag.store.engine import Store
|
||||
from haiku.rag.store.exceptions import SourceUnavailableError
|
||||
from haiku.rag.store.exceptions import (
|
||||
MigrationRequiredError,
|
||||
ReadOnlyError,
|
||||
SourceUnavailableError,
|
||||
)
|
||||
from haiku.rag.store.models.chunk import Chunk, SearchResult, SearchType
|
||||
from haiku.rag.store.models.document import Document
|
||||
from haiku.rag.store.models.document_item import extract_items
|
||||
from haiku.rag.store.repositories.chunk import ChunkRepository
|
||||
from haiku.rag.store.repositories.document import DocumentRepository
|
||||
from haiku.rag.store.repositories.document_item import DocumentItemRepository
|
||||
from haiku.rag.store.repositories.settings import SettingsRepository
|
||||
from haiku.rag.utils import escape_sql_string
|
||||
from haiku.rag.store.repositories.settings import (
|
||||
ConfigMismatchError,
|
||||
SettingsRepository,
|
||||
)
|
||||
from haiku.rag.utils import escape_sql_string, locate_database
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
|
|
@ -49,6 +56,20 @@ logger = logging.getLogger(__name__)
|
|||
_VACUUM_MIN_INTERVAL_S = 300.0
|
||||
|
||||
|
||||
# Failures whose message names the remedy and never the location, so the failing
|
||||
# database is named alongside it instead of in place of it.
|
||||
_NAMEABLE_FAILURES = (MigrationRequiredError, ConfigMismatchError, ReadOnlyError)
|
||||
|
||||
|
||||
def _without_repeats(names: list[str]) -> list[str]:
|
||||
"""`names` in order, without repeats.
|
||||
|
||||
A database named twice would be searched twice and fused as two rank lists,
|
||||
which counts it double.
|
||||
"""
|
||||
return list(dict.fromkeys(names))
|
||||
|
||||
|
||||
class RebuildMode(Enum):
|
||||
"""Mode for rebuilding the database."""
|
||||
|
||||
|
|
@ -84,10 +105,11 @@ class HaikuRAG:
|
|||
create: Whether to create the database if it doesn't exist.
|
||||
read_only: Whether to open the database in read-only mode.
|
||||
sources: Names from ``config.lancedb.databases`` this client covers.
|
||||
None means all of them. Ignored when a single ``uri`` is
|
||||
configured.
|
||||
None means all of them. Ignored when a single ``uri`` or an
|
||||
explicit ``db_path`` is given.
|
||||
"""
|
||||
self._config = config if config is not None else get_config()
|
||||
self._db_path_given = db_path is not None
|
||||
if db_path is None:
|
||||
db_path = self._config.storage.data_dir / "haiku.rag.lancedb"
|
||||
|
||||
|
|
@ -123,23 +145,14 @@ class HaikuRAG:
|
|||
"""
|
||||
return get_reranker(config=self._config)
|
||||
|
||||
@staticmethod
|
||||
def _locate(location: str) -> tuple[str, "Path | None"]:
|
||||
"""Split a configured location into (uri, db_path).
|
||||
|
||||
A value with a scheme is a `lancedb.uri`; anything else is a local path.
|
||||
Routing a local path through `uri` would have `ConnectionMode` classify it
|
||||
as object storage, which opens it without the existence check a local
|
||||
database gets.
|
||||
"""
|
||||
if "://" in location:
|
||||
return location, None
|
||||
return "", Path(location)
|
||||
|
||||
def _selected(self) -> dict[str, str]:
|
||||
"""The configured databases this client covers, name to location."""
|
||||
"""The configured databases this client covers, name to location.
|
||||
|
||||
Empty when the caller named a database itself: an explicit `db_path` says
|
||||
which one to open, so it is not overridden by a configured set.
|
||||
"""
|
||||
declared = self._config.lancedb.databases
|
||||
if not declared:
|
||||
if not declared or self._db_path_given:
|
||||
return {}
|
||||
if self._requested_sources is not None and not self._requested_sources:
|
||||
raise ValueError(
|
||||
|
|
@ -172,7 +185,7 @@ class HaikuRAG:
|
|||
return self
|
||||
if selected:
|
||||
[(self._source, location)] = selected.items()
|
||||
uri, db_path = self._locate(location)
|
||||
uri, db_path = locate_database(location)
|
||||
self._config = self._config.model_copy(deep=True)
|
||||
self._config.lancedb.databases = {}
|
||||
self._config.lancedb.uri = uri
|
||||
|
|
@ -196,6 +209,13 @@ class HaikuRAG:
|
|||
except BaseException:
|
||||
self.store.close()
|
||||
raise
|
||||
except _NAMEABLE_FAILURES as error:
|
||||
# These say what to run and never where the database is, so the name
|
||||
# is added to the message rather than replacing it: the operator needs
|
||||
# both which database failed and what to do about it.
|
||||
if self._source is None:
|
||||
raise
|
||||
raise type(error)(f"database {self._source!r}: {error}") from error
|
||||
except Exception as error:
|
||||
# A legacy `uri` or `db_path` client has no name to report instead, so
|
||||
# its error passes through as it always has.
|
||||
|
|
@ -221,7 +241,11 @@ class HaikuRAG:
|
|||
Opening is per query rather than at entry: a set of 25 configured
|
||||
databases is typically queried a few at a time, and a database nobody
|
||||
asked for must not be able to fail a query, or be opened for nothing.
|
||||
|
||||
Missing ones open together: on object storage a serial loop makes the
|
||||
first query cost the sum of the opens.
|
||||
"""
|
||||
names = _without_repeats(names)
|
||||
unknown = [n for n in names if n not in self._federated]
|
||||
if unknown:
|
||||
raise KeyError(
|
||||
|
|
@ -229,15 +253,27 @@ class HaikuRAG:
|
|||
f"{', '.join(sorted(self._federated))}"
|
||||
)
|
||||
async with self._clients_lock:
|
||||
for name in names:
|
||||
if name not in self._clients:
|
||||
self._clients[name] = await self._open_client(
|
||||
name, self._federated[name]
|
||||
)
|
||||
missing = [n for n in names if n not in self._clients]
|
||||
if missing:
|
||||
opened = await asyncio.gather(
|
||||
*(self._open_client(n, self._federated[n]) for n in missing),
|
||||
return_exceptions=True,
|
||||
)
|
||||
# Whatever opened is tracked before the failure is reported, so
|
||||
# `__aexit__` closes it: `gather` does not cancel the siblings of
|
||||
# the one that raised, and an untracked connection leaks.
|
||||
failure: BaseException | None = None
|
||||
for name, result in zip(missing, opened, strict=True):
|
||||
if isinstance(result, BaseException):
|
||||
failure = failure or result
|
||||
else:
|
||||
self._clients[name] = result
|
||||
if failure is not None:
|
||||
raise failure
|
||||
return [self._clients[n] for n in names]
|
||||
|
||||
async def _open_client(self, name: str, location: str) -> "HaikuRAG":
|
||||
uri, db_path = self._locate(location)
|
||||
uri, db_path = locate_database(location)
|
||||
config = self._config.model_copy(deep=True)
|
||||
config.lancedb.databases = {}
|
||||
config.lancedb.uri = uri
|
||||
|
|
@ -483,6 +519,32 @@ class HaikuRAG:
|
|||
"""
|
||||
return await self.chunk_repository.get_by_id(chunk_id)
|
||||
|
||||
async def get_picture_bytes(
|
||||
self, document_id: str, self_ref: str, source: str | None = None
|
||||
) -> bytes | None:
|
||||
"""Get a picture's bytes, from the database named by `source`.
|
||||
|
||||
Args:
|
||||
document_id: The document holding the picture.
|
||||
self_ref: The picture's `self_ref`.
|
||||
source: The database it came from. Required when federating.
|
||||
|
||||
Returns:
|
||||
The picture bytes if found, None otherwise.
|
||||
"""
|
||||
if not self._federated:
|
||||
return await self.document_item_repository.get_picture_bytes(
|
||||
document_id, self_ref
|
||||
)
|
||||
if source is None:
|
||||
raise ValueError(
|
||||
"a picture lookup across databases needs the source it came from"
|
||||
)
|
||||
(owner,) = await self.clients_for([source])
|
||||
return await owner.document_item_repository.get_picture_bytes(
|
||||
document_id, self_ref
|
||||
)
|
||||
|
||||
async def get_document_by_uri(self, uri: str) -> Document | None:
|
||||
"""Get a document by its URI.
|
||||
|
||||
|
|
@ -593,6 +655,32 @@ class HaikuRAG:
|
|||
"""
|
||||
return await self.document_repository.count(filter=filter)
|
||||
|
||||
async def clients_covering(
|
||||
self, sources: list[str] | None = None
|
||||
) -> list["HaikuRAG"]:
|
||||
"""The clients covering this selection.
|
||||
|
||||
The named subset for a client covering a set, or this one where it covers
|
||||
a single database. Empty for a selection of none, which is not the same as
|
||||
`None` for all of them. Every read honouring `sources` decides through
|
||||
this, so the rule cannot differ between one operation and another.
|
||||
"""
|
||||
if self._federated:
|
||||
return await self.clients_for(
|
||||
list(self._federated) if sources is None else sources
|
||||
)
|
||||
if sources is None:
|
||||
return [self]
|
||||
sources = _without_repeats(sources)
|
||||
if not sources:
|
||||
return []
|
||||
if sources != [self._source]:
|
||||
raise KeyError(
|
||||
f"unknown database(s) {', '.join(sources) or '(none)'}; this "
|
||||
f"client covers {self._source or 'a single unnamed database'}"
|
||||
)
|
||||
return [self]
|
||||
|
||||
async def search(
|
||||
self,
|
||||
query: "str | bytes | PILImage.Image",
|
||||
|
|
@ -608,13 +696,8 @@ class HaikuRAG:
|
|||
return await search_sources(
|
||||
self, query, limit, search_type, filter, include_images, sources
|
||||
)
|
||||
if sources is not None and not sources:
|
||||
if not await self.clients_covering(sources):
|
||||
return []
|
||||
if sources is not None and sources != [self._source]:
|
||||
raise KeyError(
|
||||
f"unknown database(s) {', '.join(sources) or '(none)'}; this "
|
||||
f"client covers {self._source or 'a single unnamed database'}"
|
||||
)
|
||||
results = await search(self, query, limit, search_type, filter, include_images)
|
||||
# A database named in config keeps its name even when it is the only one
|
||||
# this client covers. Only a legacy single `uri` leaves source unset.
|
||||
|
|
@ -646,10 +729,11 @@ class HaikuRAG:
|
|||
question: str,
|
||||
filter: str | None = None,
|
||||
images: Sequence[bytes] | None = None,
|
||||
sources: list[str] | None = None,
|
||||
) -> "AnalysisResult":
|
||||
from haiku.rag.client.agents import analyze
|
||||
|
||||
return await analyze(self, question, filter, images)
|
||||
return await analyze(self, question, filter, images, sources)
|
||||
|
||||
async def visualize_chunk(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -50,6 +50,8 @@ async def ask(
|
|||
filter: SQL WHERE clause to filter documents.
|
||||
images: Raw image bytes attached to the question (requires a
|
||||
vision-capable QA model).
|
||||
sources: Names of the databases to ask across. None asks across every
|
||||
configured database.
|
||||
|
||||
Returns:
|
||||
Tuple of (answer text, list of resolved citations).
|
||||
|
|
@ -97,6 +99,7 @@ async def analyze(
|
|||
question: str,
|
||||
filter: str | None = None,
|
||||
images: Sequence[bytes] | None = None,
|
||||
sources: list[str] | None = None,
|
||||
) -> "AnalysisResult":
|
||||
"""Answer a question using the analysis capability.
|
||||
|
||||
|
|
@ -110,6 +113,8 @@ async def analyze(
|
|||
filter: SQL WHERE clause to filter documents during searches.
|
||||
images: Raw image bytes attached to the question (requires a
|
||||
vision-capable analysis model).
|
||||
sources: Names of the databases to analyze across. None covers every
|
||||
configured database.
|
||||
|
||||
Returns:
|
||||
AnalysisResult with the answer and resolved citations.
|
||||
|
|
@ -119,14 +124,16 @@ async def analyze(
|
|||
from haiku.rag.utils import get_model
|
||||
|
||||
capability = create_capability(
|
||||
db_path=client.store.db_path,
|
||||
db_path=None if client._federated else client.store.db_path,
|
||||
config=client._config,
|
||||
rag=client,
|
||||
defer_loading=False,
|
||||
)
|
||||
deps = _AgentDeps(
|
||||
state={
|
||||
"analysis": AnalysisState(document_filter=filter).model_dump(mode="json")
|
||||
"analysis": AnalysisState(
|
||||
document_filter=filter, sources=sources
|
||||
).model_dump(mode="json")
|
||||
}
|
||||
)
|
||||
model_config = client._config.analysis.model or client._config.qa.model
|
||||
|
|
|
|||
|
|
@ -269,6 +269,9 @@ async def _attach_picture_data(client: "HaikuRAG", chunks: list[Chunk]) -> None:
|
|||
def _dedup_picture_chunks(results: list[SearchResult]) -> list[SearchResult]:
|
||||
"""Collapse duplicate picture-only chunks to one result per ``self_ref``.
|
||||
|
||||
Keyed by database as well, since a database copied from another holds the same
|
||||
document id: collapsing across them would drop one of two real results.
|
||||
|
||||
A single picture can produce two chunks for the same self_ref: one whose
|
||||
vector is the text embedding of the picture's description, and one whose
|
||||
vector is the image embedding of the picture's bytes. Both can rank for
|
||||
|
|
@ -276,13 +279,13 @@ def _dedup_picture_chunks(results: list[SearchResult]) -> list[SearchResult]:
|
|||
their only ref, keep the higher-scoring one. Wider chunks that span the
|
||||
picture plus surrounding items pass through untouched.
|
||||
"""
|
||||
seen: dict[tuple[str | None, str], int] = {}
|
||||
seen: dict[tuple[str | None, str | None, str], int] = {}
|
||||
keep: list[bool] = [True] * len(results)
|
||||
for i, r in enumerate(results):
|
||||
if len(r.doc_item_refs) == 1 and r.doc_item_refs[0].startswith(
|
||||
PICTURE_REF_PREFIX
|
||||
):
|
||||
key = (r.document_id, r.doc_item_refs[0])
|
||||
key = (r.source, r.document_id, r.doc_item_refs[0])
|
||||
prior = seen.get(key)
|
||||
if prior is None:
|
||||
seen[key] = i
|
||||
|
|
@ -397,7 +400,25 @@ async def expand_context(
|
|||
)
|
||||
)
|
||||
merged = unsourced + [r for group in expanded_groups for r in group]
|
||||
merged.sort(key=lambda r: r.score, reverse=True)
|
||||
# Grouping by database must not become the tiebreak: fused scores tie
|
||||
# often, so equal scores keep the order they were fused in.
|
||||
arrival = {
|
||||
result.chunk_id: rank
|
||||
for rank, result in enumerate(search_results)
|
||||
if result.chunk_id
|
||||
}
|
||||
|
||||
def fused_rank(result: SearchResult) -> int:
|
||||
return min(
|
||||
(
|
||||
arrival[cid]
|
||||
for cid in (result.chunk_id, *result.chunk_ids)
|
||||
if cid in arrival
|
||||
),
|
||||
default=len(arrival),
|
||||
)
|
||||
|
||||
merged.sort(key=lambda r: (-r.score, fused_rank(r)))
|
||||
return merged
|
||||
|
||||
from haiku.rag.context import expand_with_items, window_for
|
||||
|
|
|
|||
|
|
@ -405,6 +405,7 @@ def _build_result(
|
|||
return SearchResult(
|
||||
content=expanded_content,
|
||||
score=max(r.score for r in original_results),
|
||||
source=first.source,
|
||||
chunk_id=first.chunk_id,
|
||||
chunk_ids=chunk_ids,
|
||||
chunk_meta=first.chunk_meta,
|
||||
|
|
|
|||
|
|
@ -6,3 +6,4 @@ class AnalysisContext:
|
|||
"""Mutable context accumulating data during analysis execution."""
|
||||
|
||||
filter: str | None = None
|
||||
sources: list[str] | None = None
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import os
|
|||
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
|
||||
|
||||
|
|
@ -137,10 +138,11 @@ class Sandbox:
|
|||
each read opens an ephemeral read-only connection.
|
||||
"""
|
||||
|
||||
_db_path: Path
|
||||
_db_path: Path | None
|
||||
_config: AppConfig
|
||||
_context: AnalysisContext
|
||||
_rag: "HaikuRAG | None"
|
||||
_owners: dict[str, "HaikuRAG"]
|
||||
_lock: "asyncio.Lock | None"
|
||||
_search_results: "list[SearchResult]"
|
||||
_doc_items: dict[str, list["DocumentItem"]]
|
||||
|
|
@ -155,7 +157,7 @@ class Sandbox:
|
|||
|
||||
def __init__(
|
||||
self,
|
||||
db_path: Path,
|
||||
db_path: Path | None,
|
||||
config: AppConfig,
|
||||
context: AnalysisContext,
|
||||
rag: "HaikuRAG | None" = None,
|
||||
|
|
@ -165,6 +167,7 @@ class Sandbox:
|
|||
self._config = config
|
||||
self._context = context
|
||||
self._rag = rag
|
||||
self._owners = {}
|
||||
self._lock = lock
|
||||
self._search_results = []
|
||||
self._doc_items = {}
|
||||
|
|
@ -178,22 +181,82 @@ class Sandbox:
|
|||
self._deadline = None
|
||||
|
||||
@asynccontextmanager
|
||||
async def _connection(self) -> "AsyncIterator[HaikuRAG]":
|
||||
async def _connection(
|
||||
self, owner: "HaikuRAG | None" = None
|
||||
) -> "AsyncIterator[HaikuRAG]":
|
||||
"""Yield the shared connection (serialized by the lock), or an ephemeral
|
||||
read-only one. The lock guards the whole block so a read's awaits cannot
|
||||
interleave with another task's operation on the same connection."""
|
||||
if self._rag is not None:
|
||||
interleave with another task's operation on the same connection.
|
||||
|
||||
`owner` is the client holding one document, for the reads addressed to a
|
||||
single document. The shared connection covers a set of databases and has
|
||||
no repositories of its own, so those reads have to name their owner.
|
||||
"""
|
||||
connection = owner if owner is not None else self._rag
|
||||
if connection is not None:
|
||||
if self._lock is not None:
|
||||
async with self._lock:
|
||||
yield self._rag
|
||||
yield connection
|
||||
else:
|
||||
yield self._rag
|
||||
yield connection
|
||||
return
|
||||
from haiku.rag.client import HaikuRAG
|
||||
|
||||
async with HaikuRAG(self._db_path, config=self._config, read_only=True) as rag:
|
||||
yield rag
|
||||
|
||||
async def _documents(self) -> "tuple[list[Any], dict[str, HaikuRAG]]":
|
||||
"""Every document in scope, and the client holding each of them.
|
||||
|
||||
The owners are empty where one connection serves every read: a single
|
||||
database, or the ephemeral connection opened per read when no client was
|
||||
supplied. The selection is resolved the same way a search resolves it, so
|
||||
a database the question excluded cannot be mounted.
|
||||
"""
|
||||
async with self._connection() as rag:
|
||||
if not rag._federated:
|
||||
if not await rag.clients_covering(self._context.sources):
|
||||
return [], {}
|
||||
docs = await rag.list_documents(filter=self._context.filter)
|
||||
return docs, {}
|
||||
owners = await rag.clients_covering(self._context.sources)
|
||||
groups = await asyncio.gather(
|
||||
*(owner.list_documents(filter=self._context.filter) for owner in owners)
|
||||
)
|
||||
# Interleaved, not concatenated: code that prints the listing is read
|
||||
# through a truncated output, and concatenating shows one database's
|
||||
# documents until the truncation, hiding that there are others.
|
||||
docs = [doc for row in zip_longest(*groups) for doc in row if doc is not None]
|
||||
return docs, self._holders(owners, groups)
|
||||
|
||||
@staticmethod
|
||||
def _holders(
|
||||
owners: "list[HaikuRAG]", groups: "list[list[Any]]"
|
||||
) -> "dict[str, HaikuRAG]":
|
||||
"""Map each document id to the database holding it.
|
||||
|
||||
Document ids are UUID4, so the flat `/documents/{id}/` namespace is
|
||||
unambiguous for databases that were filled independently — but not for one
|
||||
copied from another, where the same id is in both. Duplicate results are
|
||||
merely redundant in a search; here they would be two documents claiming one
|
||||
path, and whichever arrived last would answer for both. Rejected rather
|
||||
than resolved, since either answer would be wrong half the time.
|
||||
"""
|
||||
holders: dict[str, HaikuRAG] = {}
|
||||
held_by: dict[str, str | None] = {}
|
||||
for owner, group in zip(owners, groups, strict=True):
|
||||
for doc in group:
|
||||
if not doc.id: # pragma: no cover - stored rows always carry an id
|
||||
continue
|
||||
if doc.id in holders:
|
||||
raise ValueError(
|
||||
f"document {doc.id} is in databases {held_by[doc.id]!r} and "
|
||||
f"{owner._source!r}; analysis mounts one document per id"
|
||||
)
|
||||
holders[doc.id] = owner
|
||||
held_by[doc.id] = owner._source
|
||||
return holders
|
||||
|
||||
def _run_on_loop(self, coro: Coroutine[Any, Any, Any]) -> Any:
|
||||
"""Run a coroutine on the execute() loop from a synchronous callback.
|
||||
|
||||
|
|
@ -252,7 +315,12 @@ class Sandbox:
|
|||
# gets figures through the top-level `search` tool when the
|
||||
# question is visual; in-code search is for structural work.
|
||||
async with self._connection() as rag:
|
||||
results = await rag.search(query, limit=limit, filter=context.filter)
|
||||
results = await rag.search(
|
||||
query,
|
||||
limit=limit,
|
||||
filter=context.filter,
|
||||
sources=context.sources,
|
||||
)
|
||||
expanded = await rag.expand_context(results)
|
||||
self._search_results.extend(expanded)
|
||||
out: list[dict[str, Any]] = []
|
||||
|
|
@ -264,6 +332,7 @@ class Sandbox:
|
|||
{
|
||||
"chunk_id": r.chunk_id,
|
||||
"content": r.content,
|
||||
"source": r.source,
|
||||
"document_id": r.document_id,
|
||||
"document_title": r.document_title,
|
||||
"document_uri": r.document_uri,
|
||||
|
|
@ -278,14 +347,14 @@ class Sandbox:
|
|||
return out
|
||||
|
||||
async def list_documents() -> list[dict[str, Any]]:
|
||||
async with self._connection() as rag:
|
||||
docs = await rag.list_documents(filter=context.filter)
|
||||
docs, owners = await self._documents()
|
||||
return [
|
||||
{
|
||||
"id": d.id,
|
||||
"title": d.title,
|
||||
"uri": d.uri,
|
||||
"created_at": str(d.created_at),
|
||||
"source": owners[d.id]._source if d.id in owners else None,
|
||||
}
|
||||
for d in docs
|
||||
]
|
||||
|
|
@ -309,8 +378,7 @@ class Sandbox:
|
|||
def _deny_write(_path: "PurePosixPath", _content: str | bytes) -> None:
|
||||
raise PermissionError(f"Document files are read-only: {_path}")
|
||||
|
||||
async with self._connection() as rag:
|
||||
docs = await rag.list_documents(filter=self._context.filter)
|
||||
docs, self._owners = await self._documents()
|
||||
|
||||
doc_titles = {doc.id: doc.title for doc in docs if doc.id}
|
||||
|
||||
|
|
@ -323,7 +391,7 @@ class Sandbox:
|
|||
return cached
|
||||
|
||||
async def _fetch() -> list[DocumentItem]:
|
||||
async with sandbox._connection() as rag:
|
||||
async with sandbox._connection(sandbox._owners.get(did)) as rag:
|
||||
return await rag.document_item_repository.get_all_items(did)
|
||||
|
||||
items = sandbox._run_on_loop(_fetch())
|
||||
|
|
@ -337,7 +405,7 @@ class Sandbox:
|
|||
return cached
|
||||
|
||||
async def _fetch() -> dict[str, list[str]]:
|
||||
async with sandbox._connection() as rag:
|
||||
async with sandbox._connection(sandbox._owners.get(did)) as rag:
|
||||
index = (
|
||||
await rag.chunk_repository.get_chunk_ids_by_self_ref_grouped(
|
||||
[did]
|
||||
|
|
@ -430,7 +498,7 @@ class Sandbox:
|
|||
) -> Callable[["PurePosixPath"], str]:
|
||||
def read_content(_path: "PurePosixPath") -> str:
|
||||
async def _fetch() -> str:
|
||||
async with sandbox._connection() as rag:
|
||||
async with sandbox._connection(sandbox._owners.get(did)) as rag:
|
||||
content = await rag.document_repository.get_content(did)
|
||||
return content or ""
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
from .exceptions import (
|
||||
AmbiguousDatabaseError,
|
||||
MigrationRequiredError,
|
||||
ReadOnlyError,
|
||||
SourceUnavailableError,
|
||||
|
|
@ -10,5 +11,6 @@ __all__ = [
|
|||
"Document",
|
||||
"MigrationRequiredError",
|
||||
"ReadOnlyError",
|
||||
"AmbiguousDatabaseError",
|
||||
"SourceUnavailableError",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -10,6 +10,10 @@ class MigrationRequiredError(Exception):
|
|||
pass
|
||||
|
||||
|
||||
class AmbiguousDatabaseError(Exception):
|
||||
"""A command that works on one database was run against a configured set."""
|
||||
|
||||
|
||||
class SourceUnavailableError(Exception):
|
||||
"""A configured database could not be opened.
|
||||
|
||||
|
|
|
|||
|
|
@ -137,9 +137,10 @@ class SearchResult(BaseModel):
|
|||
include the metadata of any other chunks merged with it. Never part of
|
||||
``format_for_agent`` output.
|
||||
|
||||
``source`` names the configured database a result came from, and is None when
|
||||
only one is configured. It is the name from ``lancedb.databases``, never a
|
||||
path or URI, so a location cannot travel in a result, a citation or a log.
|
||||
``source`` names the configured database a result came from: the name from
|
||||
``lancedb.databases``, never a path or URI, so a location cannot travel in a
|
||||
result, a citation or a log. It is None only where no database is named, as
|
||||
with the single ``lancedb.uri``.
|
||||
"""
|
||||
|
||||
content: str
|
||||
|
|
@ -198,6 +199,9 @@ class SearchResult(BaseModel):
|
|||
Produces a structured format with metadata that helps LLMs understand
|
||||
the source and nature of the content. When rank is provided, shows
|
||||
position instead of raw score to avoid confusing LLMs with low RRF scores.
|
||||
|
||||
The database is named only where one is named at all, so a single
|
||||
unnamed database renders exactly as before.
|
||||
"""
|
||||
if rank is not None and total is not None:
|
||||
parts = [f"[{self.chunk_id}] [rank {rank} of {total}]"]
|
||||
|
|
@ -206,6 +210,9 @@ class SearchResult(BaseModel):
|
|||
else:
|
||||
parts = [f"[{self.chunk_id}] (score: {self.score:.2f})"]
|
||||
|
||||
if self.source:
|
||||
parts.append(f"Database: {self.source}")
|
||||
|
||||
# Document source info
|
||||
source_parts = []
|
||||
if self.document_title:
|
||||
|
|
|
|||
|
|
@ -16,15 +16,15 @@ class Citation(BaseModel):
|
|||
|
||||
``picture_refs`` lists the ``self_ref`` values of picture items in the
|
||||
cited chunk. Empty for text-only citations. UIs can fetch the picture
|
||||
bytes via ``DocumentItemRepository.get_picture_bytes(document_id, ref)``
|
||||
and render them alongside the text content.
|
||||
bytes via ``HaikuRAG.get_picture_bytes(document_id, ref, source)`` and
|
||||
render them alongside the text content.
|
||||
|
||||
``chunk_ids`` lists the ids of all chunks whose expansion ranges merged
|
||||
into the cited result (always includes ``chunk_id``).
|
||||
|
||||
``source`` names the configured database the cited chunk came from, and is
|
||||
None when only one is configured. It is the name from ``lancedb.databases``,
|
||||
never a path or URI.
|
||||
``source`` names the configured database the cited chunk came from: the name
|
||||
from ``lancedb.databases``, never a path or URI. It is None only where no
|
||||
database is named, as with the single ``lancedb.uri``.
|
||||
|
||||
``doc_item_refs`` are the ``self_ref`` values of every item in the cited
|
||||
content — the exact items the model saw. Visual grounding resolves bounding
|
||||
|
|
|
|||
|
|
@ -394,6 +394,8 @@ async def format_citations_rich(
|
|||
idx = c.index if c.index is not None else (i + 1)
|
||||
|
||||
header_parts: list[str] = [f"[{idx}] {_citation_label(c)}"]
|
||||
if c.source and client is not None and client._federated:
|
||||
header_parts.append(c.source)
|
||||
pages = _citation_pages(c)
|
||||
if pages:
|
||||
header_parts.append(pages)
|
||||
|
|
@ -404,7 +406,9 @@ async def format_citations_rich(
|
|||
|
||||
body: list[RenderableType] = []
|
||||
for ref in c.picture_refs:
|
||||
image_renderable = await _render_picture(client, c.document_id, ref)
|
||||
image_renderable = await _render_picture(
|
||||
client, c.document_id, ref, c.source
|
||||
)
|
||||
body.append(
|
||||
image_renderable
|
||||
if image_renderable
|
||||
|
|
@ -436,7 +440,7 @@ async def format_citations_rich(
|
|||
|
||||
|
||||
async def _render_picture(
|
||||
client: "HaikuRAG | None", document_id: str, ref: str
|
||||
client: "HaikuRAG | None", document_id: str, ref: str, source: str | None = None
|
||||
) -> "RenderableType | None":
|
||||
"""Fetch a picture and return a Rich renderable, or None on failure/no client."""
|
||||
if client is None:
|
||||
|
|
@ -446,7 +450,7 @@ async def _render_picture(
|
|||
from PIL import Image as PILImage
|
||||
from textual_image.renderable import Image as RichImage
|
||||
|
||||
data = await client.document_item_repository.get_picture_bytes(document_id, ref)
|
||||
data = await client.get_picture_bytes(document_id, ref, source)
|
||||
if not data:
|
||||
return None
|
||||
try:
|
||||
|
|
@ -472,6 +476,19 @@ def raise_missing_extra(module: str, extra: str, exc: ModuleNotFoundError) -> No
|
|||
) from exc
|
||||
|
||||
|
||||
def locate_database(location: str) -> tuple[str, Path | None]:
|
||||
"""Split a configured location into (uri, db_path).
|
||||
|
||||
A value with a scheme is a `lancedb.uri`; anything else is a local path.
|
||||
Routing a local path through `uri` would have `ConnectionMode` classify it as
|
||||
object storage, which opens it without the existence check a local database
|
||||
gets.
|
||||
"""
|
||||
if "://" in location:
|
||||
return location, None
|
||||
return "", Path(location)
|
||||
|
||||
|
||||
def get_default_data_dir() -> Path:
|
||||
"""Get the user data directory for the current system platform.
|
||||
|
||||
|
|
|
|||
|
|
@ -141,6 +141,20 @@ def test_domain_preamble_is_added_to_capability_instructions(temp_db_path):
|
|||
)
|
||||
|
||||
|
||||
def _single_database_client() -> AsyncMock:
|
||||
"""A stand-in for a client covering one unnamed database.
|
||||
|
||||
A bare AsyncMock answers every attribute with a truthy Mock, so `_federated`
|
||||
would read as a set of databases, `_source` would reach a validated field, and
|
||||
`clients_covering` would return a Mock where the code iterates clients.
|
||||
"""
|
||||
client = AsyncMock()
|
||||
client._federated = {}
|
||||
client._source = None
|
||||
client.clients_covering.return_value = [client]
|
||||
return client
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("factory", "agent_instructions", "heading"),
|
||||
|
|
@ -315,11 +329,7 @@ async def test_capability_isolated_per_run_and_round_trips_state(temp_db_path):
|
|||
@pytest.mark.asyncio
|
||||
async def test_run_error_closes_resources_and_propagates(temp_db_path):
|
||||
capability = create_rag(db_path=temp_db_path, config=AppConfig())
|
||||
client = AsyncMock()
|
||||
# Stands in for a single-database client: a bare AsyncMock's auto
|
||||
# attributes are truthy Mocks, and `_source` reaches a validated field.
|
||||
client._federated = {}
|
||||
client._source = None
|
||||
client = _single_database_client()
|
||||
capability.rag = client
|
||||
error = RuntimeError("model failed")
|
||||
|
||||
|
|
@ -388,11 +398,7 @@ async def test_a_narrower_repeat_keeps_what_the_wider_search_returned(temp_db_pa
|
|||
async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db_path):
|
||||
capability = create_rag(db_path=temp_db_path, config=AppConfig())
|
||||
capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
|
||||
client = AsyncMock()
|
||||
# Stands in for a single-database client: a bare AsyncMock's auto
|
||||
# attributes are truthy Mocks, and `_source` reaches a validated field.
|
||||
client._federated = {}
|
||||
client._source = None
|
||||
client = _single_database_client()
|
||||
client.get_chunk_by_id.side_effect = [
|
||||
Chunk(id="chunk-1", document_id="doc-1", content="first"),
|
||||
Chunk(id="chunk-2", document_id="doc-1", content="second"),
|
||||
|
|
@ -418,11 +424,7 @@ async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db
|
|||
async def test_cite_reports_unresolved_ids_on_partial_success(temp_db_path):
|
||||
capability = create_rag(db_path=temp_db_path, config=AppConfig())
|
||||
capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
|
||||
client = AsyncMock()
|
||||
# Stands in for a single-database client: a bare AsyncMock's auto
|
||||
# attributes are truthy Mocks, and `_source` reaches a validated field.
|
||||
client._federated = {}
|
||||
client._source = None
|
||||
client = _single_database_client()
|
||||
client.get_chunk_by_id.side_effect = [
|
||||
Chunk(id="chunk-1", document_id="doc-1", content="first"),
|
||||
None,
|
||||
|
|
@ -466,11 +468,7 @@ async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path):
|
|||
]
|
||||
},
|
||||
)
|
||||
client = AsyncMock()
|
||||
# Stands in for a single-database client: a bare AsyncMock's auto
|
||||
# attributes are truthy Mocks, and `_source` reaches a validated field.
|
||||
client._federated = {}
|
||||
client._source = None
|
||||
client = _single_database_client()
|
||||
client.get_chunk_by_id.return_value = None
|
||||
capability.rag = client
|
||||
|
||||
|
|
@ -1229,11 +1227,7 @@ async def test_citing_without_searching_grounds_the_question(temp_db_path):
|
|||
async def model(_messages, _info):
|
||||
return ModelResponse(parts=next(calls))
|
||||
|
||||
client = AsyncMock()
|
||||
# Stands in for a single-database client: a bare AsyncMock's auto
|
||||
# attributes are truthy Mocks, and `_source` reaches a validated field.
|
||||
client._federated = {}
|
||||
client._source = None
|
||||
client = _single_database_client()
|
||||
client.get_chunk_by_id.return_value = Chunk(
|
||||
id="chunk-1", document_id="doc-1", content="evidence"
|
||||
)
|
||||
|
|
@ -1335,19 +1329,13 @@ async def test_a_capability_fetches_its_own_evidences_pictures(temp_db_path):
|
|||
"""Compaction rehydrates through the owner, which already holds the connection."""
|
||||
capability = create_rag(db_path=temp_db_path, config=AppConfig())
|
||||
client = AsyncMock()
|
||||
# Stands in for a single-database client: a bare AsyncMock's auto
|
||||
# attributes are truthy Mocks, and `_source` reaches a validated field.
|
||||
client._federated = {}
|
||||
client._source = None
|
||||
client.document_item_repository.get_picture_bytes.return_value = b"picture-bytes"
|
||||
client.get_picture_bytes.return_value = b"picture-bytes"
|
||||
capability.rag = client
|
||||
|
||||
data = await capability.get_picture_bytes("doc-1", "#/pictures/0")
|
||||
data = await capability.get_picture_bytes("doc-1", "#/pictures/0", "beta")
|
||||
|
||||
assert data == b"picture-bytes"
|
||||
client.document_item_repository.get_picture_bytes.assert_awaited_once_with(
|
||||
"doc-1", "#/pictures/0"
|
||||
)
|
||||
client.get_picture_bytes.assert_awaited_once_with("doc-1", "#/pictures/0", "beta")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1377,11 +1365,7 @@ async def test_citing_nothing_after_citing_something_keeps_it_grounded(temp_db_p
|
|||
capability = create_rag(db_path=temp_db_path, config=AppConfig())
|
||||
capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
|
||||
capability.epoch = 5
|
||||
client = AsyncMock()
|
||||
# Stands in for a single-database client: a bare AsyncMock's auto
|
||||
# attributes are truthy Mocks, and `_source` reaches a validated field.
|
||||
client._federated = {}
|
||||
client._source = None
|
||||
client = _single_database_client()
|
||||
client.get_chunk_by_id.return_value = Chunk(
|
||||
id="chunk-1", document_id="doc-1", content="evidence"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -82,6 +82,21 @@ def discovered(
|
|||
)
|
||||
|
||||
|
||||
def test_a_retained_picture_carries_the_source_it_came_from():
|
||||
"""Compaction re-fetches cited pictures later, so the capsule has to remember
|
||||
which database each came from."""
|
||||
found = discovered(cited={"c1": [2]}, pictures={"c1": ["#/pictures/0"]})
|
||||
found = replace(
|
||||
found,
|
||||
citations={"c1": replace_citation(found.citations["c1"], source="beta")},
|
||||
)
|
||||
|
||||
capsule = build_capsule([found])
|
||||
|
||||
[picture] = capsule.pictures
|
||||
assert picture.source == "beta"
|
||||
|
||||
|
||||
def test_nothing_cited_produces_no_capsule():
|
||||
capsule = build_capsule([discovered()])
|
||||
|
||||
|
|
|
|||
|
|
@ -897,20 +897,28 @@ _STORED_RAG_STATE = {
|
|||
}
|
||||
|
||||
|
||||
# Fields declared since `_STORED_RAG_STATE` was captured, with the value an
|
||||
# older dict loads as. A field may be added here; renaming or re-nesting one of
|
||||
# the stored keys is what this test exists to catch.
|
||||
_ADDED_SINCE = {"sources": None}
|
||||
|
||||
|
||||
def test_stored_state_shape_is_unchanged():
|
||||
"""A dict stored by an older version still loads, and dumps to the same keys.
|
||||
"""A dict stored by an older version still loads, and keeps every stored key.
|
||||
|
||||
Compatibility here is semantic JSON-object equivalence: the same keys, the
|
||||
same nesting, the same values. Key *order* is not part of the contract —
|
||||
deriving both states from a shared base reordered `AnalysisState`'s fields,
|
||||
and nothing serializes, hashes or string-compares this state; every carry
|
||||
point re-validates it by key.
|
||||
same nesting, the same values, plus whatever optional fields were added since.
|
||||
Key *order* is not part of the contract — deriving both states from a shared
|
||||
base reordered `AnalysisState`'s fields, and nothing serializes, hashes or
|
||||
string-compares this state; every carry point re-validates it by key.
|
||||
"""
|
||||
from haiku.rag.capabilities.rag import RAGState
|
||||
|
||||
state = RAGState.model_validate(_STORED_RAG_STATE)
|
||||
|
||||
assert state.model_dump(mode="json") == _STORED_RAG_STATE
|
||||
dumped = state.model_dump(mode="json")
|
||||
assert dumped == _STORED_RAG_STATE | _ADDED_SINCE
|
||||
assert {k: dumped[k] for k in _STORED_RAG_STATE} == _STORED_RAG_STATE
|
||||
|
||||
|
||||
def _populated(state_type):
|
||||
|
|
|
|||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
246
tests/sandbox/test_sandbox_multi_db.py
Normal file
246
tests/sandbox/test_sandbox_multi_db.py
Normal file
|
|
@ -0,0 +1,246 @@
|
|||
import shutil
|
||||
|
||||
import pytest
|
||||
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
||||
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,
|
||||
config=rag._config,
|
||||
context=AnalysisContext(sources=sources),
|
||||
rag=rag,
|
||||
)
|
||||
docs, owners = await sandbox._documents()
|
||||
return sandbox, docs, owners
|
||||
|
||||
|
||||
class TestDocumentsAcrossDatabases:
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_corpus_covers_every_configured_database(self, tmp_path):
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
_, docs, owners = await _mounted(rag)
|
||||
|
||||
assert {d.uri for d in docs} == {
|
||||
"test://alpha/alpha document about cats",
|
||||
"test://beta/beta document about cats",
|
||||
}
|
||||
assert {owner._source for owner in owners.values()} == {"alpha", "beta"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_selected_databases_bound_the_corpus(self, tmp_path):
|
||||
"""A question scoped to one database must not mount another's documents."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
_, docs, owners = await _mounted(rag, sources=["alpha"])
|
||||
|
||||
assert [d.uri for d in docs] == ["test://alpha/alpha document about cats"]
|
||||
assert {owner._source for owner in owners.values()} == {"alpha"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_database_needs_no_owners(self, tmp_path, temp_db_path):
|
||||
"""A single connection serves every read, so nothing has to be routed."""
|
||||
config = _config(tmp_path, ["alpha"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
_, docs, owners = await _mounted(rag)
|
||||
|
||||
assert len(docs) == 1
|
||||
assert owners == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_document_is_read_from_the_database_holding_it(self, tmp_path):
|
||||
"""Reads addressed to one document go through its owner, which is the
|
||||
only client that can answer them."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
sandbox, docs, owners = await _mounted(rag)
|
||||
sandbox._owners = owners
|
||||
for doc in docs:
|
||||
assert doc.id is not None
|
||||
async with sandbox._connection(owners[doc.id]) as owner:
|
||||
content = await owner.document_repository.get_content(doc.id)
|
||||
assert content is not None
|
||||
assert owners[doc.id]._source is not None
|
||||
assert owners[doc.id]._source in content
|
||||
|
||||
|
||||
class TestExecutingAcrossDatabases:
|
||||
@pytest.mark.asyncio
|
||||
async def test_code_reads_documents_from_every_database(self, tmp_path):
|
||||
"""The virtual filesystem is one flat namespace over the whole selected
|
||||
set, so code reads a document without knowing which database holds it."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
sandbox = Sandbox(
|
||||
db_path=rag._db_path,
|
||||
config=rag._config,
|
||||
context=AnalysisContext(),
|
||||
rag=rag,
|
||||
)
|
||||
try:
|
||||
result = await sandbox.execute(
|
||||
"docs = await list_documents()\n"
|
||||
"for d in sorted(docs, key=lambda d: d['uri']):\n"
|
||||
" with open('/documents/' + d['id'] + '/content.txt') as f:\n"
|
||||
" print(f.read())"
|
||||
)
|
||||
finally:
|
||||
await sandbox.close()
|
||||
|
||||
assert result.success, result.stderr
|
||||
assert "alpha document about cats" in result.stdout
|
||||
assert "beta document about cats" in result.stdout
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_code_cannot_read_an_unselected_database(self, tmp_path):
|
||||
"""Scoping the question scopes the filesystem."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
beta = (await rag.clients_for(["beta"]))[0]
|
||||
[outside] = await beta.document_repository.list_all(limit=1)
|
||||
sandbox = Sandbox(
|
||||
db_path=rag._db_path,
|
||||
config=rag._config,
|
||||
context=AnalysisContext(sources=["alpha"]),
|
||||
rag=rag,
|
||||
)
|
||||
try:
|
||||
result = await sandbox.execute(
|
||||
"docs = await list_documents()\n"
|
||||
"print(len(docs))\n"
|
||||
f"print(open('/documents/{outside.id}/content.txt').read())"
|
||||
)
|
||||
finally:
|
||||
await sandbox.close()
|
||||
|
||||
assert not result.success
|
||||
assert "beta document" not in result.stdout
|
||||
|
||||
|
||||
class TestSelectionOnOneDatabase:
|
||||
"""A client covering a single named database answers a selection the same way
|
||||
a search does, or the sandbox would mount what a search would refuse."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_selecting_no_database_mounts_nothing(self, tmp_path):
|
||||
config = _config(tmp_path, ["alpha"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
_, docs, owners = await _mounted(rag, sources=[])
|
||||
|
||||
assert docs == []
|
||||
assert owners == {}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_selecting_another_database_is_refused(self, tmp_path):
|
||||
config = _config(tmp_path, ["alpha"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
with pytest.raises(KeyError, match="beta"):
|
||||
await _mounted(rag, sources=["beta"])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_selecting_it_by_name_mounts_it(self, tmp_path):
|
||||
config = _config(tmp_path, ["alpha"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
_, docs, _ = await _mounted(rag, sources=["alpha"])
|
||||
|
||||
assert len(docs) == 1
|
||||
|
||||
|
||||
class TestCopiedDatabases:
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_document_in_two_databases_is_refused(self, tmp_path):
|
||||
"""Ids are unique per database, not across a copy of one: two documents
|
||||
would claim one path and the last would answer for both."""
|
||||
config = _config(tmp_path, ["alpha", "clone"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
shutil.rmtree(tmp_path / "clone.lancedb", ignore_errors=True)
|
||||
shutil.copytree(tmp_path / "alpha.lancedb", tmp_path / "clone.lancedb")
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
with pytest.raises(ValueError, match="one document per id"):
|
||||
await _mounted(rag)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_refusal_names_the_databases(self, tmp_path):
|
||||
config = _config(tmp_path, ["alpha", "clone"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
shutil.rmtree(tmp_path / "clone.lancedb", ignore_errors=True)
|
||||
shutil.copytree(tmp_path / "alpha.lancedb", tmp_path / "clone.lancedb")
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
with pytest.raises(ValueError) as raised:
|
||||
await _mounted(rag)
|
||||
|
||||
assert "alpha" in str(raised.value)
|
||||
assert "clone" in str(raised.value)
|
||||
assert str(tmp_path) not in str(raised.value)
|
||||
|
||||
|
||||
class TestListingOrder:
|
||||
@pytest.mark.asyncio
|
||||
async def test_the_listing_interleaves_the_databases(self, tmp_path):
|
||||
"""Code reads the listing through a truncated output, so concatenating
|
||||
shows one database's documents until the truncation and hides the rest."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", [f"alpha {i}" for i in range(5)])
|
||||
await _seed(config, "beta", ["beta one"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
sandbox = Sandbox(
|
||||
db_path=None,
|
||||
config=config,
|
||||
context=AnalysisContext(),
|
||||
rag=rag,
|
||||
)
|
||||
docs, _ = await sandbox._documents()
|
||||
|
||||
assert len(docs) == 6
|
||||
# The head has to reveal both databases.
|
||||
assert {(d.uri or "").split("/")[2] for d in docs[:2]} == {"alpha", "beta"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_in_code_list_documents_names_the_database(self, tmp_path):
|
||||
"""`source` is what lets code group the corpus by database."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha one"])
|
||||
await _seed(config, "beta", ["beta one"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
sandbox = Sandbox(
|
||||
db_path=None,
|
||||
config=config,
|
||||
context=AnalysisContext(),
|
||||
rag=rag,
|
||||
)
|
||||
rows = await sandbox._build_external_functions()["list_documents"]()
|
||||
|
||||
assert "source" in rows[0]
|
||||
assert {r["source"] for r in rows} == {"alpha", "beta"}
|
||||
|
|
@ -1,5 +1,6 @@
|
|||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -7,9 +8,19 @@ from click.exceptions import BadParameter
|
|||
from typer.testing import CliRunner
|
||||
|
||||
from haiku.rag.cli import _cli as cli
|
||||
from haiku.rag.cli import _parse_meta_options
|
||||
from haiku.rag.cli import (
|
||||
_parse_meta_options,
|
||||
require_one_database,
|
||||
resolve_db_path,
|
||||
select_database,
|
||||
)
|
||||
from haiku.rag.cli import cli as cli_wrapper
|
||||
from haiku.rag.store.exceptions import MigrationRequiredError
|
||||
from haiku.rag.config import get_config, set_config
|
||||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
from haiku.rag.store.exceptions import (
|
||||
AmbiguousDatabaseError,
|
||||
MigrationRequiredError,
|
||||
)
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
|
@ -94,6 +105,229 @@ class TestRebuildValidation:
|
|||
assert "mutually exclusive" in result.output
|
||||
|
||||
|
||||
class TestOneDatabaseCommands:
|
||||
"""`lancedb.databases` names a set; most commands work on one database."""
|
||||
|
||||
@staticmethod
|
||||
def _config(**databases):
|
||||
return AppConfig(lancedb=LanceDBConfig(databases=databases))
|
||||
|
||||
def test_a_configured_set_refuses_a_one_database_command(self):
|
||||
with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"):
|
||||
require_one_database(
|
||||
self._config(alpha="/db/a.lancedb", beta="/db/b.lancedb"),
|
||||
None,
|
||||
federated=False,
|
||||
)
|
||||
|
||||
def test_the_refusal_names_the_databases_and_not_their_locations(self):
|
||||
"""A location in an error message travels into logs and terminals; the
|
||||
names exist so it does not have to."""
|
||||
with pytest.raises(AmbiguousDatabaseError) as raised:
|
||||
require_one_database(
|
||||
self._config(alpha="s3://bucket/prefix/a.lancedb"),
|
||||
None,
|
||||
federated=False,
|
||||
)
|
||||
|
||||
assert "alpha" in str(raised.value)
|
||||
assert "s3://bucket/prefix/a.lancedb" not in str(raised.value)
|
||||
assert "bucket" not in str(raised.value)
|
||||
|
||||
def test_a_command_covering_the_set_is_allowed(self):
|
||||
require_one_database(
|
||||
self._config(alpha="/db/a.lancedb", beta="/db/b.lancedb"),
|
||||
None,
|
||||
federated=True,
|
||||
)
|
||||
|
||||
def test_naming_a_path_is_allowed(self):
|
||||
require_one_database(
|
||||
self._config(alpha="/db/a.lancedb"),
|
||||
Path("/db/other.lancedb"),
|
||||
federated=False,
|
||||
)
|
||||
|
||||
def test_no_configured_databases_is_allowed(self):
|
||||
require_one_database(AppConfig(), None, federated=False)
|
||||
|
||||
def test_the_refusal_exits_with_an_error(self):
|
||||
with patch("haiku.rag.cli._cli") as mock_cli:
|
||||
mock_cli.side_effect = AmbiguousDatabaseError("names alpha, beta")
|
||||
|
||||
with pytest.raises(SystemExit) as exc_info:
|
||||
cli_wrapper()
|
||||
assert exc_info.value.code == 1
|
||||
|
||||
|
||||
class TestSelectingADatabaseByName:
|
||||
"""`--database NAME` is the only way to reach a configured database whose
|
||||
location is a URI, since `--db` takes a path."""
|
||||
|
||||
@staticmethod
|
||||
def _install(monkeypatch, **databases):
|
||||
import haiku.rag.config as config_module
|
||||
|
||||
monkeypatch.setattr(config_module, "_config", None)
|
||||
set_config(AppConfig(lancedb=LanceDBConfig(databases=databases)))
|
||||
|
||||
def test_a_uri_location_becomes_the_configured_uri(self, monkeypatch):
|
||||
self._install(monkeypatch, medic="s3://bucket/prefix/medic.lancedb")
|
||||
|
||||
db_path = select_database("medic")
|
||||
|
||||
assert db_path is None
|
||||
config = get_config()
|
||||
assert config.lancedb.uri == "s3://bucket/prefix/medic.lancedb"
|
||||
assert config.lancedb.databases == {}
|
||||
|
||||
def test_a_local_location_becomes_the_database_path(self, monkeypatch):
|
||||
self._install(monkeypatch, st="/data/st.lancedb")
|
||||
|
||||
db_path = select_database("st")
|
||||
|
||||
assert db_path == Path("/data/st.lancedb")
|
||||
assert get_config().lancedb.uri == ""
|
||||
|
||||
def test_an_unknown_name_names_the_configured_ones(self, monkeypatch):
|
||||
self._install(monkeypatch, alpha="/data/a.lancedb", beta="/data/b.lancedb")
|
||||
|
||||
with pytest.raises(AmbiguousDatabaseError, match="alpha, beta"):
|
||||
select_database("gamma")
|
||||
|
||||
def test_an_unknown_name_does_not_leak_locations(self, monkeypatch):
|
||||
self._install(monkeypatch, medic="s3://bucket/prefix/medic.lancedb")
|
||||
|
||||
with pytest.raises(AmbiguousDatabaseError) as raised:
|
||||
select_database("gamma")
|
||||
|
||||
assert "bucket" not in str(raised.value)
|
||||
|
||||
def test_selecting_nothing_reports_an_empty_mapping(self, monkeypatch):
|
||||
import haiku.rag.config as config_module
|
||||
|
||||
monkeypatch.setattr(config_module, "_config", None)
|
||||
set_config(AppConfig())
|
||||
|
||||
with pytest.raises(AmbiguousDatabaseError, match="nothing"):
|
||||
select_database("medic")
|
||||
|
||||
def test_the_callback_selects_before_a_command_runs(self, tmp_path, monkeypatch):
|
||||
"""`--database` is resolved once the config is loaded, so every command
|
||||
and both TUIs see the selected database."""
|
||||
import haiku.rag.config as config_module
|
||||
|
||||
monkeypatch.setattr(config_module, "_config", None)
|
||||
config_file = tmp_path / "haiku.rag.yaml"
|
||||
config_file.write_text("lancedb:\n databases:\n alpha: /data/a.lancedb\n")
|
||||
|
||||
result = runner.invoke(
|
||||
cli, ["--config", str(config_file), "--database", "nope", "list"]
|
||||
)
|
||||
|
||||
assert result.exit_code != 0
|
||||
assert isinstance(result.exception, AmbiguousDatabaseError)
|
||||
assert "nope" in str(result.exception)
|
||||
|
||||
def test_a_selection_does_not_outlive_its_invocation(self, tmp_path, monkeypatch):
|
||||
"""The selector is module state, so a second invocation without
|
||||
`--database` must not inherit the first one's database."""
|
||||
import haiku.rag.cli as cli_module
|
||||
import haiku.rag.config as config_module
|
||||
|
||||
monkeypatch.setattr(config_module, "_config", None)
|
||||
monkeypatch.setattr(cli_module, "_database", None)
|
||||
monkeypatch.setattr(cli_module, "_database_path", None)
|
||||
config_file = tmp_path / "haiku.rag.yaml"
|
||||
selected = tmp_path / "alpha.lancedb"
|
||||
config_file.write_text(f"lancedb:\n databases:\n alpha: {selected}\n")
|
||||
|
||||
runner.invoke(
|
||||
cli, ["--config", str(config_file), "--database", "alpha", "info"]
|
||||
)
|
||||
assert cli_module._database_path == selected
|
||||
|
||||
runner.invoke(cli, ["--config", str(config_file), "settings"])
|
||||
|
||||
assert cli_module._database_path is None
|
||||
assert cli_module._database is None
|
||||
|
||||
def test_a_selection_does_not_outlive_its_invocation_in_process(
|
||||
self, tmp_path, monkeypatch
|
||||
):
|
||||
"""Selecting rewrites the configuration, so a second invocation has to
|
||||
start from a freshly loaded one rather than inherit the rewrite."""
|
||||
import haiku.rag.cli as cli_module
|
||||
import haiku.rag.config as config_module
|
||||
|
||||
monkeypatch.setattr(config_module, "_config", None)
|
||||
monkeypatch.setattr(cli_module, "_database", None)
|
||||
monkeypatch.setattr(cli_module, "_database_path", None)
|
||||
config_file = tmp_path / "haiku.rag.yaml"
|
||||
config_file.write_text(
|
||||
f"lancedb:\n databases:\n alpha: {tmp_path / 'alpha.lancedb'}\n"
|
||||
f" beta: {tmp_path / 'beta.lancedb'}\n"
|
||||
)
|
||||
|
||||
runner.invoke(
|
||||
cli, ["--config", str(config_file), "--database", "alpha", "settings"]
|
||||
)
|
||||
assert get_config().lancedb.uri == ""
|
||||
assert get_config().lancedb.databases == {}
|
||||
|
||||
runner.invoke(cli, ["--config", str(config_file), "settings"])
|
||||
|
||||
assert set(get_config().lancedb.databases) == {"alpha", "beta"}
|
||||
|
||||
def test_a_selection_does_not_outlive_an_invocation_without_a_config_file(
|
||||
self, tmp_path, monkeypatch
|
||||
):
|
||||
"""No config file is still a load: the previous invocation's selected URI
|
||||
must not be what the next one talks to."""
|
||||
import haiku.rag.cli as cli_module
|
||||
import haiku.rag.config as config_module
|
||||
|
||||
monkeypatch.setattr(config_module, "_config", None)
|
||||
monkeypatch.setattr(cli_module, "_database", None)
|
||||
monkeypatch.setattr(cli_module, "_database_path", None)
|
||||
monkeypatch.delenv("HAIKU_RAG_CONFIG_PATH", raising=False)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
config_file = tmp_path / "selected.yaml"
|
||||
config_file.write_text(
|
||||
"lancedb:\n databases:\n medic: s3://bucket/medic.lancedb\n"
|
||||
)
|
||||
|
||||
runner.invoke(
|
||||
cli, ["--config", str(config_file), "--database", "medic", "settings"]
|
||||
)
|
||||
assert get_config().lancedb.uri == "s3://bucket/medic.lancedb"
|
||||
|
||||
runner.invoke(cli, ["settings"])
|
||||
|
||||
assert get_config().lancedb.uri == ""
|
||||
|
||||
|
||||
class TestResolvingTheDatabasePath:
|
||||
def test_a_path_wins_when_nothing_is_selected(self, monkeypatch):
|
||||
monkeypatch.setattr("haiku.rag.cli._database", None)
|
||||
monkeypatch.setattr("haiku.rag.cli._database_path", None)
|
||||
|
||||
assert resolve_db_path(Path("/data/one.lancedb")) == Path("/data/one.lancedb")
|
||||
|
||||
def test_the_selected_database_is_used_without_a_path(self, monkeypatch):
|
||||
monkeypatch.setattr("haiku.rag.cli._database", "st")
|
||||
monkeypatch.setattr("haiku.rag.cli._database_path", Path("/data/st.lancedb"))
|
||||
|
||||
assert resolve_db_path(None) == Path("/data/st.lancedb")
|
||||
|
||||
def test_naming_a_database_twice_is_refused(self, monkeypatch):
|
||||
monkeypatch.setattr("haiku.rag.cli._database", "st")
|
||||
monkeypatch.setattr("haiku.rag.cli._database_path", Path("/data/st.lancedb"))
|
||||
|
||||
with pytest.raises(AmbiguousDatabaseError, match="not both"):
|
||||
resolve_db_path(Path("/data/other.lancedb"))
|
||||
|
||||
|
||||
class TestCliMigrationError:
|
||||
def test_catches_migration_required_error(self):
|
||||
with patch("haiku.rag.cli._cli") as mock_cli:
|
||||
|
|
@ -351,6 +585,55 @@ class TestAskAnalyzeImageOption:
|
|||
assert mock_ask.call_args.kwargs["images"] == [buffer.getvalue()]
|
||||
|
||||
|
||||
class TestRenderingTheDatabase:
|
||||
"""Across databases a result has to say which one it came from. One database
|
||||
needs no such label, so single-database output is unchanged."""
|
||||
|
||||
@staticmethod
|
||||
def _app(tmp_path, **databases):
|
||||
from haiku.rag.app import HaikuRAGApp
|
||||
|
||||
return HaikuRAGApp(
|
||||
db_path=tmp_path / "unused",
|
||||
config=AppConfig(lancedb=LanceDBConfig(databases=databases)),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _rendered(app, result) -> str:
|
||||
from rich.console import Console
|
||||
|
||||
app.console = Console(record=True, width=200)
|
||||
app._rich_print_search_result(result)
|
||||
return app.console.export_text()
|
||||
|
||||
@staticmethod
|
||||
def _result():
|
||||
from haiku.rag.store.models import SearchResult
|
||||
|
||||
return SearchResult(
|
||||
content="a body",
|
||||
score=0.9,
|
||||
source="alpha",
|
||||
chunk_id="c1",
|
||||
document_id="d1",
|
||||
document_uri="test://alpha/one",
|
||||
)
|
||||
|
||||
def test_a_set_labels_each_result(self, tmp_path):
|
||||
app = self._app(
|
||||
tmp_path,
|
||||
alpha=str(tmp_path / "alpha.lancedb"),
|
||||
beta=str(tmp_path / "beta.lancedb"),
|
||||
)
|
||||
|
||||
assert "database: alpha" in self._rendered(app, self._result())
|
||||
|
||||
def test_one_database_is_not_labelled(self, tmp_path):
|
||||
app = self._app(tmp_path, alpha=str(tmp_path / "alpha.lancedb"))
|
||||
|
||||
assert "database:" not in self._rendered(app, self._result())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def app_stub(monkeypatch, tmp_path):
|
||||
"""Stand in for HaikuRAGApp so a command's wiring can be checked without a
|
||||
|
|
@ -358,7 +641,10 @@ def app_stub(monkeypatch, tmp_path):
|
|||
the application layer renders."""
|
||||
# AsyncMock so every command's `asyncio.run(app.x(...))` gets a coroutine.
|
||||
stub = AsyncMock()
|
||||
monkeypatch.setattr("haiku.rag.cli.create_app", lambda db=None: stub)
|
||||
monkeypatch.setattr(
|
||||
"haiku.rag.cli.create_app",
|
||||
lambda db=None, *, federated=False: stub,
|
||||
)
|
||||
return stub
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,3 +1,5 @@
|
|||
import asyncio
|
||||
|
||||
import pytest
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
from docling_core.types.doc.labels import DocItemLabel
|
||||
|
|
@ -8,6 +10,7 @@ from haiku.rag.config import get_config
|
|||
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
||||
from haiku.rag.store.exceptions import SourceUnavailableError
|
||||
from haiku.rag.store.models import Chunk, DocumentItem
|
||||
from haiku.rag.utils import locate_database
|
||||
|
||||
|
||||
class TestConfig:
|
||||
|
|
@ -48,6 +51,112 @@ async def _seed(config, name, contents):
|
|||
)
|
||||
|
||||
|
||||
class TestNamingADatabaseDirectly:
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_explicit_db_path_wins_over_the_configured_set(
|
||||
self, tmp_path, temp_db_path
|
||||
):
|
||||
"""A caller that names a path means that database, not the configured
|
||||
set: the CLI resolves `--db` to one and must not fan out instead."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
async with HaikuRAG(temp_db_path, config=config, create=True) as rag:
|
||||
assert rag._federated == {}
|
||||
assert rag._source is None
|
||||
assert rag.store.db_path == temp_db_path
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_configured_database_is_opened_by_name(self, tmp_path):
|
||||
"""A set of one is not federated, and the client resolves it."""
|
||||
config = _config(tmp_path, ["alpha"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
assert rag._federated == {}
|
||||
assert rag._source == "alpha"
|
||||
results = await rag.search("cats", search_type="fts", limit=10)
|
||||
|
||||
assert [r.source for r in results] == ["alpha"]
|
||||
|
||||
|
||||
class TestOpeningDatabases:
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_databases_open_together(self, tmp_path):
|
||||
"""A cold fan-out costs one open, not their sum. On object storage a
|
||||
serial loop is the difference between one round trip and N."""
|
||||
names = ["alpha", "beta", "gamma"]
|
||||
config = _config(tmp_path, names)
|
||||
for name in names:
|
||||
await _seed(config, name, [f"{name} document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
barrier = asyncio.Barrier(len(names))
|
||||
open_one = rag._open_client
|
||||
|
||||
async def gated(name: str, location: str):
|
||||
# Every open has to be in flight before any of them finishes, so
|
||||
# a serial loop cannot get past this and the wait times out.
|
||||
await barrier.wait()
|
||||
return await open_one(name, location)
|
||||
|
||||
rag._open_client = gated
|
||||
clients = await asyncio.wait_for(rag.clients_for(names), timeout=15)
|
||||
|
||||
assert {client._source for client in clients} == set(names)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_failed_open_does_not_leak_the_ones_that_worked(self, tmp_path):
|
||||
"""Opening together means a failure has siblings already open. They are
|
||||
tracked before it is reported, so closing the set closes them."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
config.lancedb.databases["beta"] = str(tmp_path / "absent.lancedb")
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
with pytest.raises(SourceUnavailableError, match="beta"):
|
||||
await rag.clients_for(["alpha", "beta"])
|
||||
|
||||
assert set(rag._clients) == {"alpha"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_database_named_twice_is_opened_once(self, tmp_path):
|
||||
"""Fusion would count a repeated database as two rank lists."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
clients = await rag.clients_for(["alpha", "alpha", "beta"])
|
||||
|
||||
assert [client._source for client in clients] == ["alpha", "beta"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_database_named_twice_returns_each_result_once(self, tmp_path):
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
results = await rag.search(
|
||||
"cats", limit=10, search_type="fts", sources=["alpha", "alpha"]
|
||||
)
|
||||
|
||||
assert [r.source for r in results] == ["alpha"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_one_database_named_twice_is_still_that_database(self, tmp_path):
|
||||
"""A client covering a single named database compares the selection
|
||||
against its own name, so repeats have to collapse first."""
|
||||
config = _config(tmp_path, ["alpha"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
covering = await rag.clients_covering(["alpha", "alpha"])
|
||||
|
||||
assert [client._source for client in covering] == ["alpha"]
|
||||
|
||||
|
||||
class TestFederatedSearch:
|
||||
@pytest.mark.asyncio
|
||||
async def test_results_carry_their_source(self, tmp_path):
|
||||
|
|
@ -120,13 +229,13 @@ class TestSingleDatabaseUnchanged:
|
|||
|
||||
class TestLocate:
|
||||
def test_a_scheme_is_a_uri(self):
|
||||
assert HaikuRAG._locate("s3://bucket/one.lancedb") == (
|
||||
assert locate_database("s3://bucket/one.lancedb") == (
|
||||
"s3://bucket/one.lancedb",
|
||||
None,
|
||||
)
|
||||
|
||||
def test_anything_else_is_a_local_path(self):
|
||||
uri, db_path = HaikuRAG._locate("/data/one.lancedb")
|
||||
uri, db_path = locate_database("/data/one.lancedb")
|
||||
assert uri == ""
|
||||
assert db_path is not None and str(db_path) == "/data/one.lancedb"
|
||||
|
||||
|
|
|
|||
|
|
@ -1,13 +1,39 @@
|
|||
import pytest
|
||||
from docling_core.types.doc.document import DoclingDocument
|
||||
from docling_core.types.doc.labels import DocItemLabel
|
||||
from pydantic_ai import ModelRetry
|
||||
|
||||
from haiku.rag.capabilities.rag import RAGState, create_capability
|
||||
from haiku.rag.client import HaikuRAG
|
||||
from haiku.rag.store.models import SearchResult
|
||||
from haiku.rag.config import get_config
|
||||
from haiku.rag.store.models import Chunk, Document, DocumentItem, SearchResult
|
||||
from haiku.rag.store.models.citation import resolve_citations
|
||||
from tests.test_multi_db import _config, _seed
|
||||
|
||||
|
||||
async def _seed_expandable(config, name, sentences):
|
||||
"""One document whose chunk covers a single item, so expansion has
|
||||
neighbours to pull in and rebuilds the result rather than passing it
|
||||
through."""
|
||||
dim = get_config().embeddings.model.vector_dim
|
||||
doc = DoclingDocument(name=name)
|
||||
for sentence in sentences:
|
||||
doc.add_text(label=DocItemLabel.TEXT, text=sentence)
|
||||
async with HaikuRAG(config=config, create=True, sources=[name]) as rag:
|
||||
await rag.import_document(
|
||||
doc,
|
||||
[
|
||||
Chunk(
|
||||
content=sentences[0],
|
||||
embedding=[0.1] * dim,
|
||||
order=0,
|
||||
metadata={"doc_item_refs": ["#/texts/0"]},
|
||||
)
|
||||
],
|
||||
uri=f"test://{name}/expandable",
|
||||
)
|
||||
|
||||
|
||||
class TestExpansionRouting:
|
||||
@pytest.mark.asyncio
|
||||
async def test_expansion_routes_each_result_to_its_database(self, tmp_path):
|
||||
|
|
@ -26,6 +52,45 @@ class TestExpansionRouting:
|
|||
assert r.source is not None
|
||||
assert r.source in r.content
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_expanded_result_keeps_its_source(self, tmp_path):
|
||||
"""Expansion rebuilds the result, and the rebuilt one has to name the
|
||||
database it was expanded through."""
|
||||
config = _config(tmp_path, ["alpha"])
|
||||
await _seed_expandable(
|
||||
config, "alpha", ["cats sleep often", "cats also hunt", "cats purr"]
|
||||
)
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
results = await rag.search("cats", search_type="fts", limit=10)
|
||||
expanded = await rag.expand_context(results)
|
||||
|
||||
assert len(expanded) == 1
|
||||
assert "cats also hunt" in expanded[0].content, "expansion did not run"
|
||||
assert expanded[0].source == "alpha"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_expansion_keeps_tied_results_in_fused_order(self, tmp_path):
|
||||
"""Fused scores tie often, so grouping by database must not reorder
|
||||
them: the tiebreak is the order they arrived in."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha one about cats", "alpha two about cats"])
|
||||
await _seed(config, "beta", ["beta one about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
found = await rag.search("cats", search_type="fts", limit=10)
|
||||
by_source: dict[str, list[SearchResult]] = {}
|
||||
for result in found:
|
||||
by_source.setdefault(result.source or "", []).append(result)
|
||||
# Interleaved, so grouping by database is visible as a reordering.
|
||||
fused = [by_source["alpha"][0], by_source["beta"][0], by_source["alpha"][1]]
|
||||
for result in fused:
|
||||
result.score = 0.5
|
||||
|
||||
expanded = await rag.expand_context(fused)
|
||||
|
||||
assert [r.chunk_id for r in expanded] == [r.chunk_id for r in fused]
|
||||
|
||||
|
||||
class TestCitationSource:
|
||||
def test_a_citation_carries_the_result_source(self):
|
||||
|
|
@ -58,6 +123,7 @@ class TestCitationSource:
|
|||
|
||||
class TestAskAcrossDatabases:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_the_capability_searches_the_selected_databases(self, tmp_path):
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
|
@ -74,6 +140,7 @@ class TestAskAcrossDatabases:
|
|||
assert "beta document" not in formatted
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_searching_all_databases_reaches_both(self, tmp_path):
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
|
@ -92,6 +159,7 @@ class TestAskAcrossDatabases:
|
|||
|
||||
class TestCiteFallback:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_an_id_from_a_selected_database_resolves_with_its_source(
|
||||
self, tmp_path
|
||||
):
|
||||
|
|
@ -128,6 +196,7 @@ class TestCiteFallback:
|
|||
assert citation.source == "alpha"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_an_id_outside_the_selected_databases_does_not_resolve(
|
||||
self, tmp_path
|
||||
):
|
||||
|
|
@ -154,6 +223,297 @@ class TestCiteFallback:
|
|||
with pytest.raises(ModelRetry):
|
||||
await run._cite([outside.id])
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_selecting_no_databases_cites_nothing(self, tmp_path):
|
||||
"""`sources=[]` selected nothing, which is not the same as everything:
|
||||
the fallback must not go looking where the question never looked."""
|
||||
from tests.capabilities.test_capabilities import Deps, make_context
|
||||
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
alpha = (await rag.clients_for(["alpha"]))[0]
|
||||
[chunk] = await alpha.chunk_repository.list_all(limit=1)
|
||||
assert chunk.id is not None
|
||||
|
||||
capability = create_capability(config=config, rag=rag, defer_loading=False)
|
||||
deps = Deps(state={"rag": RAGState(sources=[]).model_dump(mode="json")})
|
||||
run = await capability.for_run(make_context(deps))
|
||||
|
||||
with pytest.raises(ModelRetry):
|
||||
await run._cite([chunk.id])
|
||||
|
||||
|
||||
class TestStandaloneCapabilities:
|
||||
"""A capability nobody hands a client opens its own. It has to reach the
|
||||
configured set, or a host that only registers capabilities gets one
|
||||
database while the configuration names several."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_a_rag_capability_opens_the_configured_set(self, tmp_path):
|
||||
from tests.capabilities.test_capabilities import Deps, make_context
|
||||
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
capability = create_capability(config=config, defer_loading=False)
|
||||
assert capability.db_path is None
|
||||
run = await capability.for_run(make_context(Deps()))
|
||||
try:
|
||||
formatted = await run._search("cats", limit=10)
|
||||
finally:
|
||||
await run._close()
|
||||
|
||||
assert isinstance(formatted, str)
|
||||
assert "alpha document" in formatted
|
||||
assert "beta document" in formatted
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_an_analysis_capability_mounts_the_configured_set(self, tmp_path):
|
||||
from haiku.rag.capabilities.analysis import (
|
||||
create_capability as create_analysis,
|
||||
)
|
||||
from tests.capabilities.test_capabilities import Deps, make_context
|
||||
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
capability = create_analysis(config=config, defer_loading=False)
|
||||
run = await capability.for_run(make_context(Deps()))
|
||||
try:
|
||||
sandbox = await run._ensure_sandbox()
|
||||
docs, owners = await sandbox._documents()
|
||||
finally:
|
||||
await run._close()
|
||||
|
||||
assert len(docs) == 2
|
||||
assert {owner._source for owner in owners.values()} == {"alpha", "beta"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_single_configured_database_is_still_opened(self, tmp_path):
|
||||
"""One named database is a set of one, not a path to guess."""
|
||||
config = _config(tmp_path, ["alpha"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
capability = create_capability(config=config, defer_loading=False)
|
||||
rag = await capability._ensure_rag()
|
||||
try:
|
||||
assert rag._source == "alpha"
|
||||
finally:
|
||||
await capability._close()
|
||||
|
||||
|
||||
class TestAnalyzeAcrossDatabases:
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_the_capability_searches_the_selected_databases(self, tmp_path):
|
||||
"""`analysis_search` is the same tool as the RAG one, and the sandbox is
|
||||
scoped by the same selection."""
|
||||
from haiku.rag.capabilities.analysis import AnalysisState
|
||||
from haiku.rag.capabilities.analysis import (
|
||||
create_capability as create_analysis,
|
||||
)
|
||||
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
capability = create_analysis(config=config, rag=rag, defer_loading=False)
|
||||
capability.state = AnalysisState(sources=["alpha"])
|
||||
|
||||
formatted = await capability._search("cats", limit=10)
|
||||
sandbox = await capability._ensure_sandbox()
|
||||
await capability._close()
|
||||
|
||||
assert isinstance(formatted, str)
|
||||
assert "alpha document" in formatted
|
||||
assert "beta document" not in formatted
|
||||
assert sandbox._context.sources == ["alpha"]
|
||||
|
||||
|
||||
class TestDatabaseIdentityForTheModel:
|
||||
def test_a_result_names_its_database(self):
|
||||
"""The model has to attribute and compare evidence by database while it
|
||||
composes the answer, not only afterwards through the citations."""
|
||||
result = SearchResult(content="body", score=0.9, source="alpha", chunk_id="c1")
|
||||
|
||||
assert "Database: alpha" in result.format_for_agent()
|
||||
|
||||
def test_an_unnamed_database_is_not_mentioned(self):
|
||||
"""A single unnamed database renders as it always has."""
|
||||
result = SearchResult(content="body", score=0.9, chunk_id="c1")
|
||||
|
||||
assert "Database" not in result.format_for_agent()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.vcr()
|
||||
async def test_in_code_search_names_the_database(self, tmp_path):
|
||||
from haiku.rag.sandbox import AnalysisContext, Sandbox
|
||||
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
sandbox = Sandbox(
|
||||
db_path=None,
|
||||
config=config,
|
||||
context=AnalysisContext(),
|
||||
rag=rag,
|
||||
)
|
||||
try:
|
||||
result = await sandbox.execute(
|
||||
"rows = await search('cats', limit=10)\n"
|
||||
"print(sorted(r['source'] for r in rows))\n"
|
||||
"docs = await list_documents()\n"
|
||||
"print(sorted(d['source'] for d in docs))"
|
||||
)
|
||||
finally:
|
||||
await sandbox.close()
|
||||
|
||||
assert result.success, result.stderr
|
||||
assert "['alpha', 'beta']" in result.stdout
|
||||
assert result.stdout.count("['alpha', 'beta']") == 2
|
||||
|
||||
|
||||
class TestActionableFailures:
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_migration_error_survives_being_named(self, tmp_path, temp_db_path):
|
||||
"""The remedy is the whole value of the message, and it names no location,
|
||||
so it is not replaced by the database's name."""
|
||||
from haiku.rag.store.exceptions import MigrationRequiredError
|
||||
|
||||
config = _config(tmp_path, ["alpha"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config, sources=["alpha"]) as rag:
|
||||
await rag.store.set_haiku_version("0.20.0")
|
||||
|
||||
with pytest.raises(MigrationRequiredError) as raised:
|
||||
async with HaikuRAG(config=config, sources=["alpha"]):
|
||||
pass
|
||||
|
||||
# Both halves: which database failed, and what to run about it.
|
||||
assert "haiku-rag migrate" in str(raised.value)
|
||||
assert "alpha" in str(raised.value)
|
||||
assert str(tmp_path) not in str(raised.value)
|
||||
|
||||
|
||||
class TestPictureDeduplication:
|
||||
"""One picture yields two chunks — a text-embedded one and an image-embedded
|
||||
one — that collapse to the best. Two databases holding the same picture are
|
||||
two results, not a duplicate."""
|
||||
|
||||
@staticmethod
|
||||
def _picture(source, score):
|
||||
return SearchResult(
|
||||
content="a figure",
|
||||
score=score,
|
||||
source=source,
|
||||
chunk_id=f"{source}-c",
|
||||
document_id="doc-1",
|
||||
doc_item_refs=["#/pictures/0"],
|
||||
)
|
||||
|
||||
def test_the_same_picture_in_two_databases_survives(self):
|
||||
from haiku.rag.client.search import _dedup_picture_chunks
|
||||
|
||||
kept = _dedup_picture_chunks(
|
||||
[self._picture("alpha", 0.9), self._picture("clone", 0.5)]
|
||||
)
|
||||
|
||||
assert [r.source for r in kept] == ["alpha", "clone"]
|
||||
|
||||
def test_duplicates_within_one_database_still_collapse(self):
|
||||
from haiku.rag.client.search import _dedup_picture_chunks
|
||||
|
||||
lower = self._picture("alpha", 0.5)
|
||||
higher = self._picture("alpha", 0.9)
|
||||
|
||||
kept = _dedup_picture_chunks([lower, higher])
|
||||
|
||||
assert kept == [higher]
|
||||
|
||||
|
||||
class TestPictureRouting:
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_picture_is_fetched_from_the_database_that_holds_it(self, tmp_path):
|
||||
"""A `self_ref` repeats across databases, so the citation's source is
|
||||
what decides where the bytes come from."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
await _seed(config, "beta", ["beta document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
beta = (await rag.clients_for(["beta"]))[0]
|
||||
[document] = await beta.document_repository.list_all(limit=1)
|
||||
assert document.id is not None
|
||||
await beta.document_item_repository.create_all(
|
||||
[
|
||||
DocumentItem(
|
||||
document_id=document.id,
|
||||
self_ref="#/pictures/0",
|
||||
position=99,
|
||||
label="picture",
|
||||
text="",
|
||||
picture_data=b"beta-picture",
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
await rag.get_picture_bytes(document.id, "#/pictures/0", "beta")
|
||||
== b"beta-picture"
|
||||
)
|
||||
assert (
|
||||
await rag.get_picture_bytes(document.id, "#/pictures/0", "alpha")
|
||||
is None
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_single_database_needs_no_source(self, temp_db_path):
|
||||
"""One database is where the picture is, named or not."""
|
||||
async with HaikuRAG(temp_db_path, create=True) as rag:
|
||||
document = await rag.document_repository.create(
|
||||
Document(content="body", uri="test://one")
|
||||
)
|
||||
assert document.id is not None
|
||||
await rag.document_item_repository.create_all(
|
||||
[
|
||||
DocumentItem(
|
||||
document_id=document.id,
|
||||
self_ref="#/pictures/0",
|
||||
position=0,
|
||||
label="picture",
|
||||
text="",
|
||||
picture_data=b"the-picture",
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
assert (
|
||||
await rag.get_picture_bytes(document.id, "#/pictures/0")
|
||||
== b"the-picture"
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_picture_lookup_without_a_source_is_refused(self, tmp_path):
|
||||
"""Federating, nothing can say which database holds an unqualified
|
||||
reference."""
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
async with HaikuRAG(config=config) as rag:
|
||||
with pytest.raises(ValueError, match="source"):
|
||||
await rag.get_picture_bytes("doc-1", "#/pictures/0")
|
||||
|
||||
|
||||
class TestFederatedEdges:
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -182,14 +542,16 @@ class TestFederatedEdges:
|
|||
|
||||
from haiku.rag.capabilities.rag import RAGCapability
|
||||
from haiku.rag.store.models import Chunk
|
||||
from tests.capabilities.test_capabilities import Deps, make_context
|
||||
from tests.capabilities.test_capabilities import (
|
||||
Deps,
|
||||
_single_database_client,
|
||||
make_context,
|
||||
)
|
||||
|
||||
config = _config(tmp_path, ["alpha", "beta"])
|
||||
await _seed(config, "alpha", ["alpha document about cats"])
|
||||
|
||||
orphan = AsyncMock()
|
||||
orphan._federated = {}
|
||||
orphan._source = None
|
||||
orphan = _single_database_client()
|
||||
orphan.get_chunk_by_id.return_value = Chunk(
|
||||
id="orphan", document_id=None, content="no document"
|
||||
)
|
||||
|
|
|
|||
|
|
@ -769,6 +769,52 @@ async def test_format_citations_rich_header_and_footer():
|
|||
assert "chunk: chunk-uuid-1" in output
|
||||
|
||||
|
||||
async def test_format_citations_rich_names_the_database_when_federating():
|
||||
"""Across databases, a citation has to say which one it came from."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import format_citations_rich
|
||||
|
||||
citation = Citation(
|
||||
document_id="doc-uuid-1",
|
||||
chunk_id="chunk-uuid-1",
|
||||
document_uri="test://doc",
|
||||
document_title="Test Doc",
|
||||
content="Body",
|
||||
source="medic",
|
||||
)
|
||||
client = AsyncMock()
|
||||
client._federated = {"medic": "/data/medic.lancedb", "st": "/data/st.lancedb"}
|
||||
|
||||
output = _render_rich(await format_citations_rich([citation], client))
|
||||
|
||||
assert "medic" in output
|
||||
|
||||
|
||||
async def test_format_citations_rich_omits_the_database_for_one_database():
|
||||
"""A single database is not worth naming on every citation."""
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import format_citations_rich
|
||||
|
||||
citation = Citation(
|
||||
document_id="doc-uuid-1",
|
||||
chunk_id="chunk-uuid-1",
|
||||
document_uri="test://doc",
|
||||
document_title="Test Doc",
|
||||
content="Body",
|
||||
source="medic",
|
||||
)
|
||||
client = AsyncMock()
|
||||
client._federated = {}
|
||||
|
||||
output = _render_rich(await format_citations_rich([citation], client))
|
||||
|
||||
assert "medic" not in output
|
||||
|
||||
|
||||
async def test_format_citations_rich_truncates_long_content():
|
||||
from haiku.rag.store.models.citation import Citation
|
||||
from haiku.rag.utils import CITATION_PREVIEW_CHARS, format_citations_rich
|
||||
|
|
@ -937,7 +983,7 @@ async def test_render_picture_handles_stored_bytes(stored, renders):
|
|||
stored = buf.getvalue()
|
||||
|
||||
client = AsyncMock()
|
||||
client.document_item_repository.get_picture_bytes = AsyncMock(return_value=stored)
|
||||
client.get_picture_bytes = AsyncMock(return_value=stored)
|
||||
|
||||
result = await _render_picture(client, "doc1", "#/pictures/0")
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue