Merge pull request #335 from ggozad/feat/s3-object-store

add S3/object storage support for LanceDB connections
This commit is contained in:
Yiorgis Gozadinos 2026-04-09 15:27:18 +03:00 committed by GitHub
commit 4c1eec4ae7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 650 additions and 84 deletions

View file

@ -1,6 +1,11 @@
# Changelog # Changelog
## [Unreleased] ## [Unreleased]
### Added
- **S3/Object storage support**: Connect to LanceDB on S3, GCS, Azure Blob, or HDFS via `lancedb.uri` and `storage_options` config. Supports S3-compatible stores with custom endpoints.
- **Remote skill generation**: `create-skill` now supports remote databases — omit `--db` and provide `--config-file` to generate skills that connect to object storage at runtime instead of bundling the database.
## [0.38.0] - 2026-04-07 ## [0.38.0] - 2026-04-07
### Added ### Added

View file

@ -32,26 +32,47 @@ lancedb:
# Amazon S3 # Amazon S3
lancedb: lancedb:
uri: s3://my-bucket/my-table uri: s3://my-bucket/my-table
# Use AWS credentials or IAM roles storage_options:
region: us-east-1
# Amazon S3 with explicit credentials
lancedb:
uri: s3://my-bucket/my-table
storage_options:
aws_access_key_id: YOUR_ACCESS_KEY
aws_secret_access_key: YOUR_SECRET_KEY
region: us-east-1
# S3-compatible (SeaweedFS, Tigris, etc.)
lancedb:
uri: s3://my-bucket/my-table
storage_options:
endpoint: http://localhost:8333
aws_access_key_id: YOUR_ACCESS_KEY
aws_secret_access_key: YOUR_SECRET_KEY
region: us-east-1
allow_http: "true"
# Azure Blob Storage # Azure Blob Storage
lancedb: lancedb:
uri: az://my-container/my-table uri: az://my-container/my-table
# Use Azure credentials
# Google Cloud Storage # Google Cloud Storage
lancedb: lancedb:
uri: gs://my-bucket/my-table uri: gs://my-bucket/my-table
# Use GCP credentials
# HDFS # HDFS
lancedb: lancedb:
uri: hdfs://namenode:port/path/to/table uri: hdfs://namenode:port/path/to/table
``` ```
Authentication is handled through standard cloud provider credentials (AWS CLI, Azure CLI, gcloud, etc.) or by setting `api_key` for LanceDB Cloud. - **LanceDB Cloud** (`db://`): Requires `api_key` and `region`. Table optimization and indexing are managed server-side.
- **Object storage** (`s3://`, `gs://`, `az://`, `hdfs://`): Uses `storage_options` for credentials and endpoint configuration. Authentication can also be provided via environment variables (`AWS_ACCESS_KEY_ID`, `AWS_SECRET_ACCESS_KEY`, etc.) or cloud provider SDK defaults (AWS CLI, Azure CLI, gcloud).
- **S3-compatible stores** (MinIO, Tigris, etc.): Set `endpoint` in `storage_options`. When using `http://` endpoints, also set `allow_http: "true"`.
**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. The `storage_options` keys are case-insensitive and passed directly to the underlying object store library. Available keys depend on the backend — see the [LanceDB storage docs](https://lancedb.com/docs/storage/) for details.
**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 and vector indexing are still performed normally.
## Database Creation ## Database Creation

View file

@ -45,9 +45,14 @@ class HaikuRAGApp: # pragma: no cover
self.before = before self.before = before
self.console = Console() self.console = Console()
from haiku.rag.store.engine import ConnectionMode
self._is_local = ConnectionMode.from_config(self.config) == ConnectionMode.LOCAL
self._display_path = self.db_path if self._is_local else self.config.lancedb.uri
async def init(self): async def init(self):
"""Initialize a new database.""" """Initialize a new database."""
if self.db_path.exists(): if self._is_local and self.db_path.exists():
self.console.print( self.console.print(
f"[yellow]Database already exists at {self.db_path}[/yellow]" f"[yellow]Database already exists at {self.db_path}[/yellow]"
) )
@ -57,31 +62,35 @@ class HaikuRAGApp: # pragma: no cover
client = HaikuRAG(db_path=self.db_path, config=self.config, create=True) client = HaikuRAG(db_path=self.db_path, config=self.config, create=True)
client.close() client.close()
self.console.print( self.console.print(
f"[bold green]Database initialized at {self.db_path}[/bold green]" f"[bold green]Database initialized at {self._display_path}[/bold green]"
) )
async def info(self): async def info(self):
"""Display read-only information about the database without modifying it.""" """Display read-only information about the database without modifying it."""
import lancedb from haiku.rag.store.engine import Store, connect_lancedb
# Basic: show path # Basic: show path/URI
self.console.print("[bold]haiku.rag database info[/bold]") self.console.print("[bold]haiku.rag database info[/bold]")
self.console.print( self.console.print(
f" [repr.attrib_name]path[/repr.attrib_name]: {self.db_path}" f" [repr.attrib_name]path[/repr.attrib_name]: {self._display_path}"
) )
if not self.db_path.exists(): if self._is_local and not self.db_path.exists():
self.console.print("[red]Database path does not exist.[/red]") self.console.print("[red]Database path does not exist.[/red]")
return return
# Connect without going through Store to avoid upgrades/validation writes # Connect without going through Store to avoid upgrades/validation writes
db = lancedb.connect(self.db_path) db = connect_lancedb(self.config, self.db_path)
if not db.list_tables().tables:
self.console.print(
"[red]Database is empty. Use 'haiku-rag init' to initialize.[/red]"
)
return
versions = get_package_versions() versions = get_package_versions()
from haiku.rag.store.engine import Store
store = Store( store = Store(
self.db_path, self.db_path,
config=self.config, config=self.config,
@ -201,7 +210,7 @@ class HaikuRAGApp: # pragma: no cover
""" """
from haiku.rag.store.engine import Store from haiku.rag.store.engine import Store
if not self.db_path.exists(): if self._is_local and not self.db_path.exists():
self.console.print("[red]Database path does not exist.[/red]") self.console.print("[red]Database path does not exist.[/red]")
return return

View file

@ -718,7 +718,7 @@ def serve(
@_cli.command( @_cli.command(
"create-skill", "create-skill",
help="Generate a standalone skill package with an embedded database", help="Generate a standalone skill package with an embedded or remote database",
) )
def create_skill_cmd( # pragma: no cover def create_skill_cmd( # pragma: no cover
name: str = typer.Option( name: str = typer.Option(
@ -726,10 +726,10 @@ def create_skill_cmd( # pragma: no cover
"--name", "--name",
help="Skill name (lowercase alphanumeric and hyphens)", help="Skill name (lowercase alphanumeric and hyphens)",
), ),
db: Path = typer.Option( db: Path | None = typer.Option(
..., None,
"--db", "--db",
help="Path to the LanceDB database to embed", help="Path to the LanceDB database to embed (omit for remote storage)",
), ),
description: str | None = typer.Option( description: str | None = typer.Option(
None, None,

View file

@ -60,6 +60,7 @@ class LanceDBConfig(BaseModel):
uri: str = "" uri: str = ""
api_key: str = "" api_key: str = ""
region: str = "" region: str = ""
storage_options: dict[str, str] = Field(default_factory=dict)
class EmbeddingsConfig(BaseModel): class EmbeddingsConfig(BaseModel):

View file

@ -64,21 +64,23 @@ class InfoModal(ModalScreen):
async def on_mount(self) -> None: async def on_mount(self) -> None:
"""Load and display database info.""" """Load and display database info."""
import lancedb from haiku.rag.store.engine import ConnectionMode, connect_lancedb
lines: list[str] = [] lines: list[str] = []
# Path # Path
lines.append(f"[bold $accent]path[/bold $accent]: {self.db_path}") lines.append(f"[bold $accent]path[/bold $accent]: {self.db_path}")
if not self.db_path.exists(): is_local = self.client.store._connection_mode == ConnectionMode.LOCAL
if is_local and not self.db_path.exists():
lines.append("[red]Database path does not exist.[/red]") lines.append("[red]Database path does not exist.[/red]")
self._content_widget.update("\n".join(lines)) self._content_widget.update("\n".join(lines))
return return
# Connect to get table info # Connect to get table info
config = self.client.store._config
try: try:
db = lancedb.connect(self.db_path) db = connect_lancedb(config, self.db_path)
table_names = set(db.list_tables().tables) table_names = set(db.list_tables().tables)
except Exception as e: except Exception as e:
lines.append(f"[red]Failed to open database: {e}[/red]") lines.append(f"[red]Failed to open database: {e}[/red]")

View file

@ -74,6 +74,7 @@ def render_templates(
description: str, description: str,
tool_names: list[str], tool_names: list[str],
preamble: str | None = None, preamble: str | None = None,
remote: bool = False,
) -> pathlib.Path: ) -> pathlib.Path:
if preamble is None: if preamble is None:
preamble = DEFAULT_PREAMBLE preamble = DEFAULT_PREAMBLE
@ -88,6 +89,7 @@ def render_templates(
"tool_names": tool_names, "tool_names": tool_names,
"preamble": preamble, "preamble": preamble,
"rag_version": rag_version, "rag_version": rag_version,
"remote": remote,
} }
result_dir = output_dir / f"{name}-skill" result_dir = output_dir / f"{name}-skill"
@ -115,7 +117,7 @@ def render_templates(
def generate_skill( def generate_skill(
db_path: pathlib.Path, db_path: pathlib.Path | None,
output_dir: pathlib.Path, output_dir: pathlib.Path,
name: str, name: str,
description: str, description: str,
@ -125,7 +127,16 @@ def generate_skill(
) -> pathlib.Path: ) -> pathlib.Path:
validate_metadata(name, description) validate_metadata(name, description)
validate_tools(tool_names) validate_tools(tool_names)
validate_db_path(db_path)
if db_path is None:
if config_path is None:
raise ValueError(
"config_path is required when db_path is not provided "
"(remote storage needs connection config)"
)
else:
validate_db_path(db_path)
validate_output_dir(output_dir, name) validate_output_dir(output_dir, name)
result = render_templates( result = render_templates(
@ -134,11 +145,14 @@ def generate_skill(
description=description, description=description,
tool_names=tool_names, tool_names=tool_names,
preamble=preamble, preamble=preamble,
remote=db_path is None,
) )
pkg_name = name.replace("-", "_") pkg_name = name.replace("-", "_")
assets_dir = result / f"{pkg_name}_skill" / "assets" assets_dir = result / f"{pkg_name}_skill" / "assets"
shutil.copytree(db_path, assets_dir / f"{name}.lancedb")
if db_path is not None:
shutil.copytree(db_path, assets_dir / f"{name}.lancedb")
if config_path is not None: if config_path is not None:
shutil.copy2(config_path, assets_dir / "haiku.rag.yaml") shutil.copy2(config_path, assets_dir / "haiku.rag.yaml")

View file

@ -27,7 +27,11 @@ from haiku.rag.skills._tools import AnalysisEntry
_TOOL_NAMES = {{ tool_names | tojson }} _TOOL_NAMES = {{ tool_names | tojson }}
_ASSETS_DIR = Path(__file__).resolve().parent / "assets" _ASSETS_DIR = Path(__file__).resolve().parent / "assets"
{% if remote %}
_DB_PATH = None
{% else %}
_DB_PATH = _ASSETS_DIR / "{{ name }}.lancedb" _DB_PATH = _ASSETS_DIR / "{{ name }}.lancedb"
{% endif %}
_CONFIG_PATH = _ASSETS_DIR / "haiku.rag.yaml" _CONFIG_PATH = _ASSETS_DIR / "haiku.rag.yaml"

View file

@ -2,6 +2,7 @@ import asyncio
import json import json
import logging import logging
from datetime import datetime, timedelta from datetime import datetime, timedelta
from enum import Enum
from importlib import metadata from importlib import metadata
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any
@ -19,6 +20,40 @@ from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError
logger = logging.getLogger(__name__) logger = logging.getLogger(__name__)
class ConnectionMode(Enum):
LOCAL = "local"
CLOUD = "cloud"
OBJECT_STORAGE = "object_storage"
@staticmethod
def from_config(config: AppConfig) -> "ConnectionMode":
uri = config.lancedb.uri
if not uri:
return ConnectionMode.LOCAL
if uri.startswith("db://"):
return ConnectionMode.CLOUD
return ConnectionMode.OBJECT_STORAGE
def connect_lancedb(config: AppConfig, db_path: Path | None = None):
mode = ConnectionMode.from_config(config)
if mode == ConnectionMode.CLOUD:
return lancedb.connect(
uri=config.lancedb.uri,
api_key=config.lancedb.api_key,
region=config.lancedb.region,
)
elif mode == ConnectionMode.OBJECT_STORAGE:
kwargs: dict[str, Any] = {"uri": config.lancedb.uri}
if config.lancedb.storage_options:
kwargs["storage_options"] = config.lancedb.storage_options
return lancedb.connect(**kwargs)
else:
if db_path is None:
raise ValueError("No lancedb.uri configured and no db_path provided")
return lancedb.connect(db_path)
class DocumentRecord(LanceModel): class DocumentRecord(LanceModel):
id: str = Field(default_factory=lambda: str(uuid4())) id: str = Field(default_factory=lambda: str(uuid4()))
content: str content: str
@ -97,7 +132,7 @@ class Store:
# Check if database exists (for local filesystem only) # Check if database exists (for local filesystem only)
is_new_db = False is_new_db = False
if not self._has_cloud_config(): if self._connection_mode == ConnectionMode.LOCAL:
if not db_path.exists(): if not db_path.exists():
if not create: if not create:
raise FileNotFoundError( raise FileNotFoundError(
@ -110,7 +145,13 @@ class Store:
Path.mkdir(db_path.parent, parents=True) Path.mkdir(db_path.parent, parents=True)
# Connect to LanceDB # Connect to LanceDB
self.db = self._connect_to_lancedb(db_path) self.db = connect_lancedb(self._config, db_path)
# For remote stores, detect new DB by checking if tables exist
if not is_new_db and self._connection_mode != ConnectionMode.LOCAL:
existing_tables = self.db.list_tables().tables
if not existing_tables:
is_new_db = True
# For existing databases, read stored vector dimension to create ChunkRecord # For existing databases, read stored vector dimension to create ChunkRecord
# that can read existing chunks. For new databases, use config's dimension. # that can read existing chunks. For new databases, use config's dimension.
@ -198,9 +239,7 @@ class Store:
""" """
self._assert_writable() self._assert_writable()
if self._has_cloud_config() and str(self._config.lancedb.uri).startswith( if self._connection_mode == ConnectionMode.CLOUD:
"db://"
):
return return
# Skip if already running (non-blocking) # Skip if already running (non-blocking)
@ -224,26 +263,9 @@ class Store:
# Handle resource errors gracefully # Handle resource errors gracefully
logger.debug(f"Vacuum skipped due to resource constraints: {e}") logger.debug(f"Vacuum skipped due to resource constraints: {e}")
def _connect_to_lancedb(self, db_path: Path): @property
"""Establish connection to LanceDB (local, cloud, or object storage).""" def _connection_mode(self) -> ConnectionMode:
# Check if we have cloud configuration return ConnectionMode.from_config(self._config)
if self._has_cloud_config():
return lancedb.connect(
uri=self._config.lancedb.uri,
api_key=self._config.lancedb.api_key,
region=self._config.lancedb.region,
)
else:
# Local file system connection
return lancedb.connect(db_path)
def _has_cloud_config(self) -> bool:
"""Check if cloud configuration is complete."""
return bool(
self._config.lancedb.uri
and self._config.lancedb.api_key
and self._config.lancedb.region
)
def get_stats(self) -> dict: def get_stats(self) -> dict:
"""Get comprehensive table statistics. """Get comprehensive table statistics.
@ -298,7 +320,7 @@ class Store:
it will be replaced (using replace=True parameter). it will be replaced (using replace=True parameter).
Note: Index creation requires sufficient training data. Note: Index creation requires sufficient training data.
""" """
if self._has_cloud_config(): if self._connection_mode == ConnectionMode.CLOUD:
return return
try: try:

View file

@ -0,0 +1,18 @@
services:
seaweedfs:
image: chrislusf/seaweedfs
ports:
- "8333:8333"
command: server -s3 -s3.config=/etc/seaweedfs/s3-config.json
volumes:
- ./s3-config.json:/etc/seaweedfs/s3-config.json:ro
createbucket:
image: chrislusf/seaweedfs
depends_on:
- seaweedfs
entrypoint: >
/bin/sh -c "
sleep 3 &&
echo 's3.bucket.create -name test-bucket' | weed shell -master seaweedfs:9333
"

View file

@ -0,0 +1,14 @@
{
"identities": [
{
"name": "admin",
"credentials": [
{
"accessKey": "testkey",
"secretKey": "testsecret"
}
],
"actions": ["Admin", "Read", "Write", "List", "Tagging"]
}
]
}

View file

@ -1,8 +1,10 @@
import json import json
from unittest.mock import patch
import pytest import pytest
from haiku.rag.app import HaikuRAGApp from haiku.rag.app import HaikuRAGApp
from haiku.rag.config.models import AppConfig, LanceDBConfig
@pytest.mark.asyncio @pytest.mark.asyncio
@ -153,3 +155,93 @@ async def test_app_info_with_vector_index(temp_db_path, capsys):
# Check basic info still present # Check basic info still present
assert "documents: 1" in out assert "documents: 1" in out
assert "chunks: 512" in out assert "chunks: 512" in out
@pytest.mark.asyncio
async def test_app_info_uses_connect_lancedb_for_remote(tmp_path, capsys):
"""info() should use connect_lancedb() instead of direct lancedb.connect() for remote URIs."""
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path",
storage_options={"endpoint": "http://localhost:9000"},
)
)
app = HaikuRAGApp(db_path=nonexistent, config=config)
with patch("haiku.rag.store.engine.connect_lancedb") as mock_connect:
# Make the mock return something that lets info() proceed minimally
mock_db = mock_connect.return_value
mock_table = mock_db.open_table.return_value
mock_table.search.return_value.where.return_value.limit.return_value.to_arrow.return_value.to_pylist.return_value = [
{
"settings": json.dumps(
{
"version": "1.0.0",
"embeddings": {
"model": {
"provider": "test",
"name": "test",
"vector_dim": 3,
}
},
}
)
}
]
with patch("haiku.rag.store.engine.Store") as mock_store_cls:
mock_store = mock_store_cls.return_value
mock_store.get_stats.return_value = {
"documents": {"exists": True, "num_rows": 0, "total_bytes": 0},
"chunks": {
"exists": True,
"num_rows": 0,
"total_bytes": 0,
"has_vector_index": False,
"num_indexed_rows": 0,
"num_unindexed_rows": 0,
},
}
await app.info()
mock_connect.assert_called_once_with(config, nonexistent)
@pytest.mark.asyncio
async def test_app_init_skips_exists_check_for_remote(tmp_path):
"""init() should not check db_path.exists() for remote URIs."""
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path",
storage_options={"endpoint": "http://localhost:9000"},
)
)
app = HaikuRAGApp(db_path=nonexistent, config=config)
with patch("haiku.rag.app.HaikuRAG") as mock_client_cls:
await app.init()
# Should have called HaikuRAG to create, not returned early
mock_client_cls.assert_called_once()
@pytest.mark.asyncio
async def test_app_history_skips_exists_check_for_remote(tmp_path):
"""history() should not check db_path.exists() for remote URIs."""
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path",
storage_options={"endpoint": "http://localhost:9000"},
)
)
app = HaikuRAGApp(db_path=nonexistent, config=config)
with patch("haiku.rag.store.engine.Store") as mock_store_cls:
mock_store = mock_store_cls.return_value
mock_store.documents_table.list_versions.return_value = []
mock_store.chunks_table.list_versions.return_value = []
mock_store.settings_table.list_versions.return_value = []
await app.history()
mock_store_cls.assert_called_once()

View file

@ -3,46 +3,211 @@ from unittest.mock import patch
import pytest import pytest
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.store.engine import Store from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.engine import ConnectionMode, Store, connect_lancedb
@pytest.mark.asyncio class TestConnectionMode:
async def test_lancedb_cloud_skips_optimization(temp_db_path): def test_local_when_uri_empty(self):
"""Test that vacuum is skipped when using LanceDB Cloud (db:// URI).""" config = AppConfig(lancedb=LanceDBConfig(uri=""))
# Create a store assert ConnectionMode.from_config(config) == ConnectionMode.LOCAL
store = Store(temp_db_path, create=True)
# Mock all cloud config to simulate LanceDB Cloud usage def test_cloud_when_db_uri(self):
with ( config = AppConfig(
patch.object(Config.lancedb, "uri", "db://test-database"), lancedb=LanceDBConfig(
patch.object(Config.lancedb, "api_key", "test-api-key"), uri="db://my-database", api_key="key", region="us-east-1"
patch.object(Config.lancedb, "region", "us-east-1"), )
): )
# Mock the optimize method to track if it's called assert ConnectionMode.from_config(config) == ConnectionMode.CLOUD
with patch.object(store.chunks_table, "optimize") as mock_optimize:
# Call vacuum - this should skip optimization for LanceDB Cloud
await store.vacuum()
# The optimize method should NOT have been called for LanceDB Cloud def test_object_storage_s3(self):
mock_optimize.assert_not_called() config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
store.close() def test_object_storage_gs(self):
config = AppConfig(lancedb=LanceDBConfig(uri="gs://bucket/path"))
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
def test_object_storage_az(self):
config = AppConfig(lancedb=LanceDBConfig(uri="az://container/path"))
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
def test_object_storage_hdfs(self):
config = AppConfig(lancedb=LanceDBConfig(uri="hdfs://namenode/path"))
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
def test_unknown_uri_treated_as_object_storage(self):
config = AppConfig(lancedb=LanceDBConfig(uri="custom://something"))
assert ConnectionMode.from_config(config) == ConnectionMode.OBJECT_STORAGE
@pytest.mark.asyncio class TestConnectLancedb:
async def test_local_storage_calls_optimization(temp_db_path): def test_local_passes_db_path(self, temp_db_path):
"""Test that vacuum calls optimization for local storage.""" config = AppConfig(lancedb=LanceDBConfig(uri=""))
# Create a store with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect:
store = Store(temp_db_path, create=True) connect_lancedb(config, db_path=temp_db_path)
mock_connect.assert_called_once_with(temp_db_path)
# Ensure uri is empty (local storage) def test_cloud_passes_uri_api_key_region(self):
with patch.object(Config.lancedb, "uri", ""): config = AppConfig(
# Mock the optimize method to track if it's called lancedb=LanceDBConfig(
with patch.object(store.chunks_table, "optimize") as mock_optimize: uri="db://my-database", api_key="test-key", region="us-west-2"
# Call vacuum - this should optimize all tables for local storage )
await store.vacuum() )
with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect:
connect_lancedb(config)
mock_connect.assert_called_once_with(
uri="db://my-database", api_key="test-key", region="us-west-2"
)
# The optimize method SHOULD have been called for local storage def test_object_storage_passes_uri_and_storage_options(self):
mock_optimize.assert_called() config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path",
storage_options={
"endpoint": "http://minio:9000",
"region": "us-east-1",
},
)
)
with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect:
connect_lancedb(config)
mock_connect.assert_called_once_with(
uri="s3://bucket/path",
storage_options={
"endpoint": "http://minio:9000",
"region": "us-east-1",
},
)
store.close() def test_object_storage_without_storage_options(self):
config = AppConfig(lancedb=LanceDBConfig(uri="s3://bucket/path"))
with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect:
connect_lancedb(config)
mock_connect.assert_called_once_with(uri="s3://bucket/path")
def test_local_without_db_path_raises(self):
config = AppConfig(lancedb=LanceDBConfig(uri=""))
with pytest.raises(
ValueError, match="No lancedb.uri configured and no db_path provided"
):
connect_lancedb(config)
class TestStoreConnectionMode:
def test_store_connection_mode_local(self, temp_db_path):
store = Store(temp_db_path, create=True)
assert store._connection_mode == ConnectionMode.LOCAL
store.close()
def test_store_connection_mode_cloud(self, temp_db_path):
store = Store(temp_db_path, create=True)
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"),
):
assert store._connection_mode == ConnectionMode.CLOUD
store.close()
def test_store_connection_mode_object_storage(self, temp_db_path):
store = Store(temp_db_path, create=True)
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
assert store._connection_mode == ConnectionMode.OBJECT_STORAGE
store.close()
class TestVacuumByConnectionMode:
@pytest.mark.asyncio
async def test_cloud_skips_vacuum(self, temp_db_path):
store = Store(temp_db_path, create=True)
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"),
):
with patch.object(store.chunks_table, "optimize") as mock_optimize:
await store.vacuum()
mock_optimize.assert_not_called()
store.close()
@pytest.mark.asyncio
async def test_object_storage_runs_vacuum(self, temp_db_path):
store = Store(temp_db_path, create=True)
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
with patch.object(store.chunks_table, "optimize") as mock_optimize:
await store.vacuum()
mock_optimize.assert_called()
store.close()
@pytest.mark.asyncio
async def test_local_runs_vacuum(self, temp_db_path):
store = Store(temp_db_path, create=True)
with patch.object(Config.lancedb, "uri", ""):
with patch.object(store.chunks_table, "optimize") as mock_optimize:
await store.vacuum()
mock_optimize.assert_called()
store.close()
class TestVectorIndexByConnectionMode:
def test_cloud_skips_index_creation(self, temp_db_path):
store = Store(temp_db_path, create=True)
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"),
):
with patch.object(store.chunks_table, "count_rows") as mock_count:
store._ensure_vector_index()
mock_count.assert_not_called()
store.close()
def test_object_storage_runs_index_creation(self, temp_db_path):
store = Store(temp_db_path, create=True)
with patch.object(Config.lancedb, "uri", "s3://bucket/path"):
with patch.object(
store.chunks_table, "count_rows", return_value=0
) as mock_count:
store._ensure_vector_index()
mock_count.assert_called()
store.close()
class TestStoreSkipsPathValidationForRemote:
def test_skips_path_check_for_cloud(self, tmp_path):
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
config = AppConfig(
lancedb=LanceDBConfig(
uri="db://test-database", api_key="key", region="us-east-1"
)
)
with patch("haiku.rag.store.engine.lancedb.connect"):
with patch.object(Store, "_init_tables"):
store = Store(
nonexistent,
config=config,
create=True,
skip_validation=True,
skip_migration_check=True,
)
store.close()
def test_skips_path_check_for_object_storage(self, tmp_path):
nonexistent = tmp_path / "does_not_exist" / "db.lancedb"
config = AppConfig(
lancedb=LanceDBConfig(
uri="s3://bucket/path",
storage_options={"endpoint": "http://localhost:9000"},
)
)
with patch("haiku.rag.store.engine.lancedb.connect"):
with patch.object(Store, "_init_tables"):
store = Store(
nonexistent,
config=config,
create=True,
skip_validation=True,
skip_migration_check=True,
)
store.close()

