Add unified RAG skill with session context and reuse existing tools;

This commit is contained in:
Yiorgis Gozadinos 2026-02-19 15:52:49 +02:00
parent 4855051936
commit 524647c501
No known key found for this signature in database
12 changed files with 876 additions and 35 deletions

View file

@ -0,0 +1,336 @@
import os
from pathlib import Path
from typing import Any
from pydantic import BaseModel
from pydantic_ai import RunContext
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.rag.tools.qa import QAHistoryEntry
from haiku.skills.models import Skill, SkillSource
from haiku.skills.parser import parse_skill_md
from haiku.skills.state import SkillRunDeps
class ResearchEntry(BaseModel):
question: str
title: str
executive_summary: str
class RAGState(BaseModel):
citations: list[Any] = []
qa_history: list[QAHistoryEntry] = []
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = {}
documents: list[DocumentInfo] = []
reports: list[ResearchEntry] = []
def create_skill(
db_path: Path | None = None,
config: Any = None,
) -> Skill:
"""Create a RAG skill for searching and analyzing documents.
Args:
db_path: Path to the LanceDB database. Resolved from:
1. This argument
2. HAIKU_RAG_DB environment variable
3. haiku.rag default (config.storage.data_dir / "haiku.rag.lancedb")
config: haiku.rag AppConfig instance. If None, uses get_config().
"""
from haiku.rag.config import get_config
if config is None:
config = get_config()
if db_path is None:
env_db = os.environ.get("HAIKU_RAG_DB")
if env_db:
db_path = Path(env_db).expanduser()
else:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
path = Path(__file__).parent / "rag"
metadata, instructions = parse_skill_md(path / "SKILL.md")
async def search(
ctx: RunContext[SkillRunDeps], query: str, limit: int | None = None
) -> str:
"""Search the knowledge base using hybrid search (vector + full-text).
Returns ranked results with content and metadata.
Args:
query: The search query.
limit: Maximum number of results.
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
results = await rag.search(query, limit=limit)
results = await rag.expand_context(results)
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
ctx.deps.state.searches[query] = list(results)
return "\n\n---\n\n".join(
r.format_for_agent(rank=i + 1, total=len(results))
for i, r in enumerate(results)
)
async def list_documents(
ctx: RunContext[SkillRunDeps],
limit: int | None = None,
offset: int | None = None,
filter: str | None = None,
) -> list[dict[str, Any]]:
"""List documents in the knowledge base with optional pagination and filtering.
Args:
limit: Maximum number of documents to return.
offset: Number of documents to skip.
filter: Optional SQL WHERE clause to filter documents.
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
documents = await rag.list_documents(limit, offset, filter)
result = [
{
"id": doc.id,
"title": doc.title,
"uri": doc.uri,
"metadata": doc.metadata,
"created_at": str(doc.created_at),
"updated_at": str(doc.updated_at),
}
for doc in documents
]
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
for doc_dict in result:
doc_info = DocumentInfo(
id=str(doc_dict["id"]),
title=doc_dict["title"] or "Untitled",
uri=doc_dict.get("uri") or "",
created=doc_dict.get("created_at", ""),
)
if not any(d.id == doc_info.id for d in ctx.deps.state.documents):
ctx.deps.state.documents.append(doc_info)
return result
async def get_document(
ctx: RunContext[SkillRunDeps], query: str
) -> dict[str, Any] | None:
"""Retrieve a document by ID, title, or URI.
Args:
query: Document ID, title, or URI to look up.
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
document = await rag.resolve_document(query)
if document is None:
return None
result = {
"id": document.id,
"content": document.content,
"title": document.title,
"uri": document.uri,
"metadata": document.metadata,
"created_at": str(document.created_at),
"updated_at": str(document.updated_at),
}
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
doc_info = DocumentInfo(
id=str(result["id"]),
title=result["title"] or "Untitled",
uri=result.get("uri") or "",
created=result.get("created_at", ""),
)
if not any(d.id == doc_info.id for d in ctx.deps.state.documents):
ctx.deps.state.documents.append(doc_info)
return result
async def ask(ctx: RunContext[SkillRunDeps], question: str) -> str:
"""Ask a question and get an answer with citations from the knowledge base.
Args:
question: The question to ask.
"""
from haiku.rag.client import HaikuRAG
from haiku.rag.utils import format_citations
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
answer, citations = await rag.ask(question)
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
ctx.deps.state.citations.extend(citations)
ctx.deps.state.qa_history.append(
QAHistoryEntry(question=question, answer=answer, citations=citations)
)
if citations:
answer += "\n\n" + format_citations(citations)
return answer
async def analyze(
ctx: RunContext[SkillRunDeps],
question: str,
document: str | None = None,
filter: str | None = None,
) -> str:
"""Answer complex analytical questions using code execution.
Use this for questions requiring computation, aggregation, or
data traversal across documents.
Args:
question: The question to answer.
document: Optional document ID or title to pre-load for analysis.
filter: Optional SQL WHERE clause to filter documents.
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
documents = [document] if document else None
result = await rag.rlm(question, documents=documents, filter=filter)
output = result.answer
if result.program:
output += f"\n\nProgram:\n{result.program}"
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
ctx.deps.state.qa_history.append(
QAHistoryEntry(question=question, answer=output)
)
return output
async def get_session_context(ctx: RunContext[SkillRunDeps], query: str) -> str:
"""Retrieve relevant prior Q&A from the current session.
Call this before other tools when there may be prior questions
in the session that are relevant to the current query.
Args:
query: The current question or topic to find relevant context for.
"""
from haiku.rag.embeddings import get_embedder
from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD
from haiku.rag.utils import cosine_similarity
state = (
ctx.deps.state
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState)
else None
)
if state is None or not state.qa_history:
return "No prior questions in this session."
embedder = get_embedder(config)
query_embedding = await embedder.embed_query(query)
to_embed = []
to_embed_indices = []
for i, qa in enumerate(state.qa_history):
if qa.question_embedding is None:
to_embed.append(qa.question)
to_embed_indices.append(i)
if to_embed:
new_embeddings = await embedder.embed_documents(to_embed)
for i, idx in enumerate(to_embed_indices):
state.qa_history[idx].question_embedding = new_embeddings[i]
matches = []
for qa in state.qa_history:
if qa.question_embedding is not None:
similarity = cosine_similarity(query_embedding, qa.question_embedding)
if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:
matches.append(qa)
if not matches:
return "No relevant prior questions found for this query."
parts = []
for qa in matches:
parts.append(f"Q: {qa.question}\nA: {qa.answer}")
return "Relevant prior Q&A:\n\n" + "\n\n---\n\n".join(parts)
async def research(ctx: RunContext[SkillRunDeps], question: str) -> str:
"""Conduct deep multi-agent research on a question.
Iteratively searches, analyzes, and synthesizes information from the
knowledge base to produce a comprehensive research report.
Only use when the user explicitly requests deep research.
Args:
question: The research question to investigate.
"""
from haiku.rag.client import HaikuRAG
async with HaikuRAG(db_path, config=config, read_only=True) as rag:
report = await rag.research(question)
if ctx.deps and ctx.deps.state and isinstance(ctx.deps.state, RAGState):
ctx.deps.state.reports.append(
ResearchEntry(
question=question,
title=report.title,
executive_summary=report.executive_summary,
)
)
ctx.deps.state.qa_history.append(
QAHistoryEntry(question=question, answer=report.executive_summary)
)
parts = [
f"# {report.title}",
f"\n## Executive Summary\n{report.executive_summary}",
]
if report.main_findings:
parts.append("\n## Main Findings")
for finding in report.main_findings:
parts.append(f"- {finding}")
if report.conclusions:
parts.append("\n## Conclusions")
for conclusion in report.conclusions:
parts.append(f"- {conclusion}")
if report.limitations:
parts.append("\n## Limitations")
for limitation in report.limitations:
parts.append(f"- {limitation}")
if report.recommendations:
parts.append("\n## Recommendations")
for rec in report.recommendations:
parts.append(f"- {rec}")
parts.append(f"\n## Sources\n{report.sources_summary}")
return "\n".join(parts)
return Skill(
metadata=metadata,
source=SkillSource.ENTRYPOINT,
path=path,
instructions=instructions,
tools=[
search,
list_documents,
get_document,
ask,
analyze,
research,
get_session_context,
],
state_type=RAGState,
state_namespace="rag",
)

