fix: include 'db_path'/'config' in 'extras'
fix: 'skills.rlm.create_skill' includes extras Closes #329. Closes #330.
This commit is contained in:
parent
13a71e34a7
commit
ed69bf4262
5 changed files with 61 additions and 29 deletions
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue