diff --git a/haiku_rag_slim/haiku/rag/skills/__init__.py b/haiku_rag_slim/haiku/rag/skills/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/haiku_rag_slim/haiku/rag/skills/rag.py b/haiku_rag_slim/haiku/rag/skills/rag.py new file mode 100644 index 00000000..ca76fb9f --- /dev/null +++ b/haiku_rag_slim/haiku/rag/skills/rag.py @@ -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", + ) diff --git a/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md b/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md new file mode 100644 index 00000000..374c8a0b --- /dev/null +++ b/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md @@ -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?" diff --git a/haiku_rag_slim/haiku/rag/tools/document.py b/haiku_rag_slim/haiku/rag/tools/document.py index f053f137..8b189bb1 100644 --- a/haiku_rag_slim/haiku/rag/tools/document.py +++ b/haiku_rag_slim/haiku/rag/tools/document.py @@ -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"), diff --git a/haiku_rag_slim/haiku/rag/tools/qa.py b/haiku_rag_slim/haiku/rag/tools/qa.py index 76d301ac..7f7b48e2 100644 --- a/haiku_rag_slim/haiku/rag/tools/qa.py +++ b/haiku_rag_slim/haiku/rag/tools/qa.py @@ -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: diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 79e4f9c6..0c89f49c 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -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. diff --git a/haiku_rag_slim/pyproject.toml b/haiku_rag_slim/pyproject.toml index 84f34710..c65c186b 100644 --- a/haiku_rag_slim/pyproject.toml +++ b/haiku_rag_slim/pyproject.toml @@ -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" diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index 659ad1ff..9f986c4e 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/tests/agents/chat/test_chat_agent.py @@ -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 diff --git a/tests/skills/__init__.py b/tests/skills/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/skills/conftest.py b/tests/skills/conftest.py new file mode 100644 index 00000000..d6def468 --- /dev/null +++ b/tests/skills/conftest.py @@ -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 diff --git a/tests/skills/test_rag.py b/tests/skills/test_rag.py new file mode 100644 index 00000000..17c64f0c --- /dev/null +++ b/tests/skills/test_rag.py @@ -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) diff --git a/uv.lock b/uv.lock index 17c4041f..a6f08a61 100644 --- a/uv.lock +++ b/uv.lock @@ -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"