Update agent.py

agent.py

@@ -7,267 +7,289 @@ from pathlib import Path
 from typing import Optional, Union, Dict, Any, List
 from dotenv import load_dotenv
 
-from langgraph.graph import StateGraph, MessagesState
-from langgraph.prebuilt import create_react_agent
-from langchain_core.messages import HumanMessage, SystemMessage
-from langchain_core.tools import tool
-from langchain_openai import ChatOpenAI
-
 load_dotenv()
 
-
-class OpenRouterLLM(ChatOpenAI):
-    """Custom OpenRouter LLM wrapper for LangGraph"""
+# Simple tool-based agent without LangGraph for now
+class SimpleAgent:
+    """Simple agent with tool capabilities"""
     
-    def __init__(self, model: str = "deepseek/deepseek-v3.1-terminus", **kwargs):
-        api_key = os.getenv("OPENROUTER_API_KEY") or os.getenv("my_key")
-        
-        super().__init__(
-            model=model,
-            openai_api_key=api_key,
-            openai_api_base="https://openrouter.ai/api/v1",
-            **kwargs
-        )
+    def __init__(self, llm):
+        self.llm = llm
+        self.tools = {
+            'search_web': self.search_web,
+            'search_wikipedia': self.search_wikipedia,
+            'execute_python': self.execute_python,
+            'read_excel_file': self.read_excel_file,
+            'read_text_file': self.read_text_file,
+        }
+    
+    def search_web(self, query: str) -> str:
+        """Search the web using DuckDuckGo for current information."""
+        try:
+            search_url = f"https://api.duckduckgo.com/?q={query}&format=json&no_html=1&skip_disambig=1"
+            response = requests.get(search_url, timeout=10)
+            
+            if response.status_code == 200:
+                data = response.json()
+                results = []
+                if data.get("AbstractText"):
+                    results.append(f"Abstract: {data['AbstractText']}")
+                
+                if data.get("RelatedTopics"):
+                    for topic in data["RelatedTopics"][:3]:
+                        if isinstance(topic, dict) and topic.get("Text"):
+                            results.append(f"Related: {topic['Text']}")
+                
+                if results:
+                    return "\n".join(results)
+                else:
+                    return f"Search performed for '{query}' but no specific results found."
+            else:
+                return f"Search failed with status code {response.status_code}"
+        except Exception as e:
+            return f"Search error: {str(e)}"
 
+    def search_wikipedia(self, query: str) -> str:
+        """Search Wikipedia for factual information."""
+        try:
+            search_url = "https://en.wikipedia.org/api/rest_v1/page/summary/" + query.replace(" ", "_")
+            response = requests.get(search_url, timeout=10)
+            
+            if response.status_code == 200:
+                data = response.json()
+                extract = data.get("extract", "")
+                if extract:
+                    return f"Wikipedia: {extract[:500]}..."
+                else:
+                    return f"Wikipedia page found for '{query}' but no extract available."
+            else:
+                return f"Wikipedia search failed for '{query}'"
+        except Exception as e:
+            return f"Wikipedia search error: {str(e)}"
 
-@tool
-def search_web(query: str) -> str:
-    """Search the web using DuckDuckGo for current information."""
-    try:
-        # Simple web search using DuckDuckGo
-        search_url = f"https://api.duckduckgo.com/?q={query}&format=json&no_html=1&skip_disambig=1"
-        response = requests.get(search_url, timeout=10)
-        
-        if response.status_code == 200:
-            data = response.json()
+    def execute_python(self, code: str) -> str:
+        """Execute Python code and return the result."""
+        try:
+            import io
+            import sys
             
-            # Extract results
-            results = []
-            if data.get("AbstractText"):
-                results.append(f"Abstract: {data['AbstractText']}")
+            safe_globals = {
+                '__builtins__': {
+                    'print': print, 'len': len, 'str': str, 'int': int, 'float': float,
+                    'bool': bool, 'list': list, 'dict': dict, 'tuple': tuple, 'set': set,
+                    'range': range, 'sum': sum, 'max': max, 'min': min, 'abs': abs,
+                    'round': round, 'sorted': sorted, 'enumerate': enumerate, 'zip': zip,
+                },
+                'math': __import__('math'),
+                'json': __import__('json'),
+            }
             
