Upload app.py

app.py

@@ -3,12 +3,10 @@ import gradio as gr
 from typing import List
 import logging
 import logging.handlers
-import time
-import random
+import json
 from langchain_openai import ChatOpenAI
 from langchain_core.tools import tool
-from langgraph.prebuilt import create_react_agent
-from langchain_core.messages import HumanMessage
+from langchain_core.messages import HumanMessage, AIMessage, ToolMessage
 from langchain_tavily import TavilySearch
 
 # Configuration - set to False to disable detailed logging
@@ -58,27 +56,27 @@ if ENABLE_DETAILED_LOGGING:
     else:
         logger.warning("No Tavily API key found in environment variables")
 
-# Tavily search tool integration
-
-class ReactAgentChat:
+# Tool calling agent implementation
+class ToolCallingAgentChat:
     def __init__(self, ip: str, port: str, api_key: str, model: str):
         self.ip = ip
         self.port = port
         self.api_key = api_key
         self.model = model
-        self.agent = None
+        self.llm = None
+        self.tools = []
         self._setup_agent()
     
     def _setup_agent(self):
-        """Initialize the LangGraph ReAct agent"""
+        """Initialize the tool calling agent"""
         try:
             if ENABLE_DETAILED_LOGGING:
-                logger.info(f"=== SETTING UP AGENT ===")
+                logger.info(f"=== SETTING UP TOOL CALLING AGENT ===")
                 logger.info(f"LLM URL: http://{self.ip}:{self.port}/v1")
                 logger.info(f"Model: {self.model}")
             
             # Create OpenAI-compatible model
