Merge pull request #148 from ggozad/feat/vector-index
Vector indexing & Search accuracy improvements
This commit is contained in:
commit
28088aa31d
12 changed files with 464 additions and 75 deletions
21
CHANGELOG.md
21
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
|
||||
|
|
|
|||
28
docs/cli.md
28
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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]")
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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