Merge pull request #335 from ggozad/feat/s3-object-store
add S3/object storage support for LanceDB connections
This commit is contained in:
commit
4c1eec4ae7
15 changed files with 650 additions and 84 deletions
|
|
@ -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
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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,
|
||||||
|
|
|
||||||
|
|
@ -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):
|
||||||
|
|
|
||||||
|
|
@ -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]")
|
||||||
|
|
|
||||||
|
|
@ -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")
|
||||||
|
|
|
||||||
|
|
@ -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"
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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:
|
||||||
|
|
|
||||||
18
tests/docker/docker-compose.s3.yml
Normal file
18
tests/docker/docker-compose.s3.yml
Normal 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
14
tests/docker/s3-config.json
Normal file
|
|
@ -0,0 +1,14 @@
|
||||||
|
{
|
||||||
|
"identities": [
|
||||||
|
{
|
||||||
|
"name": "admin",
|
||||||
|
"credentials": [
|
||||||
|
{
|
||||||
|
"accessKey": "testkey",
|
||||||
|
"secretKey": "testsecret"
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"actions": ["Admin", "Read", "Write", "List", "Tagging"]
|
||||||
|
}
|
||||||
|
]
|
||||||
|
}
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
|
|
@ -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()
|
||||||
|
|
|
||||||
151
tests/test_s3_integration.py
Normal file
151
tests/test_s3_integration.py
Normal 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
|
||||||
|
|
@ -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"],
|
||||||
|
)
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue