diff --git a/CHANGELOG.md b/CHANGELOG.md index 004bf3bb..e50da252 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,12 +1,30 @@ # Changelog ## [Unreleased] +### Added + +- **Manual Vector Indexing**: New `create-index` CLI command for explicit vector index creation + - Creates IVF_PQ indexes + - Requires minimum 256 chunks (LanceDB training data requirement) + - New `search.vector_index_metric` config option: `cosine` (default), `l2`, or `dot` + - New `search.vector_refine_factor` config option (default: 30) for accuracy/speed tradeoff + - Indexes not created automatically during ingestion to avoid performance degradation + - Manual rebuilding required after adding significant new data +- **Enhanced Info Command**: `haiku-rag info` now shows storage sizes and vector index statistics + - Displays storage size for documents and chunks tables in human-readable format + - Shows vector index status (exists/not created) + - Shows indexed and unindexed chunk counts for monitoring index staleness + ### Changed -- **Evaluations**: Improved evaluation dataset naming and simplified evaluator configuration +- **Evaluations**: Improved evaluation dataset naming and simplified evaluator +- configuration - `EvalDataset` now accepts dataset name for better organization in Logfire - Added `--name` CLI parameter to override evaluation run names - Removed `IsInstance` evaluator, using only `LLMJudge` for QA evaluation +- **Search Accuracy**: Applied `refine_factor` to vector and hybrid searches for improved accuracy + - Retrieves `refine_factor * limit` candidates and re-ranks in memory + - Higher values increase accuracy but slow down queries ### Fixed @@ -16,6 +34,7 @@ - `get_model()` utility function accepts `config` parameter (defaults to global Config) - Allows creating multiple graphs with different configurations in the same application + ## [0.17.2] - 2025-11-19 ### Added diff --git a/docs/cli.md b/docs/cli.md index 53d629b0..ec7a164b 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -215,14 +215,38 @@ Shows: - path to the database - stored haiku.rag version (from settings) - embeddings provider/model and vector dimension -- number of documents +- number of documents and chunks (with storage sizes) +- vector index status (exists/not created, indexed/unindexed chunks) - table versions per table (documents, chunks) -At the end, a separate “Versions” section lists runtime package versions: +At the end, a separate "Versions" section lists runtime package versions: - haiku.rag - lancedb - docling +### Create Vector Index + +Create a vector index on the chunks table for fast approximate nearest neighbor search: + +```bash +haiku-rag create-index [--db /path/to/your.lancedb] +``` + +**Requirements:** +- Minimum 256 chunks required for index creation (LanceDB training data requirement) +- Creates an IVF_PQ index using the configured `search.vector_index_metric` (cosine/l2/dot) + +**When to use:** +- After ingesting documents (indexes are not created automatically) +- After adding significant new data to rebuild the index +- Use `haiku-rag info` to check index status and see how many chunks are indexed/unindexed + +**Search behavior:** +- Without index: Brute-force kNN search (exact nearest neighbors, slower for large datasets) +- With index: Fast ANN (approximate nearest neighbors) using IVF_PQ +- With stale index: LanceDB combines indexed results (fast ANN) + brute-force kNN on unindexed rows +- Performance degrades as more unindexed data accumulates + ### Vacuum (Optimize and Cleanup) Reduce disk usage by optimizing and pruning old table versions across all tables: diff --git a/docs/configuration.md b/docs/configuration.md index 9e50ee4d..8a1fdb23 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -86,6 +86,10 @@ research: confidence_threshold: 0.8 max_concurrency: 1 +search: + vector_index_metric: cosine # cosine, l2, or dot + vector_refine_factor: 30 + agui: host: "0.0.0.0" port: 8000 @@ -721,6 +725,51 @@ haiku.rag intelligently handles database creation based on operation type: This prevents the common mistake where a search query accidentally creates an empty database. To initialize your database, simply add your first document using `haiku-rag add` or `haiku-rag add-src`. +### Vector Indexing + +Configure vector indexing behavior for efficient similarity search: + +```yaml +search: + vector_index_metric: cosine # cosine, l2, or dot + vector_refine_factor: 30 # Re-ranking factor for accuracy +``` + +- **vector_index_metric**: Distance metric for vector similarity: + - `cosine`: Cosine similarity (default, best for most embeddings) + - `l2`: Euclidean distance + - `dot`: Dot product similarity +- **vector_refine_factor**: Improves accuracy when using a vector index by retrieving `refine_factor * limit` candidates (using approximate search) and re-ranking them with exact distances. Higher values increase accuracy but slow down queries. Default: 30 + - **Only applies with a vector index** - has no effect on brute-force search, which already returns exact results + +!!! note + Vector indexes are only necessary for large datasets with over 100,000 chunks. For smaller datasets, LanceDB's brute-force kNN search provides exact results with good performance. Only create an index if you notice search performance degradation on large datasets. + +**Index creation:** + +Vector indexes are **not created automatically** during document ingestion to avoid slowing down the process. After you've added documents (at least 256 chunks required), create the index manually: + +```bash +haiku-rag create-index +``` + +This command: +- Checks if you have enough data (minimum 256 chunks) +- Creates an IVF_PQ index for fast approximate nearest neighbor (ANN) search +- Uses LanceDB's automatic parameter calculation based on your dataset size and vector dimensions + +**Re-indexing:** + +Indexes are not automatically updated when you add new documents. After adding a significant amount of new data: + +```bash +haiku-rag create-index # Rebuilds the index with all data +``` + +Searches still work with stale indexes - LanceDB uses the index for old data (fast ANN) and brute-force kNN for new unindexed rows, then combines the results. However, performance degrades as more unindexed data accumulates. + +For datasets with fewer than 256 chunks, searches use brute-force kNN scans (exact nearest neighbors, 100% recall) which work well for small datasets but don't scale beyond a few hundred thousand vectors. + ### Document Processing ```yaml diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index 01bf2401..7e75a575 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -18,6 +18,7 @@ from haiku.rag.mcp import create_mcp_server from haiku.rag.monitor import FileWatcher from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document +from haiku.rag.utils import format_bytes logger = logging.getLogger(__name__) @@ -83,10 +84,24 @@ class HaikuRAGApp: embed_model = embeddings.get("model") vector_dim = embeddings.get("vector_dim") - num_docs = 0 - if "documents" in table_names: - docs_tbl = db.open_table("documents") - num_docs = int(docs_tbl.count_rows()) # type: ignore[attr-defined] + # Get comprehensive table statistics + from haiku.rag.store.engine import Store + + store = Store( + self.db_path, config=self.config, skip_validation=True, read_only=True + ) + table_stats = store.get_stats() + store.close() + + num_docs = table_stats["documents"].get("num_rows", 0) + doc_bytes = table_stats["documents"].get("total_bytes", 0) + + num_chunks = table_stats["chunks"].get("num_rows", 0) + chunk_bytes = table_stats["chunks"].get("total_bytes", 0) + + has_vector_index = table_stats["chunks"].get("has_vector_index", False) + num_indexed_rows = table_stats["chunks"].get("num_indexed_rows", 0) + num_unindexed_rows = table_stats["chunks"].get("num_unindexed_rows", 0) # Table versions per table (direct API) doc_versions = ( @@ -116,8 +131,43 @@ class HaikuRAGApp: " [repr.attrib_name]embeddings[/repr.attrib_name]: unknown" ) self.console.print( - f" [repr.attrib_name]documents[/repr.attrib_name]: {num_docs}" + f" [repr.attrib_name]documents[/repr.attrib_name]: {num_docs} " + f"({format_bytes(doc_bytes)})" ) + self.console.print( + f" [repr.attrib_name]chunks[/repr.attrib_name]: {num_chunks} " + f"({format_bytes(chunk_bytes)})" + ) + + # Vector index information + if has_vector_index: + 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}" + ) + if num_unindexed_rows > 0: + self.console.print( + f" [repr.attrib_name]unindexed chunks[/repr.attrib_name]: [yellow]{num_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}" + ) + else: + if num_chunks >= 256: + self.console.print( + " [repr.attrib_name]vector index[/repr.attrib_name]: [yellow]✗ not created[/yellow] " + "(run: haiku-rag create-index)" + ) + else: + self.console.print( + f" [repr.attrib_name]vector index[/repr.attrib_name]: ✗ not created " + f"(need {256 - num_chunks} more chunks)" + ) + self.console.print( f" [repr.attrib_name]versions (documents)[/repr.attrib_name]: {doc_versions}" ) @@ -138,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: @@ -173,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: @@ -195,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: @@ -220,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: @@ -267,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]") @@ -397,6 +447,39 @@ class HaikuRAGApp: except Exception as e: self.console.print(f"[red]Error during vacuum: {e}[/red]") + async def create_index(self): + """Create vector index on the chunks table.""" + try: + async with HaikuRAG( + db_path=self.db_path, config=self.config, skip_validation=True + ) as client: + row_count = client.store.chunks_table.count_rows() + self.console.print(f"Chunks in database: {row_count}") + + if row_count < 256: + self.console.print( + f"[yellow]Warning: Need at least 256 chunks to create an index (have {row_count})[/yellow]" + ) + return + + # Check if index already exists + indices = client.store.chunks_table.list_indices() + has_vector_index = any("vector" in str(idx).lower() for idx in indices) + + if has_vector_index: + self.console.print( + "[yellow]Rebuilding existing vector index...[/yellow]" + ) + else: + self.console.print("[bold]Creating vector index...[/bold]") + + client.store._ensure_vector_index() + self.console.print( + "[bold green]Vector index created successfully.[/bold green]" + ) + except Exception as e: + self.console.print(f"[red]Error creating index: {e}[/red]") + def show_settings(self): """Display current configuration settings.""" self.console.print("[bold]haiku.rag configuration[/bold]") diff --git a/haiku_rag_slim/haiku/rag/cli.py b/haiku_rag_slim/haiku/rag/cli.py index dc99ec2a..e086179d 100644 --- a/haiku_rag_slim/haiku/rag/cli.py +++ b/haiku_rag_slim/haiku/rag/cli.py @@ -375,6 +375,18 @@ def vacuum( asyncio.run(app.vacuum()) +@cli.command("create-index", help="Create vector index for efficient similarity search") +def create_index( + db: Path | None = typer.Option( + None, + "--db", + help="Path to the LanceDB database file", + ), +): + app = create_app(db) + asyncio.run(app.create_index()) + + @cli.command("info", help="Show read-only database info (no upgrades or writes)") def info( db: Path | None = typer.Option( diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index 0ecff0fe..8b2d3e25 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/config/models.py b/haiku_rag_slim/haiku/rag/config/models.py index f71fc03b..59807bae 100644 --- a/haiku_rag_slim/haiku/rag/config/models.py +++ b/haiku_rag_slim/haiku/rag/config/models.py @@ -81,6 +81,11 @@ class ProcessingConfig(BaseModel): conversion_options: ConversionOptions = Field(default_factory=ConversionOptions) +class SearchConfig(BaseModel): + vector_index_metric: Literal["cosine", "l2", "dot"] = "cosine" + vector_refine_factor: int = 30 + + class OllamaConfig(BaseModel): base_url: str = Field( default_factory=lambda: __import__("os").environ.get( @@ -127,5 +132,6 @@ class AppConfig(BaseModel): qa: QAConfig = Field(default_factory=QAConfig) research: ResearchConfig = Field(default_factory=ResearchConfig) processing: ProcessingConfig = Field(default_factory=ProcessingConfig) + search: SearchConfig = Field(default_factory=SearchConfig) providers: ProvidersConfig = Field(default_factory=ProvidersConfig) agui: AGUIConfig = Field(default_factory=AGUIConfig) diff --git a/haiku_rag_slim/haiku/rag/store/engine.py b/haiku_rag_slim/haiku/rag/store/engine.py index 4c2a54e3..455d2e14 100644 --- a/haiku_rag_slim/haiku/rag/store/engine.py +++ b/haiku_rag_slim/haiku/rag/store/engine.py @@ -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( @@ -145,6 +146,87 @@ class Store: and self._config.lancedb.region ) + def get_stats(self) -> dict: + """Get comprehensive table statistics. + + Returns: + Dictionary with statistics for documents and chunks tables including: + - Row counts + - Storage sizes + - Vector index status and statistics + """ + stats_dict: dict = { + "documents": {"exists": False}, + "chunks": {"exists": False}, + } + + # Documents table stats + doc_stats: dict = self.documents_table.stats() # type: ignore[assignment] + stats_dict["documents"] = { + "exists": True, + "num_rows": doc_stats.get("num_rows", 0), + "total_bytes": doc_stats.get("total_bytes", 0), + } + + # Chunks table stats + chunk_stats: dict = self.chunks_table.stats() # type: ignore[assignment] + stats_dict["chunks"] = { + "exists": True, + "num_rows": chunk_stats.get("num_rows", 0), + "total_bytes": chunk_stats.get("total_bytes", 0), + } + + # Vector index stats + indices = self.chunks_table.list_indices() + has_vector_index = any("vector" in str(idx).lower() for idx in indices) + stats_dict["chunks"]["has_vector_index"] = has_vector_index + + if has_vector_index: + index_stats = self.chunks_table.index_stats("vector_idx") + if index_stats is not None: + stats_dict["chunks"]["num_indexed_rows"] = index_stats.num_indexed_rows + stats_dict["chunks"]["num_unindexed_rows"] = ( + index_stats.num_unindexed_rows + ) + + return stats_dict + + def _ensure_vector_index(self) -> None: + """Create or rebuild vector index on chunks table. + + Cloud deployments auto-create indexes, so we skip for those. + For self-hosted, creates an IVF_PQ index. If an index exists, + it will be replaced (using replace=True parameter). + Note: Index creation requires sufficient training data. + """ + if self._has_cloud_config(): + return + + try: + # Check if table has enough data (indexes require training data) + row_count = self.chunks_table.count_rows() + if row_count < 256: + logger.debug( + f"Skipping vector index creation: need at least 256 rows, have {row_count}" + ) + return + + # Create or replace index (replace=True is the default) + logger.info("Creating vector index on chunks table...") + self.chunks_table.create_index( + metric=self._config.search.vector_index_metric, + index_type="IVF_PQ", + replace=True, # Explicit: replace existing index + ) + + # Wait for index creation to complete + # Index name is column_name + "_idx" + self.chunks_table.wait_for_index(["vector_idx"], timeout=timedelta(hours=1)) + + logger.info("Vector index created successfully") + except Exception as e: + logger.warning(f"Could not create vector index: {e}") + def _validate_configuration(self) -> None: """Validate that the configuration is compatible with the database.""" from haiku.rag.store.repositories.settings import SettingsRepository @@ -190,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.""" diff --git a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py index d5083808..dd251e6a 100644 --- a/haiku_rag_slim/haiku/rag/store/repositories/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/repositories/chunk.py @@ -1,12 +1,16 @@ import inspect import json import logging -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from uuid import uuid4 if TYPE_CHECKING: import pandas as pd - from lancedb.query import LanceQueryBuilder + from lancedb.query import ( + LanceHybridQueryBuilder, + LanceQueryBuilder, + LanceVectorQueryBuilder, + ) from lancedb.rerankers import RRFReranker @@ -319,8 +323,14 @@ class ChunkRepository: # Prepare search query based on search type if search_type == "vector": query_embedding = await self.embedder.embed(query) - results = self.store.chunks_table.search( - query_embedding, query_type="vector", vector_column_name="vector" + vector_query = cast( + "LanceVectorQueryBuilder", + self.store.chunks_table.search( + query_embedding, query_type="vector", vector_column_name="vector" + ), + ) + results = vector_query.refine_factor( + self.store._config.search.vector_refine_factor ) elif search_type == "fts": @@ -331,12 +341,15 @@ class ChunkRepository: # Create RRF reranker reranker = RRFReranker() # Perform native hybrid search with RRF reranking - results = ( + hybrid_query = cast( + "LanceHybridQueryBuilder", self.store.chunks_table.search(query_type="hybrid") .vector(query_embedding) - .text(query) - .rerank(reranker) + .text(query), ) + results = hybrid_query.refine_factor( + self.store._config.search.vector_refine_factor + ).rerank(reranker) # Apply filtering if needed (common for all search types) if filtered_doc_ids is not None: diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index c08750b7..ae9d6e04 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -8,6 +8,16 @@ from types import ModuleType from packaging.version import Version, parse +def format_bytes(num_bytes: int) -> str: + """Format bytes as human-readable string.""" + size = float(num_bytes) + for unit in ["B", "KB", "MB", "GB", "TB"]: + if size < 1024.0: + return f"{size:.1f} {unit}" + size /= 1024.0 + return f"{size:.1f} PB" + + def get_default_data_dir() -> Path: """Get the user data directory for the current system platform. diff --git a/tests/test_database_autocreate.py b/tests/test_database_autocreate.py index 75359ab5..77ef7c15 100644 --- a/tests/test_database_autocreate.py +++ b/tests/test_database_autocreate.py @@ -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() diff --git a/tests/test_info.py b/tests/test_info.py index dbc9669e..9dae5f93 100644 --- a/tests/test_info.py +++ b/tests/test_info.py @@ -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"]