-            if data.get("RelatedTopics"):
-                for topic in data["RelatedTopics"][:3]:
-                    if isinstance(topic, dict) and topic.get("Text"):
-                        results.append(f"Related: {topic['Text']}")
+            old_stdout = sys.stdout
+            sys.stdout = mystdout = io.StringIO()
             
-            if results:
-                return "\n".join(results)
-            else:
-                return f"Search performed for '{query}' but no specific results found."
-        else:
-            return f"Search failed with status code {response.status_code}"
+            try:
+                exec(code, safe_globals)
+                output = mystdout.getvalue()
+            finally:
+                sys.stdout = old_stdout
             
-    except Exception as e:
-        return f"Search error: {str(e)}"
-
+            return output if output else "Code executed successfully (no output)"
+        except Exception as e:
+            return f"Python execution error: {str(e)}"
 
-@tool
-def search_wikipedia(query: str) -> str:
-    """Search Wikipedia for factual information."""
-    try:
-        # Wikipedia API search
-        search_url = "https://en.wikipedia.org/api/rest_v1/page/summary/" + query.replace(" ", "_")
-        response = requests.get(search_url, timeout=10)
-        
-        if response.status_code == 200:
-            data = response.json()
-            extract = data.get("extract", "")
-            if extract:
-                return f"Wikipedia: {extract[:500]}..."
+    def read_excel_file(self, file_path: str, sheet_name: Optional[str] = None) -> str:
+        """Read an Excel file and return its contents."""
+        try:
+            file_path_obj = Path(file_path)
+            if not file_path_obj.exists():
+                return f"Error: File not found at {file_path}"
+            
+            if sheet_name and sheet_name.isdigit():
+                sheet_name = int(sheet_name)
+            elif sheet_name is None:
+                sheet_name = 0
+                
+            df = pd.read_excel(file_path, sheet_name=sheet_name)
+            
+            if len(df) > 20:
+                result = f"Excel file with {len(df)} rows and {len(df.columns)} columns:\n\n"
+                result += "First 10 rows:\n" + df.head(10).to_string(index=False)
+                result += f"\n\n... ({len(df) - 20} rows omitted) ...\n\n"
+                result += "Last 10 rows:\n" + df.tail(10).to_string(index=False)
             else:
-                return f"Wikipedia page found for '{query}' but no extract available."
-        else:
-            return f"Wikipedia search failed for '{query}'"
+                result = f"Excel file with {len(df)} rows and {len(df.columns)} columns:\n\n"
+                result += df.to_string(index=False)
+                
+            return result
+        except Exception as e:
+            return f"Error reading Excel file: {str(e)}"
+
+    def read_text_file(self, file_path: str) -> str:
+        """Read a text file and return its contents."""
+        try:
+            file_path_obj = Path(file_path)
+            if not file_path_obj.exists():
+                return f"Error: File not found at {file_path}"
+            
+            encodings = ['utf-8', 'utf-16', 'iso-8859-1', 'cp1252']
+            
+            for encoding in encodings:
+                try:
+                    with open(file_path_obj, 'r', encoding=encoding) as f:
+                        content = f.read()
+                    return f"File content ({encoding} encoding):\n\n{content}"
+                except UnicodeDecodeError:
+                    continue
             
-    except Exception as e:
-        return f"Wikipedia search error: {str(e)}"
+            return f"Error: Could not decode file with any standard encoding"
+        except Exception as e:
+            return f"Error reading file: {str(e)}"
 
