Update final_agent.py

final_agent.py

@@ -20,56 +20,55 @@ 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_groq import ChatGroq
-from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_groq import ChatGroq # Using Groq
 from langchain_community.tools.tavily_search import TavilySearchResults
 
 # ==============================================================================
 # Environment Setup & LLM
 # ==============================================================================
 load_dotenv()
-gemini_api_key = os.getenv("GEMINI_API_KEY")
+# Removed Gemini Key handling
 tavily_api_key = os.getenv("TAVILY_API_KEY")
-groq_api_key = os,getenv("GROQ_API_KEY")
+groq_api_key = os.getenv("GROQ_API_KEY")
 
 # --- 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 NameError: pass
 # 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 Tavily search tool.")
-
-# --- Initialize LLM ---
-# Using the model specified in the user's code block
-# llm = ChatGoogleGenerativeAI(
-#     model="gemini-2.0-flash-lite", # As per user's last provided code
-#     google_api_key=gemini_api_key,
-#     temperature=0.1 # Low temperature for factual tasks
-# )
-llm = ChatGroq(
-    model="gemma2-9b-it",
-    api_key=groq_api_key,
-    temperature = 0.1
-)
-# print(f"LLM Initialized: {llm.model}")
+    raise ValueError("TAVILY_API_KEY not found in Space secrets. Required for search.")
+if not groq_api_key:
+    raise ValueError("GROQ_API_KEY not found in Space secrets. Required for Groq LLM.")
+
+# --- Initialize LLM (Using Groq) ---
+try:
+    llm = ChatGroq(
+        model="llama3-70b-8192", # Powerful model available on Groq, good for reasoning
+        # model="gemma2-9b-it", # Alternative if Llama3 causes issues
+        api_key=groq_api_key,
+        temperature=0.1 # Low temperature for factual tasks
+    )
+    print(f"LLM Initialized: Groq - {llm.model_name}")
+except Exception as e:
+     print(f"ERROR initializing Groq LLM: {e}")
+     traceback.print_exc()
+     raise # Stop if LLM fails to init
 
 # ==============================================================================
 # State Definition
 # ==============================================================================
 class AgentState(TypedDict):
     """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
+    input_question: str
+    messages: Annotated[List[BaseMessage], add_messages]
+    error: Optional[str]
+    iterations: int
 
 # ==============================================================================
 # Tools Definitions
@@ -91,22 +90,14 @@ def web_browser(url: str) -> str:
         response = requests.get(url, headers=headers, timeout=20)
         response.raise_for_status()
         response.encoding = response.apparent_encoding or 'utf-8'
-        # Configure html2text
-        h = html2text.HTML2Text(bodywidth=0)
-        h.ignore_links = True
-        h.ignore_images = True
-        # Convert HTML to text
+        h = html2text.HTML2Text(bodywidth=0); h.ignore_links = True; h.ignore_images = True
         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]"
+        if len(clean_text) > max_length: return clean_text[:max_length] + "\n\n... [Content Truncated]"
         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)}"
+    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
@@ -114,43 +105,30 @@ 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)
         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)}"
+    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} ---")
-    # 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"
-        summary += f"\nFirst 5 rows:\n{df.head().to_string()}\n"
+        df = pd.read_csv(file_path); summary = f"CSV Analysis Report for {os.path.basename(file_path)}:\n- Shape: {df.shape[0]} rows, {df.shape[1]} columns\n- Columns: {', '.join(df.columns)}\n\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."
+        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)}"
@@ -162,17 +140,10 @@ def analyze_excel_file(file_path: str) -> str:
     print(f"--- [Tool] Analyzing Excel: {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"
-        summary += f"\nFirst 5 rows:\n{df.head().to_string()}\n"
+        df = pd.read_excel(file_path, engine='openpyxl'); summary = f"Excel Analysis Report for {os.path.basename(file_path)} (First Sheet):\n- Shape: {df.shape[0]} rows, {df.shape[1]} columns\n- Columns: {', '.join(df.columns)}\n\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."
+        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)}"
@@ -184,10 +155,8 @@ def extract_text_from_image(file_path: str) -> str:
     print(f"--- [Tool] Extracting text from image: {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."
@@ -195,24 +164,14 @@ def extract_text_from_image(file_path: str) -> str:
 
 # --- 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
+def add(a: float, b: float) -> float: """Adds two numbers (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
+def subtract(a: float, b: float) -> float: """Subtracts the second number from the first (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
+def multiply(a: float, b: float) -> float: """Multiplies two numbers (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} ---")
+    """Divides the first number by the second (a / b). Handles division by zero."""
     if b == 0: return "Error: Cannot divide by zero."
     return a / b
 
@@ -221,6 +180,9 @@ 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 ---
+# Ensure LLM is initialized before binding
+if 'llm' not in globals():
+    raise RuntimeError("LLM was not initialized successfully before binding tools.")
 llm_with_tools = llm.bind_tools(tools)
 print(f"Agent initialized with {len(tools)} tools.")
 
@@ -239,21 +201,18 @@ def call_agent_node(state: AgentState) -> dict:
         print(f"Warning: Reached max iterations ({MAX_ITERATIONS}). Stopping.")
         return {"error": f"Max iterations ({MAX_ITERATIONS}) reached."}
     try:
-        # Call the LLM
+        # Ensure LLM is bound with tools before invoking
+        if 'llm_with_tools' not in globals():
+             return {"error": "LLM tools not bound."}
         response = llm_with_tools.invoke(state['messages'])
-        print("--- [Node] AI Response/Action ---")
-        response.pretty_print() # Log the LLM's thoughts and actions
-        # Return the response message and incremented iteration count
+        print("--- [Node] AI Response/Action ---"); response.pretty_print()
         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} ---")
