diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index 560c249c..0f3d1c14 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -467,6 +467,8 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: task: A specific, actionable instruction describing what to compute document_name: Optional document to focus on """ + from haiku.rag.agents.rlm import RLMContext, RLMDeps, create_rlm_agent + client = ctx.deps.client session_state = ctx.deps.session_state @@ -479,8 +481,30 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: # 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) + # Call RLM agent directly to access code executions + rlm_context = RLMContext(filter=filter_clause) + deps = RLMDeps( + client=client, + config=ctx.deps.config, + context=rlm_context, + ) + + rlm_agent = create_rlm_agent(ctx.deps.config) + result = await rlm_agent.run(task, deps=deps) + + # Format response with code executions + answer = result.output.answer + code_executions = rlm_context.code_executions + + if code_executions: + code_section = "\n\n---\n**Code executed:**\n" + for i, execution in enumerate(code_executions, 1): + code_section += f"\n```python\n# Execution {i}\n{execution.code}\n```\n" + if execution.stdout.strip(): + code_section += f"Output:\n```\n{execution.stdout.strip()}\n```\n" + if execution.stderr.strip(): + code_section += f"Errors:\n```\n{execution.stderr.strip()}\n```\n" + return answer + code_section return answer diff --git a/haiku_rag_slim/haiku/rag/agents/chat/prompts.py b/haiku_rag_slim/haiku/rag/agents/chat/prompts.py index 91ab46cb..6589690f 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/prompts.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/prompts.py @@ -27,6 +27,8 @@ CRITICAL - When using "analyze", reformulate the user's question into a specific - 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" +When "analyze" returns results, include both the answer AND the "Code executed" section in your response to the user. This shows transparency about how the computation was performed. + IMPORTANT - When user mentions a document in search/ask: - If user says "search in ", "find in ", "answer from ", or " in ": - Extract the TOPIC as `query`/`question` diff --git a/tests/cassettes/test_chat_agent/test_analyze_tool.yaml b/tests/cassettes/test_chat_agent/test_analyze_tool.yaml index fdda78f1..0bdc6fd4 100644 --- a/tests/cassettes/test_chat_agent/test_analyze_tool.yaml +++ b/tests/cassettes/test_chat_agent/test_analyze_tool.yaml @@ -157,7 +157,7 @@ interactions: connection: - keep-alive content-length: - - '7744' + - '7066' content-type: - application/json host: @@ -246,12 +246,10 @@ interactions: type: function - function: description: |- - Answer CONTENT questions by retrieving and synthesizing from documents. + Answer a specific question using the knowledge base. - Use this for questions about WHAT documents say - retrieval and synthesis. - Examples: "What does X say about Y?", "What are the main findings?", "Explain concept Z" - - Do NOT use for counting/aggregation questions like "How many documents mention X?" - use analyze instead. + Use this for direct questions that need a focused answer with citations. + Uses a research graph for planning, searching, and synthesis. name: ask parameters: additionalProperties: false @@ -263,7 +261,7 @@ interactions: default: null description: Optional document name/title to search within (e.g., "tbmed593", "army manual") question: - description: The content question to answer + description: The question to answer type: string required: - question @@ -322,23 +320,13 @@ interactions: description: |- Execute a computational task via code execution. - IMPORTANT: Do NOT pass the user's question directly. Instead, provide a - clear, specific task instruction that describes exactly what to compute. + 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: - - User asks "How many documents are there?" → - task="Count the total number of documents in the database using list_documents()" - - User asks "What's the average word count?" → - task="Calculate the average word count across all documents by getting each document's content and counting words" - - User asks "Which documents mention Python?" → - task="Search for 'Python' and return the titles of all matching documents" - - Use this for: - - Counting: task="Count documents matching criteria X" - - Aggregation: task="Sum/average values Y across documents" - - Extraction: task="Extract and list all Z from documents" - - Do NOT use for content questions - use ask instead. + - "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" name: analyze parameters: additionalProperties: false @@ -360,7 +348,7 @@ interactions: response: headers: content-length: - - '515' + - '539' content-type: - application/json parsed_body: @@ -369,24 +357,24 @@ interactions: index: 0 message: content: '' - reasoning: Need count. Use analyze. + reasoning: Need to count total documents. Use analyze tool. role: assistant tool_calls: - function: arguments: '{"task":"Count the total number of documents using list_documents()"}' name: analyze - id: call_6hr3gxx2 + id: call_w7qynecj index: 0 type: function - created: 1769782360 - id: chatcmpl-797 + created: 1769785117 + id: chatcmpl-187 model: gpt-oss object: chat.completion system_fingerprint: fp_ollama usage: - completion_tokens: 39 - prompt_tokens: 1508 - total_tokens: 1547 + completion_tokens: 43 + prompt_tokens: 1373 + total_tokens: 1416 status: code: 200 message: OK @@ -641,7 +629,7 @@ interactions: response: headers: content-length: - - '523' + - '527' content-type: - application/json parsed_body: @@ -650,24 +638,24 @@ interactions: index: 0 message: content: '' - reasoning: We need to run list_documents, count. + reasoning: Need to run list_documents and count. role: assistant tool_calls: - function: - arguments: '{"code":"docs=list_documents(limit=1000);print(len(docs))"}' + arguments: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))"}' name: execute_code - id: call_j3yll9cs + id: call_quotbvly index: 0 type: function - created: 1769782362 - id: chatcmpl-963 + created: 1769785119 + id: chatcmpl-208 model: gpt-oss object: chat.completion system_fingerprint: fp_ollama usage: - completion_tokens: 45 + completion_tokens: 46 prompt_tokens: 1759 - total_tokens: 1804 + total_tokens: 1805 status: code: 200 message: OK @@ -680,7 +668,7 @@ interactions: connection: - keep-alive content-length: - - '8772' + - '8780' content-type: - application/json host: @@ -849,18 +837,18 @@ interactions: role: user - content: |- - We need to run list_documents, count. + Need to run list_documents and count. role: assistant tool_calls: - function: - arguments: '{"code":"docs=list_documents(limit=1000);print(len(docs))"}' + arguments: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))"}' name: execute_code - id: call_j3yll9cs + id: call_quotbvly type: function - - content: '{"code":"docs=list_documents(limit=1000);print(len(docs))","stdout":"3\n","stderr":"","success":true}' + - content: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))","stdout":"3\n","stderr":"","success":true}' role: tool - tool_call_id: call_j3yll9cs + tool_call_id: call_quotbvly model: gpt-oss reasoning_effort: low stream: false @@ -946,15 +934,15 @@ interactions: message: content: '{"answer":"There are 3 documents in the knowledge base."}' role: assistant - created: 1769782363 - id: chatcmpl-319 + created: 1769785120 + id: chatcmpl-514 model: gpt-oss object: chat.completion system_fingerprint: fp_ollama usage: completion_tokens: 21 - prompt_tokens: 1849 - total_tokens: 1870 + prompt_tokens: 1852 + total_tokens: 1873 status: code: 200 message: OK @@ -967,7 +955,7 @@ interactions: connection: - keep-alive content-length: - - '8099' + - '7582' content-type: - application/json host: @@ -1022,18 +1010,32 @@ interactions: role: user - content: |- - Need count. Use analyze. + Need to count total documents. Use analyze tool. role: assistant tool_calls: - function: arguments: '{"task":"Count the total number of documents using list_documents()"}' name: analyze - id: call_6hr3gxx2 + id: call_w7qynecj type: function - - content: There are 3 documents in the knowledge base. + - content: | + There are 3 documents in the knowledge base. + + --- + **Code executed:** + + ```python + # Execution 1 + docs = list_documents(limit=1000) + print(len(docs)) + ``` + Output: + ``` + 3 + ``` role: tool - tool_call_id: call_6hr3gxx2 + tool_call_id: call_w7qynecj model: gpt-oss reasoning_effort: low stream: false @@ -1070,12 +1072,10 @@ interactions: type: function - function: description: |- - Answer CONTENT questions by retrieving and synthesizing from documents. + Answer a specific question using the knowledge base. - Use this for questions about WHAT documents say - retrieval and synthesis. - Examples: "What does X say about Y?", "What are the main findings?", "Explain concept Z" - - Do NOT use for counting/aggregation questions like "How many documents mention X?" - use analyze instead. + Use this for direct questions that need a focused answer with citations. + Uses a research graph for planning, searching, and synthesis. name: ask parameters: additionalProperties: false @@ -1087,7 +1087,7 @@ interactions: default: null description: Optional document name/title to search within (e.g., "tbmed593", "army manual") question: - description: The content question to answer + description: The question to answer type: string required: - question @@ -1146,23 +1146,13 @@ interactions: description: |- Execute a computational task via code execution. - IMPORTANT: Do NOT pass the user's question directly. Instead, provide a - clear, specific task instruction that describes exactly what to compute. + 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: - - User asks "How many documents are there?" → - task="Count the total number of documents in the database using list_documents()" - - User asks "What's the average word count?" → - task="Calculate the average word count across all documents by getting each document's content and counting words" - - User asks "Which documents mention Python?" → - task="Search for 'Python' and return the titles of all matching documents" - - Use this for: - - Counting: task="Count documents matching criteria X" - - Aggregation: task="Sum/average values Y across documents" - - Extraction: task="Extract and list all Z from documents" - - Do NOT use for content questions - use ask instead. + - "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" name: analyze parameters: additionalProperties: false @@ -1184,7 +1174,7 @@ interactions: response: headers: content-length: - - '341' + - '332' content-type: - application/json parsed_body: @@ -1192,17 +1182,17 @@ interactions: - finish_reason: stop index: 0 message: - content: You’ve got three documents in the database right now. + content: There are **three** documents in the database. role: assistant - created: 1769782365 - id: chatcmpl-556 + created: 1769785121 + id: chatcmpl-668 model: gpt-oss object: chat.completion system_fingerprint: fp_ollama usage: - completion_tokens: 15 - prompt_tokens: 1574 - total_tokens: 1589 + completion_tokens: 14 + prompt_tokens: 1481 + total_tokens: 1495 status: code: 200 message: OK