danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Centralize vacuuming

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-10-07 12:05:25 +03:00
parent ce11efd17c
commit 094f0e0d44
No known key found for this signature in database
7 changed files with 51 additions and 94 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
src/haiku/rag/client.py
View file

@ -617,13 +617,13 @@ class HaikuRAG:
# Final maintenance: centralized vacuum to curb disk usage
try:
self.store.vacuum()
await self.store.vacuum()
except Exception:
pass
async def vacuum(self) -> None:
"""Optimize and clean up old versions across all tables."""
self.store.vacuum()
await self.store.vacuum()
def close(self):
"""Close the underlying store connection."""

6
src/haiku/rag/migration.py
View file

@ -27,7 +27,7 @@ class SQLiteToLanceDBMigrator:
self.lancedb_path = lancedb_path
self.console = Console()
def migrate(self) -> bool:
async def migrate(self) -> bool:
"""Perform the migration."""
try:
self.console.print(
@ -94,7 +94,7 @@ class SQLiteToLanceDBMigrator:
# Optimize and cleanup using centralized vacuum
self.console.print("[cyan]Optimizing LanceDB...[/cyan]")
try:
lance_store.vacuum()
await lance_store.vacuum()
self.console.print("[green]✅ Optimization completed[/green]")
except Exception as e:
self.console.print(
@ -313,4 +313,4 @@ async def migrate_sqlite_to_lancedb(
lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
migrator = SQLiteToLanceDBMigrator(sqlite_path, lancedb_path)
return migrator.migrate()
return await migrator.migrate()

30
src/haiku/rag/store/engine.py
View file

@ -1,3 +1,4 @@
import asyncio
import json
import logging
from datetime import timedelta
@ -51,6 +52,7 @@ class Store:
def __init__(self, db_path: Path, skip_validation: bool = False):
self.db_path: Path = db_path
self.embedder = get_embedder()
self._vacuum_lock = asyncio.Lock()
# Create the ChunkRecord model with the correct vector dimension
self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
@ -78,20 +80,38 @@ class Store:
if not skip_validation:
self._validate_configuration()
def vacuum(self, retention_seconds: int = Config.VACUUM_RETENTION_SECONDS) -> None:
async def vacuum(
self, retention_seconds: int = Config.VACUUM_RETENTION_SECONDS
) -> None:
"""Optimize and clean up old versions across all tables to reduce disk usage.
Args:
retention_seconds: Retention threshold in seconds. Only versions older
than this will be removed. Defaults to Config.VACUUM_RETENTION_SECONDS.
Note:
If vacuum is already running, this method returns immediately without blocking.
"""
if self._has_cloud_config() and str(Config.LANCEDB_URI).startswith("db://"):
return
# Perform maintenance per table using optimize() with configurable retention
retention = timedelta(seconds=retention_seconds)
for table in [self.documents_table, self.chunks_table, self.settings_table]:
table.optimize(cleanup_older_than=retention)
# Skip if already running (non-blocking)
if self._vacuum_lock.locked():
return
async with self._vacuum_lock:
try:
# Perform maintenance per table using optimize() with configurable retention
retention = timedelta(seconds=retention_seconds)
for table in [
self.documents_table,
self.chunks_table,
self.settings_table,
]:
table.optimize(cleanup_older_than=retention)
except (RuntimeError, OSError) as e:
# Handle resource errors gracefully
logger.debug(f"Vacuum skipped due to resource constraints: {e}")
def _connect_to_lancedb(self, db_path: Path):
"""Establish connection to LanceDB (local, cloud, or object storage)."""

28
src/haiku/rag/store/repositories/chunk.py
View file

@ -1,4 +1,3 @@
import asyncio
import inspect
import json
import logging
@ -23,7 +22,6 @@ class ChunkRepository:
def __init__(self, store: Store) -> None:
self.store = store
self.embedder = get_embedder()
self._optimize_lock = asyncio.Lock()
def _ensure_fts_index(self) -> None:
"""Ensure FTS index exists on the content column."""
@ -35,21 +33,6 @@ class ChunkRepository:
# Log the error but don't fail - FTS might already exist
logger.debug(f"FTS index creation skipped: {e}")
async def _optimize(self) -> None:
"""Optimize the chunks table to refresh indexes."""
# Skip optimization for LanceDB Cloud as it handles this automatically
if Config.LANCEDB_URI and Config.LANCEDB_URI.startswith("db://"):
return
async with self._optimize_lock:
try:
self.store.chunks_table.optimize()
except (RuntimeError, OSError) as e:
# Handle "too many open files" and other resource errors gracefully
logger.debug(
f"Table optimization skipped due to resource constraints: {e}"
)
async def create(self, entity: Chunk) -> Chunk:
"""Create a chunk in the database."""
assert entity.document_id, "Chunk must have a document_id to be created"
@ -77,11 +60,6 @@ class ChunkRepository:
self.store.chunks_table.add([chunk_record])
entity.id = chunk_id
# Try to optimize if not currently locked (non-blocking)
if not self._optimize_lock.locked():
asyncio.create_task(self._optimize())
return entity
async def get_by_id(self, entity_id: str) -> Chunk | None:
@ -125,10 +103,6 @@ class ChunkRepository:
"vector": embedding,
},
)
# Try to optimize if not currently locked (non-blocking)
if not self._optimize_lock.locked():
asyncio.create_task(self._optimize())
return entity
async def delete(self, entity_id: str) -> bool:
@ -227,8 +201,6 @@ class ChunkRepository:
if chunk_records:
self.store.chunks_table.add(chunk_records)
# Force optimization once at the end for bulk operations
await self._optimize()
return created_chunks
async def delete_all(self) -> None:

11
src/haiku/rag/store/repositories/document.py
View file

@ -1,3 +1,4 @@
import asyncio
import json
from datetime import datetime
from typing import TYPE_CHECKING
@ -200,8 +201,9 @@ class DocumentRepository:
chunk.order = order
await self.chunk_repository.create(chunk)
# Vacuum old versions after successful creation
self.store.vacuum()
# Vacuum old versions in background (non-blocking)
asyncio.create_task(self.store.vacuum())
return created_doc
except Exception:
# Roll back to the captured versions and re-raise
@ -232,8 +234,9 @@ class DocumentRepository:
updated_doc.id, docling_document
)
# Vacuum old versions after successful update
self.store.vacuum()
# Vacuum old versions in background (non-blocking)
asyncio.create_task(self.store.vacuum())
return updated_doc
except Exception:
# Roll back to the captured versions and re-raise

