From d52f453c449f1305c869d033c3ec30350626b80d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 17 Apr 2026 18:31:11 +0300 Subject: [PATCH] flatten skill architecture: replace ask/analyze/research with direct tools --- app/frontend/components/Chat.tsx | 25 +- app/frontend/lib/sessionStorage.ts | 32 +- haiku_rag_slim/haiku/rag/chat/app.py | 25 +- .../haiku/rag/skill_generator/__init__.py | 5 +- .../rag/skill_generator/templates/SKILL.md.j2 | 62 +-- .../skill_generator/templates/__init__.py.j2 | 27 +- haiku_rag_slim/haiku/rag/skills/_tools.py | 356 ++++---------- haiku_rag_slim/haiku/rag/skills/analysis.py | 18 +- .../haiku/rag/skills/rag-analysis/SKILL.md | 67 ++- haiku_rag_slim/haiku/rag/skills/rag.py | 9 +- haiku_rag_slim/haiku/rag/skills/rag/SKILL.md | 58 ++- ...s_dict_for_document_with_docling_data.yaml | 42 -- tests/skills/test_analysis.py | 108 ++--- tests/skills/test_rag.py | 433 ++---------------- 14 files changed, 375 insertions(+), 892 deletions(-) delete mode 100644 tests/cassettes/test_sandbox/TestSandboxDoclingDocument.test_returns_dict_for_document_with_docling_data.yaml 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