fix: include 'db_path'/'config' in 'extras'

fix: 'skills.rlm.create_skill' includes extras

Closes #329.
Closes #330.
This commit is contained in:
Tres Seaver 2026-03-27 14:20:31 -04:00
parent 13a71e34a7
commit ed69bf4262
No known key found for this signature in database
GPG key ID: F2A968348913D1D8
5 changed files with 61 additions and 29 deletions

View file

@ -249,7 +249,13 @@ def create_skill_extras(
) -> dict[str, Any]:
"""Create non-tool utility functions bound to a specific database.
Returns a dict of callables that can be attached to a Skill's extras.
Returns a dict of values that can be attached to a Skill's extras:
Keys:
- 'db_path': path to the LanceDB used to configure the skill
- 'config': config passed to (or derived for) the skill
- 'list_documents': returns info for documents in the database
- 'list_documents': returns visulaizations for chunks in the database
"""
async def visualize_chunk(chunk_id: str) -> list:
@ -283,6 +289,8 @@ def create_skill_extras(
]
return {
"db_path": db_path,
"config": config,
"visualize_chunk": visualize_chunk,
"list_documents": list_documents,
}

View file

@ -54,7 +54,7 @@ def create_skill(
config: haiku.rag AppConfig instance. If None, uses get_config().
"""
from haiku.rag.config import get_config
from haiku.rag.skills._tools import create_skill_tools
from haiku.rag.skills._tools import create_skill_extras, create_skill_tools
if config is None:
config = get_config()
@ -67,6 +67,7 @@ def create_skill(
db_path = config.storage.data_dir / "haiku.rag.lancedb"
tools = create_skill_tools(db_path, config, RLMState, ["analyze"])
extras = create_skill_extras(db_path, config)
return Skill(
metadata=skill_metadata(),
@ -74,6 +75,7 @@ def create_skill(
path=_skill_path,
instructions=instructions(),
tools=list(tools.values()),
extras=extras,
state_type=STATE_TYPE,
state_namespace=STATE_NAMESPACE,
)

View file

@ -7,6 +7,7 @@ from pydantic_ai import RunContext
from haiku.rag.client import HaikuRAG
from haiku.rag.embeddings import EmbedderWrapper
from haiku.skills.state import SkillRunDeps
from haiku.rag.config.models import AppConfig
VECTOR_DIM = 2560
@ -45,6 +46,11 @@ def mock_embedder(monkeypatch):
monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)
@pytest.fixture
def test_app_config():
return AppConfig(environment="skills-test")
@pytest.fixture
async def rag_db(temp_db_path):
"""Create a test database with sample documents."""

View file

@ -42,10 +42,10 @@ class TestRAGModuleAPI:
assert isinstance(result, str)
assert len(result) > 0
def test_constants_match_create_skill(self, temp_db_path):
def test_constants_match_create_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=temp_db_path)
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.state_type is STATE_TYPE
assert skill.state_namespace == STATE_NAMESPACE
assert skill.metadata == skill_metadata()
@ -53,18 +53,18 @@ class TestRAGModuleAPI:
class TestRAGSkillCreation:
def test_create_skill_returns_valid_skill(self, temp_db_path):
def test_create_skill_returns_valid_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=temp_db_path)
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.metadata.name == "rag"
assert skill.metadata.description
assert skill.instructions
def test_create_skill_has_expected_tools(self, temp_db_path):
def test_create_skill_has_expected_tools(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=temp_db_path)
skill = create_skill(config=test_app_config, db_path=temp_db_path)
tool_names = {getattr(t, "__name__") for t in skill.tools if callable(t)}
assert tool_names == {
"search",
@ -73,18 +73,19 @@ class TestRAGSkillCreation:
"ask",
"research",
}
def test_create_skill_has_state(self, temp_db_path):
def test_create_skill_has_state(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=temp_db_path)
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill._state_type is RAGState
assert skill._state_namespace == "rag"
def test_create_skill_has_extras(self, temp_db_path):
def test_create_skill_has_extras(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=temp_db_path)
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.extras["config"] is test_app_config
assert skill.extras["db_path"] is temp_db_path
assert "visualize_chunk" in skill.extras
assert "list_documents" in skill.extras
assert callable(skill.extras["visualize_chunk"])
@ -99,33 +100,37 @@ class TestRAGSkillCreation:
class TestSkillExtras:
async def test_list_documents_returns_all(self, rag_db):
async def test_list_documents_returns_all(self, test_app_config, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
skill = create_skill(config=test_app_config, db_path=rag_db)
list_docs = skill.extras["list_documents"]
results = await list_docs()
assert len(results) == 2
assert all(k in results[0] for k in ("id", "title", "uri", "metadata"))
async def test_list_documents_with_filter(self, rag_db):
async def test_list_documents_with_filter(self, test_app_config, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
skill = create_skill(config=test_app_config, db_path=rag_db)
list_docs = skill.extras["list_documents"]
results = await list_docs(filter="title = 'AI Overview'")
assert len(results) == 1
assert results[0]["title"] == "AI Overview"
async def test_visualize_chunk_unknown_returns_empty(self, rag_db):
async def test_visualize_chunk_unknown_returns_empty(
self, test_app_config, rag_db,
):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
skill = create_skill(config=test_app_config, db_path=rag_db)
visualize = skill.extras["visualize_chunk"]
result = await visualize("nonexistent-chunk-id")
assert result == []
async def test_visualize_chunk_returns_images(self, rag_db, monkeypatch):
async def test_visualize_chunk_returns_images(
self, test_app_config, rag_db, monkeypatch,
):
from haiku.rag.client import HaikuRAG
from haiku.rag.skills.rag import create_skill
@ -133,7 +138,7 @@ class TestSkillExtras:
HaikuRAG, "visualize_chunk", AsyncMock(return_value=["img1"])
)
skill = create_skill(db_path=rag_db)
skill = create_skill(config=test_app_config, db_path=rag_db)
visualize = skill.extras["visualize_chunk"]
# Get a real chunk_id from the db

View file

@ -39,10 +39,10 @@ class TestRLMModuleAPI:
assert isinstance(result, str)
assert len(result) > 0
def test_constants_match_create_skill(self, temp_db_path):
def test_constants_match_create_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import create_skill
skill = create_skill(db_path=temp_db_path)
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.state_type is STATE_TYPE
assert skill.state_namespace == STATE_NAMESPACE
assert skill.metadata == skill_metadata()
@ -50,28 +50,39 @@ class TestRLMModuleAPI:
class TestRLMSkillCreation:
def test_create_skill_returns_valid_skill(self, temp_db_path):
def test_create_skill_returns_valid_skill(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import create_skill
skill = create_skill(db_path=temp_db_path)
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.metadata.name == "rag-rlm"
assert skill.metadata.description
assert skill.instructions
def test_create_skill_has_expected_tools(self, temp_db_path):
def test_create_skill_has_expected_tools(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import create_skill
skill = create_skill(db_path=temp_db_path)
skill = create_skill(config=test_app_config, db_path=temp_db_path)
tool_names = {getattr(t, "__name__") for t in skill.tools if callable(t)}
assert tool_names == {"analyze"}
def test_create_skill_has_state(self, temp_db_path):
def test_create_skill_has_state(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import RLMState, create_skill
skill = create_skill(db_path=temp_db_path)
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill._state_type is RLMState
assert skill._state_namespace == "rlm"
def test_create_skill_has_extras(self, test_app_config, temp_db_path):
from haiku.rag.skills.rlm import create_skill
skill = create_skill(config=test_app_config, db_path=temp_db_path)
assert skill.extras["config"] is test_app_config
assert skill.extras["db_path"] is temp_db_path
assert "visualize_chunk" in skill.extras
assert "list_documents" in skill.extras
assert callable(skill.extras["visualize_chunk"])
assert callable(skill.extras["list_documents"])
def test_create_skill_from_env(self, monkeypatch, temp_db_path):
monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path))
from haiku.rag.skills.rlm import create_skill