Update final_agent.py

final_agent.py

@@ -14,6 +14,7 @@ from urllib.parse import urlparse # For download tool
 from typing import Annotated, List, TypedDict, Optional
 from dotenv import load_dotenv
 
+# LangChain and LangGraph Imports
 from langgraph.graph import StateGraph, START, END
 from langgraph.graph.message import add_messages
 from langgraph.prebuilt import ToolNode, tools_condition
@@ -29,19 +30,27 @@ 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'
+# --- Optional: Tesseract Path ---
+# If Tesseract OCR is not in your system's PATH environment variable,
+# uncomment the following line and set the correct path to tesseract.exe
+# try:
+#     pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe' # Example path for Windows
+# except NameError: pass # Handles case where pytesseract might not be imported yet if PIL fails first
+# except Exception as e: print(f"Warning: Could not set tesseract_cmd path: {e}")
 
+
+# --- Validate API Keys ---
 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.")
+    raise ValueError("TAVILY_API_KEY not found. Required for Tavily search tool.")
 
-# LLM Choice remains as per user's last code, temperature lowered
+# --- Initialize LLM ---
+# Using the model specified in the user's code block
 llm = ChatGoogleGenerativeAI(
-    model="gemini-2.0-flash", # Keeping user's specified model
+    model="gemini-1.5-flash-latest", # As per user's last provided code
     google_api_key=gemini_api_key,
-    temperature=0.1 # Lowered temperature for GAIA task
+    temperature=0.1 # Low temperature for factual tasks
 )
 print(f"LLM Initialized: {llm.model}")
 
@@ -49,15 +58,16 @@ 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
+    """Defines the structure of the information the agent tracks during its run."""
+    input_question: str # The original question from the benchmark
+    messages: Annotated[List[BaseMessage], add_messages] # History of interactions (Human, AI, Tool)
+    error: Optional[str] # Stores any error message encountered
+    iterations: int # Counter for agent steps to prevent loops
 
 # ==============================================================================
-# Tools (Original + Integrated)
+# Tools Definitions
 # ==============================================================================
+print("Defining tools...")
 
 # --- Search Tool (Tavily) ---
 search_tool = TavilySearchResults(max_results=3, api_key=tavily_api_key)
@@ -70,16 +80,20 @@ 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'}
+        headers = {'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'}
         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
+        # Configure html2text
+        h = html2text.HTML2Text(bodywidth=0)
+        h.ignore_links = True
+        h.ignore_images = True
+        # Convert HTML to text
         clean_text = h.handle(response.text)
+        # Limit content length
         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:
@@ -93,11 +107,14 @@ 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:
+        # Generate filename if needed
         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]}"
+        # Define save path
         temp_dir = tempfile.gettempdir(); filepath = os.path.join(temp_dir, filename)
+        # Download file
         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)
@@ -113,9 +130,11 @@ def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
 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}"
+    # GAIA might provide relative paths, ensure they work or adjust logic if needed
+    if not os.path.exists(file_path): return f"Error: CSV file not found at path: {file_path}"
     try:
         df = pd.read_csv(file_path)
+        # Generate summary string
         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"
@@ -134,9 +153,10 @@ def analyze_csv_file(file_path: str) -> str:
 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}"
+    if not os.path.exists(file_path): return f"Error: Excel file not found at path: {file_path}"
     try:
         df = pd.read_excel(file_path, engine='openpyxl')
+        # Generate summary string
         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"
@@ -155,11 +175,12 @@ def analyze_excel_file(file_path: str) -> str:
 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}"
+    if not os.path.exists(file_path): return f"Error: Image file not found at path: {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 a clear message if no text found, otherwise return extracted text
         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."
@@ -188,133 +209,224 @@ def divide(a: float, b: float) -> float | str:
     if b == 0: return "Error: Cannot divide by zero."
     return a / b
 
-# --- Tool List & LLM Binding ---
+# --- Compile list of all tools ---
 tools = [ search_tool, web_browser, download_file_from_url, analyze_csv_file,
           analyze_excel_file, extract_text_from_image, add, subtract, multiply, divide ]
+
+# --- Bind tools to the LLM ---
 llm_with_tools = llm.bind_tools(tools)
 print(f"Agent initialized with {len(tools)} tools.")
 
 # ==============================================================================
 # Node Definitions
 # ==============================================================================
+print("Defining graph nodes...")
 
-# --- Agent Node (LLM Call) ---
+# --- Agent Node ---
 def call_agent_node(state: AgentState) -> dict:
