danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Return program in chat agent analyze;

Browse source
This commit is contained in:
Yiorgis Gozadinos 2026-01-30 17:10:51 +02:00
parent ed570633cd
commit 9a480cee90
No known key found for this signature in database
3 changed files with 99 additions and 83 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

28
haiku_rag_slim/haiku/rag/agents/chat/agent.py
View file

@ -467,6 +467,8 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
task: A specific, actionable instruction describing what to compute task: A specific, actionable instruction describing what to compute
document_name: Optional document to focus on document_name: Optional document to focus on
""" """
from haiku.rag.agents.rlm import RLMContext, RLMDeps, create_rlm_agent
client = ctx.deps.client client = ctx.deps.client
session_state = ctx.deps.session_state session_state = ctx.deps.session_state
@ -479,8 +481,30 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
# Combine filters: session AND tool # Combine filters: session AND tool
filter_clause = combine_filters(session_filter, tool_filter) filter_clause = combine_filters(session_filter, tool_filter)
# Call RLM agent with the task instruction # Call RLM agent directly to access code executions
answer = await client.rlm(task, filter=filter_clause) 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 return answer

2
haiku_rag_slim/haiku/rag/agents/chat/prompts.py
View file

@ -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: "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" - 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: 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>": - If user says "search in <doc>", "find in <doc>", "answer from <doc>", or "<topic> in <doc>":
- Extract the TOPIC as `query`/`question` - Extract the TOPIC as `query`/`question`

152
tests/cassettes/test_chat_agent/test_analyze_tool.yaml
View file

