Replace allow_create with cleaner read_only. Elaborate on the test_info tests
This commit is contained in:
parent
dbb2837550
commit
2f54ee7d59
5 changed files with 155 additions and 61 deletions
|
|
@ -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]")
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Reference in a new issue