View file

@ -0,0 +1,36 @@
---
name: rag
description: Search, retrieve and analyze documents using RAG (Retrieval Augmented Generation).
---
# RAG
You are a RAG (Retrieval Augmented Generation) assistant with access to a document knowledge base.
Use your tools to search and answer questions. Never make up information — always use tools to get facts from the knowledge base.
## How to decide which tool to use
- **get_session_context** — Call this first when there have been prior questions in the session. It finds relevant prior Q&A so you can avoid redundant searches and give more informed answers.
- **list_documents** — Use when the user wants to browse or see what documents are available (e.g., "what documents do you have?", "show me the documents", "list available docs").
- **get_document** — Use when the user wants the full content of a specific document (e.g., "get the paper about X", "show me document Y"). Accepts a document ID, title, or URI — partial matches work.
- **search** — Use when the user wants to find relevant passages across documents (e.g., "search for embeddings", "find mentions of transformers"). Returns matching chunks with metadata.
- **ask** — Use for questions about topics in the knowledge base (e.g., "what is DocLayNet?", "explain the methodology"). Returns an answer with citations. Always include the citations in your response.
- **analyze** — Use for any question that involves code, computation, counting, aggregation, comparison, or complex reasoning (e.g., "how many pages?", "compare the results in table 3", "write code to find the longest word", "calculate the average"). The analyze tool can write and execute Python code with full access to the knowledge base. **When in doubt between search and analyze, prefer analyze** — it can search internally and also compute over results.
- **research** — Deep multi-agent research that produces comprehensive reports. **Only use when the user explicitly requests deep research** (e.g., "do a deep research on X", "research this topic thoroughly"). Never call this tool on your own — it is slow and expensive.
## When search returns irrelevant results
If your first search returns results that clearly don't match the question, **do not keep searching with variations**. Instead:
- Use **analyze** if the question involves computation or code
- Use **ask** if the question is factual
- Report that the knowledge base doesn't contain relevant information
## When the user mentions a specific document
If the user says "search in [doc]", "find in [doc]", or "answer from [doc]":
- Extract the **topic** as the `query`/`question` parameter
- Use **get_document** or **list_documents** first to identify the document, then search/ask with a filter
Examples:
- "search for embeddings in the ML paper" -> first identify "ML paper", then search for "embeddings"
- "what does the DocLayNet paper say about annotations?" -> ask with question="what are the annotation methods?"

