diff --git a/app/backend/main.py b/app/backend/main.py index d76b33e7..cfbd7f0b 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -149,18 +149,20 @@ async def db_info(_: Request) -> JSONResponse: } ) + from haiku.rag.store.engine import get_database_stats + client = get_client() - stats = client.store.get_stats() + stats = get_database_stats(client.store.db) return JSONResponse( { "exists": True, "path": str(db_path), - "documents": stats.get("documents", {}).get("num_rows", 0), - "chunks": stats.get("chunks", {}).get("num_rows", 0), - "documents_bytes": stats.get("documents", {}).get("total_bytes", 0), - "chunks_bytes": stats.get("chunks", {}).get("total_bytes", 0), - "has_vector_index": stats.get("chunks", {}).get("has_vector_index", False), + "documents": stats["documents"].get("num_rows", 0), + "chunks": stats["chunks"].get("num_rows", 0), + "documents_bytes": stats["documents"].get("total_bytes", 0), + "chunks_bytes": stats["chunks"].get("total_bytes", 0), + "has_vector_index": stats["chunks"].get("has_vector_index", False), } ) diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index eb470c03..dba8aba6 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -68,7 +68,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 + from haiku.rag.store.engine import connect_lancedb, get_database_stats from haiku.rag.store.upgrades import get_pending_upgrades if self.before is not None: @@ -89,9 +89,9 @@ class HaikuRAGApp: # pragma: no cover # Connect directly. Don't go through Store so a database that is # missing tables (e.g. pre-migration) still reports what it can. db = connect_lancedb(self.config, self.db_path) - existing_tables = set(db.list_tables().tables) + stats = get_database_stats(db) - if not existing_tables: + if not any(entry["exists"] for entry in stats.values()): self.console.print( "[red]Database is empty. Use 'haiku-rag init' to initialize.[/red]" ) @@ -103,7 +103,7 @@ class HaikuRAGApp: # pragma: no cover embed_provider = "unknown" embed_model = "unknown" vector_dim = None - if "settings" in existing_tables: + if stats["settings"]["exists"]: settings_tbl = db.open_table("settings") rows = ( settings_tbl.search() @@ -131,16 +131,14 @@ class HaikuRAGApp: # pragma: no cover f"{embed_provider}/{embed_model} (dim: {dim_part})" ) - # Per-table row counts and sizes. Missing required tables are shown - # as "absent" rather than raising. + # Per-table row counts and sizes. Missing required tables are + # reported as "absent" rather than raising. for name in ("documents", "chunks", "document_items"): - if name in existing_tables: - stats = db.open_table(name).stats() - num_rows = stats.get("num_rows", 0) - total_bytes = stats.get("total_bytes", 0) + entry = stats[name] + if entry["exists"]: self.console.print( - f" [repr.attrib_name]{name}[/repr.attrib_name]: {num_rows} " - f"({format_bytes(total_bytes)})" + f" [repr.attrib_name]{name}[/repr.attrib_name]: {entry['num_rows']} " + f"({format_bytes(entry['total_bytes'])})" ) else: self.console.print( @@ -148,20 +146,11 @@ class HaikuRAGApp: # pragma: no cover ) # Vector index information - if "chunks" in existing_tables: - chunks_tbl = db.open_table("chunks") - num_chunks = chunks_tbl.stats().get("num_rows", 0) - indices = chunks_tbl.list_indices() - has_vector_index = any("vector" in str(idx).lower() for idx in indices) - - if has_vector_index: - index_stats = chunks_tbl.index_stats("vector_idx") - num_indexed_rows = ( - index_stats.num_indexed_rows if index_stats is not None else 0 - ) - num_unindexed_rows = ( - index_stats.num_unindexed_rows if index_stats is not None else 0 - ) + 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) self.console.print( " [repr.attrib_name]vector index[/repr.attrib_name]: ✓ exists" ) @@ -189,15 +178,15 @@ class HaikuRAGApp: # pragma: no cover f"(need {256 - num_chunks} more chunks)" ) - if "documents" in existing_tables: - doc_versions = len(list(db.open_table("documents").list_versions())) + if stats["documents"]["exists"]: self.console.print( - f" [repr.attrib_name]versions (documents)[/repr.attrib_name]: {doc_versions}" + f" [repr.attrib_name]versions (documents)[/repr.attrib_name]: " + f"{stats['documents']['num_versions']}" ) - if "chunks" in existing_tables: - chunk_versions = len(list(db.open_table("chunks").list_versions())) + if stats["chunks"]["exists"]: self.console.print( - f" [repr.attrib_name]versions (chunks)[/repr.attrib_name]: {chunk_versions}" + f" [repr.attrib_name]versions (chunks)[/repr.attrib_name]: " + f"{stats['chunks']['num_versions']}" ) # Migration status @@ -207,8 +196,8 @@ class HaikuRAGApp: # pragma: no cover self.console.rule() if pending: self.console.print( - f"[bold yellow]Pending migrations: {len(pending)}[/bold yellow] " - "(run [cyan]haiku-rag migrate[/cyan] to upgrade)" + f"[bold yellow]{len(pending)} migration(s) pending.[/bold yellow] " + "Run [cyan]haiku-rag migrate[/cyan] to upgrade." ) for step in pending: desc = step.description or "" diff --git a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py index 38eb1d76..d9fd8a31 100644 --- a/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py +++ b/haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py @@ -64,7 +64,11 @@ class InfoModal(ModalScreen): async def on_mount(self) -> None: """Load and display database info.""" - from haiku.rag.store.engine import ConnectionMode, connect_lancedb + from haiku.rag.store.engine import ( + ConnectionMode, + connect_lancedb, + get_database_stats, + ) lines: list[str] = [] @@ -81,7 +85,7 @@ class InfoModal(ModalScreen): config = self.client.store._config try: db = connect_lancedb(config, self.db_path) - table_names = set(db.list_tables().tables) + stats = get_database_stats(db) except Exception as e: lines.append(f"[red]Failed to open database: {e}[/red]") self._content_widget.update("\n".join(lines)) @@ -90,16 +94,13 @@ class InfoModal(ModalScreen): # Get versions versions = get_package_versions() - # Get stats from store - table_stats = self.client.store.get_stats() - # Read settings stored_version = "unknown" embed_provider: str | None = None embed_model: str | None = None vector_dim: int | None = None - if "settings" in table_names: + if stats["settings"]["exists"]: settings_tbl = db.open_table("settings") arrow = settings_tbl.search().where("id = 'settings'").limit(1).to_arrow() rows = arrow.to_pylist() if arrow is not None else [] @@ -113,27 +114,18 @@ class InfoModal(ModalScreen): embed_model = embed_model_obj.get("name") vector_dim = embed_model_obj.get("vector_dim") - num_docs = table_stats["documents"].get("num_rows", 0) - doc_bytes = table_stats["documents"].get("total_bytes", 0) + num_docs = stats["documents"].get("num_rows", 0) + doc_bytes = stats["documents"].get("total_bytes", 0) - num_chunks = table_stats["chunks"].get("num_rows", 0) - chunk_bytes = table_stats["chunks"].get("total_bytes", 0) + num_chunks = stats["chunks"].get("num_rows", 0) + chunk_bytes = stats["chunks"].get("total_bytes", 0) - has_vector_index = table_stats["chunks"].get("has_vector_index", False) - num_indexed_rows = table_stats["chunks"].get("num_indexed_rows", 0) - num_unindexed_rows = table_stats["chunks"].get("num_unindexed_rows", 0) + has_vector_index = stats["chunks"].get("has_vector_index", False) + num_indexed_rows = stats["chunks"].get("num_indexed_rows", 0) + num_unindexed_rows = stats["chunks"].get("num_unindexed_rows", 0) - # Table versions - doc_versions = ( - len(list(db.open_table("documents").list_versions())) - if "documents" in table_names - else 0 - ) - chunk_versions = ( - len(list(db.open_table("chunks").list_versions())) - if "chunks" in table_names - else 0 - ) + doc_versions = stats["documents"].get("num_versions", 0) + chunk_versions = stats["chunks"].get("num_versions", 0) # Build output lines.append( diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 99d0e9cf..292e6176 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -121,6 +121,50 @@ class SettingsRecord(LanceModel): settings: str = Field(default="{}") +REQUIRED_TABLES: tuple[str, ...] = ("documents", "chunks", "document_items", "settings") + + +def get_database_stats(db: lancedb.DBConnection) -> dict: + """Collect stats for every haiku.rag table on the connection. + + Missing tables return ``{"exists": False}``. Present tables include + ``num_rows``, ``total_bytes``, and ``num_versions``. The ``chunks`` + entry additionally reports vector index status and, when an index + exists, ``num_indexed_rows`` and ``num_unindexed_rows``. + """ + existing = set(db.list_tables().tables) + stats: dict = {} + tables: dict = {} + + for name in REQUIRED_TABLES: + if name not in existing: + stats[name] = {"exists": False} + continue + tbl = db.open_table(name) + tables[name] = tbl + # lancedb's .stats() stub claims TableStatistics but returns a plain dict at runtime. + tbl_stats: dict = tbl.stats() # ty: ignore[invalid-assignment] + stats[name] = { + "exists": True, + "num_rows": tbl_stats.get("num_rows", 0), + "total_bytes": tbl_stats.get("total_bytes", 0), + "num_versions": len(list(tbl.list_versions())), + } + + if stats["chunks"]["exists"]: + chunks_tbl = tables["chunks"] + indices = chunks_tbl.list_indices() + has_vector_index = any("vector" in str(idx).lower() for idx in indices) + stats["chunks"]["has_vector_index"] = has_vector_index + if has_vector_index: + index_stats = chunks_tbl.index_stats("vector_idx") + if index_stats is not None: + stats["chunks"]["num_indexed_rows"] = index_stats.num_indexed_rows + stats["chunks"]["num_unindexed_rows"] = index_stats.num_unindexed_rows + + return stats + + class Store: def __init__( self, @@ -277,51 +321,6 @@ class Store: def _connection_mode(self) -> ConnectionMode: return ConnectionMode.from_config(self._config) - def get_stats(self) -> dict: - """Get comprehensive table statistics. - - Returns: - Dictionary with statistics for documents and chunks tables including: - - Row counts - - Storage sizes - - Vector index status and statistics - """ - stats_dict: dict = { - "documents": {"exists": False}, - "chunks": {"exists": False}, - } - - # Documents table stats - doc_stats: dict = self.documents_table.stats() - stats_dict["documents"] = { - "exists": True, - "num_rows": doc_stats.get("num_rows", 0), - "total_bytes": doc_stats.get("total_bytes", 0), - } - - # Chunks table stats - chunk_stats: dict = self.chunks_table.stats() - stats_dict["chunks"] = { - "exists": True, - "num_rows": chunk_stats.get("num_rows", 0), - "total_bytes": chunk_stats.get("total_bytes", 0), - } - - # Vector index stats - indices = self.chunks_table.list_indices() - has_vector_index = any("vector" in str(idx).lower() for idx in indices) - stats_dict["chunks"]["has_vector_index"] = has_vector_index - - if has_vector_index: - index_stats = self.chunks_table.index_stats("vector_idx") - if index_stats is not None: - stats_dict["chunks"]["num_indexed_rows"] = index_stats.num_indexed_rows - stats_dict["chunks"]["num_unindexed_rows"] = ( - index_stats.num_unindexed_rows - ) - - return stats_dict - def _ensure_vector_index(self) -> None: """Create or rebuild vector index on chunks table. @@ -368,8 +367,7 @@ class Store: def _init_tables(self): """Initialize database tables (create if they don't exist).""" existing_tables = self.db.list_tables().tables - required_tables = {"documents", "chunks", "document_items", "settings"} - missing_tables = required_tables - set(existing_tables) + missing_tables = set(REQUIRED_TABLES) - set(existing_tables) if missing_tables and self._read_only: raise ReadOnlyError( diff --git a/tests/store/test_engine.py b/tests/store/test_engine.py new file mode 100644 index 00000000..e4a00256 --- /dev/null +++ b/tests/store/test_engine.py @@ -0,0 +1,103 @@ +import pytest + +from haiku.rag.store import Store +from haiku.rag.store.engine import get_database_stats + + +class TestGetDatabaseStats: + def test_empty_database_stats(self, temp_db_path): + """get_database_stats() on a fresh database reports zero rows and no vector index.""" + store = Store(temp_db_path, create=True) + + stats = get_database_stats(store.db) + + for name in ("documents", "chunks", "document_items", "settings"): + assert stats[name]["exists"] is True + assert stats[name]["num_rows"] >= 0 + assert stats[name]["total_bytes"] >= 0 + assert stats[name]["num_versions"] >= 1 + + assert stats["documents"]["num_rows"] == 0 + assert stats["chunks"]["num_rows"] == 0 + assert stats["chunks"]["has_vector_index"] is False + + store.close() + + def test_missing_tables_report_absent(self, temp_db_path): + """Tables that don't exist on the connection are reported as absent.""" + import lancedb + from lancedb.pydantic import LanceModel + from pydantic import Field + + class SettingsRecord(LanceModel): + id: str = Field(default="settings") + settings: str = Field(default="{}") + + db = lancedb.connect(temp_db_path) + db.create_table("settings", schema=SettingsRecord) + + stats = get_database_stats(db) + + assert stats["settings"]["exists"] is True + assert stats["documents"] == {"exists": False} + assert stats["chunks"] == {"exists": False} + assert stats["document_items"] == {"exists": False} + + @pytest.mark.asyncio + async def test_stats_after_adding_document(self, temp_db_path): + """get_database_stats() reflects document and chunk counts after inserts.""" + from haiku.rag.store.models import Chunk, Document + from haiku.rag.store.repositories.chunk import ChunkRepository + from haiku.rag.store.repositories.document import DocumentRepository + + store = Store(temp_db_path, create=True) + doc_repo = DocumentRepository(store) + chunk_repo = ChunkRepository(store) + + doc = await doc_repo.create(Document(content="hello world")) + assert doc.id is not None + + await chunk_repo.create( + Chunk( + content="hello world", + document_id=doc.id, + embedding=[0.0] * store.embedder._vector_dim, + ) + ) + + stats = get_database_stats(store.db) + assert stats["documents"]["num_rows"] == 1 + assert stats["chunks"]["num_rows"] == 1 + store.close() + + def test_stats_with_vector_index(self, temp_db_path): + """get_database_stats() reports vector index details once an index exists.""" + from datetime import timedelta + + store = Store(temp_db_path, create=True) + dim = store.embedder._vector_dim + + # Need >=256 rows for IVF_PQ training. + rows = [ + { + "id": f"chunk-{i}", + "document_id": "doc-1", + "content": f"content {i}", + "content_fts": "", + "metadata": "{}", + "order": i, + "vector": [float(i % 7) + 0.01 * j for j in range(dim)], + } + for i in range(256) + ] + store.chunks_table.add(rows) + store.chunks_table.create_index( + metric="cosine", index_type="IVF_PQ", replace=True + ) + store.chunks_table.wait_for_index(["vector_idx"], timeout=timedelta(minutes=1)) + + stats = get_database_stats(store.db) + assert stats["chunks"]["has_vector_index"] is True + assert stats["chunks"]["num_indexed_rows"] >= 0 + assert "num_unindexed_rows" in stats["chunks"] + store.close() diff --git a/tests/test_info.py b/tests/test_info.py index abe90307..0ef0dadc 100644 --- a/tests/test_info.py +++ b/tests/test_info.py @@ -247,7 +247,7 @@ async def test_app_info_with_missing_document_items_table(temp_db_path, capsys): assert "document_items: absent" in out # Migration status should be surfaced - assert "Pending migrations:" in out + assert "migration(s) pending" in out assert "haiku-rag migrate" in out @@ -308,7 +308,7 @@ async def test_app_info_reports_up_to_date(temp_db_path, capsys): out = capsys.readouterr().out assert "Database is up to date." in out - assert "Pending migrations" not in out + assert "migration(s) pending" not in out @pytest.mark.asyncio diff --git a/tests/test_s3_integration.py b/tests/test_s3_integration.py index b763c1df..ba482949 100644 --- a/tests/test_s3_integration.py +++ b/tests/test_s3_integration.py @@ -52,9 +52,11 @@ def _make_config() -> AppConfig: def test_store_connect_and_create(tmp_path): + from haiku.rag.store.engine import get_database_stats + config = _make_config() store = Store(tmp_path / "unused", config=config, create=True) - stats = store.get_stats() + stats = get_database_stats(store.db) assert stats["documents"]["exists"] assert stats["chunks"]["exists"] store.close() @@ -69,7 +71,7 @@ async def test_store_vacuum(tmp_path): def test_store_add_document(tmp_path): - from haiku.rag.store.engine import DocumentRecord + from haiku.rag.store.engine import DocumentRecord, get_database_stats config = _make_config() store = Store(tmp_path / "unused", config=config, create=True) @@ -77,7 +79,7 @@ def test_store_add_document(tmp_path): doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.") store.documents_table.add([doc]) - stats = store.get_stats() + stats = get_database_stats(store.db) assert stats["documents"]["num_rows"] == 1 store.close()