+    def run(self, question: str) -> str:
+        """Run the agent with tool usage"""
+        # First, try to answer directly
+        direct_response = self.llm(f"""
+Question: {question}
 
-@tool
-def execute_python(code: str) -> str:
-    """Execute Python code and return the result."""
-    try:
-        # Create a safe execution environment
-        safe_globals = {
-            '__builtins__': {
-                'print': print,
-                'len': len,
-                'str': str,
-                'int': int,
-                'float': float,
-                'bool': bool,
-                'list': list,
-                'dict': dict,
-                'tuple': tuple,
-                'set': set,
-                'range': range,
-                'sum': sum,
-                'max': max,
-                'min': min,
-                'abs': abs,
-                'round': round,
-                'sorted': sorted,
-                'enumerate': enumerate,
-                'zip': zip,
-            },
-            'math': __import__('math'),
-            'json': __import__('json'),
-            'datetime': __import__('datetime'),
-            'random': __import__('random'),
-        }
+Think step by step. If this question requires:
+- Web search for current information, say "NEED_SEARCH: <search query>"
+- Mathematical calculation, say "NEED_PYTHON: <python code>"
+- Wikipedia lookup, say "NEED_WIKI: <search term>"
+- File analysis (if file path mentioned), say "NEED_FILE: <file_path>"
+
+Otherwise, provide a direct answer.
+
+Your response:""")
+
+        # Check if tools are needed
+        if "NEED_SEARCH:" in direct_response:
+            search_query = direct_response.split("NEED_SEARCH:")[1].strip()
+            search_result = self.search_web(search_query)
+            return self.llm(f"Question: {question}\n\nSearch results: {search_result}\n\nFinal answer:")
         
-        # Capture output
-        import io
-        import sys
+        elif "NEED_PYTHON:" in direct_response:
+            code = direct_response.split("NEED_PYTHON:")[1].strip()
+            exec_result = self.execute_python(code)
+            return self.llm(f"Question: {question}\n\nCalculation result: {exec_result}\n\nFinal answer:")
         
-        old_stdout = sys.stdout
-        sys.stdout = mystdout = io.StringIO()
+        elif "NEED_WIKI:" in direct_response:
+            wiki_query = direct_response.split("NEED_WIKI:")[1].strip()
+            wiki_result = self.search_wikipedia(wiki_query)
+            return self.llm(f"Question: {question}\n\nWikipedia info: {wiki_result}\n\nFinal answer:")
         
-        try:
-            # Execute the code
-            exec(code, safe_globals)
-            output = mystdout.getvalue()
-        finally:
-            sys.stdout = old_stdout
+        elif "NEED_FILE:" in direct_response:
+            file_path = direct_response.split("NEED_FILE:")[1].strip()
+            if file_path.endswith(('.xlsx', '.xls')):
+                file_content = self.read_excel_file(file_path)
+            else:
+                file_content = self.read_text_file(file_path)
+            return self.llm(f"Question: {question}\n\nFile content: {file_content}\n\nFinal answer:")
         
-        return output if output else "Code executed successfully (no output)"
+        else:
+            return direct_response
+class OpenRouterLLM:
+    """Simple OpenRouter LLM wrapper"""
+    
+    def __init__(self, model: str = "deepseek/deepseek-v3.1-terminus"):
+        self.api_key = os.getenv("OPENROUTER_API_KEY") or os.getenv("my_key")
+        self.model = model
+        self.base_url = "https://openrouter.ai/api/v1/chat/completions"
+    
+    def __call__(self, prompt: str, max_tokens: int = 1500, temperature: float = 0.1) -> str:
+        """Make API call to OpenRouter"""
         
-    except Exception as e:
-        return f"Python execution error: {str(e)}"
-
-
-@tool
-def read_excel_file(file_path: str, sheet_name: Optional[str] = None) -> str:
-    """Read an Excel file and return its contents as a formatted string."""
-    try:
-        file_path_obj = Path(file_path)
-        if not file_path_obj.exists():
-            return f"Error: File not found at {file_path}"
+        if not self.api_key or not self.api_key.startswith('sk-or-v1-'):
+            return "Error: Invalid OpenRouter API key"
         
-        # Try to read the Excel file
-        if sheet_name and sheet_name.isdigit():
-            sheet_name = int(sheet_name)
-        elif sheet_name is None:
-            sheet_name = 0
-            
-        df = pd.read_excel(file_path, sheet_name=sheet_name)
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
         
