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

Allow multiple tool calling rounds in QA agent

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-07-08 12:51:58 +03:00
parent bea96f1b1b
commit 2f261a33f5
No known key found for this signature in database
5 changed files with 151 additions and 160 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

0
BENCHMARKS.md → docs/benchmarks.md
View file

108
src/haiku/rag/qa/anthropic.py
View file

@ -37,75 +37,69 @@ try:
messages: list[MessageParam] = [{"role": "user", "content": question}] messages: list[MessageParam] = [{"role": "user", "content": question}]
response = await anthropic_client.messages.create( max_rounds = 5 # Prevent infinite loops
model=self._model,
max_tokens=4096,
system=self._system_prompt,
messages=messages,
tools=self.tools,
temperature=0.0,
)
if response.stop_reason == "tool_use": for _ in range(max_rounds):
messages.append({"role": "assistant", "content": response.content}) response = await anthropic_client.messages.create(
model=self._model,
max_tokens=4096,
system=self._system_prompt,
messages=messages,
tools=self.tools,
temperature=0.0,
)
# Process tool calls if response.stop_reason == "tool_use":
tool_results = [] messages.append({"role": "assistant", "content": response.content})
for content_block in response.content:
if isinstance(content_block, ToolUseBlock):
if content_block.name == "search_documents":
args = content_block.input
query = (
args.get("query", question)
if isinstance(args, dict)
else question
)
limit = (
int(args.get("limit", 3))
if isinstance(args, dict)
else 3
)
search_results = await self._client.search( # Process tool calls
query, limit=limit tool_results = []
) for content_block in response.content:
if isinstance(content_block, ToolUseBlock):
context_chunks = [] if content_block.name == "search_documents":
for chunk, score in search_results: args = content_block.input
context_chunks.append( query = (
f"Content: {chunk.content}\nScore: {score:.4f}" args.get("query", question)
if isinstance(args, dict)
else question
)
limit = (
int(args.get("limit", 3))
if isinstance(args, dict)
else 3
) )
context = "\n\n".join(context_chunks) search_results = await self._client.search(
query, limit=limit
)
tool_results.append( context_chunks = []
{ for chunk, score in search_results:
"type": "tool_result", context_chunks.append(
"tool_use_id": content_block.id, f"Content: {chunk.content}\nScore: {score:.4f}"
"content": context, )
}
)
if tool_results: context = "\n\n".join(context_chunks)
messages.append({"role": "user", "content": tool_results})
final_response = await anthropic_client.messages.create( tool_results.append(
model=self._model, {
max_tokens=4096, "type": "tool_result",
system=self._system_prompt, "tool_use_id": content_block.id,
messages=messages, "content": context,
temperature=0.0, }
) )
if final_response.content:
first_content = final_response.content[0] if tool_results:
messages.append({"role": "user", "content": tool_results})
else:
# No tool use, return the response
if response.content:
first_content = response.content[0]
if isinstance(first_content, TextBlock): if isinstance(first_content, TextBlock):
return first_content.text return first_content.text
return "" return ""
if response.content: # If we've exhausted max rounds, return empty string
first_content = response.content[0]
if isinstance(first_content, TextBlock):
return first_content.text
return "" return ""
except ImportError: except ImportError:

79
src/haiku/rag/qa/ollama.py
View file

@ -14,54 +14,51 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
async def answer(self, question: str) -> str: async def answer(self, question: str) -> str:
ollama_client = AsyncClient(host=Config.OLLAMA_BASE_URL) ollama_client = AsyncClient(host=Config.OLLAMA_BASE_URL)
# Define the search tool
messages = [ messages = [
{"role": "system", "content": self._system_prompt}, {"role": "system", "content": self._system_prompt},
{"role": "user", "content": question}, {"role": "user", "content": question},
] ]
# Initial response with tool calling max_rounds = 5 # Prevent infinite loops
response = await ollama_client.chat(
model=self._model,
messages=messages,
tools=self.tools,
options=OLLAMA_OPTIONS,
think=False,
)
if response.get("message", {}).get("tool_calls"): for _ in range(max_rounds):
for tool_call in response["message"]["tool_calls"]: response = await ollama_client.chat(
if tool_call["function"]["name"] == "search_documents":
args = tool_call["function"]["arguments"]
query = args.get("query", question)
limit = int(args.get("limit", 3))
search_results = await self._client.search(query, limit=limit)
context_chunks = []
for chunk, score in search_results:
context_chunks.append(
f"Content: {chunk.content}\nScore: {score:.4f}"
)
context = "\n\n".join(context_chunks)
messages.append(response["message"])
messages.append(
{
"role": "tool",
"content": context,
"tool_call_id": tool_call.get("id", "search_tool"),
}
)
final_response = await ollama_client.chat(
model=self._model, model=self._model,
messages=messages, messages=messages,
think=False, tools=self.tools,
options=OLLAMA_OPTIONS, options=OLLAMA_OPTIONS,
think=False,
) )
return final_response["message"]["content"]
else: if response.get("message", {}).get("tool_calls"):
return response["message"]["content"] messages.append(response["message"])
for tool_call in response["message"]["tool_calls"]:
if tool_call["function"]["name"] == "search_documents":
args = tool_call["function"]["arguments"]
query = args.get("query", question)
limit = int(args.get("limit", 3))
search_results = await self._client.search(query, limit=limit)
context_chunks = []
for chunk, score in search_results:
context_chunks.append(
f"Content: {chunk.content}\nScore: {score:.4f}"
)
context = "\n\n".join(context_chunks)
messages.append(
{
"role": "tool",
"content": context,
"tool_call_id": tool_call.get("id", "search_tool"),
}
)
else:
# No tool calls, return the response
return response["message"]["content"]
# If we've exhausted max rounds, return empty string
return ""