View file

@ -0,0 +1,151 @@
# Start SeaweedFS before running:
# docker compose -f tests/docker/docker-compose.s3.yml up -d
# Stop after:
# docker compose -f tests/docker/docker-compose.s3.yml down -v
import socket
from uuid import uuid4
import pytest
from haiku.rag.app import HaikuRAGApp
from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig, LanceDBConfig
from haiku.rag.store.engine import Store
S3_ENDPOINT = "http://localhost:8333"
S3_BUCKET = "test-bucket"
S3_STORAGE_OPTIONS = {
"endpoint": S3_ENDPOINT,
"region": "us-east-1",
"allow_http": "true",
"aws_access_key_id": "testkey",
"aws_secret_access_key": "testsecret",
}
def _s3_available() -> bool:
try:
s = socket.create_connection(("localhost", 8333), timeout=1)
s.close()
return True
except OSError:
return False
pytestmark = [
pytest.mark.integration,
pytest.mark.skipif(
not _s3_available(), reason="SeaweedFS not running on localhost:8333"
),
]
def _make_config() -> AppConfig:
unique_prefix = uuid4().hex[:8]
return AppConfig(
lancedb=LanceDBConfig(
uri=f"s3://{S3_BUCKET}/test-{unique_prefix}",
storage_options=S3_STORAGE_OPTIONS,
)
)
def test_store_connect_and_create(tmp_path):
config = _make_config()
store = Store(tmp_path / "unused", config=config, create=True)
stats = store.get_stats()
assert stats["documents"]["exists"]
assert stats["chunks"]["exists"]
store.close()
@pytest.mark.asyncio
async def test_store_vacuum(tmp_path):
config = _make_config()
store = Store(tmp_path / "unused", config=config, create=True)
await store.vacuum()
store.close()
def test_store_add_document(tmp_path):
from haiku.rag.store.engine import DocumentRecord
config = _make_config()
store = Store(tmp_path / "unused", config=config, create=True)
doc = DocumentRecord(content="The quick brown fox jumps over the lazy dog.")
store.documents_table.add([doc])
stats = store.get_stats()
assert stats["documents"]["num_rows"] == 1
store.close()
@pytest.mark.asyncio
async def test_client_create_document(tmp_path):
config = _make_config()
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag:
doc = await rag.create_document(
"Python is a programming language.", uri="test://python"
)
assert doc.id
assert doc.uri == "test://python"
@pytest.mark.asyncio
async def test_client_list_documents(tmp_path):
config = _make_config()
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag:
await rag.create_document("First document.", uri="test://first")
await rag.create_document("Second document.", uri="test://second")
docs = await rag.list_documents()
assert len(docs) == 2
@pytest.mark.asyncio
async def test_client_search(tmp_path):
config = _make_config()
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag:
await rag.create_document(
"The Eiffel Tower is located in Paris, France.", uri="test://eiffel"
)
results = await rag.search("Eiffel Tower")
assert len(results) > 0
assert "Eiffel" in results[0].content
@pytest.mark.asyncio
async def test_client_delete_document(tmp_path):
config = _make_config()
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag:
doc = await rag.create_document("Temporary document.", uri="test://temp")
await rag.delete_document(doc.id)
docs = await rag.list_documents()
assert len(docs) == 0
@pytest.mark.asyncio
async def test_app_info(tmp_path, capsys):
config = _make_config()
async with HaikuRAG(tmp_path / "unused", config=config, create=True) as rag:
await rag.create_document("Info test document.", uri="test://info")
app = HaikuRAGApp(db_path=tmp_path / "unused", config=config)
await app.info()
out = capsys.readouterr().out
assert "path:" in out
assert config.lancedb.uri in out
assert "documents: 1" in out
@pytest.mark.asyncio
async def test_app_info_empty_db(tmp_path, capsys):
config = _make_config()
app = HaikuRAGApp(db_path=tmp_path / "unused", config=config)
await app.info()
out = capsys.readouterr().out
assert "Database is empty" in out