View file

@ -24,6 +24,7 @@ Document content:
class DocumentInfo(BaseModel):
"""Document info for list_documents response."""
id: str | None = None
title: str
uri: str
created: str
@ -107,6 +108,7 @@ def create_document_toolset(
return DocumentListResponse(
documents=[
DocumentInfo(
id=doc.id,
title=doc.title or "Untitled",
uri=doc.uri or "",
created=doc.created_at.strftime("%Y-%m-%d"),

View file

@ -1,4 +1,3 @@
import math
from collections.abc import Callable
from pydantic import BaseModel, Field
@ -24,20 +23,11 @@ from haiku.rag.tools.session import (
SessionState,
compute_combined_state_delta,
)
from haiku.rag.utils import cosine_similarity
PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7
def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
"""Compute cosine similarity between two vectors."""
dot_product = sum(a * b for a, b in zip(vec1, vec2))
norm1 = math.sqrt(sum(a * a for a in vec1))
norm2 = math.sqrt(sum(b * b for b in vec2))
if norm1 == 0 or norm2 == 0:
return 0.0
return dot_product / (norm1 * norm2)
class QAHistoryEntry(BaseModel):
"""A Q&A pair with optional cached embedding for similarity matching."""
@ -129,7 +119,7 @@ async def run_qa_core(
matched_answers = []
for qa in qa_session_state.qa_history:
if qa.question_embedding is not None:
similarity = _cosine_similarity(
similarity = cosine_similarity(
question_embedding, qa.question_embedding
)
if similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD:

View file

@ -1,3 +1,4 @@
import math
import sys
from datetime import UTC, datetime
from importlib import metadata
@ -14,6 +15,16 @@ if TYPE_CHECKING:
from haiku.rag.config.models import AppConfig, ModelConfig
def cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
"""Compute cosine similarity between two vectors."""
dot_product = sum(a * b for a, b in zip(vec1, vec2))
norm1 = math.sqrt(sum(a * a for a in vec1))
norm2 = math.sqrt(sum(b * b for b in vec2))
if norm1 == 0 or norm2 == 0:
return 0.0
return dot_product / (norm1 * norm2)
def parse_datetime(s: str) -> datetime:
"""Parse a datetime string into a datetime object.

View file

@ -24,7 +24,7 @@ classifiers = [
dependencies = [
"cachetools>=5.5.0",
"docling-core==2.65.1",
"haiku.skills>=0.3.0",
"haiku.skills>=0.4.0",
"httpx>=0.28.1",
"jsonpatch>=1.33",
"lancedb==0.29.2",
@ -58,6 +58,9 @@ mistral = ["pydantic-ai-slim[mistral]"]
bedrock = ["pydantic-ai-slim[bedrock]"]
vertexai = ["pydantic-ai-slim[vertexai]"]
[project.entry-points."haiku.skills"]
rag = "haiku.rag.skills.rag:create_skill"
[project.scripts]
haiku-rag = "haiku.rag.cli:cli"

View file

@ -999,41 +999,41 @@ def test_search_tool_citation_registry_logic():
# =============================================================================
def test_cosine_similarity_identical_vectors():
def testcosine_similarity_identical_vectors():
"""Test cosine similarity returns 1.0 for identical vectors."""
from haiku.rag.tools.qa import _cosine_similarity
from haiku.rag.utils import cosine_similarity
vec = [1.0, 2.0, 3.0]
assert _cosine_similarity(vec, vec) == pytest.approx(1.0)
assert cosine_similarity(vec, vec) == pytest.approx(1.0)
def test_cosine_similarity_orthogonal_vectors():
def testcosine_similarity_orthogonal_vectors():
"""Test cosine similarity returns 0.0 for orthogonal vectors."""
from haiku.rag.tools.qa import _cosine_similarity
from haiku.rag.utils import cosine_similarity
vec1 = [1.0, 0.0, 0.0]
vec2 = [0.0, 1.0, 0.0]
assert _cosine_similarity(vec1, vec2) == pytest.approx(0.0)
assert cosine_similarity(vec1, vec2) == pytest.approx(0.0)
def test_cosine_similarity_opposite_vectors():
def testcosine_similarity_opposite_vectors():
"""Test cosine similarity returns -1.0 for opposite vectors."""
from haiku.rag.tools.qa import _cosine_similarity
from haiku.rag.utils import cosine_similarity
vec1 = [1.0, 2.0, 3.0]
vec2 = [-1.0, -2.0, -3.0]
assert _cosine_similarity(vec1, vec2) == pytest.approx(-1.0)
assert cosine_similarity(vec1, vec2) == pytest.approx(-1.0)
def test_cosine_similarity_zero_vector():
def testcosine_similarity_zero_vector():
"""Test cosine similarity handles zero vectors gracefully."""
from haiku.rag.tools.qa import _cosine_similarity
from haiku.rag.utils import cosine_similarity
vec = [1.0, 2.0, 3.0]
zero = [0.0, 0.0, 0.0]
assert _cosine_similarity(vec, zero) == 0.0
assert _cosine_similarity(zero, vec) == 0.0
assert _cosine_similarity(zero, zero) == 0.0
assert cosine_similarity(vec, zero) == 0.0
assert cosine_similarity(zero, vec) == 0.0
assert cosine_similarity(zero, zero) == 0.0
def test_prior_answer_relevance_threshold_constant():
@ -1047,14 +1047,14 @@ def test_prior_answer_matching_above_threshold():
"""Test that similar questions (above threshold) are matched."""
from haiku.rag.tools.qa import (
PRIOR_ANSWER_RELEVANCE_THRESHOLD,
_cosine_similarity,
cosine_similarity,
)
# Simulate two nearly identical question embeddings
question_embedding = [0.5, 0.5, 0.5, 0.5]
prior_embedding = [0.51, 0.49, 0.5, 0.5] # Very similar
similarity = _cosine_similarity(question_embedding, prior_embedding)
similarity = cosine_similarity(question_embedding, prior_embedding)
assert similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD
@ -1062,14 +1062,14 @@ def test_prior_answer_matching_below_threshold():
"""Test that dissimilar questions (below threshold) are not matched."""
from haiku.rag.tools.qa import (
PRIOR_ANSWER_RELEVANCE_THRESHOLD,
_cosine_similarity,
cosine_similarity,
)
# Simulate two different question embeddings
question_embedding = [1.0, 0.0, 0.0, 0.0]
prior_embedding = [0.0, 1.0, 0.0, 0.0] # Orthogonal = very different
similarity = _cosine_similarity(question_embedding, prior_embedding)
similarity = cosine_similarity(question_embedding, prior_embedding)
assert similarity < PRIOR_ANSWER_RELEVANCE_THRESHOLD

0
tests/skills/__init__.py Normal file
View file

64
tests/skills/conftest.py Normal file
View file

@ -0,0 +1,64 @@
import random
from unittest.mock import MagicMock
import pytest
from pydantic_ai import RunContext
from haiku.rag.client import HaikuRAG
from haiku.rag.embeddings import EmbedderWrapper
from haiku.skills.state import SkillRunDeps
VECTOR_DIM = 2560
def _make_ctx(state=None):
"""Create a mock RunContext with SkillRunDeps."""
ctx = MagicMock(spec=RunContext)
ctx.deps = SkillRunDeps(state=state)
return ctx
def _get_tool(skill, name):
"""Get a tool function from a skill by name."""
for tool in skill.tools:
if callable(tool) and tool.__name__ == name:
return tool
raise ValueError(f"Tool {name!r} not found in skill")
@pytest.fixture(autouse=True)
def mock_embedder(monkeypatch):
"""Monkeypatch the embedder to return deterministic vectors."""
async def fake_embed_query(self, text):
random.seed(hash(text) % (2**32))
return [random.random() for _ in range(VECTOR_DIM)]
async def fake_embed_documents(self, texts):
result = []
for t in texts:
random.seed(hash(t) % (2**32))
result.append([random.random() for _ in range(VECTOR_DIM)])
return result
monkeypatch.setattr(EmbedderWrapper, "embed_query", fake_embed_query)
monkeypatch.setattr(EmbedderWrapper, "embed_documents", fake_embed_documents)
@pytest.fixture
async def rag_db(temp_db_path):
"""Create a test database with sample documents."""
async with HaikuRAG(temp_db_path, create=True) as rag:
await rag.create_document(
"Artificial intelligence is transforming industries worldwide. "
"Deep learning models are used in healthcare, finance, and transportation.",
title="AI Overview",
uri="test://ai-overview",
)
await rag.create_document(
"Machine learning is a subset of artificial intelligence. "
"It includes supervised learning, unsupervised learning, and reinforcement learning.",
title="ML Basics",
uri="test://ml-basics",
)
return temp_db_path

373
tests/skills/test_rag.py Normal file
View file

@ -0,0 +1,373 @@
from unittest.mock import AsyncMock
from haiku.rag.agents.research.models import Citation, ResearchReport
from haiku.rag.agents.rlm.models import RLMResult
from haiku.rag.client import HaikuRAG
from haiku.rag.store.models.chunk import SearchResult
from haiku.rag.tools.document import DocumentInfo
from haiku.rag.tools.qa import QAHistoryEntry
from .conftest import _get_tool, _make_ctx
class TestRAGSkillCreation:
def test_create_skill_returns_valid_skill(self, temp_db_path):
from haiku.rag.skills.rag import create_skill
skill = create_skill(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):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=temp_db_path)
tool_names = {getattr(t, "__name__") for t in skill.tools if callable(t)}
assert tool_names == {
"search",
"list_documents",
"get_document",
"ask",
"analyze",
"research",
"get_session_context",
}
def test_create_skill_has_state(self, temp_db_path):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=temp_db_path)
assert skill._state_type is RAGState
assert skill._state_namespace == "rag"
def test_create_skill_from_env(self, monkeypatch, temp_db_path):
monkeypatch.setenv("HAIKU_RAG_DB", str(temp_db_path))
from haiku.rag.skills.rag import create_skill
skill = create_skill()
assert skill.metadata.name == "rag"
class TestSearchTool:
async def test_search_returns_formatted_string(self, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
ctx = _make_ctx()
result = await search(ctx, query="artificial intelligence")
assert isinstance(result, str)
assert len(result) > 0
async def test_search_updates_state(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
state = RAGState()
ctx = _make_ctx(state)
await search(ctx, query="artificial intelligence")
assert "artificial intelligence" in state.searches
results = state.searches["artificial intelligence"]
assert len(results) > 0
assert isinstance(results[0], SearchResult)
async def test_search_without_state(self, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
search = _get_tool(skill, "search")
ctx = _make_ctx(state=None)
result = await search(ctx, query="artificial intelligence")
assert isinstance(result, str)
class TestListDocumentsTool:
async def test_list_documents_returns_results(self, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
list_docs = _get_tool(skill, "list_documents")
ctx = _make_ctx()
results = await list_docs(ctx)
assert isinstance(results, list)
assert len(results) == 2
async def test_list_documents_updates_state(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
list_docs = _get_tool(skill, "list_documents")
state = RAGState()
ctx = _make_ctx(state)
await list_docs(ctx)
assert len(state.documents) == 2
assert isinstance(state.documents[0], DocumentInfo)
assert state.documents[0].id is not None
async def test_list_documents_no_duplicates_in_state(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
list_docs = _get_tool(skill, "list_documents")
state = RAGState()
ctx = _make_ctx(state)
await list_docs(ctx)
await list_docs(ctx)
assert len(state.documents) == 2
class TestGetDocumentTool:
async def test_get_document_by_title(self, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
get_doc = _get_tool(skill, "get_document")
ctx = _make_ctx()
result = await get_doc(ctx, query="AI Overview")
assert result is not None
assert result["title"] == "AI Overview"
async def test_get_document_updates_state(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
get_doc = _get_tool(skill, "get_document")
state = RAGState()
ctx = _make_ctx(state)
await get_doc(ctx, query="AI Overview")
assert len(state.documents) == 1
assert isinstance(state.documents[0], DocumentInfo)
assert state.documents[0].title == "AI Overview"
async def test_get_document_not_found(self, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
get_doc = _get_tool(skill, "get_document")
ctx = _make_ctx()
result = await get_doc(ctx, query="nonexistent document xyz")
assert result is None
class TestAskTool:
async def test_ask_returns_answer_with_citations(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import create_skill
citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="test://ai-overview",
document_title="AI Overview",
content="AI is transforming industries.",
)
]
monkeypatch.setattr(
HaikuRAG,
"ask",
AsyncMock(return_value=("AI transforms industries worldwide.", citations)),
)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
ctx = _make_ctx()
result = await ask(ctx, question="What is AI?")
assert isinstance(result, str)
assert "AI transforms industries" in result
async def test_ask_updates_state(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import RAGState, create_skill
citations = [
Citation(
document_id="d1",
chunk_id="c1",
document_uri="test://ai-overview",
content="AI content",
)
]
monkeypatch.setattr(
HaikuRAG,
"ask",
AsyncMock(return_value=("AI transforms industries.", citations)),
)
skill = create_skill(db_path=rag_db)
ask = _get_tool(skill, "ask")
state = RAGState()
ctx = _make_ctx(state)
await ask(ctx, question="What is AI?")
assert len(state.citations) == 1
assert len(state.qa_history) == 1
assert isinstance(state.qa_history[0], QAHistoryEntry)
assert state.qa_history[0].question == "What is AI?"
assert state.qa_history[0].citations == citations
class TestAnalyzeTool:
async def test_analyze_returns_result(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import create_skill
monkeypatch.setattr(
HaikuRAG,
"rlm",
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")),
)
skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze")
ctx = _make_ctx()
result = await analyze(ctx, question="How many documents?")
assert isinstance(result, str)
assert "42" in result
async def test_analyze_updates_state(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import RAGState, create_skill
monkeypatch.setattr(
HaikuRAG,
"rlm",
AsyncMock(return_value=RLMResult(answer="42", program="print(42)")),
)
skill = create_skill(db_path=rag_db)
analyze = _get_tool(skill, "analyze")
state = RAGState()
ctx = _make_ctx(state)
await analyze(ctx, question="How many documents?")
assert len(state.qa_history) == 1
assert state.qa_history[0].question == "How many documents?"
class TestGetSessionContextTool:
async def test_no_prior_questions(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
skill = create_skill(db_path=rag_db)
get_ctx = _get_tool(skill, "get_session_context")
state = RAGState()
ctx = _make_ctx(state)
result = await get_ctx(ctx, query="What is AI?")
assert "no prior" in result.lower()
async def test_returns_relevant_prior_qa(self, rag_db):
import random
from haiku.rag.skills.rag import RAGState, create_skill
from tests.skills.conftest import VECTOR_DIM
skill = create_skill(db_path=rag_db)
get_ctx = _get_tool(skill, "get_session_context")
# Pre-compute the embedding that the fake embedder will produce
# for the query, so we can set it on the prior entry for high similarity
query_text = "Tell me about artificial intelligence"
random.seed(hash(query_text) % (2**32))
query_embedding = [random.random() for _ in range(VECTOR_DIM)]
state = RAGState(
qa_history=[
QAHistoryEntry(
question="What is artificial intelligence?",
answer="AI is the simulation of human intelligence by machines.",
question_embedding=query_embedding,
),
]
)
ctx = _make_ctx(state)
result = await get_ctx(ctx, query=query_text)
assert "artificial intelligence" in result.lower()
assert "simulation" in result.lower()
async def test_no_relevant_matches(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
from tests.skills.conftest import VECTOR_DIM
skill = create_skill(db_path=rag_db)
get_ctx = _get_tool(skill, "get_session_context")
# Use alternating ±1 embedding which is near-orthogonal to the
# all-positive vectors produced by the fake embedder
orthogonal = [1.0 if i % 2 == 0 else -1.0 for i in range(VECTOR_DIM)]
state = RAGState(
qa_history=[
QAHistoryEntry(
question="What is the weather?",
answer="It is sunny today.",
question_embedding=orthogonal,
),
]
)
ctx = _make_ctx(state)
result = await get_ctx(ctx, query="Explain quantum computing")
assert "no relevant" in result.lower()
async def test_without_state(self, rag_db):
from haiku.rag.skills.rag import create_skill
skill = create_skill(db_path=rag_db)
get_ctx = _get_tool(skill, "get_session_context")
ctx = _make_ctx(state=None)
result = await get_ctx(ctx, query="What is AI?")
assert "no prior" in result.lower()
class TestResearchTool:
async def test_research_returns_report(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import create_skill
report = ResearchReport(
title="AI Research",
executive_summary="AI is transforming industries.",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Multiple sources consulted.",
)
monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report))
skill = create_skill(db_path=rag_db)
research = _get_tool(skill, "research")
ctx = _make_ctx()
result = await research(ctx, question="What is AI?")
assert isinstance(result, str)
assert "AI Research" in result
async def test_research_updates_state(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import RAGState, create_skill
report = ResearchReport(
title="AI Research",
executive_summary="AI is transforming industries.",
main_findings=["Finding 1"],
conclusions=["Conclusion 1"],
sources_summary="Multiple sources consulted.",
)
monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report))
skill = create_skill(db_path=rag_db)
research = _get_tool(skill, "research")
state = RAGState()
ctx = _make_ctx(state)
await research(ctx, question="What is AI?")
assert len(state.reports) == 1
assert state.reports[0].question == "What is AI?"
assert len(state.qa_history) == 1
assert state.qa_history[0].question == "What is AI?"
assert state.qa_history[0].answer == "AI is transforming industries."
async def test_research_without_state(self, rag_db, monkeypatch):
from haiku.rag.skills.rag import create_skill
report = ResearchReport(
title="AI Research",
executive_summary="Summary.",
main_findings=["Finding"],
conclusions=["Conclusion"],
sources_summary="Sources.",
)
monkeypatch.setattr(HaikuRAG, "research", AsyncMock(return_value=report))
skill = create_skill(db_path=rag_db)
research = _get_tool(skill, "research")
ctx = _make_ctx(state=None)
result = await research(ctx, question="What is AI?")
assert isinstance(result, str)

34
uv.lock
View file

@ -1513,7 +1513,7 @@ requires-dist = [
{ name = "cohere", marker = "extra == 'cohere'", specifier = ">=5.20.1" },
{ name = "docling", marker = "extra == 'docling'", specifier = "==2.73.1" },
{ name = "docling-core", specifier = "==2.65.1" },
{ name = "haiku-skills", specifier = ">=0.3.0" },
{ name = "haiku-skills", specifier = ">=0.4.0" },
{ name = "httpx", specifier = ">=0.28.1" },
{ name = "jsonpatch", specifier = ">=1.33" },
{ name = "lancedb", specifier = "==0.29.2" },
@ -1544,16 +1544,17 @@ provides-extras = ["docling", "voyageai", "mxbai", "cohere", "zeroentropy", "jin
[[package]]
name = "haiku-skills"
version = "0.3.0"
version = "0.4.0"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "pydantic" },
{ name = "pydantic-ai-slim", extra = ["mcp"] },
{ name = "pyyaml" },
{ name = "skills-ref" },
]
sdist = { url = "https://files.pythonhosted.org/packages/74/d6/11dbef98d7b04f5aacc7da3354fbe7bbb9ab5aa5948517d135f68f457d7d/haiku_skills-0.3.0.tar.gz", hash = "sha256:191414a840653ba938aa8ce1804f10cc489426d44fdcc7f261feb2e32e1be041", size = 158600, upload-time = "2026-02-19T10:34:28.851Z" }
sdist = { url = "https://files.pythonhosted.org/packages/f7/3b/7545006f8e034926b615fdc09a61340f87a2ddb288c96af9532f490c089e/haiku_skills-0.4.0.tar.gz", hash = "sha256:98d6c454da419f31fb8275608bf5ef057c6e7fa9cfecd594d0958617633a479c", size = 159518, upload-time = "2026-02-19T13:36:22.887Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/fd/a2/279333965841a2ecda2e5cc695306998f9770ac5b92b2d6ad8c50206b162/haiku_skills-0.3.0-py3-none-any.whl", hash = "sha256:ac5ddaea07d920ffec368ea7cf8e88963e727feece01571fd40dfc31da4b3ccd", size = 20359, upload-time = "2026-02-19T10:34:27.235Z" },
{ url = "https://files.pythonhosted.org/packages/da/6e/340c6162b1de5cdefdd532c435e08eb053920c07b880c87c8d901c046192/haiku_skills-0.4.0-py3-none-any.whl", hash = "sha256:d12b6cfdfb276ba902d9e98a5582b408e0f1e6989c513bccdf45438ec80f2d38", size = 20565, upload-time = "2026-02-19T13:36:21.35Z" },
]
[[package]]
@ -4724,6 +4725,19 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" },
]
[[package]]
name = "skills-ref"
version = "0.1.1"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "click" },
{ name = "strictyaml" },
]
sdist = { url = "https://files.pythonhosted.org/packages/23/42/943d3ba8b097af7068b7178563a5062ad8a977982f4a7b4f67facfc575e9/skills_ref-0.1.1.tar.gz", hash = "sha256:6b400ca6e0049be62dca0167ff943ba2745fd67efb37fbba4d0ee341fccd2695", size = 93519, upload-time = "2026-01-10T13:23:41.423Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/af/25/36a43c3a61fb6cc3984e6ad5e556929b8ae71c95eba615dae4cf2f427964/skills_ref-0.1.1-py3-none-any.whl", hash = "sha256:d35db5bb8de71ae301daf5ca9cb71f8a555e8c6f83a6d40e46a5bc09f8f461b5", size = 12918, upload-time = "2026-01-10T13:23:40.106Z" },
]
[[package]]
name = "sniffio"
version = "1.3.1"
@ -4777,6 +4791,18 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/81/0d/13d1d239a25cbfb19e740db83143e95c772a1fe10202dda4b76792b114dd/starlette-0.52.1-py3-none-any.whl", hash = "sha256:0029d43eb3d273bc4f83a08720b4912ea4b071087a3b48db01b7c839f7954d74", size = 74272, upload-time = "2026-01-18T13:34:09.188Z" },
]
[[package]]
name = "strictyaml"
version = "1.7.3"
source = { registry = "https://pypi.org/simple" }
dependencies = [
{ name = "python-dateutil" },
]
sdist = { url = "https://files.pythonhosted.org/packages/b3/08/efd28d49162ce89c2ad61a88bd80e11fb77bc9f6c145402589112d38f8af/strictyaml-1.7.3.tar.gz", hash = "sha256:22f854a5fcab42b5ddba8030a0e4be51ca89af0267961c8d6cfa86395586c407", size = 115206, upload-time = "2023-03-10T12:50:27.062Z" }
wheels = [
{ url = "https://files.pythonhosted.org/packages/96/7c/a81ef5ef10978dd073a854e0fa93b5d8021d0594b639cc8f6453c3c78a1d/strictyaml-1.7.3-py3-none-any.whl", hash = "sha256:fb5c8a4edb43bebb765959e420f9b3978d7f1af88c80606c03fb420888f5d1c7", size = 123917, upload-time = "2023-03-10T12:50:17.242Z" },
]
[[package]]
name = "sympy"
version = "1.14.0"