119
src/haiku/rag/qa/openai.py
View file

@ -24,8 +24,6 @@ try:
async def answer(self, question: str) -> str: async def answer(self, question: str) -> str:
openai_client = AsyncOpenAI() openai_client = AsyncOpenAI()
# Define the search tool
messages: list[ChatCompletionMessageParam] = [ messages: list[ChatCompletionMessageParam] = [
ChatCompletionSystemMessageParam( ChatCompletionSystemMessageParam(
role="system", content=self._system_prompt role="system", content=self._system_prompt
@ -33,69 +31,70 @@ try:
ChatCompletionUserMessageParam(role="user", content=question), ChatCompletionUserMessageParam(role="user", content=question),
] ]
# Initial response with tool calling max_rounds = 5 # Prevent infinite loops
response = await openai_client.chat.completions.create(
model=self._model,
messages=messages,
tools=self.tools,
temperature=0.0,
)
response_message = response.choices[0].message for _ in range(max_rounds):
response = await openai_client.chat.completions.create(
if response_message.tool_calls:
messages.append(
ChatCompletionAssistantMessageParam(
role="assistant",
content=response_message.content,
tool_calls=[
{
"id": tc.id,
"type": "function",
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments,
},
}
for tc in response_message.tool_calls
],
)
)
for tool_call in response_message.tool_calls:
if tool_call.function.name == "search_documents":
import json
args = json.loads(tool_call.function.arguments)
query = args.get("query", question)
limit = int(args.get("limit", 3))
search_results = await self._client.search(query, limit=limit)
context_chunks = []
for chunk, score in search_results:
context_chunks.append(
f"Content: {chunk.content}\nScore: {score:.4f}"
)
context = "\n\n".join(context_chunks)
messages.append(
ChatCompletionToolMessageParam(
role="tool",
content=context,
tool_call_id=tool_call.id,
)
)
final_response = await openai_client.chat.completions.create(
model=self._model, model=self._model,
messages=messages, messages=messages,
tools=self.tools,
temperature=0.0, temperature=0.0,
) )
return final_response.choices[0].message.content or ""
else: response_message = response.choices[0].message
return response_message.content or ""
if response_message.tool_calls:
messages.append(
ChatCompletionAssistantMessageParam(
role="assistant",
content=response_message.content,
tool_calls=[
{
"id": tc.id,
"type": "function",
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments,
},
}
for tc in response_message.tool_calls
],
)
)
for tool_call in response_message.tool_calls:
if tool_call.function.name == "search_documents":
import json
args = json.loads(tool_call.function.arguments)
query = args.get("query", question)
limit = int(args.get("limit", 3))
search_results = await self._client.search(
query, limit=limit
)
context_chunks = []
for chunk, score in search_results:
context_chunks.append(
f"Content: {chunk.content}\nScore: {score:.4f}"
)
context = "\n\n".join(context_chunks)
messages.append(
ChatCompletionToolMessageParam(
role="tool",
content=context,
tool_call_id=tool_call.id,
)
)
else:
# No tool calls, return the response
return response_message.content or ""
# If we've exhausted max rounds, return empty string
return ""
except ImportError: except ImportError:
pass pass

5
src/haiku/rag/qa/prompts.py
View file

@ -5,7 +5,8 @@ Your process:
1. When a user asks a question, use the search_documents tool to find relevant information 1. When a user asks a question, use the search_documents tool to find relevant information
2. Search with specific keywords and phrases from the user's question 2. Search with specific keywords and phrases from the user's question
3. Review the search results and their relevance scores 3. Review the search results and their relevance scores
4. Provide a comprehensive answer based only on the retrieved documents 4. If you need additional context, perform follow-up searches with different keywords
5. Provide a comprehensive answer based only on the retrieved documents
Guidelines: Guidelines:
- Base your answers strictly on the provided document content - Base your answers strictly on the provided document content
@ -15,5 +16,5 @@ Guidelines:
- If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question." - If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question."
- For complex questions, consider breaking them down and performing multiple searches - For complex questions, consider breaking them down and performing multiple searches
Be thorough but concise, and always maintain accuracy over completeness. Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
""" """