62
tests/test_lancedb_connection.py
View file

@ -4,43 +4,24 @@ import pytest
from haiku.rag.config import Config
from haiku.rag.store.engine import Store
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
from haiku.rag.store.repositories.chunk import ChunkRepository
from haiku.rag.store.repositories.document import DocumentRepository
@pytest.mark.asyncio
async def test_lancedb_cloud_skips_optimization(temp_db_path):
"""Test that optimization is skipped when using LanceDB Cloud (db:// URI)."""
"""Test that vacuum is skipped when using LanceDB Cloud (db:// URI)."""
# Create a store
store = Store(temp_db_path)
chunk_repo = ChunkRepository(store)
doc_repo = DocumentRepository(store)
# Create a document
document = Document(content="Test document content", metadata={})
created_document = await doc_repo.create(document)
document_id = created_document.id
# Mock LANCEDB_URI to simulate LanceDB Cloud usage
with patch.object(Config, "LANCEDB_URI", "db://test-database"):
# Mock all cloud config to simulate LanceDB Cloud usage
with (
patch.object(Config, "LANCEDB_URI", "db://test-database"),
patch.object(Config, "LANCEDB_API_KEY", "test-api-key"),
patch.object(Config, "LANCEDB_REGION", "us-east-1"),
):
# Mock the optimize method to track if it's called
with patch.object(store.chunks_table, "optimize") as mock_optimize:
# Create a chunk - this should trigger optimization logic
chunk = Chunk(
document_id=document_id,
content="Test chunk content",
metadata={"test": "value"},
)
created_chunk = await chunk_repo.create(chunk)
assert created_chunk.id is not None
# Wait a moment to ensure any async optimization would complete
import asyncio
await asyncio.sleep(0.1)
# Call vacuum - this should skip optimization for LanceDB Cloud
await store.vacuum()
# The optimize method should NOT have been called for LanceDB Cloud
mock_optimize.assert_not_called()
@ -50,35 +31,16 @@ async def test_lancedb_cloud_skips_optimization(temp_db_path):
@pytest.mark.asyncio
async def test_local_storage_calls_optimization(temp_db_path):
"""Test that optimization is called for local storage."""
"""Test that vacuum calls optimization for local storage."""
# Create a store
store = Store(temp_db_path)
chunk_repo = ChunkRepository(store)
doc_repo = DocumentRepository(store)
# Create a document
document = Document(content="Test document content", metadata={})
created_document = await doc_repo.create(document)
document_id = created_document.id
# Ensure LANCEDB_URI is empty (local storage)
with patch.object(Config, "LANCEDB_URI", ""):
# Mock the optimize method to track if it's called
with patch.object(store.chunks_table, "optimize") as mock_optimize:
# Create a chunk - this should trigger optimization logic
chunk = Chunk(
document_id=document_id,
content="Test chunk content",
metadata={"test": "value"},
)
created_chunk = await chunk_repo.create(chunk)
assert created_chunk.id is not None
# Wait a moment to ensure async optimization completes
import asyncio
await asyncio.sleep(0.1)
# Call vacuum - this should optimize all tables for local storage
await store.vacuum()
# The optimize method SHOULD have been called for local storage
mock_optimize.assert_called()

4
tests/test_versioning.py
View file

@ -158,7 +158,7 @@ async def test_vacuum_with_retention_threshold(temp_db_path):
# Vacuum with default threshold (60 seconds) - should keep recent versions
# Note: vacuum may create new versions even when not cleaning up old ones
store.vacuum()
await store.vacuum()
after_default_doc_versions = len(list(store.documents_table.list_versions()))
after_default_chunk_versions = len(list(store.chunks_table.list_versions()))
@ -173,7 +173,7 @@ async def test_vacuum_with_retention_threshold(temp_db_path):
)
# Vacuum with 0 threshold - should significantly reduce versions
store.vacuum(retention_seconds=0)
await store.vacuum(retention_seconds=0)
after_zero_doc_versions = len(list(store.documents_table.list_versions()))
after_zero_chunk_versions = len(list(store.chunks_table.list_versions()))