-        # Convert to string representation
-        if len(df) > 20:
-            # Show first 10 and last 10 rows for large datasets
-            result = f"Excel file with {len(df)} rows and {len(df.columns)} columns:\n\n"
-            result += "First 10 rows:\n"
-            result += df.head(10).to_string(index=False)
-            result += f"\n\n... ({len(df) - 20} rows omitted) ...\n\n"
-            result += "Last 10 rows:\n"
-            result += df.tail(10).to_string(index=False)
-        else:
-            result = f"Excel file with {len(df)} rows and {len(df.columns)} columns:\n\n"
-            result += df.to_string(index=False)
-            
-        return result
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": "You are a helpful AI assistant. Provide direct, accurate answers. For GAIA evaluation, be precise and concise."
+                },
+                {
+                    "role": "user", 
+                    "content": prompt
+                }
+            ],
+            "temperature": temperature,
+            "max_tokens": max_tokens,
+        }
         
-    except Exception as e:
-        return f"Error reading Excel file: {str(e)}"
-
-
-@tool
-def read_text_file(file_path: str) -> str:
-    """Read a text file and return its contents."""
-    try:
-        file_path_obj = Path(file_path)
-        if not file_path_obj.exists():
-            return f"Error: File not found at {file_path}"
+        try:
+            response = requests.post(self.base_url, headers=headers, json=payload, timeout=30)
+            
+            if response.status_code != 200:
+                return f"API Error: {response.status_code}"
+            
+            result = response.json()
+            
+            if "choices" in result and len(result["choices"]) > 0:
+                answer = result["choices"][0]["message"]["content"].strip()
+                return self._clean_answer(answer)
+            else:
+                return "Error: No response content received"
+                
+        except Exception as e:
+            return f"Error: {str(e)}"
+    
+    def _clean_answer(self, answer: str) -> str:
+        """Clean the answer for GAIA evaluation"""
+        answer = answer.strip()
         
-        # Try different encodings
-        encodings = ['utf-8', 'utf-16', 'iso-8859-1', 'cp1252']
+        # Remove common prefixes
+        prefixes = [
+            "Answer:", "The answer is:", "Final answer:", "Result:", 
+            "Solution:", "Based on", "Therefore", "In conclusion"
+        ]
         
-        for encoding in encodings:
-            try:
-                with open(file_path_obj, 'r', encoding=encoding) as f:
-                    content = f.read()
-                return f"File content ({encoding} encoding):\n\n{content}"
-            except UnicodeDecodeError:
-                continue
+        for prefix in prefixes:
+            if answer.lower().startswith(prefix.lower()):
+                answer = answer[len(prefix):].strip()
+                if answer.startswith(':'):
+                    answer = answer[1:].strip()
+                break
         
-        return f"Error: Could not decode file with any standard encoding"
+        # Remove quotes and periods from short answers
+        if len(answer.split()) <= 3:
+            answer = answer.strip('"\'.')
         
-    except Exception as e:
-        return f"Error reading file: {str(e)}"
+        return answer
 
 
 class GaiaAgent:
-    """LangGraph-based agent for GAIA tasks using OpenRouter DeepSeek"""
+    """Simple tool-based agent for GAIA tasks"""
     
     def __init__(self):
-        print("Initializing GaiaAgent with LangGraph and OpenRouter DeepSeek...")
+        print("Initializing GaiaAgent with OpenRouter DeepSeek...")
         
         # Initialize the LLM
-        self.llm = OpenRouterLLM(
-            model="deepseek/deepseek-v3.1-terminus",
-            temperature=0.1,
-            max_tokens=2000
-        )
+        self.llm = OpenRouterLLM(model="deepseek/deepseek-v3.1-terminus")
         
-        # Define available tools
-        self.tools = [
-            search_web,
-            search_wikipedia,
-            execute_python,
-            read_excel_file,
-            read_text_file,
-        ]
-        
-        # Create the agent
-        self.agent = create_react_agent(
-            self.llm,
-            self.tools,
-            state_modifier=self._get_system_prompt()
-        )
+        # Initialize the agent with tools
+        self.agent = SimpleAgent(self.llm)
         
         print("GaiaAgent initialized successfully!")
     
