tool call improvements

app.py

@@ -225,7 +225,7 @@ class BasicAgent:
             model="gemini-2.5-flash-preview-05-20",
             max_tokens=8192,
             temperature=0,
-            convert_system_message_to_human=True  # Enable system message conversion
+            convert_system_message_to_human=True
         )
         
         # Create tool executor
@@ -240,7 +240,43 @@ class BasicAgent:
             "google_search": google_search
         }
         
-        # System message
+        # Define tool configurations
+        self.openai_tools = [{
+            "type": "function",
+            "function": {
+                "name": "google_search",
+                "description": "Search for information on the web. Use this tool to find specific information about the question.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "query": {
+                            "type": "string",
+                            "description": "The search query to find relevant information"
+                        }
+                    },
+                    "required": ["query"]
+                }
+            }
+        }]
+        
+        self.gemini_tools = [{
+            "function_declarations": [{
+                "name": "google_search",
+                "description": "Search for information on the web. Use this tool to find specific information about the question.",
+                "parameters": {
+                    "type": "object",
+                    "properties": {
+                        "query": {
+                            "type": "string",
+                            "description": "The search query to find relevant information"
+                        }
+                    },
+                    "required": ["query"]
+                }
+            }]
+        }]
+        
+        # System message with explicit tool usage instructions
         self.sys_msg = SystemMessage('''You are a general AI assistant. I will ask you a question. Follow these steps:
 
         1. First, use the google_search tool to find relevant information about the question.
@@ -260,6 +296,7 @@ class BasicAgent:
 
         IMPORTANT:
         - You MUST use the google_search tool before providing your final answer
+        - Format your tool calls as: {"name": "google_search", "arguments": {"query": "your search query"}}
         - Your final answer should ONLY be the requested information, no explanations
         - If you need to search again, use the tool again
         - Do not provide detailed analysis in your final answer
@@ -308,26 +345,15 @@ class BasicAgent:
                         raise ValueError("Primary LLM not initialized")
                         
                     logger.info("Attempting to use primary LLM (OpenAI)")
-                    # For OpenAI, we can use the system message directly
+                    
+                    # Add explicit tool usage prompt
+                    messages_with_tool_prompt = [self.sys_msg] + messages + [
+                        HumanMessage(content="Use the google_search tool to find the information. Format your response as a tool call.")
+                    ]
+                    
                     response = self.primary_llm.invoke(
-                        [self.sys_msg] + messages,
-                        tools=[{
-                            "type": "function",
-                            "function": {
-                                "name": "google_search",
-                                "description": "Search for information on the web",
-                                "parameters": {
-                                    "type": "object",
-                                    "properties": {
-                                        "query": {
-                                            "type": "string",
-                                            "description": "The search query"
-                                        }
-                                    },
-                                    "required": ["query"]
-                                }
-                            }
-                        }]
+                        messages_with_tool_prompt,
+                        tools=self.openai_tools
                     )
                     
                     if not response or not hasattr(response, 'content'):
@@ -339,7 +365,7 @@ class BasicAgent:
                         return {"messages": [response], "next": "tools"}
                     else:
                         # If no tool call, try without tools
-                        response = self.primary_llm.invoke([self.sys_msg] + messages)
+                        response = self.primary_llm.invoke(messages_with_tool_prompt)
                         if not response or not hasattr(response, 'content'):
                             raise ValueError("Invalid response format from OpenAI")
                         logger.info("Successfully used primary LLM without tools")
@@ -358,27 +384,14 @@ class BasicAgent:
                         else:
                             messages_with_system = [self.sys_msg] + messages
                         
-                        # Create tool configuration for Gemini
-                        genai_tool = {
-                            "function_declarations": [{
-                                "name": "google_search",
-                                "description": "Search for information on the web",
-                                "parameters": {
-                                    "type": "object",
-                                    "properties": {
-                                        "query": {
-                                            "type": "string",
-                                            "description": "The search query"
-                                        }
-                                    },
-                                    "required": ["query"]
-                                }
-                            }]
-                        }
+                        # Add explicit tool usage prompt
+                        messages_with_tool_prompt = messages_with_system + [
+                            HumanMessage(content="Use the google_search tool to find the information. Format your response as a tool call.")
+                        ]
                         
                         response = self.fallback_llm.invoke(
-                            messages_with_system,
-                            tools=[genai_tool]
+                            messages_with_tool_prompt,
+                            tools=self.gemini_tools
                         )
                         
                         if not response or not hasattr(response, 'content'):
@@ -390,7 +403,7 @@ class BasicAgent:
                             return {"messages": [response], "next": "tools"}
                         else:
                             # If no tool call, try without tools
-                            response = self.fallback_llm.invoke(messages_with_system)
+                            response = self.fallback_llm.invoke(messages_with_tool_prompt)
                             if not response or not hasattr(response, 'content'):
                                 raise ValueError("Invalid response format from Gemini")
                             logger.info("Successfully used fallback LLM without tools")
@@ -422,10 +435,7 @@ class BasicAgent:
                 if any(phrase in content.lower() for phrase in ["let me", "i'll", "i will", "sure", "okay", "alright"]):
                     logger.info("Model provided acknowledgment instead of tool call, prompting for search")
                     return {
-                        "messages": messages + [
-                            AIMessage(content="Please use the google_search tool to find the information."),
-                            HumanMessage(content="Please search for the information using the google_search tool.")
-                        ],
+                        "messages": [AIMessage(content="Please use the google_search tool to find the information.")],
                         "next": "agent"
                     }
                 
@@ -489,7 +499,13 @@ class BasicAgent:
                             result = execute_tool(tool_name, tool_args, self.tools)
                             logger.info(f"Tool result: {result}")
                             
+                            # Add tool result to messages
                             messages.append(AIMessage(content=f"Tool result: {result}"))
+                            
+                            # If this was a google_search, analyze the results
+                            if tool_name == "google_search":
+                                # Add a prompt to analyze the search results
+                                messages.append(HumanMessage(content="Based on the search results, please provide your final answer."))
                         except Exception as e:
                             logger.error(f"Error executing tool {tool_name}: {str(e)}")
                             messages.append(AIMessage(content=f"Tool error: {str(e)}"))
@@ -499,7 +515,6 @@ class BasicAgent:
                     if any(phrase in content for phrase in ["let me", "i'll", "i will", "sure", "okay", "alright"]):
                         logger.info("No tool calls found, prompting for search")
                         messages.append(AIMessage(content="Please use the google_search tool to find the information."))
-                        messages.append(HumanMessage(content="Please search for the information using the google_search tool."))
                     else:
                         logger.info("No tool calls found in AI message")
                         # If the message looks like a final answer, return it