Allow connecting to lancedb cloud, s3, gc, az etc

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

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] = []
 

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()
 

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()

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()