@ -157,7 +157,7 @@ interactions:
connection: connection:
- keep-alive - keep-alive
content-length: content-length:
- '7744' - '7066'
content-type: content-type:
- application/json - application/json
host: host:
@ -246,12 +246,10 @@ interactions:
type: function type: function
- function: - function:
description: |- 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. Use this for direct questions that need a focused answer with citations.
Examples: "What does X say about Y?", "What are the main findings?", "Explain concept Z" Uses a research graph for planning, searching, and synthesis.
Do NOT use for counting/aggregation questions like "How many documents mention X?" - use analyze instead.
name: ask name: ask
parameters: parameters:
additionalProperties: false additionalProperties: false
@ -263,7 +261,7 @@ interactions:
default: null default: null
description: Optional document name/title to search within (e.g., "tbmed593", "army manual") description: Optional document name/title to search within (e.g., "tbmed593", "army manual")
question: question:
description: The content question to answer description: The question to answer
type: string type: string
required: required:
- question - question
@ -322,23 +320,13 @@ interactions:
description: |- description: |-
Execute a computational task via code execution. Execute a computational task via code execution.
IMPORTANT: Do NOT pass the user's question directly. Instead, provide a IMPORTANT: Provide a clear, specific task instruction that describes
clear, specific task instruction that describes exactly what to compute. exactly what to compute. Do NOT pass the user's question directly.
Examples of good task instructions: Examples of good task instructions:
- User asks "How many documents are there?" → - "Count the total number of documents using list_documents()"
task="Count the total number of documents in the database using list_documents()" - "Search for 'Python' and return the titles of all matching documents"
- User asks "What's the average word count?" → - "Calculate the average word count across all documents"
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.
name: analyze name: analyze
parameters: parameters:
additionalProperties: false additionalProperties: false
@ -360,7 +348,7 @@ interactions:
response: response:
headers: headers:
content-length: content-length:
- '515' - '539'
content-type: content-type:
- application/json - application/json
parsed_body: parsed_body:
@ -369,24 +357,24 @@ interactions:
index: 0 index: 0
message: message:
content: '' content: ''
reasoning: Need count. Use analyze. reasoning: Need to count total documents. Use analyze tool.
role: assistant role: assistant
tool_calls: tool_calls:
- function: - function:
arguments: '{"task":"Count the total number of documents using list_documents()"}' arguments: '{"task":"Count the total number of documents using list_documents()"}'
name: analyze name: analyze
id: call_6hr3gxx2 id: call_w7qynecj
index: 0 index: 0
type: function type: function
created: 1769782360 created: 1769785117
id: chatcmpl-797 id: chatcmpl-187
model: gpt-oss model: gpt-oss
object: chat.completion object: chat.completion
system_fingerprint: fp_ollama system_fingerprint: fp_ollama
usage: usage:
completion_tokens: 39 completion_tokens: 43
prompt_tokens: 1508 prompt_tokens: 1373
total_tokens: 1547 total_tokens: 1416
status: status:
code: 200 code: 200
message: OK message: OK
@ -641,7 +629,7 @@ interactions:
response: response:
headers: headers:
content-length: content-length:
- '523' - '527'
content-type: content-type:
- application/json - application/json
parsed_body: parsed_body:
@ -650,24 +638,24 @@ interactions:
index: 0 index: 0
message: message:
content: '' content: ''
reasoning: We need to run list_documents, count. reasoning: Need to run list_documents and count.
role: assistant role: assistant
tool_calls: tool_calls:
- function: - function:
arguments: '{"code":"docs=list_documents(limit=1000);print(len(docs))"}' arguments: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))"}'
name: execute_code name: execute_code
id: call_j3yll9cs id: call_quotbvly
index: 0 index: 0
type: function type: function
created: 1769782362 created: 1769785119
id: chatcmpl-963 id: chatcmpl-208
model: gpt-oss model: gpt-oss
object: chat.completion object: chat.completion
system_fingerprint: fp_ollama system_fingerprint: fp_ollama
usage: usage:
completion_tokens: 45 completion_tokens: 46
prompt_tokens: 1759 prompt_tokens: 1759
total_tokens: 1804 total_tokens: 1805
status: status:
code: 200 code: 200
message: OK message: OK
@ -680,7 +668,7 @@ interactions:
connection: connection:
- keep-alive - keep-alive
content-length: content-length:
- '8772' - '8780'
content-type: content-type:
- application/json - application/json
host: host:
@ -849,18 +837,18 @@ interactions:
role: user role: user
- content: |- - content: |-
<think> <think>
We need to run list_documents, count. Need to run list_documents and count.
</think> </think>
role: assistant role: assistant
tool_calls: tool_calls:
- function: - function:
arguments: '{"code":"docs=list_documents(limit=1000);print(len(docs))"}' arguments: '{"code":"docs = list_documents(limit=1000)\nprint(len(docs))"}'
name: execute_code name: execute_code
id: call_j3yll9cs id: call_quotbvly
type: function 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 role: tool
tool_call_id: call_j3yll9cs tool_call_id: call_quotbvly
model: gpt-oss model: gpt-oss
reasoning_effort: low reasoning_effort: low
stream: false stream: false
@ -946,15 +934,15 @@ interactions:
message: message:
content: '{"answer":"There are 3 documents in the knowledge base."}' content: '{"answer":"There are 3 documents in the knowledge base."}'
role: assistant role: assistant
created: 1769782363 created: 1769785120
id: chatcmpl-319 id: chatcmpl-514
model: gpt-oss model: gpt-oss
object: chat.completion object: chat.completion
system_fingerprint: fp_ollama system_fingerprint: fp_ollama
usage: usage:
completion_tokens: 21 completion_tokens: 21
prompt_tokens: 1849 prompt_tokens: 1852
total_tokens: 1870 total_tokens: 1873
status: status:
code: 200 code: 200
message: OK message: OK
@ -967,7 +955,7 @@ interactions:
connection: connection:
- keep-alive - keep-alive
content-length: content-length:
- '8099' - '7582'
content-type: content-type:
- application/json - application/json
host: host:
@ -1022,18 +1010,32 @@ interactions:
role: user role: user
- content: |- - content: |-
<think> <think>
Need count. Use analyze. Need to count total documents. Use analyze tool.
</think> </think>
role: assistant role: assistant
tool_calls: tool_calls:
- function: - function:
arguments: '{"task":"Count the total number of documents using list_documents()"}' arguments: '{"task":"Count the total number of documents using list_documents()"}'
name: analyze name: analyze
id: call_6hr3gxx2 id: call_w7qynecj
type: function 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 role: tool
tool_call_id: call_6hr3gxx2 tool_call_id: call_w7qynecj
model: gpt-oss model: gpt-oss
reasoning_effort: low reasoning_effort: low
stream: false stream: false
@ -1070,12 +1072,10 @@ interactions:
type: function type: function
- function: - function:
description: |- 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. Use this for direct questions that need a focused answer with citations.
Examples: "What does X say about Y?", "What are the main findings?", "Explain concept Z" Uses a research graph for planning, searching, and synthesis.
Do NOT use for counting/aggregation questions like "How many documents mention X?" - use analyze instead.
name: ask name: ask
parameters: parameters:
additionalProperties: false additionalProperties: false
@ -1087,7 +1087,7 @@ interactions:
default: null default: null
description: Optional document name/title to search within (e.g., "tbmed593", "army manual") description: Optional document name/title to search within (e.g., "tbmed593", "army manual")
question: question:
description: The content question to answer description: The question to answer
type: string type: string
required: required:
- question - question
@ -1146,23 +1146,13 @@ interactions:
description: |- description: |-
Execute a computational task via code execution. Execute a computational task via code execution.
IMPORTANT: Do NOT pass the user's question directly. Instead, provide a IMPORTANT: Provide a clear, specific task instruction that describes
clear, specific task instruction that describes exactly what to compute. exactly what to compute. Do NOT pass the user's question directly.
Examples of good task instructions: Examples of good task instructions:
- User asks "How many documents are there?" → - "Count the total number of documents using list_documents()"
task="Count the total number of documents in the database using list_documents()" - "Search for 'Python' and return the titles of all matching documents"
- User asks "What's the average word count?" → - "Calculate the average word count across all documents"
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.
name: analyze name: analyze
parameters: parameters:
additionalProperties: false additionalProperties: false
@ -1184,7 +1174,7 @@ interactions:
response: response:
headers: headers:
content-length: content-length:
- '341' - '332'
content-type: content-type:
- application/json - application/json
parsed_body: parsed_body:
@ -1192,17 +1182,17 @@ interactions:
- finish_reason: stop - finish_reason: stop
index: 0 index: 0
message: message:
content: You’ve got three documents in the database right now. content: There are **three** documents in the database.
role: assistant role: assistant
created: 1769782365 created: 1769785121
id: chatcmpl-556 id: chatcmpl-668
model: gpt-oss model: gpt-oss
object: chat.completion object: chat.completion
system_fingerprint: fp_ollama system_fingerprint: fp_ollama
usage: usage:
completion_tokens: 15 completion_tokens: 14
prompt_tokens: 1574 prompt_tokens: 1481
total_tokens: 1589 total_tokens: 1495
status: status:
code: 200 code: 200
message: OK message: OK