Replace allow_create with cleaner read_only. Elaborate on the test_info tests

This commit is contained in:
Yiorgis Gozadinos 2025-11-21 14:37:33 +02:00
parent dbb2837550
commit 2f54ee7d59
No known key found for this signature in database
5 changed files with 155 additions and 61 deletions

View file

@ -87,7 +87,9 @@ class HaikuRAGApp:
# Get comprehensive table statistics
from haiku.rag.store.engine import Store
store = Store(self.db_path, config=self.config)
store = Store(
self.db_path, config=self.config, skip_validation=True, read_only=True
)
table_stats = store.get_stats()
store.close()
@ -186,7 +188,7 @@ class HaikuRAGApp:
async def list_documents(self, filter: str | None = None):
async with HaikuRAG(
db_path=self.db_path, config=self.config, allow_create=False
db_path=self.db_path, config=self.config, read_only=True
) as self.client:
documents = await self.client.list_documents(filter=filter)
for doc in documents:
@ -221,7 +223,7 @@ class HaikuRAGApp:
async def get_document(self, doc_id: str):
async with HaikuRAG(
db_path=self.db_path, config=self.config, allow_create=False
db_path=self.db_path, config=self.config, read_only=True
) as self.client:
doc = await self.client.get_document_by_id(doc_id)
if doc is None:
@ -243,7 +245,7 @@ class HaikuRAGApp:
async def search(self, query: str, limit: int = 5, filter: str | None = None):
async with HaikuRAG(
db_path=self.db_path, config=self.config, allow_create=False
db_path=self.db_path, config=self.config, read_only=True
) as self.client:
results = await self.client.search(query, limit=limit, filter=filter)
if not results:
@ -268,7 +270,7 @@ class HaikuRAGApp:
verbose: Show verbose output
"""
async with HaikuRAG(
db_path=self.db_path, config=self.config, allow_create=False
db_path=self.db_path, config=self.config, read_only=True
) as self.client:
try:
if deep:
@ -315,7 +317,7 @@ class HaikuRAGApp:
verbose: Show AG-UI event stream during execution
"""
async with HaikuRAG(
db_path=self.db_path, config=self.config, allow_create=False
db_path=self.db_path, config=self.config, read_only=True
) as client:
try:
self.console.print("[bold cyan]Starting research[/bold cyan]")

View file

@ -29,7 +29,7 @@ class HaikuRAG:
db_path: Path | None = None,
config: AppConfig = Config,
skip_validation: bool = False,
allow_create: bool = True,
read_only: bool = False,
):
"""Initialize the RAG client with a database path.
@ -37,8 +37,8 @@ class HaikuRAG:
db_path: Path to the database file. If None, uses config.storage.data_dir.
config: Configuration to use. Defaults to global Config.
skip_validation: Whether to skip configuration validation on database load.
allow_create: Whether to allow database creation. If False, will raise error
if database doesn't exist (for read operations).
read_only: Whether to open in read-only mode. If True, will raise error
if database doesn't exist and will skip upgrades.
"""
self._config = config
if db_path is None:
@ -47,7 +47,7 @@ class HaikuRAG:
db_path,
config=self._config,
skip_validation=skip_validation,
allow_create=allow_create,
read_only=read_only,
)
self.document_repository = DocumentRepository(self.store)
self.chunk_repository = ChunkRepository(self.store)

View file

