Refactor agent architecture with wrapper pattern and modular structure

This commit restructures the agent codebase to support multiple agent
implementations through a flexible wrapper pattern, improving code
organization and maintainability.

MAJOR CHANGES:

1. Agent Architecture Refactoring:
- Created MyGAIAAgents wrapper class for managing multiple agent types
- Renamed MyLangGraphAgent → LangGraphAgent for consistency
- Moved LangGraphAgent to dedicated langgraphagent.py module
- Renamed internal helper methods to private (prefixed with _)
- Added ACTIVE_AGENT config variable to switch between agent types

2. New Files:
- langgraphagent.py: Standalone LangGraphAgent implementation with
private methods (_create_llm_client, _init_questions, _assistant,
_should_continue, _build_graph)

3. Code Organization Improvements:
- agents.py: Now contains only the MyGAIAAgents wrapper (35 lines vs 337)
- Better separation of concerns between wrapper and implementation
- Cleaner import structure across the codebase

4. API Refactoring (app.py):
- Moved ResultFormatter.format_for_api() call into submit_and_score()
- submit_and_score() now accepts raw results instead of formatted payload
- run_and_submit_all() simplified - passes raw results directly
- Renamed results_for_display → logs_for_display in run_test_code()
- Better responsibility distribution between functions

5. Configuration (config.py):
- Added ACTIVE_AGENT = "LangGraph" configuration variable
- Supports future agent types: ReActLangGraph, LLamaIndex, SMOL

6. Updated Imports:
- agent_runner.py: MyLangGraphAgent → MyGAIAAgents
- app.py: Updated docstrings and removed unused import

7. Documentation (README.md):
- Updated code examples to use MyGAIAAgents wrapper
- Replaced "Change LLM Provider" section with "Change Agent Type"
- Documents ACTIVE_AGENT configuration usage

BENEFITS:
- Easier to add new agent implementations (just add to wrapper)
- Better code modularity and single responsibility principle
- Cleaner API boundaries between functions
- Improved testability with separated concerns
- Future-ready for multiple agent architectures

FILES CHANGED:
- Modified: README.md, agent_runner.py, agents.py, app.py, config.py
- Added: langgraphagent.py

README.md

