Update agent.py

agent.py

@@ -3,7 +3,6 @@ import re
 from datetime import datetime, timedelta
 from typing import TypedDict, Annotated
 import sympy as sp
-from sympy import *
 import math
 from langchain_openai import ChatOpenAI
 from langchain_community.tools.tavily_search import TavilySearchResults
@@ -137,9 +136,12 @@ class GAIAAgent:
         if not openai_key:
             raise ValueError("OPENAI_API_KEY environment variable is required")
         if not tavily_key:
-            raise ValueError("TAVILY_API_KEY environment variable is required")
+            print("⚠️ TAVILY_API_KEY not found - web search will be disabled")
+            self.has_search = False
+        else:
+            self.has_search = True
         
-        print("✅ API keys found - initializing agent...")
+        print("✅ Initializing GAIA agent...")
         
         # Initialize LLM (using OpenAI GPT-4)
         self.llm = ChatOpenAI(
@@ -148,17 +150,20 @@ class GAIAAgent:
             openai_api_key=openai_key
         )
         
-        # Initialize tools
-        self.search_tool = TavilySearchResults(
-            max_results=5,
-            tavily_api_key=tavily_key
-        )
-        
-        # Create tool list
-        self.tools = [self.search_tool]
+        # Initialize tools only if we have Tavily key
+        self.tools = []
+        if self.has_search:
+            self.search_tool = TavilySearchResults(
+                max_results=5,
+                tavily_api_key=tavily_key
+            )
+            self.tools = [self.search_tool]
         
-        # Create LLM with tools
-        self.llm_with_tools = self.llm.bind_tools(self.tools)
+        # Create LLM with tools (only if we have tools)
+        if self.tools:
+            self.llm_with_tools = self.llm.bind_tools(self.tools)
+        else:
+            self.llm_with_tools = self.llm
         
         # Build the graph
         self.graph = self._build_graph()
@@ -172,57 +177,93 @@ class GAIAAgent:
             """Main agent reasoning node"""
             messages = state["messages"]
             
-            # Add system message if not present
+            # Add system message if not present at the beginning
             if not any(isinstance(msg, SystemMessage) for msg in messages):
                 system_msg = SystemMessage(content=self.system_prompt)
                 messages = [system_msg] + messages
             
-            # Get the last human message to check if it needs special processing
-            last_human_msg = None
-            for msg in reversed(messages):
+            # Get the original question (the first HumanMessage)
+            original_question = None
+            for msg in messages:
                 if isinstance(msg, HumanMessage):
-                    last_human_msg = msg.content
+                    original_question = msg.content
                     break
             
-            # Check if this is a math problem
-            if last_human_msg and self._is_math_problem(last_human_msg):
-                math_result = math_calculator(last_human_msg)
-                enhanced_msg = f"Math calculation result: {math_result}\n\nOriginal question: {last_human_msg}\n\nProvide your final answer based on this calculation."
-                messages[-1] = HumanMessage(content=enhanced_msg)
+            # Check if this is a fresh question (not after tool calls)
+            last_msg = messages[-1]
+            is_fresh_question = isinstance(last_msg, HumanMessage)
             
-            # Check if this is a date/time problem
-            elif last_human_msg and self._is_datetime_problem(last_human_msg):
-                datetime_result = date_time_processor(last_human_msg)
-                enhanced_msg = f"Date/time processing result: {datetime_result}\n\nOriginal question: {last_human_msg}\n\nProvide your final answer based on this information."
-                messages[-1] = HumanMessage(content=enhanced_msg)
+            # Only do special processing for fresh questions
+            if is_fresh_question and original_question:
+                # Check if this is a math problem
+                if self._is_math_problem(original_question):
+                    try:
+                        math_result = math_calculator(original_question)
+                        enhanced_msg = f"Question: {original_question}\n\nMath calculation result: {math_result}\n\nBased on this calculation, provide your final answer using the format: FINAL ANSWER: [your answer]"
+                        messages[-1] = HumanMessage(content=enhanced_msg)
+                    except Exception as e:
+                        print(f"Math calculation error: {e}")
+                
+                # Check if this is a date/time problem
+                elif self._is_datetime_problem(original_question):
+                    try:
+                        datetime_result = date_time_processor(original_question)
+                        enhanced_msg = f"Question: {original_question}\n\nDate/time processing result: {datetime_result}\n\nBased on this information, provide your final answer using the format: FINAL ANSWER: [your answer]"
+                        messages[-1] = HumanMessage(content=enhanced_msg)
+                    except Exception as e:
+                        print(f"DateTime processing error: {e}")
             
-            response = self.llm_with_tools.invoke(messages)
-            return {"messages": messages + [response]}
+            try:
+                response = self.llm_with_tools.invoke(messages)
+                return {"messages": messages + [response]}
+            except Exception as e:
+                print(f"LLM invocation error: {e}")
+                # Return a simple response on error
+                error_response = HumanMessage(content=f"FINAL ANSWER: Error processing question: {str(e)}")
+                return {"messages": messages + [error_response]}
         
         def tool_node(state: AgentState):
             """Tool execution node"""
-            messages = state["messages"]
-            last_message = messages[-1]
-            
-            # Execute tool calls
-            tool_node_instance = ToolNode(self.tools)
-            result = tool_node_instance.invoke(state)
-            return result
+            try:
+                tool_node_instance = ToolNode(self.tools)
+                result = tool_node_instance.invoke(state)
+                return result
+            except Exception as e:
+                print(f"Tool execution error: {e}")
+                # Add an error message and continue
+                messages = state["messages"]
+                error_msg = HumanMessage(content=f"Tool execution failed: {str(e)}. Please provide your best answer without tools.")
+                return {"messages": messages + [error_msg]}
         
         def should_continue(state: AgentState):
             """Decide whether to continue or end"""
-            last_message = state["messages"][-1]
-            
-            # If the last message has tool calls, continue to tools
-            if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
-                return "tools"
-            
-            # If we have a final answer, end
-            if hasattr(last_message, 'content') and "FINAL ANSWER:" in last_message.content:
+            try:
+                last_message = state["messages"][-1]
+                
+                # If we don't have tools, just end
+                if not self.tools:
+                    return "end"
+                
+                # If the last message has tool calls, continue to tools
+                if hasattr(last_message, 'tool_calls') and last_message.tool_calls:
+                    return "tools"
+                
+                # If we have a final answer, end
+                if (hasattr(last_message, 'content') and 
+                    last_message.content and 
+                    "FINAL ANSWER:" in str(last_message.content)):
+                    return "end"
+                
+                # Check if we've had too many iterations (prevent infinite loops)
+                if len(state["messages"]) > 10:
+                    return "end"
+                
+                # Otherwise end (be conservative)
+                return "end"
+                
+            except Exception as e:
+                print(f"Should continue error: {e}")
                 return "end"
-            
-            # Otherwise continue
-            return "end"
         
         # Build the graph
         workflow = StateGraph(AgentState)
@@ -239,9 +280,8 @@ class GAIAAgent:
         })
         workflow.add_edge("tools", "agent")
         
-        # Compile
-        memory = MemorySaver()
-        return workflow.compile(checkpointer=memory)
+        # Compile without checkpointer to avoid state issues
+        return workflow.compile()
     
     def _is_math_problem(self, text: str) -> bool:
         """Check if the text contains mathematical expressions"""
@@ -268,14 +308,20 @@ class GAIAAgent:
         try:
             print(f"Processing question: {question[:100]}...")
             
+            # Check for file/media requirements that we can't handle
+            if any(indicator in question.lower() for indicator in [
+                'attached', 'audio', 'video', 'image', 'file', 'mp3', 'pdf', 
+                'excel', 'spreadsheet', 'listen to', 'watch', 'download'
+            ]):
+                return "Unable to process files or media attachments"
+            
             # Create initial state
             initial_state = {
                 "messages": [HumanMessage(content=question)]
             }
             
             # Run the graph
-            config = {"configurable": {"thread_id": "gaia_thread"}}
-            final_state = self.graph.invoke(initial_state, config)
+            final_state = self.graph.invoke(initial_state)
             
             # Extract the final answer
             last_message = final_state["messages"][-1]
@@ -289,7 +335,13 @@ class GAIAAgent:
             
         except Exception as e:
             print(f"Error processing question: {e}")
-            return f"Error: {str(e)}"
+            # Try to provide a meaningful fallback
+            if "api" in str(e).lower() or "key" in str(e).lower():
+                return "Error: API key configuration issue"
+            elif "tool" in str(e).lower():
+                return "Error: Tool execution issue"
+            else:
+                return f"Unable to process question due to technical error"
     
     def _extract_final_answer(self, response: str) -> str:
         """Extract the final answer from the response"""