From 9e5f11060d5f8eff837c1920e235adf001456964 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 2 Sep 2025 12:45:17 +0300 Subject: [PATCH] Allow connecting to lancedb cloud, s3, gc, az etc --- docs/configuration.md | 32 ++++++++- src/haiku/rag/config.py | 4 ++ src/haiku/rag/store/engine.py | 11 ++- src/haiku/rag/store/repositories/chunk.py | 5 ++ tests/test_lancedb_connection.py | 86 +++++++++++++++++++++++ 5 files changed, 136 insertions(+), 2 deletions(-) create mode 100644 tests/test_lancedb_connection.py diff --git a/docs/configuration.md b/docs/configuration.md index e02ae2d2..7e472bd5 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -140,11 +140,41 @@ COHERE_API_KEY="your-api-key" ### Database and Storage +By default, `haiku.rag` uses a local LanceDB database: + ```bash -# Default data directory (where SQLite database is stored) +# Default data directory (where local LanceDB is stored) DEFAULT_DATA_DIR="/path/to/data" ``` +For remote storage, use the `LANCEDB_URI` setting with various backends: + +```bash +# LanceDB Cloud +LANCEDB_URI="db://your-database-name" +LANCEDB_API_KEY="your-api-key" +LANCEDB_REGION="us-west-2" # optional + +# Amazon S3 +LANCEDB_URI="s3://my-bucket/my-table" +# Use AWS credentials or IAM roles + +# Azure Blob Storage +LANCEDB_URI="az://my-container/my-table" +# Use Azure credentials + +# Google Cloud Storage +LANCEDB_URI="gs://my-bucket/my-table" +# Use GCP credentials + +# HDFS +LANCEDB_URI="hdfs://namenode:port/path/to/table" +``` + +Authentication is handled through standard cloud provider credentials (AWS CLI, Azure CLI, gcloud, etc.) or by setting `LANCEDB_API_KEY` for LanceDB Cloud. + +**Note:** Table optimization is automatically handled by LanceDB Cloud (`db://` URIs) and is disabled for better performance. For object storage backends (S3, Azure, GCS), optimization is still performed locally. + ### Document Processing ```bash diff --git a/src/haiku/rag/config.py b/src/haiku/rag/config.py index 8910fab8..285c36f5 100644 --- a/src/haiku/rag/config.py +++ b/src/haiku/rag/config.py @@ -12,6 +12,10 @@ load_dotenv() class AppConfig(BaseModel): ENV: str = "production" + LANCEDB_API_KEY: str = "" + LANCEDB_URI: str = "" + LANCEDB_REGION: str = "" + DEFAULT_DATA_DIR: Path = get_default_data_dir() MONITOR_DIRECTORIES: list[Path] = [] diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index 87cd86a5..e23fbb49 100644 --- a/src/haiku/rag/store/engine.py +++ b/src/haiku/rag/store/engine.py @@ -46,7 +46,16 @@ class Store: # Create the ChunkRecord model with the correct vector dimension self.ChunkRecord = create_chunk_model(self.embedder._vector_dim) - self.db = lancedb.connect(db_path) + # Connect to LanceDB (local, cloud, or object storage) + if Config.LANCEDB_URI and Config.LANCEDB_API_KEY and Config.LANCEDB_REGION: + self.db = lancedb.connect( + uri=Config.LANCEDB_URI, + api_key=Config.LANCEDB_API_KEY, + region=Config.LANCEDB_REGION, + ) + else: + # Local file system connection + self.db = lancedb.connect(db_path) self.create_or_update_db() diff --git a/src/haiku/rag/store/repositories/chunk.py b/src/haiku/rag/store/repositories/chunk.py index d0287d1b..48b6857e 100644 --- a/src/haiku/rag/store/repositories/chunk.py +++ b/src/haiku/rag/store/repositories/chunk.py @@ -6,6 +6,7 @@ from docling_core.types.doc.document import DoclingDocument from lancedb.rerankers import RRFReranker from haiku.rag.chunker import chunker +from haiku.rag.config import Config from haiku.rag.embeddings import get_embedder from haiku.rag.store.engine import DocumentRecord, Store from haiku.rag.store.models.chunk import Chunk @@ -28,6 +29,10 @@ class ChunkRepository: 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() diff --git a/tests/test_lancedb_connection.py b/tests/test_lancedb_connection.py new file mode 100644 index 00000000..ae821fe9 --- /dev/null +++ b/tests/test_lancedb_connection.py @@ -0,0 +1,86 @@ +from unittest.mock import patch + +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).""" + # 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 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) + + # The optimize method should NOT have been called for LanceDB Cloud + mock_optimize.assert_not_called() + + store.close() + + +@pytest.mark.asyncio +async def test_local_storage_calls_optimization(temp_db_path): + """Test that optimization is called 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) + + # The optimize method SHOULD have been called for local storage + mock_optimize.assert_called() + + store.close()