-    """Calls the LLM to decide the next step/response, increments iteration count."""
+    """Invokes the LLM with current state to decide the next step."""
     print(f"\n--- [Node] Agent thinking... (Iteration {state['iterations']}) ---")
-    MAX_ITERATIONS = 10 # Max steps for the entire task
+    MAX_ITERATIONS = 10 # Max steps allowed for the 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:
+        # Call the LLM
         response = llm_with_tools.invoke(state['messages'])
         print("--- [Node] AI Response/Action ---")
-        response.pretty_print()
-        # Important: Return the iterations incremented
+        response.pretty_print() # Log the LLM's thoughts and actions
+        # Return the response message and incremented iteration count
         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
+        traceback.print_exc() # Print full traceback for debugging LLM errors
+        # Return an error message and set error state
+        return {"messages": [AIMessage(content=f"Sorry, I encountered an error: {error_message}")], "error": error_message, "iterations": current_iterations + 1}
 
-# --- Tool Node (Executor) ---
+# --- Tool Node ---
+# Use the prebuilt ToolNode to handle execution of the bound tools
 tool_node = ToolNode(tools)
 
 # ==============================================================================
 # Graph Construction (Non-conversational)
 # ==============================================================================
+print("Building agent graph...")
 builder = StateGraph(AgentState)
 
-# Add nodes
+# Add the agent and tool nodes
 builder.add_node("agent", call_agent_node)
 builder.add_node("tools", tool_node)
 
-# Define edges
+# Set the entry point
 builder.add_edge(START, "agent")
 
-# Conditional edge after agent thinks
+# Define the conditional logic after the agent node runs
 builder.add_conditional_edges(
     "agent",
-    tools_condition, # Check if the last message has tool calls
+    tools_condition, # Built-in function checks if the last message has tool_calls
     {
-        "tools": "tools", # If yes, execute tools
-        END: END          # If no, END the graph
+        "tools": "tools", # If tool calls exist, route to the tools node
+        END: END          # If no tool calls, the agent is done, route to END
     }
 )
 
-# Edge after tools execute
-builder.add_edge("tools", "agent") # Loop back to agent to process tool results
+# Define the edge after the tools node runs
+builder.add_edge("tools", "agent") # Always return to the agent node to process tool results
 
-# Compile the graph
-graph = builder.compile()
-print("GAIA agent graph compiled with integrated tools.")
+# Compile the graph into a runnable object
+# NOTE: This compilation happens when the script is imported by app.py
+try:
+    graph = builder.compile()
+    print("GAIA agent graph compiled successfully.")
+except Exception as e:
+    print(f"ERROR: Failed to compile LangGraph graph: {e}")
+    traceback.print_exc()
+    # Raise or handle appropriately - app might fail to start if graph doesn't compile
+    raise
 
 # ==============================================================================
-# Execution (Single run for GAIA task)
+# Main Execution Function for GAIA Benchmark <<<< WRAPPER FUNCTION >>>>
 # ==============================================================================
-
-# --- 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).
+def answer_gaia_task(question: str, file_path: Optional[str] = None) -> str:
+    """
+    Runs the compiled GAIA agent graph for a given question and optional file path.
+    This is the main entry point expected by the benchmark runner.
+
+    Args:
+        question: The question text from the GAIA benchmark.
+        file_path: Optional path to a file associated with the question.
+
+    Returns:
+        A string containing the final answer extracted by the agent, or an error message.
+    """
+    # Ensure the compiled graph is available
+    if 'graph' not in globals():
+         return "Error: Agent graph was not compiled successfully."
+
+    print(f"\n{'='*20} Running Agent for GAIA Task {'='*20}")
+    print(f"Question: {question}")
+    file_context_info = ""
+    if file_path:
+        print(f"Associated File Path: {file_path}")
+        file_context_info = f"An associated file is provided at path: '{file_path}'. Your tools should use this path if they require a file path not explicitly mentioned in the question."
+
+    # Define the initial prompt sent to the agent
+    prompt_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).
+
+{file_context_info}
 
 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.
