Extract gather_database_info shared by the info command
This commit is contained in:
Yiorgis Gozadinos
parent
d45f8f5de8
commit
0be2e5b24f
3 changed files with 267 additions and 64 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
117
tests/store/test_database_info.py
Normal file
117
tests/store/test_database_info.py
Normal file
|
|
@ -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)
|
||||
Loading…
Reference in a new issue