Allow multiple tool calling rounds in QA agent
This commit is contained in:
Yiorgis Gozadinos
parent
bea96f1b1b
commit
2f261a33f5
5 changed files with 151 additions and 160 deletions
|
|
@ -37,75 +37,69 @@ try:
|
|||
|
||||
messages: list[MessageParam] = [{"role": "user", "content": question}]
|
||||
|
||||
response = await anthropic_client.messages.create(
|
||||
model=self._model,
|
||||
max_tokens=4096,
|
||||
system=self._system_prompt,
|
||||
messages=messages,
|
||||
tools=self.tools,
|
||||
temperature=0.0,
|
||||
)
|
||||
max_rounds = 5 # Prevent infinite loops
|
||||
|
||||
if response.stop_reason == "tool_use":
|
||||
messages.append({"role": "assistant", "content": response.content})
|
||||
for _ in range(max_rounds):
|
||||
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
|
||||
tool_results = []
|
||||
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
|
||||
)
|
||||
if response.stop_reason == "tool_use":
|
||||
messages.append({"role": "assistant", "content": response.content})
|
||||
|
||||
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}"
|
||||
# Process tool calls
|
||||
tool_results = []
|
||||
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
|
||||
)
|
||||
|
||||
context = "\n\n".join(context_chunks)
|
||||
search_results = await self._client.search(
|
||||
query, limit=limit
|
||||
)
|
||||
|
||||
tool_results.append(
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": content_block.id,
|
||||
"content": context,
|
||||
}
|
||||
)
|
||||
context_chunks = []
|
||||
for chunk, score in search_results:
|
||||
context_chunks.append(
|
||||
f"Content: {chunk.content}\nScore: {score:.4f}"
|
||||
)
|
||||
|
||||
if tool_results:
|
||||
messages.append({"role": "user", "content": tool_results})
|
||||
context = "\n\n".join(context_chunks)
|
||||
|
||||
final_response = await anthropic_client.messages.create(
|
||||
model=self._model,
|
||||
max_tokens=4096,
|
||||
system=self._system_prompt,
|
||||
messages=messages,
|
||||
temperature=0.0,
|
||||
)
|
||||
if final_response.content:
|
||||
first_content = final_response.content[0]
|
||||
tool_results.append(
|
||||
{
|
||||
"type": "tool_result",
|
||||
"tool_use_id": content_block.id,
|
||||
"content": context,
|
||||
}
|
||||
)
|
||||
|
||||
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):
|
||||
return first_content.text
|
||||
return ""
|
||||
|
||||
if response.content:
|
||||
first_content = response.content[0]
|
||||
if isinstance(first_content, TextBlock):
|
||||
return first_content.text
|
||||
# If we've exhausted max rounds, return empty string
|
||||
return ""
|
||||
|
||||
except ImportError:
|
||||
|
|
|
|||
|
|
@ -14,54 +14,51 @@ class QuestionAnswerOllamaAgent(QuestionAnswerAgentBase):
|
|||
async def answer(self, question: str) -> str:
|
||||
ollama_client = AsyncClient(host=Config.OLLAMA_BASE_URL)
|
||||
|
||||
# Define the search tool
|
||||
|
||||
messages = [
|
||||
{"role": "system", "content": self._system_prompt},
|
||||
{"role": "user", "content": question},
|
||||
]
|
||||
|
||||
# Initial response with tool calling
|
||||
response = await ollama_client.chat(
|
||||
model=self._model,
|
||||
messages=messages,
|
||||
tools=self.tools,
|
||||
options=OLLAMA_OPTIONS,
|
||||
think=False,
|
||||
)
|
||||
max_rounds = 5 # Prevent infinite loops
|
||||
|
||||
if response.get("message", {}).get("tool_calls"):
|
||||
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(response["message"])
|
||||
messages.append(
|
||||
{
|
||||
"role": "tool",
|
||||
"content": context,
|
||||
"tool_call_id": tool_call.get("id", "search_tool"),
|
||||
}
|
||||
)
|
||||
|
||||
final_response = await ollama_client.chat(
|
||||
for _ in range(max_rounds):
|
||||
response = await ollama_client.chat(
|
||||
model=self._model,
|
||||
messages=messages,
|
||||
think=False,
|
||||
tools=self.tools,
|
||||
options=OLLAMA_OPTIONS,
|
||||
think=False,
|
||||
)
|
||||
return final_response["message"]["content"]
|
||||
else:
|
||||
return response["message"]["content"]
|
||||
|
||||
if response.get("message", {}).get("tool_calls"):
|
||||
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 ""
|
||||
|
|
|
|||
|
|
@ -24,8 +24,6 @@ try:
|
|||
async def answer(self, question: str) -> str:
|
||||
openai_client = AsyncOpenAI()
|
||||
|
||||
# Define the search tool
|
||||
|
||||
messages: list[ChatCompletionMessageParam] = [
|
||||
ChatCompletionSystemMessageParam(
|
||||
role="system", content=self._system_prompt
|
||||
|
|
@ -33,69 +31,70 @@ try:
|
|||
ChatCompletionUserMessageParam(role="user", content=question),
|
||||
]
|
||||
|
||||
# Initial response with tool calling
|
||||
response = await openai_client.chat.completions.create(
|
||||
model=self._model,
|
||||
messages=messages,
|
||||
tools=self.tools,
|
||||
temperature=0.0,
|
||||
)
|
||||
max_rounds = 5 # Prevent infinite loops
|
||||
|
||||
response_message = response.choices[0].message
|
||||
|
||||
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(
|
||||
for _ in range(max_rounds):
|
||||
response = await openai_client.chat.completions.create(
|
||||
model=self._model,
|
||||
messages=messages,
|
||||
tools=self.tools,
|
||||
temperature=0.0,
|
||||
)
|
||||
return final_response.choices[0].message.content or ""
|
||||
else:
|
||||
return response_message.content or ""
|
||||
|
||||
response_message = response.choices[0].message
|
||||
|
||||
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:
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -5,7 +5,8 @@ Your process:
|
|||
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
|
||||
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:
|
||||
- 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."
|
||||
- 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.
|
||||
"""
|
||||
|
|
|
|||
Loading…
Reference in a new issue