View file

@ -1,8 +1,10 @@
import shutil import shutil
import subprocess import subprocess
import zipfile import zipfile
from pathlib import Path
import pytest import pytest
import yaml
from haiku.rag.skill_generator import ( from haiku.rag.skill_generator import (
AVAILABLE_TOOLS, AVAILABLE_TOOLS,
@ -449,3 +451,49 @@ class TestGenerateSkill:
assert any(n.endswith("SKILL.md") for n in names) assert any(n.endswith("SKILL.md") for n in names)
assert any("assets/" in n and n.endswith("data.lance") for n in names) assert any("assets/" in n and n.endswith("data.lance") for n in names)
assert any(n.endswith("haiku.rag.yaml") for n in names) assert any(n.endswith("haiku.rag.yaml") for n in names)
def _make_remote_config(tmp_path: Path) -> Path:
config_file = tmp_path / "haiku.rag.yaml"
config_file.write_text(
yaml.dump(
{
"lancedb": {
"uri": "s3://my-bucket/haiku-rag",
"storage_options": {
"endpoint": "http://minio:9000",
"region": "us-east-1",
},
}
}
)
)
return config_file
class TestGenerateSkillRemote:
def test_remote_skips_copytree(self, tmp_path):
config_file = _make_remote_config(tmp_path)
result = generate_skill(
db_path=None,
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search", "ask"],
config_path=config_file,
)
assets = result / "recipes_skill" / "assets"
# No bundled database
assert not (assets / "recipes.lancedb").exists()
# Config must be copied
assert (assets / "haiku.rag.yaml").is_file()
def test_remote_requires_config_path(self, tmp_path):
with pytest.raises(ValueError, match="config_path.*required"):
generate_skill(
db_path=None,
output_dir=tmp_path,
name="recipes",
description="A recipe skill.",
tool_names=["search"],
)