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

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

Browse source 
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
Show stats Download patch file Download diff file Expand all files Collapse all files

5
CHANGELOG.md
View file

@ -1,6 +1,11 @@
# Changelog
## [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
### Added

31
docs/configuration/storage.md
View file

@ -32,26 +32,47 @@ lancedb:
# Amazon S3
lancedb:
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
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 `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

29
haiku_rag_slim/haiku/rag/app.py
View file

@ -45,9 +45,14 @@ class HaikuRAGApp: # pragma: no cover
self.before = before
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):
"""Initialize a new database."""
if self.db_path.exists():
if self._is_local and self.db_path.exists():
self.console.print(
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.close()
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):
"""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(
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]")
return
# 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()
from haiku.rag.store.engine import Store
store = Store(
self.db_path,
config=self.config,
@ -201,7 +210,7 @@ class HaikuRAGApp: # pragma: no cover
"""
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]")
return

8
haiku_rag_slim/haiku/rag/cli.py
View file

@ -718,7 +718,7 @@ def serve(
@_cli.command(
"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
name: str = typer.Option(
@ -726,10 +726,10 @@ def create_skill_cmd( # pragma: no cover
"--name",
help="Skill name (lowercase alphanumeric and hyphens)",
),
db: Path = typer.Option(
...,
db: Path | None = typer.Option(
None,
"--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(
None,

1
haiku_rag_slim/haiku/rag/config/models.py
View file

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

8
haiku_rag_slim/haiku/rag/inspector/widgets/info_modal.py
View file

@ -64,21 +64,23 @@ class InfoModal(ModalScreen):
async def on_mount(self) -> None:
"""Load and display database info."""
import lancedb
from haiku.rag.store.engine import ConnectionMode, connect_lancedb
lines: list[str] = []
# 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]")
self._content_widget.update("\n".join(lines))
return
# Connect to get table info
config = self.client.store._config
try:
db = lancedb.connect(self.db_path)
db = connect_lancedb(config, self.db_path)
table_names = set(db.list_tables().tables)
except Exception as e:
lines.append(f"[red]Failed to open database: {e}[/red]")

20
haiku_rag_slim/haiku/rag/skill_generator/__init__.py
View file

@ -74,6 +74,7 @@ def render_templates(
description: str,
tool_names: list[str],
preamble: str | None = None,
remote: bool = False,
) -> pathlib.Path:
if preamble is None:
preamble = DEFAULT_PREAMBLE
@ -88,6 +89,7 @@ def render_templates(
"tool_names": tool_names,
"preamble": preamble,
"rag_version": rag_version,
"remote": remote,
}
result_dir = output_dir / f"{name}-skill"
@ -115,7 +117,7 @@ def render_templates(
def generate_skill(
db_path: pathlib.Path,
db_path: pathlib.Path | None,
output_dir: pathlib.Path,
name: str,
description: str,
@ -125,7 +127,16 @@ def generate_skill(
) -> pathlib.Path:
validate_metadata(name, description)
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)
result = render_templates(
@ -134,11 +145,14 @@ def generate_skill(
description=description,
tool_names=tool_names,
preamble=preamble,
remote=db_path is None,
)
pkg_name = name.replace("-", "_")
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:
shutil.copy2(config_path, assets_dir / "haiku.rag.yaml")

4
haiku_rag_slim/haiku/rag/skill_generator/templates/__init__.py.j2
View file

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

74
haiku_rag_slim/haiku/rag/store/engine.py
View file

@ -2,6 +2,7 @@ import asyncio
import json
import logging
from datetime import datetime, timedelta
from enum import Enum
from importlib import metadata
from pathlib import Path
from typing import Any
@ -19,6 +20,40 @@ from haiku.rag.store.exceptions import MigrationRequiredError, ReadOnlyError
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):
id: str = Field(default_factory=lambda: str(uuid4()))
content: str
@ -97,7 +132,7 @@ class Store:
# Check if database exists (for local filesystem only)
is_new_db = False
if not self._has_cloud_config():
if self._connection_mode == ConnectionMode.LOCAL:
if not db_path.exists():
if not create:
raise FileNotFoundError(
@ -110,7 +145,13 @@ class Store:
Path.mkdir(db_path.parent, parents=True)
# 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
# that can read existing chunks. For new databases, use config's dimension.
@ -198,9 +239,7 @@ class Store:
"""
self._assert_writable()
if self._has_cloud_config() and str(self._config.lancedb.uri).startswith(
"db://"
):
if self._connection_mode == ConnectionMode.CLOUD:
return
# Skip if already running (non-blocking)
@ -224,26 +263,9 @@ class Store:
# 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)."""
# Check if we have cloud configuration
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
)
@property
def _connection_mode(self) -> ConnectionMode:
return ConnectionMode.from_config(self._config)
def get_stats(self) -> dict:
"""Get comprehensive table statistics.
@ -298,7 +320,7 @@ class Store:
it will be replaced (using replace=True parameter).
Note: Index creation requires sufficient training data.
"""
if self._has_cloud_config():
if self._connection_mode == ConnectionMode.CLOUD:
return
try:

18
tests/docker/docker-compose.s3.yml Normal file
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
"

14
tests/docker/s3-config.json Normal file
View file

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

92
tests/test_info.py
View file

@ -1,8 +1,10 @@
import json
from unittest.mock import patch
import pytest
from haiku.rag.app import HaikuRAGApp
from haiku.rag.config.models import AppConfig, LanceDBConfig
@pytest.mark.asyncio
@ -153,3 +155,93 @@ async def test_app_info_with_vector_index(temp_db_path, capsys):
# Check basic info still present
assert "documents: 1" 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()

231
tests/test_lancedb_connection.py
View file

@ -3,46 +3,211 @@ from unittest.mock import patch
import pytest
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
async def test_lancedb_cloud_skips_optimization(temp_db_path):
"""Test that vacuum is skipped when using LanceDB Cloud (db:// URI)."""
# Create a store
store = Store(temp_db_path, create=True)
class TestConnectionMode:
def test_local_when_uri_empty(self):
config = AppConfig(lancedb=LanceDBConfig(uri=""))
assert ConnectionMode.from_config(config) == ConnectionMode.LOCAL
# 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:
# Call vacuum - this should skip optimization for LanceDB Cloud
await store.vacuum()
def test_cloud_when_db_uri(self):
config = AppConfig(
lancedb=LanceDBConfig(
uri="db://my-database", api_key="key", region="us-east-1"
)
)
assert ConnectionMode.from_config(config) == ConnectionMode.CLOUD
# The optimize method should NOT have been called for LanceDB Cloud
mock_optimize.assert_not_called()
def test_object_storage_s3(self):
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
async def test_local_storage_calls_optimization(temp_db_path):
"""Test that vacuum calls optimization for local storage."""
# Create a store
store = Store(temp_db_path, create=True)
class TestConnectLancedb:
def test_local_passes_db_path(self, temp_db_path):
config = AppConfig(lancedb=LanceDBConfig(uri=""))
with patch("haiku.rag.store.engine.lancedb.connect") as mock_connect:
connect_lancedb(config, db_path=temp_db_path)
mock_connect.assert_called_once_with(temp_db_path)
# Ensure 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:
# Call vacuum - this should optimize all tables for local storage
await store.vacuum()
def test_cloud_passes_uri_api_key_region(self):
config = AppConfig(
lancedb=LanceDBConfig(
uri="db://my-database", api_key="test-key", region="us-west-2"
)
)
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
mock_optimize.assert_called()
def test_object_storage_passes_uri_and_storage_options(self):
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()

151
tests/test_s3_integration.py Normal file
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

48
tests/test_skill_generator.py
View file

@ -1,8 +1,10 @@
import shutil
import subprocess
import zipfile
from pathlib import Path
import pytest
import yaml
from haiku.rag.skill_generator import (
AVAILABLE_TOOLS,
@ -449,3 +451,49 @@ class TestGenerateSkill:
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(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"],
)