9a3e2b86ca
`Store` already parses the settings blob on open, so keep it: the inspector's info modal was opening the settings table, querying it and parsing the JSON a second time, and `doctor` was asking `SettingsRepository` for it on a store it already held. Creating a database refreshes it, so a new one reports the version init wrote. `gather_database_info` keeps its own parse: it goes around `Store` so a pre-migration database still reports what it can.
131 lines
5.1 KiB
Python
131 lines
5.1 KiB
Python
import pytest
|
|
|
|
from haiku.rag.store.engine import Store
|
|
from haiku.rag.store.info import get_database_stats
|
|
|
|
|
|
class TestStoredSettings:
|
|
@pytest.mark.asyncio
|
|
async def test_a_new_database_carries_the_version_it_was_created_with(
|
|
self, temp_db_path
|
|
):
|
|
"""Creating writes the settings row, so the store has to report it
|
|
rather than the emptiness it opened on."""
|
|
from importlib import metadata
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
assert store.stored_settings["version"] == metadata.version(
|
|
"haiku.rag-slim"
|
|
)
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_an_existing_database_carries_its_stored_settings(self, temp_db_path):
|
|
"""Read once on open: reporting on a database reads them from here
|
|
instead of querying the settings table again."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
written = store.stored_settings
|
|
|
|
async with Store(temp_db_path) as store:
|
|
assert store.stored_settings == written
|
|
assert store.stored_settings["embeddings"]["model"]["vector_dim"] > 0
|
|
|
|
|
|
class TestGetDatabaseStats:
|
|
@pytest.mark.asyncio
|
|
async def test_empty_database_stats(self, temp_db_path):
|
|
"""get_database_stats() on a fresh database reports zero rows and no vector index."""
|
|
async with Store(temp_db_path, create=True) as store:
|
|
stats = await 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
|
|
|
|
@pytest.mark.asyncio
|
|
async 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 = await lancedb.connect_async(temp_db_path)
|
|
await db.create_table("settings", schema=SettingsRecord)
|
|
|
|
stats = await 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
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
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 = await get_database_stats(store.db)
|
|
assert stats["documents"]["num_rows"] == 1
|
|
assert stats["chunks"]["num_rows"] == 1
|
|
|
|
@pytest.mark.asyncio
|
|
async 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
|
|
|
|
from lancedb.index import IvfPq
|
|
|
|
async with Store(temp_db_path, create=True) as store:
|
|
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)
|
|
]
|
|
await store.chunks_table.add(rows)
|
|
await store.chunks_table.create_index(
|
|
"vector", config=IvfPq(distance_type="cosine"), replace=True
|
|
)
|
|
await store.chunks_table.wait_for_index(
|
|
["vector_idx"], timeout=timedelta(minutes=1)
|
|
)
|
|
|
|
stats = await 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"]
|