From ed69bf4262417ebfee7aa26755be2e57dd92a1ad Mon Sep 17 00:00:00 2001 From: Tres Seaver Date: Fri, 27 Mar 2026 14:20:31 -0400 Subject: [PATCH] fix: include 'db_path'/'config' in 'extras' fix: 'skills.rlm.create_skill' includes extras Closes #329. Closes #330. --- haiku_rag_slim/haiku/rag/skills/_tools.py | 10 +++++- haiku_rag_slim/haiku/rag/skills/rlm.py | 4 ++- tests/skills/conftest.py | 6 ++++ tests/skills/test_rag.py | 43 +++++++++++++---------- tests/skills/test_rlm.py | 27 +++++++++----- 5 files changed, 61 insertions(+), 29 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py index 79b1d3e8..1998797a 100644 --- a/haiku_rag_slim/haiku/rag/skills/_tools.py +++ b/haiku_rag_slim/haiku/rag/skills/_tools.py @@ -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, } diff --git a/haiku_rag_slim/haiku/rag/skills/rlm.py b/haiku_rag_slim/haiku/rag/skills/rlm.py index 7b48bbad..7efa403e 100644 --- a/haiku_rag_slim/haiku/rag/skills/rlm.py +++ b/haiku_rag_slim/haiku/rag/skills/rlm.py @@ -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, ) diff --git a/tests/skills/conftest.py b/tests/skills/conftest.py index d6def468..f25caf03 100644 --- a/tests/skills/conftest.py +++ b/tests/skills/conftest.py @@ -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.""" diff --git a/tests/skills/test_rag.py b/tests/skills/test_rag.py index 8cd49a8c..c7acca69 100644 --- a/tests/skills/test_rag.py +++ b/tests/skills/test_rag.py @@ -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 diff --git a/tests/skills/test_rlm.py b/tests/skills/test_rlm.py index 6e3317d6..c106222a 100644 --- a/tests/skills/test_rlm.py +++ b/tests/skills/test_rlm.py @@ -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