@ -54,19 +54,20 @@ class Store:
db_path: Path,
config: AppConfig = Config,
skip_validation: bool = False,
allow_create: bool = True,
read_only: bool = False,
):
self.db_path: Path = db_path
self._config = config
self.embedder = get_embedder(config=self._config)
self._vacuum_lock = asyncio.Lock()
self._read_only = read_only
# Create the ChunkRecord model with the correct vector dimension
self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
# Local filesystem handling for DB directory
if not self._has_cloud_config():
if not allow_create:
if read_only:
# Read operations should not create the database
if not db_path.exists():
raise FileNotFoundError(
@ -271,41 +272,43 @@ class Store:
)
# Run pending upgrades based on stored version and package version
try:
from haiku.rag.store.upgrades import run_pending_upgrades
current_version = metadata.version("haiku.rag-slim")
db_version = self.get_haiku_version()
if db_version != "0.0.0":
run_pending_upgrades(self, db_version, current_version)
# After upgrades complete (or if none), set stored version
# to the greater of the installed package version and the
# highest available upgrade step version in code.
# Skip in read-only mode to avoid modifying the database
if not self._read_only:
try:
from packaging.version import parse as _v
from haiku.rag.store.upgrades import run_pending_upgrades
from haiku.rag.store.upgrades import upgrades as _steps
current_version = metadata.version("haiku.rag-slim")
db_version = self.get_haiku_version()
highest_step = max((_v(u.version) for u in _steps), default=None)
effective_version = (
str(max(_v(current_version), highest_step))
if highest_step is not None
else current_version
if db_version != "0.0.0":
run_pending_upgrades(self, db_version, current_version)
# After upgrades complete (or if none), set stored version
# to the greater of the installed package version and the
# highest available upgrade step version in code.
try:
from packaging.version import parse as _v
from haiku.rag.store.upgrades import upgrades as _steps
highest_step = max((_v(u.version) for u in _steps), default=None)
effective_version = (
str(max(_v(current_version), highest_step))
if highest_step is not None
else current_version
)
except Exception:
effective_version = current_version
self.set_haiku_version(effective_version)
except Exception as e:
# Avoid hard failure on initial connection; log and continue so CLI remains usable.
logger.warning(
"Skipping upgrade due to error (db=%s -> pkg=%s): %s",
self.get_haiku_version(),
metadata.version("haiku.rag-slim"),
e,
)
except Exception:
effective_version = current_version
self.set_haiku_version(effective_version)
except Exception as e:
# Avoid hard failure on initial connection; log and continue so CLI remains usable.
logger.warning(
"Skipping upgrade due to error (db=%s -> pkg=%s): %s",
self.get_haiku_version(),
metadata.version("haiku.rag-slim"),
e,
)
def get_haiku_version(self) -> str:
"""Returns the user version stored in settings."""

View file

@ -14,12 +14,12 @@ def test_read_operations_do_not_create_database():
config = AppConfig()
# Read operation with allow_create=False should fail
# Read operation with read_only=True should fail
with pytest.raises(
FileNotFoundError,
match="Database does not exist.*Use a write operation",
):
HaikuRAG(db_path=db_path, config=config, allow_create=False)
HaikuRAG(db_path=db_path, config=config, read_only=True)
def test_write_operations_create_database():
@ -29,8 +29,8 @@ def test_write_operations_create_database():
config = AppConfig()
# Write operation with allow_create=True (default) should succeed
client = HaikuRAG(db_path=db_path, config=config, allow_create=True)
# Write operation with read_only=False (default) should succeed
client = HaikuRAG(db_path=db_path, config=config, read_only=False)
assert db_path.exists()
client.close()
@ -43,9 +43,7 @@ async def test_add_document_creates_database():
config = AppConfig()
# Create a document (write operation) should work and create DB
async with HaikuRAG(
db_path=db_path, config=config, allow_create=True
) as client:
async with HaikuRAG(db_path=db_path, config=config, read_only=False) as client:
doc = await client.create_document("Test content")
assert doc.id is not None
assert doc.content == "Test content"
@ -65,7 +63,7 @@ async def test_search_fails_if_database_does_not_exist():
match="Database does not exist.*Use a write operation",
):
async with HaikuRAG(
db_path=db_path, config=config, allow_create=False
db_path=db_path, config=config, read_only=True
) as client:
await client.search("test query")
@ -78,28 +76,24 @@ async def test_read_operations_work_after_database_created():
config = AppConfig()
# First, create DB via write operation
async with HaikuRAG(
db_path=db_path, config=config, allow_create=True
) as client:
async with HaikuRAG(db_path=db_path, config=config, read_only=False) as client:
await client.create_document("Test content", uri="test://doc1")
# Now read operations should work since DB exists
async with HaikuRAG(
db_path=db_path, config=config, allow_create=False
) as client:
async with HaikuRAG(db_path=db_path, config=config, read_only=True) as client:
docs = await client.list_documents()
assert len(docs) == 1
assert docs[0].content == "Test content"
def test_default_allow_create_is_true():
"""Test that allow_create defaults to True for backward compatibility."""
def test_default_read_only_is_false():
"""Test that read_only defaults to False for backward compatibility."""
with tempfile.TemporaryDirectory() as tmpdir:
db_path = Path(tmpdir) / "test.lancedb"
config = AppConfig()
# Without specifying allow_create, it should default to True
# Without specifying read_only, it should default to False (allow creation)
client = HaikuRAG(db_path=db_path, config=config)
assert db_path.exists()
client.close()

View file

@ -70,8 +70,20 @@ async def test_app_info_outputs_and_read_only(temp_db_path, capsys):
assert f"path: \n{temp_db_path}" in out
assert "haiku.rag version (db): 1.2.3" in out
assert "embeddings: openai/text-embedding-3-small (dim: 3)" in out
assert "lancedb:" 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
# Verify no versions changed (read-only)
# Re-open to ensure fresh view
@ -79,3 +91,86 @@ async def test_app_info_outputs_and_read_only(temp_db_path, capsys):
assert int(db2.open_table("settings").version) == before_versions["settings"]
assert int(db2.open_table("documents").version) == before_versions["documents"]
assert int(db2.open_table("chunks").version) == before_versions["chunks"]
@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.pydantic import LanceModel, Vector
from pydantic import Field
db = lancedb.connect(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 = db.create_table("settings", schema=SettingsRecord)
docs_tbl = db.create_table("documents", schema=DocumentRecord)
chunks_tbl = db.create_table("chunks", schema=ChunkRecord)
# Insert settings
settings_tbl.add(
[
SettingsRecord(
id="settings",
settings='{"version": "1.0.0", "embeddings": {"provider": "ollama", "model": "test", "vector_dim": 3}}',
)
]
)
# Insert document
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)
]
chunks_tbl.add(chunks)
# Create vector index
chunks_tbl.create_index(metric="cosine", index_type="IVF_PQ")
# Capture versions before
before_versions = {
"settings": int(settings_tbl.version),
"documents": int(docs_tbl.version),
"chunks": int(chunks_tbl.version),
}
app = HaikuRAGApp(db_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
# Verify no versions changed (read-only)
db2 = lancedb.connect(temp_db_path)
assert int(db2.open_table("settings").version) == before_versions["settings"]
assert int(db2.open_table("documents").version) == before_versions["documents"]
assert int(db2.open_table("chunks").version) == before_versions["chunks"]