From 0be2e5b24fd5b25bc28e102e7d4ef155d088fce4 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 9 Jun 2026 11:09:28 +0300 Subject: [PATCH] Extract gather_database_info shared by the info command --- haiku_rag_slim/haiku/rag/app.py | 102 ++++++++------------ haiku_rag_slim/haiku/rag/store/engine.py | 112 +++++++++++++++++++++- tests/store/test_database_info.py | 117 +++++++++++++++++++++++ 3 files changed, 267 insertions(+), 64 deletions(-) create mode 100644 tests/store/test_database_info.py diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index a9887ad1..98be2d3f 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -1,4 +1,3 @@ -import json import logging from datetime import datetime from pathlib import Path @@ -24,7 +23,7 @@ from haiku.rag.store.models.document import Document if TYPE_CHECKING: from haiku.rag.store.models import SearchResult -from haiku.rag.utils import format_bytes, format_citations_rich, get_package_versions +from haiku.rag.utils import format_bytes, format_citations_rich logger = logging.getLogger(__name__) @@ -66,8 +65,7 @@ class HaikuRAGApp: # pragma: no cover async def info(self): """Display read-only information about the database without modifying it.""" - from haiku.rag.store.engine import connect_lancedb, get_database_stats - from haiku.rag.store.upgrades import get_pending_upgrades + from haiku.rag.store.engine import gather_database_info if self.before is not None: self.console.print( @@ -84,55 +82,37 @@ class HaikuRAGApp: # pragma: no cover self.console.print("[red]Database path does not exist.[/red]") return - # Connect directly. Don't go through Store so a database that is - # missing tables (e.g. pre-migration) still reports what it can. - db = await connect_lancedb(self.config, self.db_path) - stats = await get_database_stats(db) + info = await gather_database_info(self.config, self.db_path) - if not any(entry["exists"] for entry in stats.values()): + if not info.exists: self.console.print( "[red]Database is empty. Use 'haiku-rag init' to initialize.[/red]" ) return - versions = get_package_versions() - - stored_version = "unknown" - embed_provider = "unknown" - embed_model = "unknown" - vector_dim = None - if stats["settings"]["exists"]: - settings_tbl = await db.open_table("settings") - rows = ( - await settings_tbl.query().where("id = 'settings'").limit(1).to_arrow() - ).to_pylist() - if rows: - raw = rows[0].get("settings") or "{}" - data = json.loads(raw) if isinstance(raw, str) else (raw or {}) - stored_version = str(data.get("version", "unknown")) - embeddings = data.get("embeddings", {}) - embed_model_obj = embeddings.get("model", {}) - embed_provider = embed_model_obj.get("provider", "unknown") - embed_model = embed_model_obj.get("name", "unknown") - vector_dim = embed_model_obj.get("vector_dim") - self.console.print( - f" [repr.attrib_name]haiku.rag version (db)[/repr.attrib_name]: {stored_version}" + f" [repr.attrib_name]haiku.rag version (db)[/repr.attrib_name]: {info.stored_version}" + ) + dim_part = ( + f"{info.embeddings.vector_dim}" + if info.embeddings.vector_dim is not None + else "unknown" ) - dim_part = f"{vector_dim}" if vector_dim is not None else "unknown" self.console.print( " [repr.attrib_name]embeddings[/repr.attrib_name]: " - f"{embed_provider}/{embed_model} (dim: {dim_part})" + f"{info.embeddings.provider}/{info.embeddings.name} (dim: {dim_part})" ) + tables = {t.name: t for t in info.tables} + # Per-table row counts and sizes. Missing required tables are # reported as "absent" rather than raising. for name in ("documents", "chunks", "document_items"): - entry = stats[name] - if entry["exists"]: + entry = tables[name] + if entry.exists: self.console.print( - f" [repr.attrib_name]{name}[/repr.attrib_name]: {entry['num_rows']} " - f"({format_bytes(entry['total_bytes'])})" + f" [repr.attrib_name]{name}[/repr.attrib_name]: {entry.num_rows} " + f"({format_bytes(entry.total_bytes)})" ) else: self.console.print( @@ -140,25 +120,23 @@ class HaikuRAGApp: # pragma: no cover ) # Vector index information - if stats["chunks"]["exists"]: - num_chunks = stats["chunks"]["num_rows"] - if stats["chunks"].get("has_vector_index"): - num_indexed_rows = stats["chunks"].get("num_indexed_rows", 0) - num_unindexed_rows = stats["chunks"].get("num_unindexed_rows", 0) + if tables["chunks"].exists: + num_chunks = tables["chunks"].num_rows + if info.vector_index.exists: self.console.print( " [repr.attrib_name]vector index[/repr.attrib_name]: ✓ exists" ) self.console.print( - f" [repr.attrib_name]indexed chunks[/repr.attrib_name]: {num_indexed_rows}" + f" [repr.attrib_name]indexed chunks[/repr.attrib_name]: {info.vector_index.indexed_rows}" ) - if num_unindexed_rows > 0: + if info.vector_index.unindexed_rows > 0: self.console.print( - f" [repr.attrib_name]unindexed chunks[/repr.attrib_name]: [yellow]{num_unindexed_rows}[/yellow] " + f" [repr.attrib_name]unindexed chunks[/repr.attrib_name]: [yellow]{info.vector_index.unindexed_rows}[/yellow] " "(consider running: haiku-rag create-index)" ) else: self.console.print( - f" [repr.attrib_name]unindexed chunks[/repr.attrib_name]: {num_unindexed_rows}" + f" [repr.attrib_name]unindexed chunks[/repr.attrib_name]: {info.vector_index.unindexed_rows}" ) else: if num_chunks >= 256: @@ -172,49 +150,47 @@ class HaikuRAGApp: # pragma: no cover f"(need {256 - num_chunks} more chunks)" ) - if stats["documents"]["exists"]: + if tables["documents"].exists: self.console.print( f" [repr.attrib_name]versions (documents)[/repr.attrib_name]: " - f"{stats['documents']['num_versions']}" + f"{tables['documents'].num_versions}" ) - if stats["chunks"]["exists"]: + if tables["chunks"].exists: self.console.print( f" [repr.attrib_name]versions (chunks)[/repr.attrib_name]: " - f"{stats['chunks']['num_versions']}" + f"{tables['chunks'].num_versions}" ) # Migration status - pending = ( - get_pending_upgrades(stored_version) if stored_version != "unknown" else [] - ) self.console.rule() - if pending: + if info.pending_migrations: self.console.print( - f"[bold yellow]{len(pending)} migration(s) pending.[/bold yellow] " + f"[bold yellow]{len(info.pending_migrations)} migration(s) pending.[/bold yellow] " "Run [cyan]haiku-rag migrate[/cyan] to upgrade." ) - for step in pending: - desc = step.description or "" - self.console.print(f" [yellow]→[/yellow] {step.version}: {desc}") + for step in info.pending_migrations: + self.console.print( + f" [yellow]→[/yellow] {step.version}: {step.description}" + ) else: self.console.print("[green]Database is up to date.[/green]") self.console.rule() self.console.print("[bold]Versions[/bold]") self.console.print( - f" [repr.attrib_name]haiku.rag[/repr.attrib_name]: {versions['haiku_rag']}" + f" [repr.attrib_name]haiku.rag[/repr.attrib_name]: {info.packages['haiku_rag']}" ) self.console.print( - f" [repr.attrib_name]lancedb[/repr.attrib_name]: {versions['lancedb']}" + f" [repr.attrib_name]lancedb[/repr.attrib_name]: {info.packages['lancedb']}" ) self.console.print( - f" [repr.attrib_name]docling[/repr.attrib_name]: {versions['docling']}" + f" [repr.attrib_name]docling[/repr.attrib_name]: {info.packages['docling']}" ) self.console.print( - f" [repr.attrib_name]pydantic-ai[/repr.attrib_name]: {versions['pydantic_ai']}" + f" [repr.attrib_name]pydantic-ai[/repr.attrib_name]: {info.packages['pydantic_ai']}" ) self.console.print( - f" [repr.attrib_name]docling-document schema[/repr.attrib_name]: {versions['docling_document_schema']}" + f" [repr.attrib_name]docling-document schema[/repr.attrib_name]: {info.packages['docling_document_schema']}" ) async def history(self, table: str | None = None, limit: int | None = None): diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 2e158a5c..98bcfac4 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -13,7 +13,7 @@ import pyarrow as pa from lancedb.index import FTS, BTree, IvfPq from lancedb.pydantic import LanceModel, Vector from packaging.version import parse -from pydantic import Field +from pydantic import BaseModel, Field from haiku.rag.config import AppConfig, Config from haiku.rag.embeddings import get_embedder @@ -215,6 +215,116 @@ async def get_database_stats(db: lancedb.AsyncConnection) -> dict: return stats +class EmbeddingsInfo(BaseModel): + provider: str = "unknown" + name: str = "unknown" + vector_dim: int | None = None + + +class TableInfo(BaseModel): + name: str + exists: bool + num_rows: int = 0 + total_bytes: int = 0 + num_versions: int = 0 + + +class VectorIndexInfo(BaseModel): + exists: bool = False + indexed_rows: int = 0 + unindexed_rows: int = 0 + + +class PendingMigration(BaseModel): + version: str + description: str = "" + + +class DatabaseInfo(BaseModel): + """Structured snapshot of a haiku.rag database, shared by the `info` CLI + command and the ingester control plane. Read-only; gathered without + opening a Store.""" + + path: str + exists: bool + stored_version: str = "unknown" + embeddings: EmbeddingsInfo = Field(default_factory=EmbeddingsInfo) + tables: list[TableInfo] = Field(default_factory=list) + vector_index: VectorIndexInfo = Field(default_factory=VectorIndexInfo) + pending_migrations: list[PendingMigration] = Field(default_factory=list) + packages: dict[str, str] = Field(default_factory=dict) + + +async def gather_database_info(config: AppConfig, db_path: Path) -> DatabaseInfo: + """Collect read-only database state without going through Store, so a + database missing tables (e.g. pre-migration) still reports what it can.""" + from haiku.rag.store.upgrades import get_pending_upgrades + from haiku.rag.utils import get_package_versions + + display_path = config.lancedb.uri or str(db_path) + + db = await connect_lancedb(config, db_path) + stats = await get_database_stats(db) + + if not any(entry["exists"] for entry in stats.values()): + return DatabaseInfo(path=display_path, exists=False) + + stored_version = "unknown" + embeddings = EmbeddingsInfo() + if stats["settings"]["exists"]: + settings_tbl = await db.open_table("settings") + rows = ( + await settings_tbl.query().where("id = 'settings'").limit(1).to_arrow() + ).to_pylist() + if rows: + raw = rows[0].get("settings") or "{}" + data = json.loads(raw) if isinstance(raw, str) else (raw or {}) + stored_version = str(data.get("version", "unknown")) + model = data.get("embeddings", {}).get("model", {}) + embeddings = EmbeddingsInfo( + provider=model.get("provider", "unknown"), + name=model.get("name", "unknown"), + vector_dim=model.get("vector_dim"), + ) + + tables = [ + TableInfo( + name=name, + exists=stats[name]["exists"], + num_rows=stats[name].get("num_rows", 0), + total_bytes=stats[name].get("total_bytes", 0), + num_versions=stats[name].get("num_versions", 0), + ) + for name in ("documents", "chunks", "document_items") + ] + + vector_index = VectorIndexInfo() + if stats["chunks"]["exists"] and stats["chunks"].get("has_vector_index"): + vector_index = VectorIndexInfo( + exists=True, + indexed_rows=stats["chunks"].get("num_indexed_rows", 0), + unindexed_rows=stats["chunks"].get("num_unindexed_rows", 0), + ) + + pending = ( + get_pending_upgrades(stored_version) if stored_version != "unknown" else [] + ) + + return DatabaseInfo( + path=display_path, + exists=True, + stored_version=stored_version, + embeddings=embeddings, + tables=tables, + vector_index=vector_index, + pending_migrations=[ + PendingMigration(version=step.version, description=step.description or "") + for step in pending + ], + packages=get_package_versions(), + ) + + class Store: def __init__( self, diff --git a/tests/store/test_database_info.py b/tests/store/test_database_info.py new file mode 100644 index 00000000..21bd7bb1 --- /dev/null +++ b/tests/store/test_database_info.py @@ -0,0 +1,117 @@ +import json + +import pytest +from lancedb.pydantic import LanceModel, Vector +from pydantic import Field + +from haiku.rag.config.models import AppConfig +from haiku.rag.store.engine import DocumentItemRecord, gather_database_info + + +class _SettingsRecord(LanceModel): + id: str = Field(default="settings") + settings: str = Field(default="{}") + + +class _DocumentRecord(LanceModel): + id: str + content: str + + +class _ChunkRecord(LanceModel): + id: str + document_id: str + content: str + vector: Vector(3) # type: ignore + + +async def _seed(temp_db_path, *, version: str, with_items: bool = True): + import lancedb + + db = await lancedb.connect_async(temp_db_path) + settings_tbl = await db.create_table("settings", schema=_SettingsRecord) + docs_tbl = await db.create_table("documents", schema=_DocumentRecord) + chunks_tbl = await db.create_table("chunks", schema=_ChunkRecord) + if with_items: + await db.create_table("document_items", schema=DocumentItemRecord) + + await settings_tbl.add( + [ + _SettingsRecord( + id="settings", + settings=json.dumps( + { + "version": version, + "embeddings": { + "model": { + "provider": "openai", + "name": "text-embedding-3-small", + "vector_dim": 3, + } + }, + } + ), + ) + ] + ) + await docs_tbl.add([_DocumentRecord(id="doc-1", content="hello")]) + await chunks_tbl.add( + [ + _ChunkRecord( + id="c1", document_id="doc-1", content="c", vector=[0.1, 0.2, 0.3] + ) + ] + ) + + +@pytest.mark.asyncio +async def test_gather_database_info_reports_tables_and_settings(temp_db_path): + await _seed(temp_db_path, version="1.2.3") + + info = await gather_database_info(AppConfig(), temp_db_path) + + assert info.exists is True + assert info.path == str(temp_db_path) + assert info.stored_version == "1.2.3" + assert info.embeddings.provider == "openai" + assert info.embeddings.name == "text-embedding-3-small" + assert info.embeddings.vector_dim == 3 + + tables = {t.name: t for t in info.tables} + assert tables["documents"].exists and tables["documents"].num_rows == 1 + assert tables["chunks"].exists and tables["chunks"].num_rows == 1 + assert tables["document_items"].exists + assert tables["documents"].num_versions >= 1 + assert tables["chunks"].num_versions >= 1 + + # Only one chunk: no vector index. + assert info.vector_index.exists is False + + assert "haiku_rag" in info.packages + assert "lancedb" in info.packages + + +@pytest.mark.asyncio +async def test_gather_database_info_flags_missing_table_and_pending_migrations( + temp_db_path, +): + await _seed(temp_db_path, version="0.39.0", with_items=False) + + info = await gather_database_info(AppConfig(), temp_db_path) + + tables = {t.name: t for t in info.tables} + assert tables["document_items"].exists is False + assert info.pending_migrations # 0.39.0 is behind current schema + assert all(m.version and m.description is not None for m in info.pending_migrations) + + +@pytest.mark.asyncio +async def test_gather_database_info_empty_database(temp_db_path): + import lancedb + + await lancedb.connect_async(temp_db_path) # creates the dir, no tables + + info = await gather_database_info(AppConfig(), temp_db_path) + + assert info.exists is False + assert info.path == str(temp_db_path)