@@ -113,10 +113,10 @@ Edit the question indices in [app.py:196](app.py#L196) to customize which questi
 ### Using the Agent Programmatically
 
 ```python
-from agents import MyLangGraphAgent
+from agents import MyGAIAAgents
 
-# Initialize agent
-agent = MyLangGraphAgent()
+# Initialize agent (automatically uses ACTIVE_AGENT from config)
+agent = MyGAIAAgents()
 
 # Ask a question
 answer = agent("What is the capital of France?")
@@ -174,19 +174,17 @@ The agent follows strict output formatting rules defined in [system_prompt.py](s
 
 ## Configuration
 
-### Change LLM Provider
+### Change Agent Type
 
-Edit [agents.py:52](agents.py#L52) in the `create_llm_client` method:
+Edit the `ACTIVE_AGENT` variable in [config.py:32](config.py#L32):
 
 ```python
-# Use Google Gemini (default)
-agent = MyLangGraphAgent()
-
-# Use Hugging Face models
-def create_llm_client(self, model_provider: str = "huggingface"):
-    # ...
+# Valid values: "LangGraph", "ReActLangGraph", "LLamaIndex", "SMOL"
+ACTIVE_AGENT = "LangGraph"  # Currently only LangGraph is implemented
 ```
 
+The `MyGAIAAgents` wrapper class will automatically instantiate the correct agent based on this configuration.
+
 ### Adjust Step Limits
 
 Modify the maximum iteration count in [agents.py:169](agents.py#L169):

agent_runner.py

@@ -2,7 +2,7 @@
 
 from typing import Optional, Tuple, List, Dict
 from colorama import Fore, Style
-from agents import MyLangGraphAgent
+from agents import MyGAIAAgents
 import config
 
 
@@ -16,7 +16,7 @@ class AgentRunner:
     def initialize_agent(self) -> bool:
         """Initialize the agent. Returns True if successful."""
         try:
-            self.agent = MyLangGraphAgent()
+            self.agent = MyGAIAAgents()
             return True
         except Exception as e:
             print(f"{Fore.RED}Error instantiating agent: {e}{Style.RESET_ALL}")

agents.py

@@ -1,305 +1,35 @@
-import os
-import logging
-import warnings
-import re
-import time
+"""Agent wrapper module for GAIA Benchmark."""
 
-# Suppress TensorFlow/Keras warnings
-os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
-logging.getLogger('tensorflow').setLevel(logging.ERROR)
-warnings.filterwarnings('ignore', module='tensorflow')
-warnings.filterwarnings('ignore', module='tf_keras')
-
-from typing import TypedDict, Optional, List, Annotated
-from langchain_core.messages import HumanMessage, SystemMessage
-from langgraph.graph import MessagesState, StateGraph, START, END
-from langgraph.graph.message import add_messages
-from langgraph.prebuilt import tools_condition
-from langgraph.prebuilt import ToolNode
-from langchain_google_genai import ChatGoogleGenerativeAI
-from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
-
-from custom_tools import get_custom_tools_list
-from system_prompt import SYSTEM_PROMPT
 import config
-
-# Suppress BeautifulSoup GuessedAtParserWarning
-try:
-    from bs4 import GuessedAtParserWarning
-    warnings.filterwarnings('ignore', category=GuessedAtParserWarning)
-except ImportError:
-    pass
+from langgraphagent import LangGraphAgent
 
 
-class AgentState(TypedDict):
-    question: str
-    messages: Annotated[list , add_messages]   # for LangGraph
-    answer: str
-    step_count: int  # Track number of iterations to prevent infinite loops
-    file_name: str  # Optional file name for questions that reference files
+class MyGAIAAgents:
+    """Wrapper class to manage multiple agent implementations.
 
-
-class MyLangGraphAgent:
+    This class provides a unified interface for different agent types.
+    The active agent is determined by the ACTIVE_AGENT configuration.
+    """
 
     def __init__(self):
-        # Validate API keys
-        if not os.getenv("GOOGLE_API_KEY"):
-            print("WARNING: GOOGLE_API_KEY not found - analyze_youtube_video will fail")
-
-        self.tools = get_custom_tools_list()
-        self.llm_client_with_tools = self.create_llm_client()
-        self.graph = self.build_graph()
-
-    def create_llm_client(self, model_provider: str = "google"):
-        """Create and return the LLM client with tools bound based on the model provider."""
-
-        if model_provider == "google":
-            apikey = os.getenv("GOOGLE_API_KEY")
-
-            return ChatGoogleGenerativeAI(
-                model="gemini-2.5-flash",  # Changed from gemini-2.5-flash-lite - better tool calling
-                temperature=0,
-                api_key=apikey,
-                timeout=60  # Add timeout to prevent hanging
-                ).bind_tools(self.tools)
-
-        elif model_provider == "huggingface":
-            LLM_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
-            apikey = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-
-            llmObject = HuggingFaceEndpoint(
-                repo_id=LLM_MODEL,
-                task="text-generation",
-                max_new_tokens=512,
-                temperature=0.7,
-                do_sample=False,
-                repetition_penalty=1.03,
-                huggingfacehub_api_token=apikey
-            )
-            return ChatHuggingFace(llm=llmObject).bind_tools(self.tools)
-
-    # Nodes
-    def init_questions(self, state: AgentState):
-        """Initialize the messages in the state with system prompt and user question."""
-
-        # Build the question message, including file name if available
-        question_content = state["question"]
-        if state.get("file_name"):
-            question_content += f'\n\nNote: This question references a file: {state["file_name"]}'
-
-        return {
-            "messages": [
-                    SystemMessage(content=SYSTEM_PROMPT),
-                    HumanMessage(content=question_content)
-                    ],
-            "step_count": 0  # Initialize step counter
-                }
-
-    def assistant(self, state: AgentState):
-        """Assistant node which calls the LLM with tools"""
-
-        # Track and log current step
-        current_step = state.get("step_count", 0) + 1
-        print(f"[STEP {current_step}] Calling assistant with {len(state['messages'])} messages")
-
-        # Invoke LLM with tools enabled, with retry logic for 504 errors
-        max_retries = config.MAX_RETRIES
-        delay = config.INITIAL_RETRY_DELAY
-
-        for attempt in range(max_retries + 1):
-            try:
-                response = self.llm_client_with_tools.invoke(state["messages"])
-                # Success - break out of retry loop
-                break
-            except Exception as e:
-                error_msg = str(e)
-
-                # Check if this is a 504 DEADLINE_EXCEEDED error
-                if "504" in error_msg and "DEADLINE_EXCEEDED" in error_msg:
-                    if attempt < max_retries:
-                        print(f"[RETRY] Attempt {attempt + 1}/{max_retries} failed with 504 DEADLINE_EXCEEDED")
-                        print(f"[RETRY] Retrying in {delay:.1f} seconds...")
-                        time.sleep(delay)
-                        delay *= config.RETRY_BACKOFF_FACTOR
-                        continue
-                    else:
-                        print(f"[RETRY] All {max_retries} retries exhausted for 504 error")
-                        print(f"[ERROR] LLM invocation failed after retries: {e}")
-                        return {
-                            "messages": [],
-                            "answer": f"Error: LLM failed after {max_retries} retries - {str(e)[:100]}",
-                            "step_count": current_step
-                        }
-                else:
-                    # Not a 504 error - fail immediately without retry
-                    print(f"[ERROR] LLM invocation failed: {e}")
-                    return {
-                        "messages": [],
-                        "answer": f"Error: LLM failed - {str(e)[:100]}",
-                        "step_count": current_step
-                    }
-
-        # If no tool calls, set the final answer
-        if not response.tool_calls:
-            content = response.content
-            print(f"[FINAL ANSWER] Agent produced answer (no tool calls)")
-
-            # Handle case where content is a list (e.g. mixed content from Gemini)
-            if isinstance(content, list):
-                # Extract text from list of content parts
-                text_parts = []
-                for item in content:
-                    if isinstance(item, dict) and 'text' in item:
-                        text_parts.append(item['text'])
-                    elif hasattr(item, 'text'):
-                        text_parts.append(item.text)
-                    else:
-                        text_parts.append(str(item))
-                content = " ".join(text_parts)
-            elif isinstance(content, dict) and 'text' in content:
-                # Handle single dict with 'text' field
-                content = content['text']
-            elif hasattr(content, 'text'):
-                # Handle object with text attribute
-                content = content.text
-            else:
-                # Fallback to string conversion
-                content = str(content)
-
-            # Clean up any remaining noise
-            content = content.strip()
-            print(f"[EXTRACTED TEXT] {content[:100]}{'...' if len(content) > 100 else ''}")
-
-            return {
-                "messages": [response],
-                "answer": content,
-                "step_count": current_step
-            }
-
-        # Has tool calls, log them
-        print(f"[TOOL CALLS] Agent requesting {len(response.tool_calls)} tool(s):")
-        for tc in response.tool_calls:
-            print(f"  - {tc['name']}")
+        """Initialize the wrapper with the active agent based on config."""
+        active_agent = config.ACTIVE_AGENT
 
-        return {
-            "messages": [response],
-            "step_count": current_step
-        }
-
-
-    def should_continue(self, state: AgentState):
-        """Check if we should continue or stop based on step count and other conditions."""
-
-        step_count = state.get("step_count", 0)
-
-        # Stop if we've exceeded maximum steps
-        if step_count >= 40:  # Increased from 25 to handle complex multi-step reasoning
-            print(f"[WARNING] Max steps (40) reached, forcing termination")
-            # Force a final answer if we don't have one
-            if not state.get("answer"):
-                state["answer"] = "Error: Maximum iteration limit reached"
-            return END
-
-        # Otherwise use the default tools_condition
-        return tools_condition(state)
-
-
-    def build_graph(self):
-        """Build and return the Compiled Graph for the agent."""
-
-        graph = StateGraph(AgentState)
-
-        # Build graph
-        graph.add_node("init", self.init_questions)
-        graph.add_node("assistant", self.assistant)
-        graph.add_node("tools", ToolNode(self.tools))
-        graph.add_edge(START, "init")
-        graph.add_edge("init", "assistant")
-        graph.add_conditional_edges(
-            "assistant",
-            # Use custom should_continue instead of tools_condition
-            self.should_continue,
-        )
-        graph.add_edge("tools", "assistant")
-        # Compile graph
-        return graph.compile()
+        if active_agent == "LangGraph":
+            self.agent = LangGraphAgent()
+        else:
+            # Default to LangGraph if unknown agent type
+            print(f"[WARNING] Unknown agent type '{active_agent}', defaulting to LangGraph")
+            self.agent = LangGraphAgent()
 
     def __call__(self, question: str, file_name: str = None) -> str:
-        """Invoke the agent graph with the given question and return the final answer.
+        """Invoke the active agent with the given question.
 
         Args:
             question: The question to answer
             file_name: Optional file name if the question references a file
-        """
-
-        print(f"\n{'='*60}")
-        print(f"[AGENT START] Question: {question}")
-        if file_name:
-            print(f"[FILE] {file_name}")
-        print(f"{'='*60}")
 
-        start_time = time.time()
-
-        try:
-            response = self.graph.invoke(
-                {"question": question, "messages": [], "answer": None, "step_count": 0, "file_name": file_name or ""},
-                config={"recursion_limit": 80}  # Must be >= 2x step limit (40 * 2 = 80)
-            )
-
-            elapsed_time = time.time() - start_time
-            print(f"[AGENT COMPLETE] Time: {elapsed_time:.2f}s")
-            print(f"{'='*60}\n")
-
-            answer = response.get("answer")
-            if answer is None:
-                print("[WARNING] Agent completed but returned None as answer")
-                return "Error: No answer generated"
-
-            # Final safety check: ensure answer is plain text string
-            if isinstance(answer, dict):
-                # If it's a dict, try to extract text field
-                if 'text' in answer:
-                    answer = answer['text']
-                else:
-                    answer = str(answer)
-                print(f"[WARNING] Answer was dict, extracted: {answer[:100]}")
-            elif isinstance(answer, list):
-                # If it's a list, extract text from each item
-                text_parts = []
-                for item in answer:
-                    if isinstance(item, dict) and 'text' in item:
-                        text_parts.append(item['text'])
-                    else:
-                        text_parts.append(str(item))
-                answer = " ".join(text_parts)
-                print(f"[WARNING] Answer was list, extracted: {answer[:100]}")
-            elif not isinstance(answer, str):
-                # Convert to string if it's any other type
-                answer = str(answer)
-                print(f"[WARNING] Answer was {type(answer)}, converted to string")
-
-            answer = answer.strip()
-
-            # Additional validation for numerical answers
-            # Remove common formatting issues that break exact matching
-            if answer:
-                # Remove comma separators from numbers (e.g., "1,000" -> "1000")
-                if ',' in answer and answer.replace(',', '').replace('.', '').isdigit():
-                    answer = answer.replace(',', '')
-                    print(f"[VALIDATION] Removed comma separators from answer")
-
-                # Ensure no trailing/leading whitespace or punctuation
-                answer = answer.strip().rstrip('.')
-
-                # Log if answer looks suspicious (for debugging)
-                if any(char in answer for char in ['{', '}', '[', ']', '`', '*', '#']):
-                    print(f"[WARNING] Answer contains suspicious formatting characters: {answer[:100]}")
-
-            print(f"[FINAL ANSWER] {answer}")
-            return answer
-
-        except Exception as e:
-            elapsed_time = time.time() - start_time
-            print(f"[AGENT ERROR] Failed after {elapsed_time:.2f}s: {e}")
-            print(f"{'='*60}\n")
-            return f"Error: {str(e)[:100]}"
+        Returns:
+            The agent's answer as a string
+        """
+        return self.agent(question, file_name)

app.py

@@ -13,8 +13,7 @@ init(autoreset=True)
 # Import configuration
 import config
 
-# Import agent-related code from agents module
-from agents import MyLangGraphAgent
+# Agent-related code is imported via agent_runner module
 # Import Gradio UI creation function
 from gradioapp import create_ui
 # Import scoring function for answer verification
@@ -40,13 +39,13 @@ def _submit_to_server(submit_url: str, submission_data: dict) -> dict:
     response.raise_for_status()
     return response.json()
 
-def submit_and_score(username: str, answers_payload: list) -> str:
+def submit_and_score(username: str, results: list) -> str:
     """
     Submit answers to the GAIA scoring server and return status message.
 
     Args:
         username: Hugging Face username for submission
-        answers_payload: List of dicts with {"task_id": str, "submitted_answer": str}
+        results: List of tuples (task_id, question_text, answer)
 
     Returns:
         str: Status message (success or error details)
@@ -59,6 +58,14 @@ def submit_and_score(username: str, answers_payload: list) -> str:
         print(error_msg)
         return error_msg
 
+    # Format results for API submission
+    answers_payload = ResultFormatter.format_for_api(results)
+
+    if not answers_payload:
+        error_msg = "No answers to submit."
+        print(error_msg)
+        return error_msg
+
     space_id = config.SPACE_ID
     submit_url = f"{config.DEFAULT_API_URL}/submit"
     agent_code = f"https://huggingface.co/spaces/{space_id}/tree/main"
@@ -118,7 +125,7 @@ def submit_and_score(username: str, answers_payload: list) -> str:
 
 def run_and_submit_all(username: str) -> tuple:
     """
-    Fetches all questions, runs the MyLangGraphAgent on them, submits all answers,
+    Fetches all questions, runs the GAIA agent on them, submits all answers,
     and displays the results.
 
     Returns:
@@ -142,16 +149,11 @@ def run_and_submit_all(username: str) -> tuple:
     if results is None:
         return "Error initializing agent.", None
 
-    # Format data structures: one for API submission, one for UI display
-    answers_for_api = ResultFormatter.format_for_api(results)
-    results_for_display = ResultFormatter.format_for_display(results)
-
-    if not answers_for_api:
-        print("Agent did not produce any answers to submit.")
-        return "Agent did not produce any answers to submit.", pd.DataFrame(results_for_display)
+    # Submit answers and get score (formatting happens inside submit_and_score)
+    status_message = submit_and_score(username, results)
 
-    # Submit answers and get score
-    status_message = submit_and_score(username, answers_for_api)
+    # Format results for UI display
+    results_for_display = ResultFormatter.format_for_display(results)
     results_df = pd.DataFrame(results_for_display)
     return status_message, results_df
 
@@ -239,8 +241,8 @@ def run_test_code(filter=None) -> pd.DataFrame:
         pd.DataFrame: Results and verification output
     """
     start_time = time.time()
-    results_for_display = []
-    results_for_display.append("=== Processing Example Questions One by One ===")
+    logs_for_display = []
+    logs_for_display.append("=== Processing Example Questions One by One ===")
 
     # Fetch questions (OFFLINE for testing)
     try:
@@ -263,10 +265,10 @@ def run_test_code(filter=None) -> pd.DataFrame:
     # Apply filter or use all questions
     if filter is not None:
         questions_to_process = [questions_data[i] for i in filter]
-        results_for_display.append(f"Testing {len(questions_to_process)} selected questions (indices: {filter})")
+        logs_for_display.append(f"Testing {len(questions_to_process)} selected questions (indices: {filter})")
     else:
         questions_to_process = questions_data
-        results_for_display.append(f"Testing all {len(questions_to_process)} questions")
+        logs_for_display.append(f"Testing all {len(questions_to_process)} questions")
 
     # Run agent on selected questions
     results = AgentRunner().run_on_questions(questions_to_process)
@@ -274,15 +276,15 @@ def run_test_code(filter=None) -> pd.DataFrame:
     if results is None:
         return pd.DataFrame(["Error initializing agent."])
 
-    results_for_display.append("\n=== Completed Example Questions ===")
+    logs_for_display.append("\n=== Completed Example Questions ===")
 
     # Calculate runtime
     elapsed_time = time.time() - start_time
     minutes = int(elapsed_time // 60)
     seconds = int(elapsed_time % 60)
 
-    verify_answers(results, results_for_display, runtime=(minutes, seconds))
-    return pd.DataFrame(results_for_display)
+    verify_answers(results, logs_for_display, runtime=(minutes, seconds))
+    return pd.DataFrame(logs_for_display)
 
 
 def main() -> None:

config.py

@@ -29,6 +29,8 @@ SPACE_HOST = os.getenv("SPACE_HOST")
 SPACE_ID = os.getenv("SPACE_ID")
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 
+ACTIVE_AGENT = "LangGraph" # Valid vales are ReActLangGraph, LLamaIndex, LangGraph, SMOL
+
 # Model Configuration
 GEMINI_MODEL = "gemini-2.5-flash"
 GEMINI_TEMPERATURE = 0

langgraphagent.py

@@ -0,0 +1,305 @@
+import os
+import logging
+import warnings
+import re
+import time
+
+# Suppress TensorFlow/Keras warnings
+os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
+logging.getLogger('tensorflow').setLevel(logging.ERROR)
+warnings.filterwarnings('ignore', module='tensorflow')
+warnings.filterwarnings('ignore', module='tf_keras')
+
+from typing import TypedDict, Optional, List, Annotated
+from langchain_core.messages import HumanMessage, SystemMessage
+from langgraph.graph import MessagesState, StateGraph, START, END
+from langgraph.graph.message import add_messages
+from langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_huggingface import ChatHuggingFace, HuggingFaceEndpoint
+
+from custom_tools import get_custom_tools_list
+from system_prompt import SYSTEM_PROMPT
+import config
+
+# Suppress BeautifulSoup GuessedAtParserWarning
+try:
+    from bs4 import GuessedAtParserWarning
+    warnings.filterwarnings('ignore', category=GuessedAtParserWarning)
+except ImportError:
+    pass
+
+
+class AgentState(TypedDict):
+    question: str
+    messages: Annotated[list , add_messages]   # for LangGraph
+    answer: str
+    step_count: int  # Track number of iterations to prevent infinite loops
+    file_name: str  # Optional file name for questions that reference files
+
+
+class LangGraphAgent:
+
+    def __init__(self):
+        # Validate API keys
+        if not os.getenv("GOOGLE_API_KEY"):
+            print("WARNING: GOOGLE_API_KEY not found - analyze_youtube_video will fail")
+
+        self.tools = get_custom_tools_list()
+        self.llm_client_with_tools = self._create_llm_client()
+        self.graph = self._build_graph()
+
+    def _create_llm_client(self, model_provider: str = "google"):
+        """Create and return the LLM client with tools bound based on the model provider."""
+
+        if model_provider == "google":
+            apikey = os.getenv("GOOGLE_API_KEY")
+
+            return ChatGoogleGenerativeAI(
+                model="gemini-2.5-flash",  # Changed from gemini-2.5-flash-lite - better tool calling
+                temperature=0,
+                api_key=apikey,
+                timeout=60  # Add timeout to prevent hanging
+                ).bind_tools(self.tools)
+
+        elif model_provider == "huggingface":
+            LLM_MODEL = "meta-llama/Llama-3.1-8B-Instruct"
+            apikey = os.getenv("HUGGINGFACEHUB_API_TOKEN")
+
+            llmObject = HuggingFaceEndpoint(
+                repo_id=LLM_MODEL,
+                task="text-generation",
+                max_new_tokens=512,
+                temperature=0.7,
+                do_sample=False,
+                repetition_penalty=1.03,
+                huggingfacehub_api_token=apikey
+            )
+            return ChatHuggingFace(llm=llmObject).bind_tools(self.tools)
+
+    # Nodes
+    def _init_questions(self, state: AgentState):
+        """Initialize the messages in the state with system prompt and user question."""
+
+        # Build the question message, including file name if available
+        question_content = state["question"]
+        if state.get("file_name"):
+            question_content += f'\n\nNote: This question references a file: {state["file_name"]}'
+
+        return {
+            "messages": [
+                    SystemMessage(content=SYSTEM_PROMPT),
+                    HumanMessage(content=question_content)
+                    ],
+            "step_count": 0  # Initialize step counter
+                }
+
+    def _assistant(self, state: AgentState):
+        """Assistant node which calls the LLM with tools"""
+
+        # Track and log current step
+        current_step = state.get("step_count", 0) + 1
+        print(f"[STEP {current_step}] Calling assistant with {len(state['messages'])} messages")
+
+        # Invoke LLM with tools enabled, with retry logic for 504 errors
+        max_retries = config.MAX_RETRIES
+        delay = config.INITIAL_RETRY_DELAY
+
+        for attempt in range(max_retries + 1):
+            try:
+                response = self.llm_client_with_tools.invoke(state["messages"])
+                # Success - break out of retry loop
+                break
+            except Exception as e:
+                error_msg = str(e)
+
+                # Check if this is a 504 DEADLINE_EXCEEDED error
+                if "504" in error_msg and "DEADLINE_EXCEEDED" in error_msg:
+                    if attempt < max_retries:
+                        print(f"[RETRY] Attempt {attempt + 1}/{max_retries} failed with 504 DEADLINE_EXCEEDED")
+                        print(f"[RETRY] Retrying in {delay:.1f} seconds...")
+                        time.sleep(delay)
+                        delay *= config.RETRY_BACKOFF_FACTOR
+                        continue
+                    else:
+                        print(f"[RETRY] All {max_retries} retries exhausted for 504 error")
+                        print(f"[ERROR] LLM invocation failed after retries: {e}")
+                        return {
+                            "messages": [],
+                            "answer": f"Error: LLM failed after {max_retries} retries - {str(e)[:100]}",
+                            "step_count": current_step
+                        }
+                else:
+                    # Not a 504 error - fail immediately without retry
+                    print(f"[ERROR] LLM invocation failed: {e}")
+                    return {
+                        "messages": [],
+                        "answer": f"Error: LLM failed - {str(e)[:100]}",
+                        "step_count": current_step
+                    }
+
+        # If no tool calls, set the final answer
+        if not response.tool_calls:
+            content = response.content
+            print(f"[FINAL ANSWER] Agent produced answer (no tool calls)")
+
+            # Handle case where content is a list (e.g. mixed content from Gemini)
+            if isinstance(content, list):
+                # Extract text from list of content parts
+                text_parts = []
+                for item in content:
+                    if isinstance(item, dict) and 'text' in item:
+                        text_parts.append(item['text'])
+                    elif hasattr(item, 'text'):
+                        text_parts.append(item.text)
+                    else:
+                        text_parts.append(str(item))
+                content = " ".join(text_parts)
+            elif isinstance(content, dict) and 'text' in content:
+                # Handle single dict with 'text' field
+                content = content['text']
+            elif hasattr(content, 'text'):
+                # Handle object with text attribute
+                content = content.text
+            else:
+                # Fallback to string conversion
+                content = str(content)
+
+            # Clean up any remaining noise
+            content = content.strip()
+            print(f"[EXTRACTED TEXT] {content[:100]}{'...' if len(content) > 100 else ''}")
+
+            return {
+                "messages": [response],
+                "answer": content,
+                "step_count": current_step
+            }
+
+        # Has tool calls, log them
+        print(f"[TOOL CALLS] Agent requesting {len(response.tool_calls)} tool(s):")
+        for tc in response.tool_calls:
+            print(f"  - {tc['name']}")
+
+        return {
+            "messages": [response],
+            "step_count": current_step
+        }
+
+
+    def _should_continue(self, state: AgentState):
+        """Check if we should continue or stop based on step count and other conditions."""
+
+        step_count = state.get("step_count", 0)
+
+        # Stop if we've exceeded maximum steps
+        if step_count >= 40:  # Increased from 25 to handle complex multi-step reasoning
+            print(f"[WARNING] Max steps (40) reached, forcing termination")
+            # Force a final answer if we don't have one
+            if not state.get("answer"):
+                state["answer"] = "Error: Maximum iteration limit reached"
+            return END
+
+        # Otherwise use the default tools_condition
+        return tools_condition(state)
+
+
+    def _build_graph(self):
+        """Build and return the Compiled Graph for the agent."""
+
+        graph = StateGraph(AgentState)
+
+        # Build graph
+        graph.add_node("init", self._init_questions)
+        graph.add_node("assistant", self._assistant)
+        graph.add_node("tools", ToolNode(self.tools))
+        graph.add_edge(START, "init")
+        graph.add_edge("init", "assistant")
+        graph.add_conditional_edges(
+            "assistant",
+            # Use custom should_continue instead of tools_condition
+            self._should_continue,
+        )
+        graph.add_edge("tools", "assistant")
+        # Compile graph
+        return graph.compile()
+
+    def __call__(self, question: str, file_name: str = None) -> str:
+        """Invoke the agent graph with the given question and return the final answer.
+
+        Args:
+            question: The question to answer
+            file_name: Optional file name if the question references a file
+        """
+
+        print(f"\n{'='*60}")
+        print(f"[AGENT START] Question: {question}")
+        if file_name:
+            print(f"[FILE] {file_name}")
+        print(f"{'='*60}")
+
+        start_time = time.time()
+
+        try:
+            response = self.graph.invoke(
+                {"question": question, "messages": [], "answer": None, "step_count": 0, "file_name": file_name or ""},
+                config={"recursion_limit": 80}  # Must be >= 2x step limit (40 * 2 = 80)
+            )
+
+            elapsed_time = time.time() - start_time
+            print(f"[AGENT COMPLETE] Time: {elapsed_time:.2f}s")
+            print(f"{'='*60}\n")
+
+            answer = response.get("answer")
+            if answer is None:
+                print("[WARNING] Agent completed but returned None as answer")
+                return "Error: No answer generated"
+
+            # Final safety check: ensure answer is plain text string
+            if isinstance(answer, dict):
+                # If it's a dict, try to extract text field
+                if 'text' in answer:
+                    answer = answer['text']
+                else:
+                    answer = str(answer)
+                print(f"[WARNING] Answer was dict, extracted: {answer[:100]}")
+            elif isinstance(answer, list):
+                # If it's a list, extract text from each item
+                text_parts = []
+                for item in answer:
+                    if isinstance(item, dict) and 'text' in item:
+                        text_parts.append(item['text'])
+                    else:
+                        text_parts.append(str(item))
+                answer = " ".join(text_parts)
+                print(f"[WARNING] Answer was list, extracted: {answer[:100]}")
+            elif not isinstance(answer, str):
+                # Convert to string if it's any other type
+                answer = str(answer)
+                print(f"[WARNING] Answer was {type(answer)}, converted to string")
+
+            answer = answer.strip()
+
+            # Additional validation for numerical answers
+            # Remove common formatting issues that break exact matching
+            if answer:
+                # Remove comma separators from numbers (e.g., "1,000" -> "1000")
+                if ',' in answer and answer.replace(',', '').replace('.', '').isdigit():
+                    answer = answer.replace(',', '')
+                    print(f"[VALIDATION] Removed comma separators from answer")
+
+                # Ensure no trailing/leading whitespace or punctuation
+                answer = answer.strip().rstrip('.')
+
+                # Log if answer looks suspicious (for debugging)
+                if any(char in answer for char in ['{', '}', '[', ']', '`', '*', '#']):
+                    print(f"[WARNING] Answer contains suspicious formatting characters: {answer[:100]}")
+
+            print(f"[FINAL ANSWER] {answer}")
+            return answer
+
+        except Exception as e:
+            elapsed_time = time.time() - start_time
+            print(f"[AGENT ERROR] Failed after {elapsed_time:.2f}s: {e}")
+            print(f"{'='*60}\n")
+            return f"Error: {str(e)[:100]}"