Coverage for skill extras
This commit is contained in:
parent
f4dd2dafc3
commit
61a38fc0a9
2 changed files with 49 additions and 1 deletions
|
|
@ -64,7 +64,6 @@ members = ["haiku_rag_slim", "evaluations"]
|
||||||
"haiku.rag-slim" = { workspace = true }
|
"haiku.rag-slim" = { workspace = true }
|
||||||
"haiku.rag-evals" = { workspace = true }
|
"haiku.rag-evals" = { workspace = true }
|
||||||
|
|
||||||
|
|
||||||
[dependency-groups]
|
[dependency-groups]
|
||||||
dev = [
|
dev = [
|
||||||
"haiku.rag-evals",
|
"haiku.rag-evals",
|
||||||
|
|
|
||||||
|
|
@ -98,6 +98,55 @@ class TestRAGSkillCreation:
|
||||||
assert skill.metadata.name == "rag"
|
assert skill.metadata.name == "rag"
|
||||||
|
|
||||||
|
|
||||||
|
class TestSkillExtras:
|
||||||
|
async def test_list_documents_returns_all(self, rag_db):
|
||||||
|
from haiku.rag.skills.rag import create_skill
|
||||||
|
|
||||||
|
skill = create_skill(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):
|
||||||
|
from haiku.rag.skills.rag import create_skill
|
||||||
|
|
||||||
|
skill = create_skill(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):
|
||||||
|
from haiku.rag.skills.rag import create_skill
|
||||||
|
|
||||||
|
skill = create_skill(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):
|
||||||
|
from haiku.rag.client import HaikuRAG
|
||||||
|
from haiku.rag.skills.rag import create_skill
|
||||||
|
|
||||||
|
monkeypatch.setattr(
|
||||||
|
HaikuRAG, "visualize_chunk", AsyncMock(return_value=["img1"])
|
||||||
|
)
|
||||||
|
|
||||||
|
skill = create_skill(db_path=rag_db)
|
||||||
|
visualize = skill.extras["visualize_chunk"]
|
||||||
|
|
||||||
|
# Get a real chunk_id from the db
|
||||||
|
async with HaikuRAG(rag_db, read_only=True) as rag:
|
||||||
|
docs = await rag.list_documents()
|
||||||
|
doc = await rag.get_document_by_id(docs[0].id)
|
||||||
|
chunks = await rag.chunk_repository.get_by_document_id(doc.id)
|
||||||
|
chunk_id = str(chunks[0].id)
|
||||||
|
|
||||||
|
result = await visualize(chunk_id)
|
||||||
|
assert result == ["img1"]
|
||||||
|
|
||||||
|
|
||||||
class TestSearchTool:
|
class TestSearchTool:
|
||||||
async def test_search_returns_formatted_string(self, rag_db):
|
async def test_search_returns_formatted_string(self, rag_db):
|
||||||
from haiku.rag.skills.rag import create_skill
|
from haiku.rag.skills.rag import create_skill
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue