Update agent.py

agent.py

@@ -7,10 +7,6 @@ import math
 from langchain_openai import ChatOpenAI
 from langchain_community.tools.tavily_search import TavilySearchResults
 from langchain_core.messages import HumanMessage, SystemMessage
-from langgraph.graph import StateGraph, MessagesState, START, END
-from langgraph.prebuilt import ToolNode
-from langgraph.checkpoint.memory import MemorySaver
-import json
 
 # Load environment variables
 from dotenv import load_dotenv
@@ -123,9 +119,7 @@ def date_time_processor(query: str) -> str:
     except Exception as e:
         return f"Error processing date/time query: {str(e)}"
 
-# Define the agent state
-class AgentState(TypedDict):
-    messages: Annotated[list, "The messages in the conversation"]
+# Removed LangGraph dependencies - using simpler approach
 
 class GAIAAgent:
     def __init__(self):
@@ -150,138 +144,41 @@ class GAIAAgent:
             openai_api_key=openai_key
         )
         
-        # Initialize tools only if we have Tavily key
-        self.tools = []
+        # Initialize search tool if available
         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 (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()
+            self.search_tool = None
         
         self.system_prompt = read_system_prompt()
 
-    def _build_graph(self):
-        """Build the LangGraph workflow"""
-        
-        def agent_node(state: AgentState):
-            """Main agent reasoning node"""
-            messages = state["messages"]
-            
-            # 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 original question (the first HumanMessage)
-            original_question = None
-            for msg in messages:
-                if isinstance(msg, HumanMessage):
-                    original_question = msg.content
-                    break
-            
-            # Check if this is a fresh question (not after tool calls)
-            last_msg = messages[-1]
-            is_fresh_question = isinstance(last_msg, HumanMessage)
-            
-            # 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}")
-            
-            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 _search_web(self, query: str) -> str:
+        """Perform web search if available"""
+        if not self.search_tool:
+            return "Web search not available (no Tavily API key)"
         
-        def tool_node(state: AgentState):
-            """Tool execution node"""
-            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"""
-            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"
+        try:
+            results = self.search_tool.invoke({"query": query})
+            if results and len(results) > 0:
+                # Format the results nicely
+                formatted_results = []
+                for i, result in enumerate(results[:3], 1):  # Top 3 results
+                    if isinstance(result, dict):
+                        title = result.get('title', 'No title')
+                        content = result.get('content', 'No content')
+                        url = result.get('url', 'No URL')
+                        formatted_results.append(f"{i}. {title}\n   {content}\n   Source: {url}")
+                    else:
+                        formatted_results.append(f"{i}. {str(result)}")
                 
-            except Exception as e:
-                print(f"Should continue error: {e}")
-                return "end"
-        
-        # Build the graph
-        workflow = StateGraph(AgentState)
-        
-        # Add nodes
-        workflow.add_node("agent", agent_node)
-        workflow.add_node("tools", tool_node)
-        
-        # Add edges
-        workflow.add_edge(START, "agent")
-        workflow.add_conditional_edges("agent", should_continue, {
-            "tools": "tools",
-            "end": END
-        })
-        workflow.add_edge("tools", "agent")
-        
-        # Compile without checkpointer to avoid state issues
-        return workflow.compile()
+                return "\n\n".join(formatted_results)
+            else:
+                return "No search results found"
+        except Exception as e:
+            return f"Search error: {str(e)}"
     
     def _is_math_problem(self, text: str) -> bool:
         """Check if the text contains mathematical expressions"""
@@ -315,19 +212,20 @@ class GAIAAgent:
             ]):
                 return "Unable to process files or media attachments"
             
-            # Create initial state
-            initial_state = {
-                "messages": [HumanMessage(content=question)]
-            }
+            # Build the prompt based on question type
+            enhanced_question = self._enhance_question(question)
             
-            # Run the graph
-            final_state = self.graph.invoke(initial_state)
+            # Create messages
+            messages = [
+                SystemMessage(content=self.system_prompt),
+                HumanMessage(content=enhanced_question)
+            ]
             
-            # Extract the final answer
-            last_message = final_state["messages"][-1]
-            response_content = last_message.content if hasattr(last_message, 'content') else str(last_message)
+            # Get response from LLM
+            response = self.llm.invoke(messages)
+            response_content = response.content if hasattr(response, 'content') else str(response)
             
-            # Extract just the final answer part
+            # Extract the final answer
             final_answer = self._extract_final_answer(response_content)
             
             print(f"Final answer: {final_answer}")
@@ -343,6 +241,52 @@ class GAIAAgent:
             else:
                 return f"Unable to process question due to technical error"
     
+    def _enhance_question(self, question: str) -> str:
+        """Enhance the question with relevant context and tools"""
+        try:
+            # Check if this is a math problem
+            if self._is_math_problem(question):
+                try:
+                    math_result = math_calculator(question)
+                    return f"Question: {question}\n\nMath calculation result: {math_result}\n\nBased on this calculation, provide your final answer using the format: FINAL ANSWER: [your answer]"
+                except Exception as e:
+                    print(f"Math calculation error: {e}")
+            
+            # Check if this is a date/time problem
+            elif self._is_datetime_problem(question):
+                try:
+                    datetime_result = date_time_processor(question)
+                    return f"Question: {question}\n\nDate/time processing result: {datetime_result}\n\nBased on this information, provide your final answer using the format: FINAL ANSWER: [your answer]"
+                except Exception as e:
+                    print(f"DateTime processing error: {e}")
+            
+            # Check if this needs web search
+            elif self._needs_web_search(question):
+                try:
+                    search_result = self._search_web(question)
+                    return f"Question: {question}\n\nWeb search results:\n{search_result}\n\nBased on this information, provide your final answer using the format: FINAL ANSWER: [your answer]"
+                except Exception as e:
+                    print(f"Web search error: {e}")
+            
+            # For other questions, just add the format instruction
+            return f"Question: {question}\n\nProvide your final answer using the format: FINAL ANSWER: [your answer]"
+            
+        except Exception as e:
+            print(f"Question enhancement error: {e}")
+            return f"Question: {question}\n\nProvide your final answer using the format: FINAL ANSWER: [your answer]"
+    
+    def _needs_web_search(self, text: str) -> bool:
+        """Check if the question likely needs web search"""
+        search_indicators = [
+            'who', 'what', 'when', 'where', 'which', 'published', 'article',
+            'wikipedia', 'latest', 'recent', 'current', 'news', 'website',
+            'url', 'http', 'www', 'competition', 'olympics', 'award', 
+            'winner', 'recipient', 'author', 'published in', 'paper',
+            'study', 'research', 'species', 'city', 'country'
+        ]
+        text_lower = text.lower()
+        return any(indicator in text_lower for indicator in search_indicators)
+    
     def _extract_final_answer(self, response: str) -> str:
         """Extract the final answer from the response"""
         if "FINAL ANSWER:" in response: