9baa213b34
Store, connect_lancedb, gather_database_info and run_doctor take a location, a path or a URI, and classify it with ConnectionMode.of. SingleDatabaseSession owns the resolved DatabaseRef and passes its location down. This removes DatabaseRef.connection(), default_db_path, the placeholder path for URI-backed databases and the per-database config copies, so the configuration a client holds is the one the caller gave it. The chat hands its capabilities the scope it opened along with the client it lends, and the v0.58.0 migration no longer checks local free disk for a database behind a URI.
474 lines
16 KiB
Python
474 lines
16 KiB
Python
import json
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from haiku.rag.app import HaikuRAGApp
|
|
from haiku.rag.config.models import AppConfig, LanceDBConfig
|
|
from haiku.rag.store.schema import DocumentItemRecord
|
|
from tests.conftest import for_path
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_info_outputs(temp_db_path, capsys):
|
|
# Build a minimal LanceDB with settings, documents, and chunks without using Store
|
|
import lancedb
|
|
from lancedb.pydantic import LanceModel, Vector
|
|
from pydantic import Field
|
|
|
|
db = await lancedb.connect_async(temp_db_path)
|
|
|
|
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
|
|
|
|
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)
|
|
await db.create_table("document_items", schema=DocumentItemRecord)
|
|
|
|
# Insert one of each - using the new config format
|
|
await settings_tbl.add(
|
|
[
|
|
SettingsRecord(
|
|
id="settings",
|
|
settings=json.dumps(
|
|
{
|
|
"version": "1.2.3",
|
|
"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])]
|
|
)
|
|
|
|
app = HaikuRAGApp(scope=for_path(temp_db_path))
|
|
await app.info()
|
|
|
|
out = capsys.readouterr().out
|
|
# Validate expected content substrings
|
|
# Note: Rich console may wrap long paths to new lines, so check separately
|
|
assert "path:" in out
|
|
# Rich may wrap long paths across lines — check with newlines stripped
|
|
out_no_wrap = out.replace("\n", "")
|
|
assert str(temp_db_path) in out_no_wrap
|
|
assert "haiku.rag version (db):" in out
|
|
assert "embeddings: openai/text-embedding-3-small (dim: 3)" in out
|
|
assert "documents: 1" in out
|
|
assert "chunks: 1" in out
|
|
|
|
# Vector index should not exist (only 1 chunk, need 256)
|
|
assert "vector index: ✗ not created" in out
|
|
assert "need 255 more chunks" in out
|
|
|
|
# Table versions should be shown
|
|
assert "versions (documents):" in out
|
|
assert "versions (chunks):" in out
|
|
|
|
# Package versions section
|
|
assert "lancedb:" in out
|
|
assert "haiku.rag:" in out
|
|
assert "docling-document schema:" in out
|
|
assert "pydantic-ai:" in out
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_info_with_vector_index(temp_db_path, capsys):
|
|
# Build a database with enough chunks to create a vector index
|
|
import lancedb
|
|
from lancedb.index import IvfPq
|
|
from lancedb.pydantic import LanceModel, Vector
|
|
from pydantic import Field
|
|
|
|
db = await lancedb.connect_async(temp_db_path)
|
|
|
|
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
|
|
|
|
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)
|
|
await db.create_table("document_items", schema=DocumentItemRecord)
|
|
|
|
# Insert settings
|
|
await settings_tbl.add(
|
|
[
|
|
SettingsRecord(
|
|
id="settings",
|
|
settings='{"version": "1.0.0", "embeddings": {"model": {"provider": "ollama", "name": "test", "vector_dim": 3}}}',
|
|
)
|
|
]
|
|
)
|
|
|
|
# Insert document
|
|
await docs_tbl.add([DocumentRecord(id="doc-1", content="test")])
|
|
|
|
# Insert 512 chunks to allow index creation (PQ needs more than 256 for training)
|
|
chunks = [
|
|
ChunkRecord(
|
|
id=f"chunk-{i}",
|
|
document_id="doc-1",
|
|
content=f"content {i}",
|
|
vector=[0.1 * i, 0.2 * i, 0.3 * i],
|
|
)
|
|
for i in range(512)
|
|
]
|
|
await chunks_tbl.add(chunks)
|
|
|
|
# Create vector index
|
|
await chunks_tbl.create_index("vector", config=IvfPq(distance_type="cosine"))
|
|
|
|
app = HaikuRAGApp(scope=for_path(temp_db_path))
|
|
await app.info()
|
|
|
|
out = capsys.readouterr().out
|
|
|
|
# Check vector index exists
|
|
assert "vector index: ✓ exists" in out
|
|
assert "indexed chunks: 512" in out
|
|
assert "unindexed chunks: 0" in out
|
|
|
|
# Check basic info still present
|
|
assert "documents: 1" in out
|
|
assert "chunks: 512" in out
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_info_opens_a_named_remote_database(tmp_path):
|
|
"""A database named in `lancedb.databases` can sit behind a URI while the
|
|
configuration's own `uri` is empty; info derives and opens the URI."""
|
|
from haiku.rag.client.scope import DatabaseScope
|
|
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(databases={"papers": "s3://bucket/papers.lancedb"})
|
|
)
|
|
scope = DatabaseScope.resolve(config, database_name="papers")
|
|
app = HaikuRAGApp(scope=scope, config=config)
|
|
|
|
assert app._is_local is False
|
|
assert app._location == "s3://bucket/papers.lancedb"
|
|
|
|
with patch(
|
|
"haiku.rag.store.info.connect_lancedb", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
mock_db = mock_connect.return_value
|
|
mock_list_result = MagicMock()
|
|
mock_list_result.tables = []
|
|
mock_db.list_tables = AsyncMock(return_value=mock_list_result)
|
|
await app.info()
|
|
|
|
assert mock_connect.call_args.args[0] == "s3://bucket/papers.lancedb"
|
|
|
|
|
|
async def test_app_doctor_opens_a_named_remote_database():
|
|
"""`run_doctor` is handed the database's location, not the configuration
|
|
naming the set."""
|
|
from haiku.rag.client.scope import DatabaseScope
|
|
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(databases={"papers": "s3://bucket/papers.lancedb"})
|
|
)
|
|
app = HaikuRAGApp(scope=DatabaseScope.resolve(config, database_name="papers"))
|
|
|
|
with patch("haiku.rag.doctor.run_doctor", new_callable=AsyncMock) as run:
|
|
run.return_value = MagicMock(checks=[], ok=True, duplicates=None)
|
|
await app.doctor()
|
|
|
|
assert run.call_args.args[1] == "s3://bucket/papers.lancedb"
|
|
|
|
|
|
async def test_app_info_uses_connect_lancedb_for_remote(tmp_path):
|
|
"""info() should use connect_lancedb() instead of direct lancedb.connect() for remote URIs."""
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path",
|
|
storage_options={"endpoint": "http://localhost:9000"},
|
|
)
|
|
)
|
|
app = HaikuRAGApp(scope=for_path(None, config), config=config)
|
|
|
|
with patch(
|
|
"haiku.rag.store.info.connect_lancedb", new_callable=AsyncMock
|
|
) as mock_connect:
|
|
# Empty DB triggers the early-return path - enough to prove connect_lancedb was used
|
|
mock_db = mock_connect.return_value
|
|
mock_list_result = MagicMock()
|
|
mock_list_result.tables = []
|
|
mock_db.list_tables = AsyncMock(return_value=mock_list_result)
|
|
await app.info()
|
|
|
|
mock_connect.assert_called_once()
|
|
assert mock_connect.call_args.args[0] == "s3://bucket/path"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_info_with_missing_document_items_table(temp_db_path, capsys):
|
|
"""info() should still output database info and report pending migrations
|
|
when a required table is absent (as for a DB created before 0.40.0)."""
|
|
import lancedb
|
|
from lancedb.pydantic import LanceModel, Vector
|
|
from pydantic import Field
|
|
|
|
db = await lancedb.connect_async(temp_db_path)
|
|
|
|
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
|
|
|
|
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)
|
|
# Intentionally omit document_items (added in 0.40.0)
|
|
|
|
await settings_tbl.add(
|
|
[
|
|
SettingsRecord(
|
|
id="settings",
|
|
settings=json.dumps(
|
|
{
|
|
"version": "0.39.0",
|
|
"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])]
|
|
)
|
|
|
|
app = HaikuRAGApp(scope=for_path(temp_db_path))
|
|
await app.info()
|
|
|
|
out = capsys.readouterr().out
|
|
|
|
# Core stats should still be reported
|
|
assert "haiku.rag version (db): 0.39.0" in out
|
|
assert "documents: 1" in out
|
|
assert "chunks: 1" in out
|
|
|
|
# Missing table should be flagged, not cause a crash
|
|
assert "document_items: absent" in out
|
|
|
|
# Migration status should be surfaced
|
|
assert "migration(s) pending" in out
|
|
assert "haiku-rag migrate" in out
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_info_reports_up_to_date(temp_db_path, capsys):
|
|
"""info() should report the database is up to date when no migrations
|
|
are pending."""
|
|
from importlib import metadata
|
|
|
|
import lancedb
|
|
from lancedb.pydantic import LanceModel, Vector
|
|
from pydantic import Field
|
|
|
|
db = await lancedb.connect_async(temp_db_path)
|
|
|
|
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
|
|
|
|
settings_tbl = await db.create_table("settings", schema=SettingsRecord)
|
|
await db.create_table("documents", schema=DocumentRecord)
|
|
await db.create_table("chunks", schema=ChunkRecord)
|
|
await db.create_table("document_items", schema=DocumentItemRecord)
|
|
|
|
current_version = metadata.version("haiku.rag-slim")
|
|
await settings_tbl.add(
|
|
[
|
|
SettingsRecord(
|
|
id="settings",
|
|
settings=json.dumps(
|
|
{
|
|
"version": current_version,
|
|
"embeddings": {
|
|
"model": {
|
|
"provider": "openai",
|
|
"name": "text-embedding-3-small",
|
|
"vector_dim": 3,
|
|
}
|
|
},
|
|
}
|
|
),
|
|
)
|
|
]
|
|
)
|
|
|
|
app = HaikuRAGApp(scope=for_path(temp_db_path))
|
|
await app.info()
|
|
|
|
out = capsys.readouterr().out
|
|
assert "Database is up to date." in out
|
|
assert "migration(s) pending" not in out
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_init_skips_exists_check_for_remote(tmp_path):
|
|
"""init() should not check db_path.exists() for remote URIs."""
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path",
|
|
storage_options={"endpoint": "http://localhost:9000"},
|
|
)
|
|
)
|
|
app = HaikuRAGApp(scope=for_path(None, config), config=config)
|
|
|
|
with patch("haiku.rag.app.HaikuRAG") as mock_client_cls:
|
|
mock_client = AsyncMock()
|
|
covering = mock_client_cls._covering.return_value
|
|
covering.__aenter__ = AsyncMock(return_value=mock_client)
|
|
covering.__aexit__ = AsyncMock(return_value=False)
|
|
await app.init()
|
|
# A missing local path is opened to create, not returned early on.
|
|
mock_client_cls._covering.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_history_skips_exists_check_for_remote(tmp_path):
|
|
"""history() should not check db_path.exists() for remote URIs."""
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path",
|
|
storage_options={"endpoint": "http://localhost:9000"},
|
|
)
|
|
)
|
|
app = HaikuRAGApp(scope=for_path(None, config), config=config)
|
|
|
|
with patch("haiku.rag.store.engine.Store") as mock_store_cls:
|
|
mock_store = AsyncMock()
|
|
mock_store.list_table_versions = AsyncMock(return_value=[])
|
|
mock_store.list_tags = AsyncMock(return_value={})
|
|
mock_store_cls.return_value.__aenter__ = AsyncMock(return_value=mock_store)
|
|
mock_store_cls.return_value.__aexit__ = AsyncMock(return_value=False)
|
|
await app.history()
|
|
mock_store_cls.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_tag_rendering_escapes_markup(tmp_path):
|
|
"""lance forbids markup characters in ref names, but externally created
|
|
tags are rendered defensively: markup-looking names must come out as
|
|
literal text in tag list and history, not be interpreted by Rich."""
|
|
from rich.console import Console
|
|
|
|
from haiku.rag.store.engine import TagInfo
|
|
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path",
|
|
storage_options={"endpoint": "http://localhost:9000"},
|
|
)
|
|
)
|
|
app = HaikuRAGApp(scope=for_path(None, config), config=config)
|
|
app.console = Console(record=True, width=200)
|
|
|
|
hostile = "[red]release[/red]"
|
|
tags = {hostile: TagInfo(tables={"documents": 1}, missing_tables=[])}
|
|
|
|
with patch("haiku.rag.store.engine.Store") as mock_store_cls:
|
|
mock_store = AsyncMock()
|
|
mock_store.list_tags = AsyncMock(return_value=tags)
|
|
mock_store.list_table_versions = AsyncMock(
|
|
return_value=[{"version": 1, "timestamp": "2026-07-14 10:00:00"}]
|
|
)
|
|
mock_store_cls.return_value.__aenter__ = AsyncMock(return_value=mock_store)
|
|
mock_store_cls.return_value.__aexit__ = AsyncMock(return_value=False)
|
|
|
|
await app.list_tags()
|
|
await app.history(table="documents")
|
|
|
|
output = app.console.export_text()
|
|
assert output.count(hostile) == 2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_app_history_survives_tag_annotation_failure(tmp_path):
|
|
"""history degrades to version history without annotations, with a
|
|
warning, when aggregate tag loading fails."""
|
|
from rich.console import Console
|
|
|
|
config = AppConfig(
|
|
lancedb=LanceDBConfig(
|
|
uri="s3://bucket/path",
|
|
storage_options={"endpoint": "http://localhost:9000"},
|
|
)
|
|
)
|
|
app = HaikuRAGApp(scope=for_path(None, config), config=config)
|
|
app.console = Console(record=True, width=200)
|
|
|
|
with patch("haiku.rag.store.engine.Store") as mock_store_cls:
|
|
mock_store = AsyncMock()
|
|
mock_store.list_tags = AsyncMock(side_effect=RuntimeError("tags boom"))
|
|
mock_store.list_table_versions = AsyncMock(
|
|
return_value=[{"version": 1, "timestamp": "2026-07-15 10:00:00"}]
|
|
)
|
|
mock_store_cls.return_value.__aenter__ = AsyncMock(return_value=mock_store)
|
|
mock_store_cls.return_value.__aexit__ = AsyncMock(return_value=False)
|
|
|
|
await app.history(table="documents")
|
|
|
|
output = app.console.export_text()
|
|
assert "v1" in output
|
|
assert "2026-07-15 10:00:00" in output
|
|
assert "tags boom" in output
|