+2. If external files are mentioned (e.g., 'data.csv', 'image.png'), use the appropriate analysis tool directly on the provided file path/name. Assume files are accessible in the current directory unless a URL or the separate file path is given.
+3. If a URL is given for a file, use 'download_file_from_url' first, then analyze the downloaded file using its returned 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
-)
+Question: {question}
+"""
+
+    # Create the initial state for the graph run
+    initial_state = AgentState(
+        input_question=question,
+        messages=[HumanMessage(content=prompt_content)],
+        error=None,
+        iterations=0
+    )
+
+    final_answer = "Error: Agent execution did not complete successfully." # Default fallback
+
+    try:
+        # Invoke the compiled graph
+        final_state = graph.invoke(initial_state, config={"recursion_limit": 15}) # Set recursion limit
+
+        # Process the final state to extract the answer
+        if final_state:
+            if final_state.get("error"):
+                print(f"--- Agent stopped due to ERROR: {final_state['error']} ---")
+                final_answer = f"Error: {final_state['error']}"
+            # Check if the last message is an AIMessage and capture its content
+            elif final_state.get('messages') and isinstance(final_state['messages'][-1], AIMessage):
+                # Extract content from the last AI message - relies on prompt working
+                potential_answer = final_state['messages'][-1].content
+                print(f"--- Final Answer (from AI): {potential_answer} ---")
+                final_answer = potential_answer
+            else:
+                 print("--- Could not determine final answer (last message not AI or missing). Check logs. ---")
+                 # Log final state details for debugging
+                 print(f"Final State: Error={final_state.get('error')}, Iterations={final_state.get('iterations')}")
+
+    except Exception as e:
+        print(f"--- Graph execution failed unexpectedly: {e} ---")
+        traceback.print_exc()
+        final_answer = f"Error: Graph execution failed - {str(e)}"
+
+    print(f"{'='*20} Agent Run Finished {'='*20}")
+    # Return the final answer string
+    return str(final_answer)
 
-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
+# Local Testing Block (Optional)
 # ==============================================================================
-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'
+# This block allows you to test the agent by running final_agent.py directly.
+# It will not run when the script is imported by app.py in the Space.
+if __name__ == "__main__":
+    print("\n--- Running Local Test ---")
+    # --- Define Test Question ---
+    test_question = "What is the result of multiplying the number of rows (excluding the header) in 'data.csv' by the number found after the phrase 'total items:' in 'image.png'?"
+
+    # --- Create Dummy Files for Local Test ---
+    print("Creating dummy files for local test...")
+    dummy_files_created = True
+    try:
+        # Dummy CSV with 3 data rows + header
+        with open("data.csv", "w") as f:
+            f.write("Header1,Header2\nRow1Val1,Row1Val2\nRow2Val1,Row2Val2\nRow3Val1,Row3Val2")
+
+        # Dummy Image containing the required text
+        try:
+            img = Image.new('RGB', (300, 50), color = (255, 255, 255)) # White background
+            from PIL import ImageDraw, ImageFont # Import drawing tools locally
+            draw = ImageDraw.Draw(img)
+            # Use a basic font if specific ones aren't found
+            try: font = ImageFont.truetype("arial.ttf", 15)
+            except IOError: font = ImageFont.load_default()
+            draw.text((10,10), "Some random info... total items: 7 ... more text", fill=(0,0,0), font=font) # Black text
+            img.save("image.png")
+            print("Dummy data.csv and image.png created successfully.")
+        except ImportError:
+            print("Pillow/ImageDraw/ImageFont not installed. Cannot create dummy image file.")
+            dummy_files_created = False
+        except Exception as img_e:
+            print(f"Error creating dummy image: {img_e}")
+            dummy_files_created = False
+
+    except Exception as file_e:
+        print(f"Error creating dummy files: {file_e}")
+        dummy_files_created = False
+    # ---------------------------------------------
+
+    # --- Run the Test ---
+    if dummy_files_created:
+        # Call the main function, simulating how the benchmark runner would call it.
+        # For this specific question, file_path argument is None as paths are in the question text.
+        result = answer_gaia_task(question=test_question, file_path=None)
+
+        print(f"\n--- Local Test Result ---")
+        # Expected answer for dummy files: 3 data rows * 7 = 21
+        print(f"Returned Answer: {result}")
+        print(f"Expected Answer (for dummy files): 21")
     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.")
+        print("Skipping test execution due to issues creating dummy files.")
+
+    # --- Clean up Dummy Files ---
+    print("\nCleaning up dummy files...")
+    for dummy_file in ["data.csv", "image.png"]:
+        if os.path.exists(dummy_file):
+            try: os.remove(dummy_file)
+            except Exception as e: print(f"Could not remove {dummy_file}: {e}")
+    print("Dummy file cleanup attempted.")