-        traceback.print_exc() # Print full traceback for debugging LLM errors
-        # Return an error message and set error state
+        error_message = f"LLM invocation failed: {str(e)}"; print(f"--- [Node] ERROR: {error_message} ---")
+        traceback.print_exc()
         return {"messages": [AIMessage(content=f"Sorry, I encountered an error: {error_message}")], "error": error_message, "iterations": current_iterations + 1}
 
 # --- Tool Node ---
-# Use the prebuilt ToolNode to handle execution of the bound tools
 tool_node = ToolNode(tools)
 
 # ==============================================================================
@@ -261,37 +220,20 @@ tool_node = ToolNode(tools)
 # ==============================================================================
 print("Building agent graph...")
 builder = StateGraph(AgentState)
-
-# Add the agent and tool nodes
 builder.add_node("agent", call_agent_node)
 builder.add_node("tools", tool_node)
-
-# Set the entry point
 builder.add_edge(START, "agent")
+builder.add_conditional_edges("agent", tools_condition, {"tools": "tools", END: END})
+builder.add_edge("tools", "agent")
 
-# Define the conditional logic after the agent node runs
-builder.add_conditional_edges(
-    "agent",
-    tools_condition, # Built-in function checks if the last message has tool_calls
-    {
-        "tools": "tools", # If tool calls exist, route to the tools node
-        END: END          # If no tool calls, the agent is done, route to END
-    }
-)
-
-# 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 into a runnable object
-# NOTE: This compilation happens when the script is imported by app.py
+# Compile the graph globally so it's ready for the function call
 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
+    graph = None # Set graph to None if compilation fails
 
 # ==============================================================================
 # Main Execution Function for GAIA Benchmark <<<< WRAPPER FUNCTION >>>>
@@ -300,40 +242,33 @@ 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."
+    # Check if graph compilation was successful
+    if graph is None:
+         return "Error: Agent graph was not compiled successfully during setup."
 
     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."
+    file_context_info = f"An associated file is provided at path: '{file_path}'. Use this path if relevant." if file_path else ""
 
-    # 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).
+    # Define the initial prompt sent to the agent, incorporating strict formatting rules
+    prompt_content = f"""You are a precise AI assistant answering a specific question based *only* on information obtained using your tools.
 
 {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., '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: {question}
+1. Analyze the question: {question}
+2. Use tools (web_search, web_browser, download_file_from_url, analyze_csv_file, analyze_excel_file, extract_text_from_image, add, subtract, multiply, divide) ONLY if necessary to gather the specific information required. Assume local file paths mentioned in the question (like 'data.csv') are accessible.
+3. Synthesize the final answer from the gathered information.
+
+**CRITICAL OUTPUT FORMATTING RULES:**
+*   Your final response MUST be ONLY the answer, without any other text, explanations, or introductions.
+*   **Numbers:** Do not use commas (e.g., 1000). Do not include units ($ , %) unless explicitly asked for.
+*   **Strings:** Do not use articles (a, an, the) unless part of a required proper noun. Do not use abbreviations (e.g., write "Saint Petersburg") unless the abbreviation is the answer. Write digits as numerals (5).
+*   **Lists:** If a list is required, provide it as comma-separated values (e.g., apple,banana,cherry). Apply number/string rules to elements.
+*   If you cannot find the answer using the tools, output only the exact phrase: Information not found
+
+Provide ONLY the final answer according to these rules.
 """
 
     # Create the initial state for the graph run
@@ -355,15 +290,17 @@ Question: {question}
             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
+                # Basic cleanup for potential quotes added by LLM
+                if isinstance(potential_answer, str):
+                     if (potential_answer.startswith('"') and potential_answer.endswith('"')) or \
+                        (potential_answer.startswith("'") and potential_answer.endswith("'")):
+                          potential_answer = potential_answer[1:-1].strip()
                 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:
@@ -380,57 +317,33 @@ Question: {question}
 # Local Testing Block (Optional)
 # ==============================================================================
 # 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
+        with open("data.csv", "w") as f: f.write("Header1,Header2\nRow1Val1,Row1Val2\nRow2Val1,Row2Val2\nRow3Val1,Row3Val2")
         try:
-            img = Image.new('RGB', (300, 50), color = (255, 255, 255)) # White background
-            from PIL import ImageDraw, ImageFont # Import drawing tools locally
+            img = Image.new('RGB', (300, 50), color = (255, 255, 255))
+            from PIL import ImageDraw, ImageFont
             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
+            draw.text((10,10), "Some random info... total items: 7 ... more text", fill=(0,0,0), font=font)
             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 ---
+        except ImportError: print("Pillow/ImageDraw/ImageFont not installed. Cannot create dummy image."); 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
+
     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:
-        print("Skipping test execution due to issues creating dummy files.")
+    else: 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):