-    def _get_system_prompt(self) -> str:
-        """Get the system prompt for the agent"""
-        return """You are an advanced AI agent designed to answer complex questions that may require:
-
-1. Web searches for current information
-2. Mathematical calculations using Python
-3. File analysis (Excel, text files)  
-4. Multi-step reasoning and problem solving
-
-For GAIA evaluation:
-- Provide EXACT, DIRECT answers
-- Use tools when necessary to gather information or perform calculations
-- For math problems, show your calculation but end with just the number
-- For yes/no questions, answer just "Yes" or "No"
-- For factual questions, provide just the fact
-
-When you encounter files:
-- Use read_excel_file for .xlsx, .xls files
-- Use read_text_file for text-based files
-- Analyze the file content to answer the question
-
-Be thorough in your analysis but concise in your final answer."""
-
     def __call__(self, task_id: str, question: str) -> str:
         """Process a question and return the answer"""
         try:
             print(f"Processing task {task_id}: {question[:100]}...")
             
-            # Create the input state
-            messages = [HumanMessage(content=question)]
+            # Check if there are file references in the question
+            enhanced_question = self._enhance_question_with_file_analysis(question)
             
             # Run the agent
-            result = self.agent.invoke({"messages": messages})
+            answer = self.agent.run(enhanced_question)
             
-            # Extract the final answer
-            final_message = result["messages"][-1]
-            answer = final_message.content
-            
-            # Clean up the answer for GAIA evaluation
-            clean_answer = self._clean_answer(answer)
+            # Clean up the answer
+            clean_answer = self._clean_final_answer(answer)
             
             print(f"Agent answer for {task_id}: {clean_answer}")
             return clean_answer
@@ -277,31 +299,47 @@ Be thorough in your analysis but concise in your final answer."""
             print(f"Error processing task {task_id}: {error_msg}")
             return error_msg
     
-    def _clean_answer(self, answer: str) -> str:
-        """Clean the answer to extract the final result"""
+    def _enhance_question_with_file_analysis(self, question: str) -> str:
+        """Check if question mentions files and enhance accordingly"""
+        # Look for file path mentions in the question
+        file_patterns = [
+            r'/tmp/gaia_cached_files/[^\s]+',
+            r'saved locally at:\s*([^\s]+)',
+            r'file.*?\.xlsx?',
+            r'file.*?\.csv',
+            r'file.*?\.txt'
+        ]
+        
+        for pattern in file_patterns:
+            matches = re.findall(pattern, question, re.IGNORECASE)
+            if matches:
+                # File found, the agent will handle it automatically
+                break
+        
+        return question
+    
+    def _clean_final_answer(self, answer: str) -> str:
+        """Final cleaning of the answer"""
         answer = answer.strip()
         
-        # Look for "Final Answer:" pattern
+        # Look for final answer pattern
         if "final answer:" in answer.lower():
-            parts = re.split(r'final answer:', answer, flags=re.IGNORECASE)
+            parts = answer.lower().split("final answer:")
             if len(parts) > 1:
-                answer = parts[-1].strip()
+                answer = answer.split(":")[-1].strip()
         
-        # Remove common prefixes
-        prefixes = [
-            "The answer is", "Answer:", "Result:", "Solution:",
-            "Based on", "Therefore", "In conclusion", "So the answer is"
+        # Remove common unnecessary phrases
+        cleanup_phrases = [
+            "based on the", "according to", "the answer is", "therefore",
+            "in conclusion", "as a result", "so the answer is"
         ]
         
-        for prefix in prefixes:
-            if answer.lower().startswith(prefix.lower()):
-                answer = answer[len(prefix):].strip()
-                if answer.startswith(':'):
-                    answer = answer[1:].strip()
+        for phrase in cleanup_phrases:
+            if answer.lower().startswith(phrase):
+                answer = answer[len(phrase):].strip()
                 break
         
-        # Remove quotes and periods from short answers
-        if len(answer.split()) <= 3:
-            answer = answer.strip('"\'.')
+        # Clean up formatting
+        answer = answer.strip('.,;:"\'')
         
         return answer