-            llm = ChatOpenAI(
+            self.llm = ChatOpenAI(
                 base_url=f"http://{self.ip}:{self.port}/v1",
                 api_key=self.api_key,
                 model=self.model,
@@ -87,15 +85,14 @@ class ReactAgentChat:
             if ENABLE_DETAILED_LOGGING:
                 logger.info("LLM created successfully")
             
-            # Define tools - use Tavily search API with graceful error handling
+            # Define web search tool
             if tavily_key:
                 if ENABLE_DETAILED_LOGGING:
                     logger.info("Setting up Tavily search tool")
                 try:
-                    # Create custom wrapper for Tavily with error handling
                     @tool
                     def web_search(query: str) -> str:
-                        """Search the web for current information about any topic."""
+                        """Search the web for current information about any topic. Use this when you need up-to-date information, current events, or real-time data."""
                         try:
                             tavily_tool = TavilySearch(
                                 tavily_api_key=tavily_key,
@@ -113,8 +110,6 @@ class ReactAgentChat:
                             if ENABLE_DETAILED_LOGGING:
                                 logger.error(f"Tavily search failed for query '{query}': {e}")
                                 logger.error(f"Exception type: {type(e).__name__}")
-                                import traceback
-                                logger.error(f"Full traceback: {traceback.format_exc()}")
                             
                             # Check for rate limit or quota issues
                             if any(keyword in error_str for keyword in ['rate limit', 'quota', 'limit exceeded', 'usage limit', 'billing']):
@@ -126,42 +121,30 @@ class ReactAgentChat:
                                     logger.error(f"Tavily API error: {e}")
                                 return f"I can't search the web right now. Error: {str(e)[:100]}"
                     
-                    search_tool = web_search
+                    self.tools = [web_search]
                     if ENABLE_DETAILED_LOGGING:
-                        logger.info("Tavily search tool wrapper created successfully")
+                        logger.info("Tavily search tool created successfully")
                 except Exception as e:
                     if ENABLE_DETAILED_LOGGING:
-                        logger.error(f"Failed to create Tavily tool wrapper: {e}")
-                    # Fallback tool
-                    @tool
-                    def no_search(query: str) -> str:
-                        """Search tool unavailable."""
-                        return "I can't search the web right now."
-                    search_tool = no_search
+                        logger.error(f"Failed to create Tavily tool: {e}")
+                    self.tools = []
             else:
                 if ENABLE_DETAILED_LOGGING:
-                    logger.warning("No Tavily API key found, creating fallback tool")
-                @tool
-                def no_search(query: str) -> str:
-                    """Search tool unavailable."""
-                    if ENABLE_DETAILED_LOGGING:
-                        logger.error("Search attempted but no Tavily API key configured")
-                    return "I can't search the web right now."
-                search_tool = no_search
-            
-            tools = [search_tool]
-            if ENABLE_DETAILED_LOGGING:
-                logger.info(f"Tools defined: {[tool.name for tool in tools]}")
+                    logger.warning("No Tavily API key found, no web search tool available")
+                self.tools = []
             
             # Bind tools to the model
-            model_with_tools = llm.bind_tools(tools)
-            if ENABLE_DETAILED_LOGGING:
-                logger.info("Tools bound to model")
+            if self.tools:
+                self.llm_with_tools = self.llm.bind_tools(self.tools)
+                if ENABLE_DETAILED_LOGGING:
+                    logger.info(f"Tools bound to model: {[tool.name for tool in self.tools]}")
+            else:
+                self.llm_with_tools = self.llm
+                if ENABLE_DETAILED_LOGGING:
+                    logger.info("No tools available, using base model")
             
-            # Create the ReAct agent
-            self.agent = create_react_agent(model_with_tools, tools)
             if ENABLE_DETAILED_LOGGING:
-                logger.info("ReAct agent created successfully")
+                logger.info("Tool calling agent created successfully")
             
         except Exception as e:
             logger.error(f"=== AGENT SETUP ERROR ===")
@@ -181,9 +164,9 @@ class ReactAgentChat:
             self._setup_agent()
     
     def chat(self, message: str, history: List[List[str]]) -> str:
-        """Generate chat response using ReAct agent"""
+        """Generate chat response using tool calling"""
         try:
-            if not self.agent:
+            if not self.llm_with_tools:
                 return "Error: Agent not initialized"
             
             if ENABLE_DETAILED_LOGGING:
@@ -191,34 +174,75 @@ class ReactAgentChat:
                 logger.info(f"Message: {message}")
                 logger.info(f"History length: {len(history)}")
             
-            # Convert history to messages for context handling
+            # Convert history to messages for context
             messages = []
             for user_msg, assistant_msg in history:
                 messages.append(HumanMessage(content=user_msg))
                 if assistant_msg:  # Only add if assistant responded
-                    from langchain_core.messages import AIMessage
                     messages.append(AIMessage(content=assistant_msg))
             
             # Add current message
             messages.append(HumanMessage(content=message))
             
-            # Invoke the agent
+            # Get initial response from LLM
             if ENABLE_DETAILED_LOGGING:
-                logger.info(f"=== INVOKING AGENT ===")
-                logger.info(f"Total messages in history: {len(messages)}")
-            response = self.agent.invoke({"messages": messages})
+                logger.info(f"=== INVOKING LLM ===")
+                logger.info(f"Total messages in context: {len(messages)}")
+            
+            response = self.llm_with_tools.invoke(messages)
             
             if ENABLE_DETAILED_LOGGING:
-                logger.info(f"=== AGENT RESPONSE ===")
-                logger.info(f"Full response: {response}")
-                logger.info(f"Number of messages: {len(response.get('messages', []))}")
+                logger.info(f"=== LLM RESPONSE ===")
+                logger.info(f"Response type: {type(response)}")
+                logger.info(f"Has tool calls: {bool(response.tool_calls if hasattr(response, 'tool_calls') else False)}")
+            
+            # Check if LLM wants to call tools
+            if hasattr(response, 'tool_calls') and response.tool_calls:
+                if ENABLE_DETAILED_LOGGING:
+                    logger.info(f"=== TOOL CALLS DETECTED ===")
+                    logger.info(f"Number of tool calls: {len(response.tool_calls)}")
                 
-                # Log each message in the response
-                for i, msg in enumerate(response.get("messages", [])):
-                    logger.info(f"Message {i}: Type={type(msg).__name__}, Content={getattr(msg, 'content', 'No content')}")
+                # Add the LLM response to messages
+                messages.append(response)
+                
+                # Execute tool calls
+                for tool_call in response.tool_calls:
+                    if ENABLE_DETAILED_LOGGING:
+                        logger.info(f"Executing tool: {tool_call['name']} with args: {tool_call['args']}")
+                    
+                    # Find and execute the tool
+                    tool_result = None
+                    for tool in self.tools:
+                        if tool.name == tool_call['name']:
+                            try:
+                                tool_result = tool.invoke(tool_call['args'])
+                                if ENABLE_DETAILED_LOGGING:
+                                    logger.info(f"Tool executed successfully: {tool_call['name']}")
+                                break
+                            except Exception as e:
+                                tool_result = f"Tool execution failed: {str(e)}"
+                                if ENABLE_DETAILED_LOGGING:
+                                    logger.error(f"Tool execution failed: {e}")
+                    
+                    if tool_result is None:
+                        tool_result = f"Tool {tool_call['name']} not found"
+                    
+                    # Add tool result to messages
+                    messages.append(ToolMessage(
+                        content=str(tool_result),
+                        tool_call_id=tool_call['id']
+                    ))
+                
+                # Get final response from LLM after tool execution
+                if ENABLE_DETAILED_LOGGING:
+                    logger.info(f"=== GETTING FINAL RESPONSE ===")
+                
+                final_response = self.llm_with_tools.invoke(messages)
+                final_message = final_response.content
+            else:
+                # No tool calls, use the direct response
+                final_message = response.content
             
-            # Extract the final response
-            final_message = response["messages"][-1].content
             if ENABLE_DETAILED_LOGGING:
                 logger.info(f"=== FINAL MESSAGE ===")
                 logger.info(f"Final message: {final_message}")
@@ -235,19 +259,19 @@ class ReactAgentChat:
             return error_msg
 
 # Global agent instance
-react_agent = ReactAgentChat(llm_ip, llm_port, llm_key, llm_model)
+tool_calling_agent = ToolCallingAgentChat(llm_ip, llm_port, llm_key, llm_model)
 
 def generate_response(message: str, history: List[List[str]], system_prompt: str, 
                      max_tokens: int, ip: str, port: str, api_key: str, model: str):
-    """Generate response using ReAct agent"""
-    global react_agent
+    """Generate response using tool calling agent"""
+    global tool_calling_agent
     
     try:
         # Update agent configuration if changed
-        react_agent.update_config(ip, port, api_key, model)
+        tool_calling_agent.update_config(ip, port, api_key, model)
         
         # Generate response
-        response = react_agent.chat(message, history)
+        response = tool_calling_agent.chat(message, history)
         
         # Stream the response word by word for better UX
         words = response.split()
@@ -273,7 +297,7 @@ chatbot = gr.ChatInterface(
     ),
     additional_inputs=[
         gr.Textbox(
-            "You are a helpful AI assistant with web search capabilities.", 
+            "You are a helpful AI assistant with web search capabilities. Use web search when you need current information, recent events, or real-time data.", 
             label="System Prompt",
             lines=2
         ),
@@ -288,8 +312,8 @@ chatbot = gr.ChatInterface(
         gr.Textbox(llm_model, label="Model Name", 
                   info="Name of the model to use"),
     ],
-    title="🤖 LangGraph ReAct Agent with Tavily Search",
-    description="Chat with a LangGraph ReAct agent that can search the web using Tavily. Ask about current events, research topics, or any questions that require up-to-date information!",
+    title="🚀 Fast Tool Calling Agent with Tavily Search",
+    description="Chat with a fast tool calling agent that can search the web using Tavily. The agent automatically decides when to search based on your query - much faster than ReAct agents!",
     theme="finlaymacklon/smooth_slate"
 )