Upload final_agent.py

final_agent.py

@@ -0,0 +1,320 @@
+# ==============================================================================
+# Imports
+# ==============================================================================
+import os
+import requests
+import traceback
+import html2text # For HTML to text conversion
+import tempfile # For file handling tools
+import pandas as pd # For CSV/Excel analysis
+import openpyxl # For Excel analysis
+from PIL import Image # For image text extraction
+import pytesseract # For image text extraction
+from urllib.parse import urlparse # For download tool
+from typing import Annotated, List, TypedDict, Optional
+from dotenv import load_dotenv
+
+from langgraph.graph import StateGraph, START, END
+from langgraph.graph.message import add_messages
+from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage
+from langchain_core.tools import tool
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_community.tools.tavily_search import TavilySearchResults
+
+# ==============================================================================
+# Environment Setup & LLM
+# ==============================================================================
+load_dotenv()
+gemini_api_key = os.getenv("GEMINI_API_KEY")
+tavily_api_key = os.getenv("TAVILY_API_KEY")
+
+# --- Optional: Tesseract Path (Ensure commented out if Tesseract is in PATH) ---
+# pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'
+
+if not gemini_api_key:
+    raise ValueError("GEMINI_API_KEY not found in environment variables.")
+if not tavily_api_key:
+    raise ValueError("TAVILY_API_KEY not found. Required for search.")
+
+# LLM Choice remains as per user's last code, temperature lowered
+llm = ChatGoogleGenerativeAI(
+    model="gemini-1.5-pro", # Keeping user's specified model
+    google_api_key=gemini_api_key,
+    temperature=0.1 # Lowered temperature for GAIA task
+)
+print(f"LLM Initialized: {llm.model}")
+
+# ==============================================================================
+# State Definition
+# ==============================================================================
+class AgentState(TypedDict):
+    """Core state for the GAIA agent."""
+    input_question: str # Added input_question back
+    messages: Annotated[List[BaseMessage], add_messages]
+    error: Optional[str]
+    iterations: int
+
+# ==============================================================================
+# Tools (Original + Integrated)
+# ==============================================================================
+
+# --- Search Tool (Tavily) ---
+search_tool = TavilySearchResults(max_results=3, api_key=tavily_api_key)
+search_tool.name = "web_search"
+search_tool.description = "Performs a web search (using Tavily) to find relevant URLs/snippets for a query."
+
+# --- Web Browser Tool (html2text) ---
+@tool
+def web_browser(url: str) -> str:
+    """Fetches text content from a webpage URL using html2text. Use after 'web_search'."""
+    print(f"--- [Tool] Browsing (html2text): {url} ---")
+    try:
+        headers = {'User-Agent': 'Mozilla/5.0'}
+        response = requests.get(url, headers=headers, timeout=20)
+        response.raise_for_status()
+        response.encoding = response.apparent_encoding or 'utf-8'
+        h = html2text.HTML2Text(bodywidth=0); h.ignore_links = True; h.ignore_images = True
+        clean_text = h.handle(response.text)
+        max_length = 6000
+        if len(clean_text) > max_length:
+            return clean_text[:max_length] + "\n\n... [Content Truncated]"
+        # Ensure we return error string if empty after strip
+        cleaned_and_stripped = clean_text.strip()
+        return cleaned_and_stripped if cleaned_and_stripped else f"Error: No meaningful content via html2text for {url}."
+    except requests.exceptions.RequestException as e:
+        return f"Error: Network request failed for URL: {url}. Reason: {e}"
+    except Exception as e:
+        return f"Error: Unexpected error processing URL with html2text: {url}. Reason: {str(e)}"
+
+# --- File Download Tool ---
+@tool
+def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
+    """Downloads a file from a URL to a temporary directory. Input: file URL. Returns: path to downloaded file or error."""
+    print(f"--- [Tool] Downloading file from: {url} ---")
+    try:
+        if not filename:
+            try: path = urlparse(url).path; filename = os.path.basename(path) if path else None
+            except Exception: filename = None
+            if not filename: import uuid; filename = f"downloaded_{uuid.uuid4().hex[:8]}"
+        temp_dir = tempfile.gettempdir(); filepath = os.path.join(temp_dir, filename)
+        response = requests.get(url, stream=True, timeout=30); response.raise_for_status()
+        with open(filepath, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192): f.write(chunk)
+        print(f"--- [Tool] File downloaded to: {filepath} ---")
+        return f"File downloaded to {filepath}. Use appropriate tools (e.g., analyze_csv_file) to process it."
+    except requests.exceptions.RequestException as e:
+        return f"Error downloading file: Network issue for {url}. Reason: {e}"
+    except Exception as e:
+        return f"Error downloading file: Unexpected error for {url}. Reason: {str(e)}"
+
+# --- CSV Analysis Tool ---
+@tool
+def analyze_csv_file(file_path: str) -> str:
+    """Analyzes a CSV file at the given path using pandas. Returns a summary of content or error."""
+    print(f"--- [Tool] Analyzing CSV: {file_path} ---")
+    if not os.path.exists(file_path): return f"Error: CSV file not found: {file_path}"
+    try:
+        df = pd.read_csv(file_path)
+        summary = f"CSV Analysis Report for {os.path.basename(file_path)}:\n"
+        summary += f"- Shape: {df.shape[0]} rows, {df.shape[1]} columns\n"
+        summary += f"- Columns: {', '.join(df.columns)}\n"
+        summary += f"\nFirst 5 rows:\n{df.head().to_string()}\n"
+        numeric_cols = df.select_dtypes(include=['number'])
+        if not numeric_cols.empty:
+             summary += f"\nBasic Stats (Numeric):\n{numeric_cols.describe().to_string()}"
+        else:
+             summary += "\nNo numeric columns for stats."
+        return summary
+    except ImportError: return "Error: 'pandas' required but not installed."
+    except Exception as e: return f"Error analyzing CSV {file_path}: {str(e)}"
+
+# --- Excel Analysis Tool ---
+@tool
+def analyze_excel_file(file_path: str) -> str:
+    """Analyzes an Excel file (.xlsx, .xls) at the given path. Returns a summary of the first sheet or error."""
+    print(f"--- [Tool] Analyzing Excel: {file_path} ---")
+    if not os.path.exists(file_path): return f"Error: Excel file not found: {file_path}"
+    try:
+        df = pd.read_excel(file_path, engine='openpyxl')
+        summary = f"Excel Analysis Report for {os.path.basename(file_path)} (First Sheet):\n"
+        summary += f"- Shape: {df.shape[0]} rows, {df.shape[1]} columns\n"
+        summary += f"- Columns: {', '.join(df.columns)}\n"
+        summary += f"\nFirst 5 rows:\n{df.head().to_string()}\n"
+        numeric_cols = df.select_dtypes(include=['number'])
+        if not numeric_cols.empty:
+             summary += f"\nBasic Stats (Numeric):\n{numeric_cols.describe().to_string()}"
+        else:
+             summary += "\nNo numeric columns for stats."
+        return summary
+    except ImportError: return "Error: 'pandas' and 'openpyxl' required but not installed."
+    except Exception as e: return f"Error analyzing Excel {file_path}: {str(e)}"
+
+# --- Image Text Extraction Tool (OCR) ---
+@tool
+def extract_text_from_image(file_path: str) -> str:
+    """Extracts text from an image file at the given path using Tesseract OCR. Returns extracted text or error."""
+    print(f"--- [Tool] Extracting text from image: {file_path} ---")
+    if not os.path.exists(file_path): return f"Error: Image file not found: {file_path}"
+    try:
+        # Need to explicitly handle potential empty string from pytesseract
+        text = pytesseract.image_to_string(Image.open(file_path))
+        text_stripped = text.strip()
+        return f"Extracted text from image '{os.path.basename(file_path)}':\n{text_stripped}" if text_stripped else "No text found in image."
+    except ImportError: return "Error: 'Pillow' or 'pytesseract' required but not installed."
+    except pytesseract.TesseractNotFoundError: return "Error: Tesseract OCR not installed or not in PATH."
+    except Exception as e: return f"Error extracting text from image {file_path}: {str(e)}"
+
+# --- Basic Math Tools ---
+@tool
+def add(a: float, b: float) -> float:
+    """Adds two numbers (a + b). Handles float inputs."""
+    print(f"--- [Tool] Calculating: {a} + {b} ---")
+    return a + b
+@tool
+def subtract(a: float, b: float) -> float:
+    """Subtracts the second number from the first (a - b). Handles float inputs."""
+    print(f"--- [Tool] Calculating: {a} - {b} ---")
+    return a - b
+@tool
+def multiply(a: float, b: float) -> float:
+    """Multiplies two numbers (a * b). Handles float inputs."""
+    print(f"--- [Tool] Calculating: {a} * {b} ---")
+    return a * b
+@tool
+def divide(a: float, b: float) -> float | str:
+    """Divides the first number by the second (a / b). Handles float inputs and division by zero."""
+    print(f"--- [Tool] Calculating: {a} / {b} ---")
+    if b == 0: return "Error: Cannot divide by zero."
+    return a / b
+
+# --- Tool List & LLM Binding ---
+tools = [ search_tool, web_browser, download_file_from_url, analyze_csv_file,
+          analyze_excel_file, extract_text_from_image, add, subtract, multiply, divide ]
+llm_with_tools = llm.bind_tools(tools)
+print(f"Agent initialized with {len(tools)} tools.")
+
+# ==============================================================================
+# Node Definitions
+# ==============================================================================
+
+# --- Agent Node (LLM Call) ---
+def call_agent_node(state: AgentState) -> dict:
+    """Calls the LLM to decide the next step/response, increments iteration count."""
+    print(f"\n--- [Node] Agent thinking... (Iteration {state['iterations']}) ---")
+    MAX_ITERATIONS = 10 # Max steps for the entire task
+    current_iterations = state.get('iterations', 0)
+    if current_iterations >= MAX_ITERATIONS:
+        print(f"Warning: Reached max iterations ({MAX_ITERATIONS}). Stopping.")
+        # Return error message in state
+        return {"error": f"Max iterations ({MAX_ITERATIONS}) reached."}
+    try:
+        response = llm_with_tools.invoke(state['messages'])
+        print("--- [Node] AI Response/Action ---")
+        response.pretty_print()
+        # Important: Return the iterations incremented
+        return {"messages": [response], "iterations": current_iterations + 1}
+    except Exception as e:
+        error_message = f"LLM invocation failed: {str(e)}"
+        print(f"--- [Node] ERROR: {error_message} ---")
+        # Add error message to history and set error state
+        return {"messages": [AIMessage(content=f"Sorry, I encountered an error: {error_message}")], "error": error_message, "iterations": current_iterations + 1} # Still increment iteration on error
+
+# --- Tool Node (Executor) ---
+tool_node = ToolNode(tools)
+
+# ==============================================================================
+# Graph Construction (Non-conversational)
+# ==============================================================================
+builder = StateGraph(AgentState)
+
+# Add nodes
+builder.add_node("agent", call_agent_node)
+builder.add_node("tools", tool_node)
+
+# Define edges
+builder.add_edge(START, "agent")
+
+# Conditional edge after agent thinks
+builder.add_conditional_edges(
+    "agent",
+    tools_condition, # Check if the last message has tool calls
+    {
+        "tools": "tools", # If yes, execute tools
+        END: END          # If no, END the graph
+    }
+)
+
+# Edge after tools execute
+builder.add_edge("tools", "agent") # Loop back to agent to process tool results
+
+# Compile the graph
+graph = builder.compile()
+print("GAIA agent graph compiled with integrated tools.")
+
+# ==============================================================================
+# Execution (Single run for GAIA task)
+# ==============================================================================
+
+# --- GAIA Task Input ---
+# IMPORTANT: Replace this with the actual question/task from the benchmark environment
+#            Also, ensure any file paths mentioned are correctly handled/accessible.
+gaia_question = "What is the result of multiplying the number of rows in the provided CSV file ('data.csv') by the number found after the phrase 'total items:' in the text extracted from the provided image file ('image.png')?"
+# Example assumption: 'data.csv' and 'image.png' are expected to be in the current working directory or provided via the GAIA framework.
+
+print(f"\n--- Running Agent for GAIA Question: {gaia_question} ---")
+
+# --- Set up Initial State ---
+initial_state = AgentState(
+    input_question=gaia_question,
+    messages=[HumanMessage(content=f"""Your task is to accurately answer the following question based *only* on information obtained using your tools (web search, web browser, file download, csv/excel analysis, image OCR, math).
+
+Follow these steps methodically:
+1. Analyze the question to understand required information and tools needed.
+2. If external files are mentioned (e.g., CSV, image paths like 'data.csv', 'image.png'), use the appropriate analysis tool directly on the provided file path. Assume files mentioned are accessible in the current directory unless a URL is given.
+3. If a URL is given for a file, use 'download_file_from_url' first, then analyze the downloaded file using its path.
+4. If web information is needed, use 'web_search' then 'web_browser' on relevant URLs.
+5. If calculations are needed, use the math tools.
+6. Synthesize the information gathered from tools to arrive at the final answer.
+7. **CRITICAL:** Your final output MUST contain ONLY the precise numerical or text answer requested by the question. Do NOT include explanations, reasoning steps, units unless explicitly asked for, context, apologies, or any introductory phrases like "The final answer is...". Just the required answer string or number itself.
+
+Question: {gaia_question}
+""")],
+    error=None,
+    iterations=0
+)
+
+try:
+    # Run the graph from start to end for the single task
+    final_state = graph.invoke(initial_state, config={"recursion_limit": 15})
+except Exception as e:
+    print(f"--- Graph execution failed unexpectedly: {e} ---")
+    traceback.print_exc()
+    final_state = None
+
+# ==============================================================================
+# Results Processing
+# ==============================================================================
+print("\n--- Agent Run Finished ---")
+if final_state:
+    if final_state.get("error"):
+        print(f"Agent stopped due to ERROR: {final_state['error']}")
+
+    # Extract final answer (expected to be ONLY the answer in the last AI message)
+    if final_state.get('messages') and isinstance(final_state['messages'][-1], AIMessage):
+        potential_answer = final_state['messages'][-1].content
+        print(f"\nFinal Answer (Submit This): {potential_answer}")
+        # For GAIA submission, programmatically extract and return/save 'potential_answer'
+    else:
+        # Handle cases where the agent errored out or didn't produce a final AIMessage
+        print("Could not determine final answer (last message not AI or missing).")
+        print("Final State:")
+        # Print relevant parts of the final state for debugging
+        print(f"  Error: {final_state.get('error')}")
+        print(f"  Iterations: {final_state.get('iterations')}")
+        print("  Last few messages:")
+        for msg in final_state.get('messages', [])[-3:]: # Print last 3 messages
+             msg.pretty_print()
+
+else:
+    print("Execution failed, no final state.")