Basic integration of RLM with chat agent

docs/agents.md

@@ -62,13 +62,14 @@ Key features:
 
 ### Tools
 
-The chat agent uses five tools:
+The chat agent uses six tools:
 
 - `list_documents` — Browse available documents in the knowledge base
 - `summarize_document` — Generate a summary of a specific document
 - `get_document` — Retrieve a specific document by title or URI
 - `search` — Hybrid search with optional document filter
 - `ask` — Answer questions using the conversational research graph (automatically recalls prior answers)
+- `analyze` — Complex analytical questions via code execution (counting, aggregation, comparison)
 
 The `ask` tool automatically checks conversation history before running research. It uses embedding similarity (0.7 cosine threshold) to find semantically matching prior answers, which are passed to the research planner as context. When prior answers are sufficient, the planner can skip searching entirely.
 

haiku_rag_slim/haiku/rag/agents/chat/agent.py

@@ -447,4 +447,41 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
 
         return f"**Summary of {doc.title or doc.uri}:**\n\n{result.output}"
 
+    @agent.tool
+    async def analyze(
+        ctx: RunContext[ChatDeps],
+        task: str,
+        document_name: str | None = None,
+    ) -> str:
+        """Execute a computational task via code execution.
+
+        IMPORTANT: Provide a clear, specific task instruction that describes
+        exactly what to compute. Do NOT pass the user's question directly.
+
+        Examples of good task instructions:
+        - "Count the total number of documents using list_documents()"
+        - "Search for 'Python' and return the titles of all matching documents"
+        - "Calculate the average word count across all documents"
+
+        Args:
+            task: A specific, actionable instruction describing what to compute
+            document_name: Optional document to focus on
+        """
+        client = ctx.deps.client
+        session_state = ctx.deps.session_state
+
+        # Build session filter from document_filter
+        session_filter = build_multi_document_filter(session_state.document_filter)
+
+        # Build tool filter from document_name parameter
+        tool_filter = build_document_filter(document_name) if document_name else None
+
+        # Combine filters: session AND tool
+        filter_clause = combine_filters(session_filter, tool_filter)
+
+        # Call RLM agent with the task instruction
+        answer = await client.rlm(task, filter=filter_clause)
+
+        return answer
+
     return agent

haiku_rag_slim/haiku/rag/agents/chat/prompts.py

@@ -13,8 +13,19 @@ How to decide which tool to use:
 - "list_documents" - Use when the user wants to browse or see what documents are available (e.g., "what documents are available?", "show me the documents", "list available docs").
 - "summarize_document" - Use when the user wants an overview or summary of a specific document (e.g., "summarize document X", "what does Y cover?", "give me an overview of Z").
 - "get_document" - Use when the user wants the FULL content of a specific document (e.g., "get the paper about Y", "fetch 2412.00566", "show me the full document").
-- "ask" - Use for questions about topics in the knowledge base. It automatically finds relevant prior answers from conversation history and searches across documents to return answers with citations.
+- "ask" - Use for CONTENT questions: "What does X say about Y?", "What are the main findings?", "Explain concept Z from the documents". This tool retrieves and synthesizes text from documents.
 - "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
+- "analyze" - Use for COMPUTATIONAL tasks. IMPORTANT: Do NOT pass the user's question directly. Instead, write a specific task instruction describing what to compute.
+
+IMPORTANT - Choosing between "ask" and "analyze":
+- "ask" answers WHAT questions about content (retrieval + synthesis)
+- "analyze" answers HOW MANY/HOW MUCH questions requiring computation
+
+CRITICAL - When using "analyze", reformulate the user's question into a specific task:
+- User: "How many documents are there?" → task="Count the total number of documents using list_documents()"
+- User: "What is the total revenue across all reports?" → task="Search for revenue figures in all documents, extract the numeric values, and calculate the sum"
+- User: "How many documents discuss climate change?" → task="Search for 'climate change' and count the number of unique documents returned"
+- User: "List all the dates mentioned" → task="Search across documents, extract all date patterns, and return a deduplicated list"
 
 IMPORTANT - When user mentions a document in search/ask:
 - If user says "search in <doc>", "find in <doc>", "answer from <doc>", or "<topic> in <doc>":

haiku_rag_slim/haiku/rag/agents/rlm/prompts.py

@@ -140,9 +140,16 @@ print(sentiment)
 
 ## Output Format
 
-After executing code and gathering information, provide:
+CRITICAL: Your final response MUST be valid JSON matching this exact schema:
+```json
+{"answer": "Your complete answer here as a string"}
+```
+
+The `answer` field should contain:
 1. A clear answer to the user's question
 2. Key findings from your analysis
 3. References to specific documents/chunks that informed your answer
 
+Do NOT return arbitrary JSON structures. Always use the exact format: {"answer": "..."}
+
 CRITICAL: You MUST call execute_code at least once before providing your answer. Never give up without trying to execute code first."""

tests/agents/chat/test_chat_agent.py

@@ -1265,3 +1265,47 @@ async def test_summarization_task_cleanup_on_completion():
 
     # Task should be cleaned up
     assert session_id not in _summarization_tasks
+
+
+# =============================================================================
+# analyze Tool Tests
+# =============================================================================
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_analyze_tool(allow_model_requests, temp_db_path):
+    """Test the analyze tool for complex analytical questions."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        # Add test documents
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+        await client.create_document(
+            content=DOCLAYNET_ANNOTATION,
+            uri="doclaynet-annotation",
+            title="DocLayNet Annotation",
+        )
+        await client.create_document(
+            content=DOCLAYNET_DATA_SOURCES,
+            uri="doclaynet-sources",
+            title="DocLayNet Sources",
+        )
+
+        agent = create_chat_agent(Config)
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+        )
+
+        # Ask an analytical question that requires computation
+        result = await agent.run(
+            "How many documents are in the database?",
+            deps=deps,
+        )
+
+        assert result.output is not None
+        # The answer should mention 3 documents
+        assert "3" in result.output or "three" in result.output.lower()