support remote storage in skill generator
This commit is contained in:
parent
343bfd7199
commit
2522d46304
4 changed files with 73 additions and 7 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
)
|
||||
|
|
|
|||
Loading…
Reference in a new issue