danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

extract get_database_stats and share it across info, inspector, and backend

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-04-17 15:33:21 +03:00
parent a0a9a3410b
commit a1181ff52d
No known key found for this signature in database
7 changed files with 202 additions and 116 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
app/backend/main.py
View file

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

57
haiku_rag_slim/haiku/rag/app.py
View file

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

40
haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py
View file

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

92
haiku_rag_slim/haiku/rag/store/engine.py
View file

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

103
tests/store/test_engine.py Normal file
View file

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

4
tests/test_info.py
View file

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

8
tests/test_s3_integration.py
View file

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