diff --git a/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx
index d0c24f48..f820da4d 100644
--- a/app/frontend/components/Chat.tsx
+++ b/app/frontend/components/Chat.tsx
@@ -21,8 +21,8 @@ import { FilterIcon } from "../lib/icons";
import type { RAGState } from "../lib/sessionStorage";
import {
createSession,
- deriveCitationsHistory,
getActiveSessionId,
+ getLatestCitations,
getSession,
normalizeRAGState,
updateSessionMessages,
@@ -292,7 +292,7 @@ function MessageViewWithCitations({
isRunning?: boolean;
}) {
const ragState = useContext(ChatStateContext);
- const citationsHistory = ragState ? deriveCitationsHistory(ragState) : [];
+ const latestCitations = ragState ? getLatestCitations(ragState) : [];
// Collect completed tool_call_ids from skill_tool_result activity messages
const completedToolCallIds = useMemo(() => {
@@ -326,7 +326,6 @@ function MessageViewWithCitations({
{({ messageElements }) => {
const result: React.ReactNode[] = [];
let elemIdx = 0;
- let citIdx = 0;
let seenToolCalls = false;
for (const msg of messages) {
@@ -368,19 +367,15 @@ function MessageViewWithCitations({
}
// After an assistant text response that followed tool calls,
- // inject the next citations entry (one per turn)
+ // show citations from the latest turn
if (msg.role === "assistant" && msg.content && seenToolCalls) {
- if (citIdx < citationsHistory.length) {
- const citations = citationsHistory[citIdx];
- if (citations?.length) {
- result.push(
- ,
- );
- }
- citIdx++;
+ if (latestCitations.length > 0) {
+ result.push(
+ ,
+ );
}
seenToolCalls = false;
}
diff --git a/app/frontend/lib/sessionStorage.ts b/app/frontend/lib/sessionStorage.ts
index bf1f39c2..47737433 100644
--- a/app/frontend/lib/sessionStorage.ts
+++ b/app/frontend/lib/sessionStorage.ts
@@ -9,12 +9,6 @@ export interface Citation {
content: string;
}
-export interface QAHistoryEntry {
- question: string;
- answer: string;
- citations: Citation[];
-}
-
export interface DocumentInfo {
id: string;
title: string;
@@ -22,20 +16,13 @@ export interface DocumentInfo {
created: string;
}
-export interface ResearchEntry {
- question: string;
- title: string;
- executive_summary: string;
-}
-
// Matches RAGState from the backend skill
export interface RAGState {
- citations: Citation[];
- qa_history: QAHistoryEntry[];
+ citation_index: Record;
+ citations: string[][];
document_filter: string | null;
searches: Record;
documents: DocumentInfo[];
- reports: ResearchEntry[];
}
export interface StoredMessage {
@@ -59,20 +46,21 @@ const ACTIVE_SESSION_KEY = "haiku.rag.activeSession";
export function normalizeRAGState(state?: Partial): RAGState {
return {
+ citation_index: state?.citation_index ?? {},
citations: state?.citations ?? [],
- qa_history: state?.qa_history ?? [],
document_filter: state?.document_filter ?? null,
searches: state?.searches ?? {},
documents: state?.documents ?? [],
- reports: state?.reports ?? [],
};
}
-// Derive per-turn citation arrays from qa_history
-export function deriveCitationsHistory(state: RAGState): Citation[][] {
- return state.qa_history
- .filter((entry) => entry.citations?.length > 0)
- .map((entry) => entry.citations);
+export function getLatestCitations(state: RAGState): Citation[] {
+ const turns = state.citations;
+ if (turns.length === 0) return [];
+ const latestIds = turns[turns.length - 1];
+ return latestIds
+ .map((id) => state.citation_index[id])
+ .filter((c): c is Citation => c !== undefined);
}
export function getAllSessions(): StoredSession[] {
diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py
index 6e0d36ca..b0f1a1ee 100644
--- a/haiku_rag_slim/haiku/rag/chat/app.py
+++ b/haiku_rag_slim/haiku/rag/chat/app.py
@@ -243,8 +243,7 @@ class ChatApp(App):
content=accumulated_text,
)
)
- # Show citations from RAG state
- await self._show_citations(chat_history)
+ await self._show_citations_and_programs(chat_history)
elif event.type == EventType.TOOL_CALL_START:
assert isinstance(event, ToolCallStartEvent)
chat_history.hide_thinking()
@@ -320,18 +319,32 @@ class ChatApp(App):
chat_input.disabled = False
chat_input.focus()
- async def _show_citations(self, chat_history: "ChatHistory") -> None:
- """Show citations from skill states after an agent response."""
+ async def _show_citations_and_programs(self, chat_history: "ChatHistory") -> None:
+ """Show citations and programs from skill states after an agent response."""
if not self._toolset:
return
citations = []
for namespace in (RAG_STATE_NAMESPACE, ANALYSIS_STATE_NAMESPACE):
state = self._toolset.get_namespace(namespace)
- if state:
- citations.extend(getattr(state, "citations", []))
+ if not state:
+ continue
+ citation_turns = getattr(state, "citations", [])
+ citation_index = getattr(state, "citation_index", {})
+ if citation_turns:
+ latest_ids = citation_turns[-1]
+ for cid in latest_ids:
+ if cid in citation_index:
+ citations.append(citation_index[cid])
if citations:
await chat_history.add_citations(citations)
+ analysis_state = self._toolset.get_namespace(ANALYSIS_STATE_NAMESPACE)
+ if analysis_state:
+ executions = getattr(analysis_state, "executions", [])
+ successful = [e for e in executions if e.success]
+ if successful:
+ await chat_history.add_program(successful[-1].code)
+
async def action_clear_chat(self) -> None:
"""Clear the chat history and reset session."""
chat_history = self.query_one(ChatHistory)
diff --git a/haiku_rag_slim/haiku/rag/skill_generator/__init__.py b/haiku_rag_slim/haiku/rag/skill_generator/__init__.py
index 50c50019..2d284bfe 100644
--- a/haiku_rag_slim/haiku/rag/skill_generator/__init__.py
+++ b/haiku_rag_slim/haiku/rag/skill_generator/__init__.py
@@ -8,9 +8,8 @@ AVAILABLE_TOOLS: set[str] = {
"list_documents",
"get_document",
"search",
- "ask",
- "research",
- "analyze",
+ "execute_code",
+ "cite",
}
DEFAULT_PREAMBLE = (
diff --git a/haiku_rag_slim/haiku/rag/skill_generator/templates/SKILL.md.j2 b/haiku_rag_slim/haiku/rag/skill_generator/templates/SKILL.md.j2
index 8bf07475..fc25efe6 100644
--- a/haiku_rag_slim/haiku/rag/skill_generator/templates/SKILL.md.j2
+++ b/haiku_rag_slim/haiku/rag/skill_generator/templates/SKILL.md.j2
@@ -7,48 +7,58 @@ description: {{ description }}
{{ preamble }}
-## How to decide which tool to use
-{% if "ask" in tool_names %}
+## Tools
+{% if "search" in tool_names %}
-**Default rule:** If the user is asking a question, use **ask**. Only use **search** when the user explicitly wants to browse or find passages.
+### search
+Search the knowledge base using hybrid search (vector + full-text). Returns ranked results with context-expanded content. Use for answering questions, finding passages, exploring topics.
{% endif %}
{% if "list_documents" in tool_names %}
-- **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").
+
+### list_documents
+List all documents in the knowledge base.
{% endif %}
{% if "get_document" in tool_names %}
-- **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.
+
+### get_document
+Retrieve a document by ID, title, or URI. Partial matches work.
{% endif %}
-{% if "search" in tool_names %}
-- **search** — Use when the user wants to browse, explore, or find specific passages across documents (e.g., "search for embeddings", "find mentions of transformers"). Returns all matching results as sources.
+{% if "execute_code" in tool_names %}
+
+### execute_code
+Execute Python code in a sandboxed interpreter. Inside the code you have access to `await search()`, `await list_documents()`, `await llm()`, and a virtual filesystem at `/documents/` with document content and structure.
{% endif %}
-{% if "ask" in tool_names %}
-- **ask** — Use for factual questions that need a synthesized answer (e.g., "what is DocLayNet?", "explain the methodology"). Searches, synthesizes, and returns only the chunks actually used as citations. Always include the citations in your response.
-{% endif %}
-{% if "research" in tool_names %}
-- **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.
-{% endif %}
-{% if "analyze" in tool_names %}
-- **analyze** — Use for complex analytical questions that require computation, aggregation, or data traversal across documents (e.g., "how many pages?", "compare table 3 across documents", "calculate average word count"). Executes Python code in a sandboxed interpreter.
+{% if "cite" in tool_names %}
+
+### cite
+Register chunk IDs as citations. Call after formulating your answer with chunk_id values from search results that support it. Do NOT include chunk IDs in your answer text.
{% endif %}
{% if "search" in tool_names %}
-## When search returns irrelevant results
+## How to answer questions
-If your first search returns results that clearly don't match the question, **do not keep searching with variations**. Instead:
-{% if "ask" in tool_names %}
-- Use **ask** if the question is factual
+1. Call `search` with relevant keywords from the question
+2. Review results — they are ordered by relevance (rank 1 = best match)
+3. If needed, search again with different keywords (up to 3-4 searches total)
+4. Synthesize a concise answer based strictly on the retrieved content
+{% if "cite" in tool_names %}
+5. Call `cite` with the chunk IDs you referenced
+{% endif %}
+
+## Guidelines
+
+- Base answers strictly on retrieved content — do not use external knowledge
+- Be concise and direct — avoid elaboration unless asked
+- If results don't match the question, report that the knowledge base lacks the information
+{% if "cite" in tool_names %}
+- Do NOT include chunk IDs or UUIDs in your answer text — use the `cite` tool separately
{% endif %}
-- Report that the knowledge base doesn't contain relevant information
{% endif %}
{% if "get_document" in tool_names %}
## 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?"
+- Use **get_document** or **list_documents** first to identify the document
+- Then search for the topic
{% endif %}
diff --git a/haiku_rag_slim/haiku/rag/skill_generator/templates/__init__.py.j2 b/haiku_rag_slim/haiku/rag/skill_generator/templates/__init__.py.j2
index a82437af..829d006b 100644
--- a/haiku_rag_slim/haiku/rag/skill_generator/templates/__init__.py.j2
+++ b/haiku_rag_slim/haiku/rag/skill_generator/templates/__init__.py.j2
@@ -5,23 +5,17 @@ from pydantic import BaseModel, Field
from haiku.rag.config.models import AppConfig
from haiku.skills.models import Skill
from haiku.skills.parser import parse_skill_md
-{% if "ask" in tool_names or "research" in tool_names %}
+{% if "cite" in tool_names %}
from haiku.rag.agents.research.models import Citation
{% endif %}
{% if "search" in tool_names or "list_documents" in tool_names or "get_document" in tool_names %}
from haiku.rag.tools.document import DocumentInfo
{% endif %}
-{% if "ask" in tool_names %}
-from haiku.rag.tools.qa import QAHistoryEntry
-{% endif %}
{% if "search" in tool_names %}
from haiku.rag.store.models.chunk import SearchResult
{% endif %}
-{% if "research" in tool_names %}
-from haiku.rag.skills._tools import ResearchEntry
-{% endif %}
-{% if "analyze" in tool_names %}
-from haiku.rag.skills._tools import AnalysisEntry
+{% if "execute_code" in tool_names %}
+from haiku.rag.skills._tools import CodeExecutionEntry
{% endif %}
_TOOL_NAMES = {{ tool_names | tojson }}
@@ -36,11 +30,9 @@ _CONFIG_PATH = _ASSETS_DIR / "haiku.rag.yaml"
class SkillState(BaseModel):
-{% if "ask" in tool_names or "research" in tool_names %}
- citations: list[Citation] = Field(default_factory=list)
-{% endif %}
-{% if "ask" in tool_names %}
- qa_history: list[QAHistoryEntry] = Field(default_factory=list)
+{% if "cite" in tool_names %}
+ citation_index: dict[str, Citation] = Field(default_factory=dict)
+ citations: list[list[str]] = Field(default_factory=list)
{% endif %}
document_filter: str | None = None
{% if "search" in tool_names %}
@@ -49,11 +41,8 @@ class SkillState(BaseModel):
{% if "search" in tool_names or "list_documents" in tool_names or "get_document" in tool_names %}
documents: list[DocumentInfo] = Field(default_factory=list)
{% endif %}
-{% if "research" in tool_names %}
- reports: list[ResearchEntry] = Field(default_factory=list)
-{% endif %}
-{% if "analyze" in tool_names %}
- analyses: list[AnalysisEntry] = Field(default_factory=list)
+{% if "execute_code" in tool_names %}
+ executions: list[CodeExecutionEntry] = Field(default_factory=list)
{% endif %}
diff --git a/haiku_rag_slim/haiku/rag/skills/_tools.py b/haiku_rag_slim/haiku/rag/skills/_tools.py
index 36e9eb32..de37cf4a 100644
--- a/haiku_rag_slim/haiku/rag/skills/_tools.py
+++ b/haiku_rag_slim/haiku/rag/skills/_tools.py
@@ -8,57 +8,14 @@ from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
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.state import SkillRunDeps
-class ResearchEntry(BaseModel):
- question: str
- title: str
- executive_summary: str
-
-
-class AnalysisEntry(BaseModel):
- question: str
- answer: str
- program: str | None = None
-
-
-async def find_relevant_prior_qa(
- qa_history: list[QAHistoryEntry],
- query: str,
- config: AppConfig,
-) -> list[QAHistoryEntry]:
- from haiku.rag.embeddings import get_embedder
- from haiku.rag.tools.qa import PRIOR_ANSWER_RELEVANCE_THRESHOLD
- from haiku.rag.utils import cosine_similarity
-
- if not qa_history:
- return []
-
- embedder = get_embedder(config)
- query_embedding = await embedder.embed_query(query)
-
- to_embed = []
- to_embed_indices = []
- for i, qa in enumerate(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):
- qa_history[idx].question_embedding = new_embeddings[i]
-
- matches = []
- for qa in 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)
-
- return matches
+class CodeExecutionEntry(BaseModel):
+ code: str
+ stdout: str
+ stderr: str = ""
+ success: bool = True
async def skill_search(
@@ -88,14 +45,12 @@ async def skill_search(
async def skill_list_documents(
db_path: Path,
config: AppConfig,
- limit: int | None = None,
- offset: int | None = None,
filter: str | None = None,
) -> list[dict[str, Any]]:
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=filter)
+ documents = await rag.list_documents(filter=filter)
return [
{
"id": doc.id,
@@ -131,100 +86,6 @@ async def skill_get_document(
}
-async def skill_ask(
- db_path: Path,
- config: AppConfig,
- question: str,
- qa_history: list[QAHistoryEntry] | None = None,
- document_filter: str | None = None,
-) -> tuple[str, list[Citation]]:
- from haiku.rag.client import HaikuRAG
- from haiku.rag.utils import format_citations
-
- ask_question = question
- if qa_history:
- matches = await find_relevant_prior_qa(qa_history, question, config)
- if matches:
- prior_parts = []
- for qa in matches:
- part = f"Q: {qa.question}\nA: {qa.answer}"
- if qa.citations:
- part += "\n" + format_citations(qa.citations)
- prior_parts.append(part)
- ask_question = (
- "Context from prior questions in this session:\n\n"
- + "\n\n---\n\n".join(prior_parts)
- + "\n\n---\n\nCurrent question: "
- + question
- )
-
- async with HaikuRAG(db_path, config=config, read_only=True) as rag:
- answer, citations = await rag.ask(
- ask_question,
- filter=document_filter,
- )
-
- return answer, citations
-
-
-async def skill_research(
- db_path: Path,
- config: AppConfig,
- question: str,
- document_filter: str | None = None,
-) -> tuple[str, str, str]:
- from haiku.rag.client import HaikuRAG
-
- async with HaikuRAG(db_path, config=config, read_only=True) as rag:
- report = await rag.research(question, filter=document_filter)
-
- 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}")
-
- formatted = "\n".join(parts)
- return formatted, report.title, report.executive_summary
-
-
-async def skill_analyze(
- db_path: Path,
- config: AppConfig,
- question: str,
- document: str | None = None,
- document_filter: str | None = None,
-) -> tuple[str, str, str | None, "list[Citation]"]:
- 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.analyze(
- question, documents=documents, filter=document_filter
- )
- output = result.answer
- if result.program:
- output += f"\n\nProgram:\n{result.program}"
-
- return output, result.answer, result.program, result.citations
-
-
def update_documents_state(
documents_state: list[DocumentInfo],
doc_dicts: list[dict[str, Any]],
@@ -246,13 +107,18 @@ def _get_state(ctx: RunContext[SkillRunDeps], state_type: type[BaseModel]) -> An
return None
-def _append_citations(state: Any, citations: "list[Citation]") -> None:
- """Index and append citations to a skill state's citations list."""
- next_index = len(state.citations) + 1
+def _register_citations(state: Any, citations: "list[Citation]") -> None:
+ """Add citations to the index and record the turn's chunk IDs."""
+ chunk_ids = []
+ next_index = len(state.citation_index) + 1
for citation in citations:
- citation.index = next_index
- next_index += 1
- state.citations.extend(citations)
+ cid = citation.chunk_id
+ if cid not in state.citation_index:
+ citation.index = next_index
+ next_index += 1
+ state.citation_index[cid] = citation
+ chunk_ids.append(cid)
+ state.citations.append(chunk_ids)
def create_skill_extras(
@@ -323,6 +189,8 @@ def create_skill_tools(
tools: dict[str, Any] = {}
if "search" in tool_names:
+ max_searches = config.qa.max_searches
+ search_counts: dict[str, int] = {}
async def search(
ctx: RunContext[SkillRunDeps], query: str, limit: int | None = None
@@ -335,6 +203,14 @@ def create_skill_tools(
query: The search query.
limit: Maximum number of results.
"""
+ rid = ctx.run_id or ""
+ search_counts[rid] = search_counts.get(rid, 0) + 1
+ if search_counts[rid] > max_searches:
+ return (
+ "Search limit reached. Answer the question using "
+ "the results you already have."
+ )
+
state = _get_state(ctx, state_type)
formatted, results = await skill_search(
db_path,
@@ -353,21 +229,12 @@ def create_skill_tools(
async def list_documents(
ctx: RunContext[SkillRunDeps],
- limit: int | None = None,
- offset: int | None = None,
) -> list[dict[str, Any]]:
- """List documents in the knowledge base with optional pagination.
-
- Args:
- limit: Maximum number of documents to return.
- offset: Number of documents to skip.
- """
+ """List all documents in the knowledge base."""
state = _get_state(ctx, state_type)
result = await skill_list_documents(
db_path,
config,
- limit,
- offset,
filter=state.document_filter if state else None,
)
if state:
@@ -395,119 +262,84 @@ def create_skill_tools(
tools["get_document"] = get_document
- if "ask" in tool_names:
+ if "execute_code" in tool_names:
- async def ask(ctx: RunContext[SkillRunDeps], question: str) -> str:
- """Ask a question and get an answer with citations from the knowledge base.
+ async def execute_code(ctx: RunContext[SkillRunDeps], code: str) -> str:
+ """Execute Python code in a sandboxed interpreter.
+
+ The code has access to search(), list_documents(), llm() functions
+ and a virtual filesystem at /documents/ with document content and
+ structure (metadata.json, content.txt, items.jsonl per document).
+
+ Use print() to output results. Each call runs in a fresh
+ interpreter — variables do not persist between calls.
Args:
- question: The question to ask.
+ code: Python code to execute.
"""
- from haiku.rag.utils import format_citations
+ from haiku.rag.agents.analysis.dependencies import AnalysisContext
+ from haiku.rag.agents.analysis.sandbox import Sandbox
+ from haiku.rag.client import HaikuRAG
state = _get_state(ctx, state_type)
- answer, citations = await skill_ask(
- db_path,
- config,
- question,
- qa_history=state.qa_history if state else None,
- document_filter=state.document_filter if state else None,
- )
+ doc_filter = state.document_filter if state else None
+ context = AnalysisContext(filter=doc_filter)
+
+ async with HaikuRAG(db_path, config=config, read_only=True) as rag:
+ sandbox = Sandbox(client=rag, config=config, context=context)
+ result = await sandbox.execute(code)
+
+ if state and sandbox._search_results:
+ existing = state.searches.get("_sandbox", [])
+ seen = {r.chunk_id for r in existing}
+ for sr in sandbox._search_results:
+ if sr.chunk_id not in seen:
+ existing.append(sr)
+ seen.add(sr.chunk_id)
+ state.searches["_sandbox"] = existing
if state:
- _append_citations(state, citations)
- state.qa_history.append(
- QAHistoryEntry(
- question=question, answer=answer, citations=citations
+ state.executions.append(
+ CodeExecutionEntry(
+ code=code,
+ stdout=result.stdout,
+ stderr=result.stderr,
+ success=result.success,
)
)
+ if result.success:
+ return result.stdout if result.stdout else "No output."
+ return f"Error: {result.stderr}\n\nOutput: {result.stdout}"
+
+ tools["execute_code"] = execute_code
+
+ if "cite" in tool_names:
+
+ async def cite(ctx: RunContext[SkillRunDeps], chunk_ids: list[str]) -> str:
+ """Register chunk IDs as citations for your answer.
+
+ Call this after searching, with the chunk_id values from search
+ results that support your answer.
+
+ Args:
+ chunk_ids: List of chunk_id values from search results.
+ """
+ from haiku.rag.agents.research.models import resolve_citations
+
+ state = _get_state(ctx, state_type)
+ if not state:
+ return "No state available."
+
+ all_results = []
+ for results_list in state.searches.values():
+ all_results.extend(results_list)
+
+ citations = resolve_citations(chunk_ids, all_results)
if citations:
- answer += "\n\n" + format_citations(citations)
+ _register_citations(state, citations)
+ return f"Registered {len(citations)} citation(s)."
- return answer
-
- tools["ask"] = ask
-
- if "research" in tool_names:
-
- 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.
- """
- state = _get_state(ctx, state_type)
- formatted, title, executive_summary = await skill_research(
- db_path,
- config,
- question,
- document_filter=state.document_filter if state else None,
- )
-
- if state:
- state.reports.append(
- ResearchEntry(
- question=question,
- title=title,
- executive_summary=executive_summary,
- )
- )
- state.qa_history.append(
- QAHistoryEntry(question=question, answer=executive_summary)
- )
-
- return formatted
-
- tools["research"] = research
-
- if "analyze" in tool_names:
-
- async def analyze(
- ctx: RunContext[SkillRunDeps],
- question: str,
- document: 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.
- """
- from haiku.rag.utils import format_citations
-
- state = _get_state(ctx, state_type)
- state_filter = state.document_filter if state else None
- output, answer, program, citations = await skill_analyze(
- db_path,
- config,
- question,
- document=document,
- document_filter=state_filter,
- )
- if state:
- state.analyses.append(
- AnalysisEntry(
- question=question,
- answer=answer,
- program=program,
- )
- )
- if citations:
- _append_citations(state, citations)
-
- if citations:
- output += "\n\n" + format_citations(citations)
-
- return output
-
- tools["analyze"] = analyze
+ tools["cite"] = cite
return tools
diff --git a/haiku_rag_slim/haiku/rag/skills/analysis.py b/haiku_rag_slim/haiku/rag/skills/analysis.py
index 022ce778..97f099cc 100644
--- a/haiku_rag_slim/haiku/rag/skills/analysis.py
+++ b/haiku_rag_slim/haiku/rag/skills/analysis.py
@@ -6,15 +6,20 @@ from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
-from haiku.rag.skills._tools import AnalysisEntry
+from haiku.rag.skills._tools import CodeExecutionEntry
+from haiku.rag.store.models.chunk import SearchResult
+from haiku.rag.tools.document import DocumentInfo
from haiku.skills.models import Skill, SkillMetadata, SkillSource, StateMetadata
from haiku.skills.parser import parse_skill_md
class AnalysisState(BaseModel):
document_filter: str | None = None
- analyses: list[AnalysisEntry] = []
- citations: list[Citation] = Field(default_factory=list)
+ executions: list[CodeExecutionEntry] = Field(default_factory=list)
+ citation_index: dict[str, Citation] = Field(default_factory=dict)
+ citations: list[list[str]] = Field(default_factory=list)
+ searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
+ documents: list[DocumentInfo] = Field(default_factory=list)
STATE_TYPE = AnalysisState
@@ -69,7 +74,12 @@ def create_skill(
else:
db_path = config.storage.data_dir / "haiku.rag.lancedb"
- tools = create_skill_tools(db_path, config, AnalysisState, ["analyze"])
+ tools = create_skill_tools(
+ db_path,
+ config,
+ AnalysisState,
+ ["search", "list_documents", "execute_code", "cite"],
+ )
extras = create_skill_extras(db_path, config)
skill_instructions = instructions()
diff --git a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md b/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md
index cf590817..31be509d 100644
--- a/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md
+++ b/haiku_rag_slim/haiku/rag/skills/rag-analysis/SKILL.md
@@ -10,4 +10,69 @@ description: >
# Analysis
-Use the `analyze` tool for complex analytical questions. It writes and executes Python code against the knowledge base in a sandboxed Python interpreter.
+You solve complex analytical questions by writing and executing Python code against the knowledge base.
+
+## Tools
+
+### execute_code
+Execute Python code in a sandboxed interpreter. Each call runs in a fresh interpreter — variables do not persist between calls. Use `print()` to output results.
+
+Inside the code, these functions are available (use `await`):
+- `await search(query, limit=10)` → list of dicts with keys: chunk_id, content, document_id, document_title, document_uri, score, page_numbers, headings, doc_item_refs, labels
+- `await list_documents()` → list of dicts with keys: id, title, uri, created_at
+- `await llm(prompt)` → string response from an LLM (for classification, summarization, extraction)
+
+Available modules: `json`, `re`, `math`, `pathlib`
+Not supported: class definitions, generators/yield, match statements, decorators, `with` statements
+
+### search
+Search the knowledge base directly (outside code execution). Use for initial exploration before writing code.
+
+### list_documents
+List available documents. Use to discover what's in the knowledge base.
+
+### cite
+Register chunk IDs as citations. Call after your analysis with chunk_id values from search results that support your answer.
+
+## Document Filesystem (inside execute_code)
+
+All documents are mounted as a virtual filesystem at `/documents/`:
+
+```
+/documents/{document_id}/
+ metadata.json # {"id", "title", "uri", "created_at"}
+ content.txt # Full document text
+ items.jsonl # Structured items (one JSON object per line)
+```
+
+### metadata.json
+Document metadata. Use `Path('/documents').iterdir()` to discover documents.
+
+### content.txt
+Full text content. Use for regex or keyword search across a whole document.
+
+### items.jsonl
+Structured document items. Each line is a JSON object with:
+- `position`: sequential position in the document
+- `self_ref`: item reference (e.g. "#/texts/5", "#/tables/0")
+- `label`: item type — "section_header", "text", "table", "list_item", "caption", "formula", "picture", "code", "footnote"
+- `text`: rendered content (tables are markdown with `|` columns)
+- `page_numbers`: list of page numbers where the item appears
+
+### Cross-referencing search results with items
+Search results include `doc_item_refs` (e.g. `["#/texts/48", "#/tables/0"]`) that correspond to `self_ref` values in items.jsonl.
+
+## Strategy
+
+1. Use `search` tool first to understand what's in the knowledge base
+2. Use `execute_code` to write analysis code
+3. Iterate: run code, examine output, refine approach
+4. Call `cite` with chunk IDs from search results you referenced
+
+## Important
+
+- Each `execute_code` call runs in a fresh interpreter (no persistent variables between calls)
+- Use `print()` to output results — the output is your only feedback
+- Always execute code to answer questions — don't just describe what code would do
+- Use `await` for all async functions inside execute_code (search, list_documents, llm)
+- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `cite` tool separately to register citations.
diff --git a/haiku_rag_slim/haiku/rag/skills/rag.py b/haiku_rag_slim/haiku/rag/skills/rag.py
index e9e52772..c1044831 100644
--- a/haiku_rag_slim/haiku/rag/skills/rag.py
+++ b/haiku_rag_slim/haiku/rag/skills/rag.py
@@ -6,10 +6,8 @@ from pydantic import BaseModel, Field
from haiku.rag.agents.research.models import Citation
from haiku.rag.config.models import AppConfig
-from haiku.rag.skills._tools import ResearchEntry
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, SkillMetadata, SkillSource, StateMetadata
from haiku.skills.parser import parse_skill_md
@@ -21,7 +19,7 @@ CRITICAL RULES:
3. When a skill returns citations, always include them in your response
"""
-_RAG_TOOLS = ["search", "list_documents", "get_document", "ask", "research"]
+_RAG_TOOLS = ["search", "list_documents", "get_document", "cite"]
def get_agent_preamble(config: AppConfig) -> str:
@@ -32,12 +30,11 @@ def get_agent_preamble(config: AppConfig) -> str:
class RAGState(BaseModel):
- citations: list[Citation] = Field(default_factory=list)
- qa_history: list[QAHistoryEntry] = Field(default_factory=list)
+ citation_index: dict[str, Citation] = Field(default_factory=dict)
+ citations: list[list[str]] = Field(default_factory=list)
document_filter: str | None = None
searches: dict[str, list[SearchResult]] = Field(default_factory=dict)
documents: list[DocumentInfo] = Field(default_factory=list)
- reports: list[ResearchEntry] = Field(default_factory=list)
STATE_TYPE = RAGState
diff --git a/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md b/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md
index a1c90206..77cb2939 100644
--- a/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md
+++ b/haiku_rag_slim/haiku/rag/skills/rag/SKILL.md
@@ -5,31 +5,55 @@ description: Search, retrieve and analyze documents using RAG (Retrieval Augment
# RAG
-You are a RAG (Retrieval Augmented Generation) assistant with access to a document knowledge base.
+You are a RAG 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
+## Tools
-**Default rule:** If the user is asking a question, use **ask**. Only use **search** when the user explicitly wants to browse or find passages.
+### search
+Search the knowledge base using hybrid search (vector + full-text). Returns ranked results with context-expanded content.
-- **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 browse, explore, or find specific passages across documents (e.g., "search for embeddings", "find mentions of transformers"). Returns all matching results as sources.
-- **ask** — Use for factual questions that need a synthesized answer (e.g., "what is DocLayNet?", "explain the methodology"). Searches, synthesizes, and returns only the chunks actually used as citations. Always include the citations in your response.
-- **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.
+Each result includes:
+- `chunk_id` in brackets and rank position (rank 1 = most relevant)
+- Source: document title and section hierarchy
+- Type: content type (paragraph, table, code, list_item)
+- Content: the actual text
-## When search returns irrelevant results
+### list_documents
+List available documents in the knowledge base. Use when the user wants to browse what's available.
-If your first search returns results that clearly don't match the question, **do not keep searching with variations**. Instead:
-- Use **ask** if the question is factual
-- Report that the knowledge base doesn't contain relevant information
+### get_document
+Retrieve a document by ID, title, or URI. Partial matches work. Use when the user wants the full content of a specific document.
+
+### cite
+Register chunk IDs as citations for your answer. Call this AFTER formulating your answer, with the `chunk_id` values from search results that support it.
+
+## How to answer questions
+
+1. Call `search` with relevant keywords from the question
+2. Review the results — they are ordered by relevance (rank 1 = best match)
+3. If needed, search again with different keywords (you have a limited number of searches)
+4. Synthesize a concise answer based strictly on the retrieved content
+5. Call `cite` with the chunk IDs of search results that informed your answer
+
+## Guidelines
+
+- Base answers strictly on retrieved content — do not use external knowledge
+- Use the Source and Type metadata to understand context
+- If multiple results are relevant, synthesize them coherently
+- Be concise and direct — avoid elaboration unless asked
+- If the search tool tells you the search limit is reached, stop searching and answer with what you have
+- If the retrieved documents do not directly address the question, say: "I cannot find enough information in the knowledge base to answer this question." Do not guess or infer from tangentially related content.
+- Do NOT include chunk IDs or UUIDs in your answer text — your answer should read naturally. Use the `cite` tool separately to register citations.
## 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
+- Use `get_document` or `list_documents` first to identify the document
+- Then search for the topic
-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?"
+## When search returns irrelevant results
+
+If your first search returns results that clearly don't match the question:
+- Try one more search with different keywords
+- If still irrelevant, report that the knowledge base doesn't contain relevant information
diff --git a/tests/cassettes/test_sandbox/TestSandboxDoclingDocument.test_returns_dict_for_document_with_docling_data.yaml b/tests/cassettes/test_sandbox/TestSandboxDoclingDocument.test_returns_dict_for_document_with_docling_data.yaml
deleted file mode 100644
index 29228880..00000000
--- a/tests/cassettes/test_sandbox/TestSandboxDoclingDocument.test_returns_dict_for_document_with_docling_data.yaml
+++ /dev/null
@@ -1,42 +0,0 @@
-interactions:
-- request:
- headers:
- accept:
- - application/json
- accept-encoding:
- - gzip, deflate, zstd
- connection:
- - keep-alive
- content-length:
- - '95'
- content-type:
- - application/json
- host:
- - localhost:11434
- method: POST
- parsed_body:
- encoding_format: base64
- input:
- - Docling processed content
- model: qwen3-embedding:4b
- uri: http://localhost:11434/v1/embeddings
- response:
- headers:
- content-type:
- - application/json
- transfer-encoding:
- - chunked
- parsed_body:
- data:
- - embedding: j5CWuf1eFDwQ3rw5qIgTPYiKkrp962g9vSmVPYLJbjwdMl08JlB8uwloKLxK5109gvwPu65eaTyNiiW8SQFjvSKh4LuAFGe8NgU1PNzi9rvtYGG8ci3XPCzoUz3IisM82ObpvGTF4rroKMS8Cz1HvVivxTzHiOs8r6cnPTpm5by5NAo9IRaPu8HbMDuQCG+8QROhvI1mbbzqAIc9H3c5vFnJvjynSsK8NZ7TPDNkkzyre607U4irvLpfJTxCl9+8+fidvJXFHLye7Mg56qa+OxDd47yyxYy84eVPPeNHazqA/q87YRStu+/q0rxAVk08q0MJvOtkWzu3ubC8XnBbvBJmpLukn8W8vvX/PLu6lLvmBgk8QJCpvArcaTtE9yE9wEoHPBpDDjxHnX+70EsNvVsKobtGTgk9c3uZvGAUuDxRlWo8gvoPvEMO1zt/FCM91PsyPFvzKT12M7Q8x+YePI+S6LxUY188vouMPDcgBz2c3HI7F7ecPIDyKbt24os7m6DXvG7GxLx030K8bW26O6RI4zp6srm7mv0gPMA9YzpgoOO7rRBmvHaw9rs/oP06kgIEu3QNKTs7thC7QCU1uPLDdrwjEVo8Ao3ou+2aP7ywoMU5dc4WPRWmtDxVuiU91R+hu8xrlTzo7o+8p60LvB2ZyTvIwLm8BoO7u6awwbtMvrc89Dq1PMQkODxAoDG8KhrGOw/Xwbzt3p68pmsKPESdPLxxqR67WwGNu14zmTw7+Hy8ZLJsu25f4zsOr7+6zJtqvFqeDL0NvpS8tolKvKSB8jun/0U7ANhAPNwBOLwgIWM8h15CPLLmEbot+Ow80r4pvB5ZAjzMDay6gCpFPEM0L7zPztI6TWOzvC13cjx0Z+07qoUDPEKjNLwW+iK7HA2tu4oYPL37lNw6jOMAvLxi+jpBU5u7jA2AvOrUNLyXTsK8ceSBu6JcE7yLghc8vyVqPCTyxTw2PW+8/bBOO6xVOjxbW907tDmVu3fwCLwtnxs8ll2aPItD8ztwA4E88JjYu+85IDw4pv670s8MvGz8erxs37o8vO+kvDLirDxxfak8DigEPF6W3jyffr27qDBrOwglLburECM8O9tFvAeTrjwXqU+8P2I+PK9n47w/Pk28d1WPugZjETx+Krs7bLXZvMyN5buPnvg8jm81PEAyTTw/Cj66+XCxu19zgDwTvYm81/81PD11bbslR5y7iIBXPFs6hrwo3go9dKOfPIY4DbxPUT27G/V2OkgaRjyTDGe8VB2ruvQYFzy1WX+8/pwnO1ePlLwKiEa8VjBePIMwLrzTpES8vkdru6pwc7zPRGy8IKu/uzqns7sKbwm8Vs9xPIZdbLx+HyW8/WcfvCrSC7xdhqO8/WqCPFinhjxb2vW7btDOvMx56rsmEJ27tSgFvPg2Y7tXipW8HQ8Ku3Rfjbs8y6i7AydjPY/YN7wso5c6aJObO9VoCTxhcAW9ah0iPCeNhTs/AEc8LnoIPS7KsbyJREg8/tAIvbEZxjsijQM7NeZGPJxzQz10kKy5lWeeO6+i3btP8X08dozMvBOEPTwdRF68wMn7vBbueDyXBZ480a7PupL5kLw2hrS8548GvOTkjjlkrO47VI6cvHOpYLut8ZM8N4qBPO66C7yXcI48xg8PPGNAEzxHLT67zR+BO00SarynkAA89lmTvLRmwjq+sfO65QVRvO9m6Lwfr1q7WmMkvX2Iyryo1pe8dJfcvNemYrxh0308zTH0PHqtsDvRxxW8+hiHvBOf9jwV1Te98GiJu/OQyztm+ZM6JTYEOxSDPj3wOqw8SRQfPFBiMbwLMgO8tfc6PJa8p7uh+Le86huSu/G2xTyln7o5kDituxth1bwWoBm9TIPDvHa7oLxgZL68GCOeuyyrQTyYHo+7jbePO+NH+DwqY9y6lEE3vJxTqry52uu6dXnxOpHUF70p7sy7spTevCsDCT1OkAM9bUxcvGn6UzxibTm8REqvPM5t2bsC72Q73bymOtCWJrw1S8s76OYavIZi57tV7088B6mqPDdBQby7mPM86SdzvArTETxdyqw6KUMlPDrQ3juaMCI7C8X/ulW0nDxw01w8mxh3OqmO6LxSmHQ9lr1yO/cLlDyizPI8lwamvFifi7xckKY7S40tvf8Oobzus2I8piX8vIG2lryPAQ49eXECOza4lTnyqwC9SNS2t+k9izzF1Qy8dCJBveCOlTxhwPc8CHX+u+/r7DuA0nG70jVjvLe4qrwokbo62garuxs4QTtiGs87KGTRu1rrszykGbO8VhnavAp7sLzh6Fc8NSeHvKKeaD24i8m8+BriPP7XrryagCe74881O3oNEL3d7PG7SAXDuw5GmDy/9BQ9fiKDvSXRjTvah1A8UnwRPQ2vDT26/5G8z6fbuy4JMTxYy4m8wEqHvNRaxLwKSEs8eyNeO65jgjxSP768PeslvLj4tb02tnI8QWbyO6yqCL3Kg/I7jzvpvNVGorzTsZS8ThSKPJBJjTpRS8e8bva+vJj/4rsa+8U7WufOujup5Dzv3Vq8Ov9tO3j+b7w+zCG8DIrhPO7ZnTzfv8I8kGguPVQZqjxWLDI9hW2ePAAR5zxZ4A686vdevFcygzkk7XG8777tvKPVZTxPmsc7GOTzPHvhID15ImO8ZVKzO73eJrx+iOg7SOLuu5AS6TyrhQ+7Tt3YvHPn2DyLxhw7vEKOvHT0k7wOMII8YYe7PLV7cDtoJD693sI7vXkRczu4/CC9yPDfu/EaALzQLQC8FS19PC13qbx8aWQ8QfwWOQsvWbyg9vI6BRxrPAAQmTw5KNO7aY9cvEPySzzEnMm8vuHWPLiauzweBAs8eK3rvOQTj7qOoaS7beL2PMlnijys93k8F+rQuvUAkDvRL5888MbVvKMmWDvRCPW7S9vFvE24CTxCM+G8YCkZvWxNKTwkFAK9W+IkPOrv8LwDCp672poaPUSvFb2eGpA8od/8O9UeRjv8xMC8HhwLPGoyiTz2mwI8lHpmOrrwm7mVqA28cG+cPCpGU7xogKM73AW0uzqqljydPLy7PhGMvG/umTz0rvQ7nLGAuA/wxbuVCjW9C+QIPNatBzyrZgk8KcaQvIRYEjvbzXI8027/O9hhbrzGuqo7NLg7PbWfszoU46G8HODCvJnq6zvAzNG7WsQnO6ZITrxY5xK9buEmvRPsJjtfP6K8Hx1bPYtbMbwE7qW89yWqutl5pryMJFg7u5krPYXYcjwk4YK8q+GAu/7jqTsXOnQ8DssPPF5ebjzDoD88h5LCvPqPJToAtdq8a3/jO62xjjo6ibi8W5tgPGPHJbxB2io7qmrJO4a9NL3WhTC9KtiQvEERLjz6l+07eYgFvJEOGb2I5CI9p7QNvBCnJTz1vn67yzTsvOQYKT1L+i87tVlcvIN2Gz1codQ8opu6PBfGhzm6aQE9RTFQvALPRr0Zxi29Igm5vMPHJDyl8wQ8HP6FPF0mwbzG7p07yRwHPHk0kLtFToc8Qga2O/gzcjy+eEo6Hl6fvFYJBzqIXOK731wPPTdXVzsEaO48qKYsvXsyJ7xJfIA8YI27u9nSvjz3bh89lZvRPCv7art3fRQ9I0hYvVT/djxKG0i8v5BkPPGKPLxE1IU8oHPMuyIM4ztt4kw8cmB7u4QBgLtMpAa8BJJYvARu6jsmLvo8+xS4vL4gNbzsIQY8S3oGPcNI+Dufh5I8yIYfPEdOljwJ1b28J4OwvLsJvrwpxTI8hB9uupWGkbyPpao7vtntO9zcHjx5HpO3x5gXPesl7jy3IoI8MzGbvD+X+Lzd9Gy72KnQu0WWJTskzbG8hNIWvPDl+bvjpMu5q4Tcu75myjtaoxa96OGwO7DxDLxuoF48/4cJPGAN6TuhJx+9MpSVPVsOerqkDBE89W1ovPyhsjvV00S8ydQRPEJf5Dvij8s8CGODvGsNMjzDdzU8d54xO3n6ajxylpi7eKxkvIw1dzwSCBK9d2wLvA+A1TzmeCy77EPwu/Z9xDz8Edc8AnRwvDmSDryFJHg8mFGgObKTeDzq2Hs7WHUsvOnbibzyHo28tZASPBwmcjqPPCm9jkwQvNC9wTsWhxY76kqqvGF7ID02DfK7viKYPLmlZrydTI68FNkAPOEYdTzyZLy7bcI5vB59JjxQZYW8WpuVvOYk2byqUhg9ATMEPKAqXjx7ofS7g4e7vFW+orvN2Iq7ANgXvHkthjvxPya7hKbYvLIn3zy4Q3o8eDaHvEOK+jq9A4I8e15gvcbx47ye5Gs6CtxcvKAiwjyO27Y61cSHOqaRljxeJ988G4DOvC6WSbzYqJg7XlHzO93t8LxZOAE8tnxLu3sPULwR1Ry8CsnzPFucS7ytyxm8VL9CPLPqQzyiMMs8J3iHuQ83nbu+5Zk7nL9cPGHbFT0PurS8D3cGPGnmNzzPyBI8voPSO0tTYzzFHCg7USkfOyTKgLvSkkE8cNrbvIUmjrwYDbu77UFEvFY0NL0+H209eEzyvM4ne7vbIO47LKs8vDs+IzuNwZu6HPqbu+j+KD0LCVs9QoEBParWYzz1jJ68sy4SPEUoND1zU4883R4YPbquRLyg0Pc8PMH1vMkyCL3nhZ06UpCmOuICwTxY+Co5ucPDvI2oYzyNpay9ge77PJU8mzvBtCE8ne5Qu5kQgDxiOlq7wwWOvKEE9bvMkvy7jvN3PMIKlTtOX/w7iiPvPDo4vzxJvt68MHNFPHUK1LxCP/y7yIYxvMMFQDzuMDe8thtiPBeb6LtESwo8X74XvGsjXbxKbWa8wgrRvH70Ujz+lhe9cdYVujNV/jvkwHM7vxEBPfEWEz0nh0Y7aSMrvE5UFbz4LmA8wzLOvBnUgrybN+S8c/dnPOManTpno9G8p2RJvGvPKDzXi3C8dKlKvcUr3LyeQqI8dvrPOziBa7yVU2E8qXL9PGxps7vR7sm8AF1qvBHoP729FrC7czaKvKpKgzxQIFI8ybvBO4ijRjyD05U89Bs3PCc4IbzIVHA85dErO/Y6QrybE5C6q2lePPLW+DxVbUm7wDe+O6eaEDx0ZZi7nPHfOwYPLTzSHeQ6X1zpvNEXR7xV6Oq7QpDMO6VT0rzZkQK7k8FbO/LW9bple9K7MZiTugPW7rsgPgC9x1LFO6Mgkru9hoM8uh5tPPpgyLzWgI284xQHPBL5v7t2Ole7ND+euxZmu7qq3tC7PFtoPHerQ7sqThM987jmPDKFpTs7zRe93dU3O2p/MDwDq6y8PEeOu+vsGb1k4U47CyqXO++aWzw2pIK8x6NJPDCxCbv1xRo9KuWaPAWNXrtJdb08csaDuyDiELtR91E8BAQ0OgXQPb3BgTE9X5iFPCg44bwLhAe8lCEgvUOz/jyMMWe8iL4QPdvYzbykNg+99xVwvHPaibslvsy8TBFqvC1iibz+IwE880IJPSxOT7vOLhE8cczWOiQ3njto1Dm8+kXAOmCIqjxoaQW8yKYCPeozGT12uvq8u5C+PCcWiDxIa388OtTwvDkOtztkr788/zsAvNKZ9LsPH8e5yI6wuxEhnrwKn0i8qaymPEKD5rypXv67DSXRO0ZULDoYK/65paGpvAM5SDv9FJ88/u6cOnjUZDxJ9da74/pHuhIZDL32p1y7QB+4PAjX6LueliQ8VWokPZfyFTzlK368Y+I4PIgYKL1lSbi7hHA1vXlFDDxVH2I7SsXqvFwGajzZrhW9BS7GPJi+IrteDAE6XvqvPGTp5LyjTN06eNfcvJ5ztrzXSS+8BywDvWyUA7wy3nW7Q3OUuxN/DzzxkZO8hofdPCFGyjumXwa8edy6PCobsLwqOhm8WcwMu0DloLwLnQU9KtQ3vAMiSb1v05S8LZcKPEsuajy2Udg7qNkPPCuNvDv6r8M74BHlvPoePrwmYZ07xKWLO+GmfrvTaJc8KPwZvYrixDtzroO8na3qPCty3TpZTTA7HdDTvKxEdrxoWtO8eSjfu9nW3Du2lhg8bG4JPFFsaDsGtRg777qGvCT11jz1dyo85VxNvB311jz5DZa8cYMEPXPmYzwNqnS8Xnr3PID3CT2qk468zqlDPKDVOj1G62q8mlWsvIWFw7vWfq07d1RMvPx7rzxIlis9bpbDvCfOTD3PlU07igmyPHIzI71LzI88GTwcPT1dbDvW5VS8h1mVu1XqpzyvDAg9gs99OqeCljz/KRC8Yli8PCtV9zqJfoe6PnF7u0/EwzyobYY7/yAdvGOIIbxCB5a7jBLmurTXGr0AziC8k/1lPKlM1DxuN4M8wbkOPVMWKrslZ4s5WZSXvJSaqjw41sy8vRrCPP7GTL0BHt68F+nZO+L9I7wz7K88EPlYPHg2pjwGyeQ8JagtPGACkjxyZK28fgOCPAwN47wctaw7L0gNvfkLZb1fxjC9pJ+wOwupD70S7Lw7Ft8qvGMBFrx4Uqm7wcYAvIePYjq34vc73wvZPEQ4PbtJQg48xnnnvPqVsroLPcO8di3Ju507Aj20D6G8qjU2O4BOg7wy2427jbICvKlQszzscg69nIM9vNFTCbxmRT88znnLPGhMLT3cxXY80t85ussMID3zcIA8LTDbPDyxqbyT/uM7jje8vJ6yJbyaj0M9GlWSPJr0aLxbUr48ah4MPTm8abvZ/1098ul6PEhDSLygEv47gDkyu/6hNbx8wBm9qQ7QO7c+1Lx5RgG8M+G5OXVijTyWRi89AksFvCVgBbziiBg909MxvRqzH7wg2/q67Qqzu+felDwaAOa7cu+0O/0DfTzSwZo76Fj/Ozo1+zul7xi8ep86PJNW5zx9c8K89EQsvDLLmbzxDla8JZISum1s1rtYedA8BMDzO58eo7trinQ80t+6u6jfuDtE9xy9tFO4PNDqQbwZGNG8YVnZPF6QmjxwNsy8afj8u8DHbzufIMG7gwElPPWW1LyYhxG902KxuuNyaLxAxXo7fo2JPD60/bvD0RE8GVs4vH2SGj1MUYg8xzILvT/SELwV1Xk7+T7Mu6FA7zo4Zuc8vbGLvHYeCLyR45e8vlj8vFn7sjvv+OY76CuAvF56hzoXPRG8lVDePIJ3sjylG0C8HvRCO77ApDxwzka8B3uEvLnngruKaKM8SVUjO19ELDxEljy8wBIbO35U5zwDn4c8SSVTPU5ABb02OTi9sVxkvFGjJL3d5SI8tBgZPQlNgjyjbzq9GpLeOzIkEDymXfi89ma/OyAjtjxzzrS72jDQObp1gzxuaao8CCkXPHDYErwckuM733bSPJ0+xDw6oMa8IZpPPYi8ibz4kt28q4RIO7HcCj29MJs8Voa6vLj0+zxFj2g8DM0JPQoEybwdeOm6SLbFPBdYcDz0nG681uEBPb+P1ryz0Ji8sX+cPJsAG7zlD0Q8OIpmvAWWujyD/BM9IKc+vT9w+DrKBsQ8BFSxOzJsBLzpgNY7p0aOvCEoIT3A1jU8OIstvfhmyjzyoKQ6pN6wO+AeOTsBmTs8N4ILO6bhM71T1lg7nI2JPNAgg7xqm+I8x8Nwu1/oBL29WDi9AYu9vLKFJj254Iy8U0c4vL7DBTwle+S7YjfXOhN7/7x4QIi7W8u5vE9DxjziJyy9bvvpu04lSTx27Nk7LIufvJaSED2qOF68s0OWO6f0FbxY6Qa8AZNMvGHLsrzWAL689DLOPIi9uru2Gui50P9qPA0U9TydyAA89gKyvHA8nLzjz6o81YQZuNiPBLw48nO79P8nPDTuZ7scaxQ7T9c4PMcKyrypOI48b2PbPCmGCT27jTo87skdPCwNyryiWyu8GFJoPE7BVbv04ug77foUPS6H6bycoZQ74HAOPLueNb1QKH28Xcc+vK8JHLx5M/a8fjM7veoRuTw3Wog83HPwuVNe1Lv4hFM8Qj05vGSPFr15XzE8aLiqvAMjaTqtXna8jc0EPbOhULy7Oia95ea7ORYxpTozMrC8Y2JbPOUomjqX6Q88N2ecvPKJ6DxsFQC9DBqdPB0m2LqsHBC9OlfQvNGAgzxPuS69fzNTPPW3R7zF85Q8uombPMwNSzzAy8e8DGdmPKnGXLr/Sg+7KFyMPByqE7xHYci7KCJrOi6/gjxTY868e3e8uUvvSTwEPQi72psBvWHHBTxKBOQ6o2kWvd1IrLp7EJq8q7IGvClodrxyKt08/5ajPCz8nzxrbsM8YzcBPT8kkrzH/L88OwpVu/NtCbxLkiU8Jt7cuwf/17vbWQ89DVv3u12Z1DtuXww9IMuQO9A8DjyudiY8SL5nvDAqKTy3Bvg7yf8MvI8mHzuc6yq8Mx76vOCm+TyLWIo8tsFkvBNpyTy8CdY8LNBsPIxzurv06ZE87/MZPCEfcj0/VdC8tDlSPIVFBDx95bM7LhCDvKDcwzz63kC8Kp3bPEA09rvd7Y26tSmsPAeY/jvW1eY6SDsNu2bEmrsIqUu7br4gvOzVhjzKyrs8oiMavLQiTDsUjTg7WaXRvN/NsLx1MxU8NK73PCo4kDuZa5C7gTuMO3mieTuZ3VS8++4/PEzASbySw6S8fE/kvGYM+Typ+rk8XnegPJspCrzMbfW7/XiUPKf1Rz3CtPe8Q10PPH5iCjyrvT+9FoV6uhMNuzu9EAc7YQs0vL67nbti7Le5nSUkPHOl7ruz/mq71PM4PDw7Ej1AQJ0897ZFPR4Q5TrAf+q7ozyivLVtX7pQ1AQ9T4GLu/t4s7pBp1Y87daYvMwhQz2KSJO8aonfuw8vWbyxygi91WmHPA/rnrzaTNY8t5LkPFD3g7y3iA09uSOWvEDAFT2wM229YXVBvKEcz7x+Zh06WrdZPAzaJDzi69g6ifYpPUXwVDzAshY9PFKZPCWCJTxL0lY8iBT5uz8PEjqMAu27fVfrO81TzTuULHK8gkc0PfCyrTtMkX68iWtYPLZLlrqP6wQ8zk8vvNnNszyXFcm8JZuuvOKoED3IS2Y8OsYUPCCQ1LuhVPq8B9LdPMnwdLzouXe8u2alOX2FC70xS3a8Z1cEPcY1D7tNCMy80XWwu3MCybxrvdw6BhRSPDHai7shMeU8mgBOvEeN0jxZS/U7eZDTO2jijTwi9MS8Hwv5u+3ruDp0JtK85mTouv0n0Tr97Z68SSHgPDsDujvghQ69BgDuvGvg17y+K4o8ZGl+PHEvXjwC2vS8CccAvSQ4Nzz3Wdw8MTPaO/al9jsfVYA8HHoQOdDGtjyDVaI8TsYlvYssDDzNo6U8Gf+DvNSDjTxdeiE83GISO3C+uTyARLw6g/+KvED2WDyFiyi8L7/OPKnlV7ztJIu8943au8JgPzp6D3m8OSQRvakj8bm5a7G8j96HO07yJLwThPQ6jLyfvL27JTxoXAg9ZsUgO+4k6Lv+8xE877N5vAzNJ7wUtOi8vlghPXR7UDuUzik7NKUhvOY9kzvmR5K7vhJHvFqOHb0mqm28FVcePJweajzYae28FRSkPKWHYbwmSTe7OxzQuxy4GT1j77S8PB4MPX9ndLwHTo+8xZKIup3q/7kLiky8T90EvQsjjLzbyly7XEMPvaerojxFszm7XK+fPA7dm7wPXpG7drAaPAIiCj3dLOW7FEs9vCyJkryD1D29LQ7tu8L4hzypzqG8QlmtOwUIJL0eXLw89vnUvM/jCr0+gTy8lNDDvFQXPT0SwgI9YgzCOhmZE7zOQ448peRIvJyhBT2bTam6yRQVOgAHHTzfS/s72FbzOholpLwrTpI8ETkdvLHKvDrzKYa87VikuxFGOr0d9rA7qBXgvHUCmbw9Ii+6g7YKOyQv/ruNNXG8NBJ9vBswIDsP4li9aZxkvX3uOjxt7xO78Br2uwebWL1+N5Q8fLhrvFmMtzuDDpq7oHwsO/ynlTy2Yr275hawPKJT0DePWTw8otoBvTGi5TwZG8M8hq5zvML0fjtHU7U7IJBIuypp+TufV/i7ybP5OelPTjvT5Vw8gGyFvJipwbz5GSU8gRCKO3aP0zzsjIC670pVPN3u0TuCP+m7cWsSvPm8SbwW2p+7bEsdu7B8qbz5RBc8hjkCPZRxjjxYH6I7z4nPPGnkDzrbNG87kLXBPCHIqLxUn+i7eq4cPT1JULua19w66s00PVydJD2+cGy8A7LGu418oTuuNWY8NJe9u6u9pLwtHU67G/ZCPI4RpTwGppu8RATrPNFA97v436W8nl7vvLxESLtuzwA885PkPOuY0zrLpdc7UoOJvCPX+Tuh+4s8pinVu4cp0bvlhUW7c0WcvCzKJz34oxG8z53Eu+GUFDuB7+28NsmwuxRIDj1x7vU82id+PHPrijsXQk88lOPpO8AvxTg//vq79GqtO60+Xjy5sMq8CBhPu7SNWDprR528hifIPElCCTo+GxA8y/yjOxpi0rv2dFC8flAEPXc7rDvNb247F7GSPE1M5bxwEp27CzwEPEUVFjuHJgm8GqvUulASj7zByQC8+B5pPFh7mbymWTk8R+hYvIk6DL3caCa83QUjPF7wnzzmE/a6pIsBvfi0zLwR/RW5FVSsOUWH8DzPKra8tneUPIWk6rydols8bOPDu25UrDsvCn+8E8fHPCsMUjyxjtm704AdOkTb27soRyy9bkNxOvipDTy3ZZK7TLsavLKlTDvY6Nu6/qD+PGF4+byFf4e8QJBsPKnTT7w814K8ft3wvJlnW7wnqEM7Yv6YvKYgxTw1ilE9UqJ3PGnNYbsSCI28ZKsFPPwR/zzxps87zayKvOz6qrxSllS8+B88u/n8R7zAED68X8L8uy/iT73KEmM8OVrkvEU2H71/ZGS7l40ZPDntGjyjrQ26fQPMOxGw87tsPn68rHbIO3r4uLsgu0Q8B7vcvIpz+zxnX526uOQCvQhBL7zHZc+7t7W2u/C2Jr3M80s8OkLjPNpucTzSbSO94AYIvMSD4TzCs0i8HMwnO3gTojy3Ib4874FHPIvMZzxhTf67flKUPH2xXDz3T1k8m7CuPEWFkbz2kPg7ZIcKvXDjUTtowBS9sFXzvC6UHT3vRhE8TDCeu49CCDx0QgK8SuzgPB4dqrtE1i29jxKePJwpNbwa5SQ7uloqum9F3jvsTl+72+oOvdfTf7pQvpm84VoWuzlEuzvaggU8SnEgvIlNXTzwKxE7QyH5O/Wja7u5Voy8HDxZO3jR9jwSQ8q7F4FPOrzq1juGAHw8k0GPu0huK72We4u8tDkVPeknwjzn1pw8gbbAPBiGMj34Hmc8PdtvvFOL0zmq7TO4QF81PC4+dTw2Cjw7WjeqPHXik7ylzpu86IkeO2YlSLxr4e46zG9TPNs7KL3Kj5a8/4ObvMCCoDyfQko8nfjgOwAAFTxWMhq9I2mFPFiBrzz01lY8UDejuxThgryZRds8S/6ePHMbabzlnV68EQYYO5YS5ryGIRe9E052vAWaCbzphqu7sS65OgIBOjvZi/W7HpWePPLyszxMcsQ7nv4+PP89pTyebAe9SgV8PGBvZ7sHixM8cFTLPLuNjjx1e4Q8wwIePdemJDzc9wG8DDrAPIgQLL0jqXC8icihvI/UT70LMwg9OlVTPFEgx7wCBIm7oVlEO9BGRLyMEZS6Iqwjvf7QkLz0K6K6JkAHvVa3Y72rDbC6IaLKvHZZLDvKZhS87NoLO855qTxbJeQ7G+qgvEJyajuWk1k8rSzMu0dUA7v0be68xWZ9O1V5srxm2Os8OZaTvD+3Ib0ggxm7efKPvEogizwcz/g88Nyeux28O7wCZR899ZLevLn7tjsfEc07gzikPHQI8rzj9hG9VAdxvH3jNb18G8Q8K0U6vJQO4TxPqQa8PNffuU27Iz3+CPe8LA2kO3vNHjwHJxm8NHmuPLNAFr3TJBY9HAoZO2QyDj0MHFI8qGVvvIzXCLxWULi7lPsoPQan0ruKKss8g/OyvIlN4TwOkfk7MMHsPDJOeTzLJ8E7BZiZPLk6GjueUzO8WAxYPJyrdrvg9xQ9PsGMPNCu8Lx9sRY5N/vbujwyoDyt7uO7lwC4u1zcD70uW1W8QpB/PCA9t7utpTG8Hql3u0eJkbvLU1G82zfvvOWOsTq6nyu8+zh9vG0njjzlbny8HBEruttcnzx7s6g8FbiJO57InjyVCZ68KSbPPMZHHT3OMue8vLlJPHa0Qj1vw5+8+Sd2vIa2jzriNaM7DLIgPLBEmrxG0ZI8jaurPEF5hDyubsm8TGWKPI3QmLvMuA88/oIKPJ7tOjswBOo87d3MuwgEjbuVUga98TrqvBg8E7zQmts6QyP2O+aW57vGkME8/5yvOsPVOzy6cdY7N6nBvIoxwDxknYG843s2OS1JJDxg+Ye8cvh1PB6rHLz1RaW8e+ANPA9HzTzULtu8tHUgPP6ogLyxc4I8iHkNOJcJ5jjiZj27Pp8JvVDirbxFk7c8FjVxugNBnDwNOVs8GshMvFZyGzyMLMQ5MiJ2u4hG8jq5p0G8VSvJvLhFPbzQoL072ngEvBAwLTyA/xC9xEe6vIDruTogAnc826XCuU1f/jwIUNw5Q/rsvDoDbjxCiIO8x77cususDTySvb08c1dFPITxvLzXAE47fi+WvHYXt7wDNR69+46ivCozyLt/e3e8BeClPC4LRLwDnYm7+6iMvPd6kjz5dhq850EkvBf+lLwRwo68JZiePPEiuDzwVyy7fvapOknfFrzG3YU7N0AhPJN1CLx/4+K6/5pDPFWEGbwScnu8JJKUPEmNkrwFgfc71G/pPB2tx7uYJdE8FCjfPLOXsTyu+wK6vn6XPDNHczxIiS28NfuBvKlUwDzu8rK8lbOOOYA1rDqeORI9xof/uiuuWbx1Ias8Mu/IPIbMQjwzQsQ8tSmNu5zYLLx2DE68bNoDPWX1DTsZeb26w+2oO8V687ustA47Kjequ3Hxczs4EJc8VSl/vIASMTxqbJy8WMspvdWMgrzuDnO8S+MaPPhFGDxjaAQ7AO+IujUKIbzp5ym72mSPvGOd+7tEqFI9/WiFvKkV1rzlHiG9/H6kvP5onDuPd/68NOAqvFUzCLySSRM9toAivTvJSrp4LAK8EJ28PMC2k7wCvrI7NBCUPHHqdjxT4T68xNYMPN+psjw7H7A85XUGPdSQ0rw0cnE8IBd+u75IyDn0wo+80QSUPEPIEbw0II67whyZPFTjv7wc42U8vh01Oy6UezwlmdC66T5TPCFamrqthEw8swtgPD/qIzygubs80f+UOzo1YrwIlTO7M3C1vFHLnTwlmJw8pvzCvNQ0ETyHdwS9/7uGuB1vqzwAanK84Gg0vCKqBD2Icvo7I3bOPD6UDLyDE+k6HYtcvGW/bbwUZ1m8jwNvvBpbtTujUxG9pGYJPUQdCbxXoQI9o3ggPKSLojwg/h06YKhcvKlQN7wYZE+7rFv4O2Fx6DtHWUI86k7RvL8zaTysKYy8L83GvMymNzq0suu8diurupgPpbuvRQC8xTERPCExwLtxir88LTQDvfUZyzuMzcA7IP6zOso1F7wGiaO8QLTbvEWnQrz895s7hblcvI4Hnzt1C9Q8nJQhO7LlirlV1ZU7rFvBOqwTNjtIpCg7QV3eu98ezrthKaK82pZRu4+8VztH5iQ8FUwNvdCsqrzMhde7eoBOvLefgbpl4se7fItAvEq8CLz4qKi8YILYvJ6zTjyLaSU7zQakuwOi5zpXJNe6g2sfPDgw4DrPrPe7nI9uvC/AUr2eLTA8uHaJug==
- index: 0
- object: embedding
- model: qwen3-embedding:4b
- object: list
- usage:
- prompt_tokens: 5
- total_tokens: 5
- status:
- code: 200
- message: OK
-version: 1
diff --git a/tests/skills/test_analysis.py b/tests/skills/test_analysis.py
index 0cada1f9..040ad023 100644
--- a/tests/skills/test_analysis.py
+++ b/tests/skills/test_analysis.py
@@ -1,7 +1,3 @@
-from unittest.mock import AsyncMock
-
-from haiku.rag.agents.analysis.models import AnalysisResult
-from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.skills.analysis import (
STATE_NAMESPACE,
@@ -64,7 +60,7 @@ class TestAnalysisSkillCreation:
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"}
+ assert tool_names == {"search", "list_documents", "execute_code", "cite"}
def test_create_skill_has_state(self, test_app_config, temp_db_path):
from haiku.rag.skills.analysis import AnalysisState, create_skill
@@ -81,8 +77,6 @@ class TestAnalysisSkillCreation:
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))
@@ -118,80 +112,50 @@ class TestDomainPreambleInAnalysisSkillInstructions:
assert base_instructions in skill.instructions
-class TestAnalyzeTool:
- async def test_analyze_returns_result(self, rag_db, monkeypatch):
+class TestExecuteCodeTool:
+ async def test_execute_code_returns_output(self, rag_db):
from haiku.rag.skills.analysis import create_skill
- monkeypatch.setattr(
- HaikuRAG,
- "analyze",
- AsyncMock(return_value=AnalysisResult(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
- assert "print(42)" in result
-
- async def test_analyze_updates_state(self, rag_db, monkeypatch):
- from haiku.rag.skills.analysis import AnalysisState, create_skill
-
- monkeypatch.setattr(
- HaikuRAG,
- "analyze",
- AsyncMock(return_value=AnalysisResult(answer="42", program="print(42)")),
- )
-
- skill = create_skill(db_path=rag_db)
- analyze = _get_tool(skill, "analyze")
+ execute_code = _get_tool(skill, "execute_code")
state = AnalysisState()
ctx = _make_ctx(state)
- await analyze(ctx, question="How many documents?")
- assert len(state.analyses) == 1
- assert state.analyses[0].question == "How many documents?"
- assert state.analyses[0].answer == "42"
- assert state.analyses[0].program == "print(42)"
+ result = await execute_code(ctx, code="print('hello')")
+ assert "hello" in result
- async def test_analyze_applies_document_filter_from_state(
- self, rag_db, monkeypatch
- ):
- from haiku.rag.skills.analysis import AnalysisState, create_skill
-
- captured_kwargs = {}
-
- async def mock_analyze(self, question, **kwargs):
- captured_kwargs.update(kwargs)
- return AnalysisResult(answer="42", program="print(42)")
-
- monkeypatch.setattr(HaikuRAG, "analyze", mock_analyze)
-
- skill = create_skill(db_path=rag_db)
- analyze = _get_tool(skill, "analyze")
- state = AnalysisState(document_filter="title = 'AI Overview'")
- ctx = _make_ctx(state)
- await analyze(ctx, question="How many documents?")
- assert captured_kwargs.get("filter") == "title = 'AI Overview'"
-
- async def test_analyze_with_document(self, rag_db, monkeypatch):
+ async def test_execute_code_updates_state(self, rag_db):
from haiku.rag.skills.analysis import create_skill
- captured_kwargs = {}
+ skill = create_skill(db_path=rag_db)
+ execute_code = _get_tool(skill, "execute_code")
+ state = AnalysisState()
+ ctx = _make_ctx(state)
+ await execute_code(ctx, code="print('hello')")
+ assert len(state.executions) == 1
+ assert state.executions[0].code == "print('hello')"
+ assert state.executions[0].success is True
+ assert "hello" in state.executions[0].stdout
- async def mock_analyze(self, question, **kwargs):
- captured_kwargs.update(kwargs)
- return AnalysisResult(answer="Result", program="code()")
-
- monkeypatch.setattr(HaikuRAG, "analyze", mock_analyze)
+ async def test_execute_code_reports_errors(self, rag_db):
+ from haiku.rag.skills.analysis import create_skill
skill = create_skill(db_path=rag_db)
- analyze = _get_tool(skill, "analyze")
- ctx = _make_ctx()
- await analyze(
- ctx,
- question="Count pages",
- document="AI Overview",
+ execute_code = _get_tool(skill, "execute_code")
+ state = AnalysisState()
+ ctx = _make_ctx(state)
+ result = await execute_code(ctx, code="x = 1/0")
+ assert "Error" in result
+ assert "ZeroDivisionError" in result
+ assert state.executions[0].success is False
+
+ async def test_execute_code_applies_document_filter(self, rag_db):
+ from haiku.rag.skills.analysis import create_skill
+
+ skill = create_skill(db_path=rag_db)
+ execute_code = _get_tool(skill, "execute_code")
+ state = AnalysisState(document_filter="title = 'AI Overview'")
+ ctx = _make_ctx(state)
+ result = await execute_code(
+ ctx, code="docs = await list_documents()\nprint(len(docs))"
)
- assert captured_kwargs.get("documents") == ["AI Overview"]
+ assert "1" in result
diff --git a/tests/skills/test_rag.py b/tests/skills/test_rag.py
index 85ea7448..10184bb6 100644
--- a/tests/skills/test_rag.py
+++ b/tests/skills/test_rag.py
@@ -1,7 +1,3 @@
-from unittest.mock import AsyncMock
-
-from haiku.rag.agents.research.models import Citation, ResearchReport
-from haiku.rag.client import HaikuRAG
from haiku.rag.config.models import AppConfig
from haiku.rag.skills.rag import (
STATE_NAMESPACE,
@@ -13,7 +9,6 @@ from haiku.rag.skills.rag import (
)
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 SkillMetadata, StateMetadata
from .conftest import _get_tool, _make_ctx
@@ -116,13 +111,7 @@ class TestRAGSkillCreation:
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",
- "list_documents",
- "get_document",
- "ask",
- "research",
- }
+ assert tool_names == {"search", "list_documents", "get_document", "cite"}
def test_create_skill_has_state(self, test_app_config, temp_db_path):
from haiku.rag.skills.rag import RAGState, create_skill
@@ -139,8 +128,6 @@ class TestRAGSkillCreation:
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))
@@ -169,44 +156,6 @@ class TestSkillExtras:
assert len(results) == 1
assert results[0]["title"] == "AI Overview"
- 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(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,
- test_app_config,
- 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(config=test_app_config, 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:
async def test_search_returns_formatted_string(self, rag_db):
@@ -274,7 +223,6 @@ class TestListDocumentsTool:
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_applies_document_filter_from_state(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
@@ -287,17 +235,6 @@ class TestListDocumentsTool:
assert len(results) == 1
assert results[0]["title"] == "AI Overview"
- 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):
@@ -310,18 +247,6 @@ class TestGetDocumentTool:
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
@@ -332,343 +257,57 @@ class TestGetDocumentTool:
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):
+class TestCiteTool:
+ async def test_cite_registers_citations(self, rag_db):
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")
+ search = _get_tool(skill, "search")
+ cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state)
- await ask(ctx, question="What is AI?")
+
+ await search(ctx, query="artificial intelligence")
+ chunk_ids = [
+ sr.chunk_id
+ for results in state.searches.values()
+ for sr in results
+ if sr.chunk_id
+ ][:2]
+
+ result = await cite(ctx, chunk_ids=chunk_ids)
+ assert "Registered" in result
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 len(state.citations[0]) == 2
+ assert all(cid in state.citation_index for cid in chunk_ids)
- async def test_ask_assigns_citation_indices(self, rag_db, monkeypatch):
+ async def test_cite_deduplicates_in_index(self, rag_db):
from haiku.rag.skills.rag import RAGState, create_skill
- first_citations = [
- Citation(
- document_id="d1",
- chunk_id="c1",
- document_uri="test://doc1",
- content="First.",
- ),
- Citation(
- document_id="d2",
- chunk_id="c2",
- document_uri="test://doc2",
- content="Second.",
- ),
- ]
- second_citations = [
- Citation(
- document_id="d3",
- chunk_id="c3",
- document_uri="test://doc3",
- content="Third.",
- ),
- ]
-
- call_count = 0
-
- async def mock_ask(self, question, **kwargs):
- nonlocal call_count
- call_count += 1
- if call_count == 1:
- return ("Answer 1", first_citations)
- return ("Answer 2", second_citations)
-
- monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
-
skill = create_skill(db_path=rag_db)
- ask = _get_tool(skill, "ask")
+ search = _get_tool(skill, "search")
+ cite = _get_tool(skill, "cite")
state = RAGState()
ctx = _make_ctx(state)
- await ask(ctx, question="First question")
- assert state.citations[0].index == 1
- assert state.citations[1].index == 2
+ await search(ctx, query="artificial intelligence")
+ chunk_ids = [
+ sr.chunk_id
+ for results in state.searches.values()
+ for sr in results
+ if sr.chunk_id
+ ][:1]
- await ask(ctx, question="Second question")
- assert state.citations[2].index == 3
+ await cite(ctx, chunk_ids=chunk_ids)
+ await cite(ctx, chunk_ids=chunk_ids)
+ assert len(state.citation_index) == 1
+ assert len(state.citations) == 2
- async def test_ask_applies_document_filter_from_state(self, rag_db, monkeypatch):
- from haiku.rag.skills.rag import RAGState, create_skill
-
- captured_kwargs = {}
-
- async def mock_ask(self, question, **kwargs):
- captured_kwargs.update(kwargs)
- return ("Answer.", [])
-
- monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
-
- skill = create_skill(db_path=rag_db)
- ask = _get_tool(skill, "ask")
- state = RAGState(document_filter="title = 'AI Overview'")
- ctx = _make_ctx(state)
- await ask(ctx, question="What is AI?")
- assert captured_kwargs.get("filter") == "title = 'AI Overview'"
-
- async def test_ask_includes_prior_qa_context(self, rag_db, monkeypatch):
- import random
-
- from haiku.rag.skills.rag import RAGState, create_skill
- from tests.skills.conftest import VECTOR_DIM
-
- captured_questions = []
-
- async def mock_ask(self, question, **kwargs):
- captured_questions.append(question)
- return ("Answer about AI.", [])
-
- monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
-
- skill = create_skill(db_path=rag_db)
- ask = _get_tool(skill, "ask")
-
- # Pre-compute the embedding the fake embedder will produce for "Tell me about AI"
- query_text = "Tell me about AI"
- random.seed(hash(query_text) % (2**32))
- query_embedding = [random.random() for _ in range(VECTOR_DIM)]
-
- prior_citations = [
- Citation(
- document_id="d1",
- chunk_id="c1",
- document_uri="test://ai-overview",
- document_title="AI Overview",
- content="AI content from source.",
- )
- ]
- state = RAGState(
- qa_history=[
- QAHistoryEntry(
- question="What is artificial intelligence?",
- answer="AI is the simulation of human intelligence by machines.",
- question_embedding=query_embedding,
- citations=prior_citations,
- ),
- ]
- )
- ctx = _make_ctx(state)
- await ask(ctx, question=query_text)
-
- # rag.ask() should receive augmented question with prior context
- assert len(captured_questions) == 1
- augmented = captured_questions[0]
- assert "Context from prior questions" in augmented
- assert "What is artificial intelligence?" in augmented
- assert "AI is the simulation" in augmented
- assert "AI Overview" in augmented
- assert query_text in augmented
-
- # State should store the original question, not the augmented one
- assert state.qa_history[-1].question == query_text
-
- async def test_ask_embeds_prior_qa_on_demand(self, rag_db, monkeypatch):
- from haiku.rag.skills.rag import RAGState, create_skill
- from tests.skills.conftest import VECTOR_DIM
-
- captured_questions = []
-
- async def mock_ask(self, question, **kwargs):
- captured_questions.append(question)
- return ("Answer about AI.", [])
-
- monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
-
- skill = create_skill(db_path=rag_db)
- ask = _get_tool(skill, "ask")
-
- # Use the same question text for the prior QA entry and query so
- # their fake embeddings are identical (cosine similarity = 1.0).
- prior_question = "Tell me about AI"
- query_text = prior_question
-
- # Leave question_embedding=None to exercise the lazy embedding path
- state = RAGState(
- qa_history=[
- QAHistoryEntry(
- question=prior_question,
- answer="AI is the simulation of human intelligence by machines.",
- question_embedding=None,
- ),
- ]
- )
- ctx = _make_ctx(state)
- await ask(ctx, question=query_text)
-
- # The lazy embedding should have populated question_embedding
- assert state.qa_history[0].question_embedding is not None
- assert len(state.qa_history[0].question_embedding) == VECTOR_DIM
-
- # The augmented question should include prior context
- assert len(captured_questions) == 1
- assert "Context from prior questions" in captured_questions[0]
- assert prior_question in captured_questions[0]
-
- async def test_ask_no_prior_qa_context_when_irrelevant(self, rag_db, monkeypatch):
- from haiku.rag.skills.rag import RAGState, create_skill
- from tests.skills.conftest import VECTOR_DIM
-
- captured_questions = []
-
- async def mock_ask(self, question, **kwargs):
- captured_questions.append(question)
- return ("Answer.", [])
-
- monkeypatch.setattr(HaikuRAG, "ask", mock_ask)
-
- skill = create_skill(db_path=rag_db)
- ask = _get_tool(skill, "ask")
-
- # Use orthogonal embedding — won't match the fake embedder's output
- 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)
- await ask(ctx, question="Explain quantum computing")
-
- # rag.ask() should receive the original question unchanged
- assert len(captured_questions) == 1
- assert captured_questions[0] == "Explain quantum computing"
-
-
-class TestResearchTool:
- async def test_research_returns_report(self, rag_db, monkeypatch):
+ async def test_cite_without_state(self, rag_db):
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_applies_document_filter_from_state(
- self, rag_db, monkeypatch
- ):
- from haiku.rag.skills.rag import RAGState, create_skill
-
- captured_kwargs = {}
-
- report = ResearchReport(
- title="AI Research",
- executive_summary="Summary.",
- main_findings=["Finding"],
- conclusions=["Conclusion"],
- sources_summary="Sources.",
- )
-
- async def mock_research(self, question, **kwargs):
- captured_kwargs.update(kwargs)
- return report
-
- monkeypatch.setattr(HaikuRAG, "research", mock_research)
-
- skill = create_skill(db_path=rag_db)
- research = _get_tool(skill, "research")
- state = RAGState(document_filter="title = 'AI Overview'")
- ctx = _make_ctx(state)
- await research(ctx, question="What is AI?")
- assert captured_kwargs.get("filter") == "title = 'AI Overview'"
-
- 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")
+ cite = _get_tool(skill, "cite")
ctx = _make_ctx(state=None)
- result = await research(ctx, question="What is AI?")
- assert isinstance(result, str)
+ result = await cite(ctx, chunk_ids=["nonexistent"])
+ assert "No state" in result