Update agent.py

agent.py

@@ -1,6 +1,6 @@
 import os
 from dotenv import load_dotenv
-from typing import List, Dict, Any, Optional, Literal
+from typing import List, Dict, Any, Optional
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.graph.message import add_messages
 from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage
@@ -10,83 +10,35 @@ from langgraph.prebuilt import tools_condition
 from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
 from langchain_core.tools import tool
 from langchain_community.document_loaders import WikipediaLoader
-from langchain_community.tools import YouTubeSearchTool
+from langchain_community.document_loaders import YoutubeLoader
 from langchain_google_genai import ChatGoogleGenerativeAI
 from langchain_tavily import TavilySearch
 import tempfile
 import pandas as pd
-import numpy as np
-import requests
-from urllib.parse import urlparse
-import uuid
-from PIL import Image, ImageDraw, ImageFont, ImageEnhance, ImageFilter
-import base64
-import io
 
 load_dotenv()
 
-# ============== SYSTEM PROMPTS FOR SPECIALIZED AGENTS ============== #
-
-COORDINATOR_SYSTEM_PROMPT = """You are a Coordinator Agent that orchestrates multiple specialized agents to solve complex tasks.
-
-Your role is to:
-1. Analyze incoming requests and determine which specialized agents are needed
-2. Break down complex tasks into subtasks for different agents
-3. Coordinate between agents when needed
-4. Synthesize final answers from multiple agent responses
-
-Available specialized agents:
-- Research Agent: Wikipedia, web search, YouTube search
-- Math Agent: Basic mathematical calculations
-- Data Analysis Agent: CSV/Excel analysis, OCR text extraction
-- Image Processing Agent: Image analysis, transformation, generation
-- File Management Agent: File operations, downloads, saves
-
-When you receive a task:
-1. THINK: What type of task is this? Which agents do I need?
-2. ROUTE: Send subtasks to appropriate agents
-3. COORDINATE: Manage dependencies between agent tasks
-4. SYNTHESIZE: Combine results into a final answer
-
-Always provide a clear, comprehensive final answer.
-"""
-
-RESEARCH_AGENT_PROMPT = """You are a Research Agent specialized in information gathering and search.
-
-Your expertise includes:
-- Wikipedia searches for encyclopedic information
-- Web searches for current information and facts
-- YouTube searches for video content
+# ReAct System Prompt
+REACT_SYSTEM_PROMPT = """You are a research assistant that uses ReAct (Reasoning + Acting) methodology. For each question, follow this systematic approach:
 
-Follow ReAct methodology:
-1. THINK: What information do I need to find?
-2. ACT: Use appropriate search tools systematically
-3. OBSERVE: Analyze and verify search results
-4. SYNTHESIZE: Provide comprehensive, accurate information
+**THINK**: First, analyze the question carefully. What type of information do you need? What tools might help?
 
-Be thorough in your research and cross-reference sources when possible.
-
-Always finish with: FINAL ANSWER: [YOUR FINAL ANSWER]
+**ACT**: Use available tools to gather information. Search thoroughly and verify facts from multiple sources when possible.
 
-Your final answer should be:
-- A number (without commas or units unless specified)
-- As few words as possible for strings (no articles, no abbreviations for cities, spell out digits)  
-- A comma-separated list following the above rules for each element
-"""
+**OBSERVE**: Analyze the results from your tools. Are they complete and reliable? Do you need more information?
 
-MATH_AGENT_PROMPT = """You are a Math Agent specialized in mathematical calculations and operations.
+**REASON**: Synthesize all information gathered. Check for consistency and identify any gaps or uncertainties.
 
-Your expertise includes:
-- Basic arithmetic operations (add, subtract, multiply, divide)
-- Mathematical reasoning and problem-solving
+**VERIFY**: Before providing your final answer, double-check your reasoning and ensure you have sufficient evidence.
 
-Follow ReAct methodology:
-1. THINK: What calculations are needed?
-2. ACT: Perform calculations systematically
-3. VERIFY: Double-check your work
-4. PROVIDE: Clear numerical answers
+For each question:
+1. Break down what you're looking for
+2. Use tools systematically to gather comprehensive information  
+3. Cross-reference information when possible
+4. Be honest about limitations - if you cannot find reliable information, say so
+5. Only provide confident answers when you have verified evidence
 
-Always show your work and verify calculations.
+When you cannot access certain content (videos, audio, images without tools), clearly state this limitation.
 
 Always finish with: FINAL ANSWER: [YOUR FINAL ANSWER]
 
@@ -94,87 +46,46 @@ Your final answer should be:
 - A number (without commas or units unless specified)
 - As few words as possible for strings (no articles, no abbreviations for cities, spell out digits)  
 - A comma-separated list following the above rules for each element
-"""
-
-DATA_ANALYSIS_AGENT_PROMPT = """You are a Data Analysis Agent specialized in processing and analyzing structured data.
-
-Your expertise includes:
-- CSV file analysis and statistics
-- Excel file processing
-- OCR text extraction from images
-- Data interpretation and insights
-
-Follow ReAct methodology:
-1. THINK: What type of data analysis is needed?
-2. ACT: Use appropriate analysis tools
-3. OBSERVE: Examine data patterns and statistics
-4. INTERPRET: Provide meaningful insights
-
-Focus on accuracy and provide clear data-driven insights.
-"""
-
-IMAGE_PROCESSING_AGENT_PROMPT = """You are an Image Processing Agent specialized in image analysis, manipulation, and generation.
-
-Your expertise includes:
-- Image analysis (properties, colors, content)
-- Image transformations (resize, rotate, crop, filters)
-- Drawing and annotation on images
-- Simple image generation
-- Combining multiple images
 
-Follow ReAct methodology:
-1. THINK: What image processing is required?
-2. ACT: Apply appropriate image operations
-3. OBSERVE: Verify results and quality
-4. DELIVER: Provide processed images with explanations
-
-Focus on quality and user requirements.
+Be thorough in your research but honest about uncertainty. Quality and accuracy are more important than speed.
 """
 
-FILE_MANAGEMENT_AGENT_PROMPT = """You are a File Management Agent specialized in file operations and data handling.
-
-Your expertise includes:
-- Saving and reading files
-- Downloading files from URLs
-- Downloading task files from APIs
-- File format handling
-
-Follow ReAct methodology:
-1. THINK: What file operations are needed?
-2. ACT: Perform file operations safely
-3. VERIFY: Confirm successful operations
-4. REPORT: Provide clear status and file paths
-
-Ensure secure and reliable file handling.
-"""
-
-# ============== TOOL DEFINITIONS (grouped by agent) ============== #
-
-# Math Agent Tools
 @tool
-def multiply(a: int, b: int) -> int:
-    """Multiply two numbers"""
+def multiply(a:int, b:int) -> int:
+    """
+    Multiply two numbers
+    """
     return a * b
 
 @tool
-def add(a: int, b: int) -> int:
-    """Add two numbers"""
+def add(a:int, b:int) -> int:
+    """
+    Add two numbers
+    """
     return a + b
 
 @tool
-def subtract(a: int, b: int) -> int:
-    """Subtract two numbers"""
+def subtract(a:int, b:int) -> int:
+    """
+    Subtract two numbers
+    """
     return a - b
 
 @tool
-def divide(a: int, b: int) -> float:
-    """Divide two numbers"""
+def divide(a:int, b:int) -> int:
+    """
+    Divide two numbers
+    """
     return a / b
 
-# Research Agent Tools
 @tool
 def wikidata_search(query: str) -> str:
-    """Search for information on Wikipedia and return maximum 2 results."""
+    """
+    Search for information on Wikipedia and return maximum 2 results.
+    
+    Args:
+        query: The search query.
+    """
     loader = WikipediaLoader(query=query, load_max_docs=2)
     docs = loader.load()
     formatted_search_docs = "\n\n---\n\n".join(
@@ -184,14 +95,49 @@ def wikidata_search(query: str) -> str:
         ])
     return {"wiki_results": formatted_search_docs}
 
-# Initialize search tools
-tavily_search_tool = TavilySearch(max_results=3, topic="general")
-youtube_search_tool = YouTubeSearchTool()
+# Initialize Tavily Search Tool
+tavily_search_tool = TavilySearch(
+    max_results=3,
+    topic="general",
+)
+
+@tool
+def load_youtube_transcript(url: str, add_video_info: bool = True, language: List[str] = ["en"], translation: str = "en") -> str:
+    """
+    Load transcript from a YouTube video URL.
+    
+    Args:
+        url: YouTube video URL
+        add_video_info: Whether to include video metadata
+        language: List of language codes in descending priority
+        translation: Language to translate transcript to
+    """
+    try:
+        loader = YoutubeLoader.from_youtube_url(
+            url, 
+            add_video_info=add_video_info,
+            language=language,
+            translation=translation
+        )
+        docs = loader.load()
+        
+        formatted_transcript = "\n\n---\n\n".join([
+            f'<Document source="{doc.metadata.get("source", "")}" title="{doc.metadata.get("title", "")}" author="{doc.metadata.get("author", "")}" length="{doc.metadata.get("length", "")}"/>\n{doc.page_content}\n</Document>'
+            for doc in docs
+        ])
+        
+        return {"youtube_transcript": formatted_transcript}
+    except Exception as e:
+        return f"Error loading YouTube transcript: {str(e)}"
 
-# File Management Agent Tools
 @tool
 def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
-    """Save content to a file and return the path."""
+    """
+    Save content to a file and return the path.
+    Args:
+        content (str): the content to save to the file
+        filename (str, optional): the name of the file. If not provided, a random name file will be created.
+    """
     temp_dir = tempfile.gettempdir()
     if filename is None:
         temp_file = tempfile.NamedTemporaryFile(delete=False, dir=temp_dir)
@@ -204,22 +150,32 @@ def save_and_read_file(content: str, filename: Optional[str] = None) -> str:
 
     return f"File saved to {filepath}. You can read this file to process its contents."
 
+
 @tool
 def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
-    """Download a file from a URL and save it to a temporary location."""
+    """
+    Download a file from a URL and save it to a temporary location.
+    Args:
+        url (str): the URL of the file to download.
+        filename (str, optional): the name of the file. If not provided, a random name file will be created.
+    """
     try:
+        # Parse URL to get filename if not provided
         if not filename:
             path = urlparse(url).path
             filename = os.path.basename(path)
             if not filename:
                 filename = f"downloaded_{uuid.uuid4().hex[:8]}"
 
+        # Create temporary file
         temp_dir = tempfile.gettempdir()
         filepath = os.path.join(temp_dir, filename)
 
+        # Download the file
         response = requests.get(url, stream=True)
         response.raise_for_status()
 
+        # Save the file
         with open(filepath, "wb") as f:
             for chunk in response.iter_content(chunk_size=8192):
                 f.write(chunk)
@@ -228,75 +184,100 @@ def download_file_from_url(url: str, filename: Optional[str] = None) -> str:
     except Exception as e:
         return f"Error downloading file: {str(e)}"
 
-@tool
-def download_task_file(task_id: str, api_url: str = "https://agents-course-unit4-scoring.hf.space") -> str:
-    """Download a file associated with a task from the evaluation API."""
-    try:
-        file_url = f"{api_url}/files/{task_id}"
-        temp_dir = tempfile.gettempdir()
-        filename = f"task_{task_id}.png"
-        filepath = os.path.join(temp_dir, filename)
-        
-        response = requests.get(file_url, stream=True)
-        response.raise_for_status()
-        
-        with open(filepath, "wb") as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                f.write(chunk)
-        
-        return f"Task file downloaded to {filepath}. You can now analyze this file."
-    except Exception as e:
-        return f"Error downloading task file: {str(e)}"
 
-# Data Analysis Agent Tools
 @tool
 def extract_text_from_image(image_path: str) -> str:
-    """Extract text from an image using OCR."""
+    """
+    Extract text from an image using OCR library pytesseract (if available).
+    Args:
+        image_path (str): the path to the image file.
+    """
     try:
-        import pytesseract
+        # Open the image
         image = Image.open(image_path)
+
+        # Extract text from the image
         text = pytesseract.image_to_string(image)
+
         return f"Extracted text from image:\n\n{text}"
     except Exception as e:
         return f"Error extracting text from image: {str(e)}"
 
+
 @tool
 def analyze_csv_file(file_path: str, query: str) -> str:
-    """Analyze a CSV file using pandas and answer a question about it."""
+    """
+    Analyze a CSV file using pandas and answer a question about it.
+    Args:
+        file_path (str): the path to the CSV file.
+        query (str): Question about the data
+    """
     try:
+        # Read the CSV file
         df = pd.read_csv(file_path)
+
+        # Run various analyses based on the query
         result = f"CSV file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
         result += f"Columns: {', '.join(df.columns)}\n\n"
+
+        # Add summary statistics
         result += "Summary statistics:\n"
         result += str(df.describe())
+
         return result
+
     except Exception as e:
         return f"Error analyzing CSV file: {str(e)}"
 
+
 @tool
 def analyze_excel_file(file_path: str, query: str) -> str:
-    """Analyze an Excel file using pandas and answer a question about it."""
+    """
+    Analyze an Excel file using pandas and answer a question about it.
+    Args:
+        file_path (str): the path to the Excel file.
+        query (str): Question about the data
+    """
     try:
+        # Read the Excel file
         df = pd.read_excel(file_path)
-        result = f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
+
+        # Run various analyses based on the query
+        result = (
+            f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
+        )
         result += f"Columns: {', '.join(df.columns)}\n\n"
+
+        # Add summary statistics
         result += "Summary statistics:\n"
         result += str(df.describe())
+
         return result
+
     except Exception as e:
         return f"Error analyzing Excel file: {str(e)}"
 
-# Image Processing Agent Tools - Helper functions
+
+### ============== IMAGE PROCESSING AND GENERATION TOOLS =============== ###
+import os
+import io
+import base64
+import uuid
+from PIL import Image, ImageDraw, ImageFont, ImageEnhance, ImageFilter
+
+# Helper functions for image processing
 def encode_image(image_path: str) -> str:
     """Convert an image file to base64 string."""
     with open(image_path, "rb") as image_file:
         return base64.b64encode(image_file.read()).decode("utf-8")
 
+
 def decode_image(base64_string: str) -> Image.Image:
     """Convert a base64 string to a PIL Image."""
     image_data = base64.b64decode(base64_string)
     return Image.open(io.BytesIO(image_data))
 
+
 def save_image(image: Image.Image, directory: str = "image_outputs") -> str:
     """Save a PIL Image to disk and return the path."""
     os.makedirs(directory, exist_ok=True)
@@ -307,7 +288,13 @@ def save_image(image: Image.Image, directory: str = "image_outputs") -> str:
 
 @tool
 def analyze_image(image_base64: str) -> Dict[str, Any]:
-    """Analyze basic properties of an image."""
+    """
+    Analyze basic properties of an image (size, mode, color analysis, thumbnail preview).
+    Args:
+        image_base64 (str): Base64 encoded image string
+    Returns:
+        Dictionary with analysis result
+    """
     try:
         img = decode_image(image_base64)
         width, height = img.size
@@ -340,29 +327,42 @@ def analyze_image(image_base64: str) -> Dict[str, Any]:
     except Exception as e:
         return {"error": str(e)}
 
+
 @tool
 def transform_image(
     image_base64: str, operation: str, params: Optional[Dict[str, Any]] = None
 ) -> Dict[str, Any]:
-    """Apply transformations: resize, rotate, crop, flip, brightness, contrast, blur, sharpen, grayscale."""
+    """
+    Apply transformations: resize, rotate, crop, flip, brightness, contrast, blur, sharpen, grayscale.
+    Args:
+        image_base64 (str): Base64 encoded input image
+        operation (str): Transformation operation
+        params (Dict[str, Any], optional): Parameters for the operation
+    Returns:
+        Dictionary with transformed image (base64)
+    """
     try:
         img = decode_image(image_base64)
         params = params or {}
 
         if operation == "resize":
-            img = img.resize((
-                params.get("width", img.width // 2),
-                params.get("height", img.height // 2),
-            ))
+            img = img.resize(
+                (
+                    params.get("width", img.width // 2),
+                    params.get("height", img.height // 2),
+                )
+            )
         elif operation == "rotate":
             img = img.rotate(params.get("angle", 90), expand=True)
         elif operation == "crop":
-            img = img.crop((
-                params.get("left", 0),
-                params.get("top", 0),
-                params.get("right", img.width),
-                params.get("bottom", img.height),
-            ))
+            img = img.crop(
+                (
+                    params.get("left", 0),
+                    params.get("top", 0),
+                    params.get("right", img.width),
+                    params.get("bottom", img.height),
+                )
+            )
         elif operation == "flip":
             if params.get("direction", "horizontal") == "horizontal":
                 img = img.transpose(Image.FLIP_LEFT_RIGHT)
@@ -388,11 +388,20 @@ def transform_image(
     except Exception as e:
         return {"error": str(e)}
 
+
 @tool
 def draw_on_image(
     image_base64: str, drawing_type: str, params: Dict[str, Any]
 ) -> Dict[str, Any]:
-    """Draw shapes (rectangle, circle, line) or text onto an image."""
+    """
+    Draw shapes (rectangle, circle, line) or text onto an image.
+    Args:
+        image_base64 (str): Base64 encoded input image
+        drawing_type (str): Drawing type
+        params (Dict[str, Any]): Drawing parameters
+    Returns:
+        Dictionary with result image (base64)
+    """
     try:
         img = decode_image(image_base64)
         draw = ImageDraw.Draw(img)
@@ -412,12 +421,16 @@ def draw_on_image(
                 width=params.get("width", 2),
             )
         elif drawing_type == "line":
-            draw.line((
-                params["start_x"],
-                params["start_y"],
-                params["end_x"],
-                params["end_y"],
-            ), fill=color, width=params.get("width", 2))
+            draw.line(
+                (
+                    params["start_x"],
+                    params["start_y"],
+                    params["end_x"],
+                    params["end_y"],
+                ),
+                fill=color,
+                width=params.get("width", 2),
+            )
         elif drawing_type == "text":
             font_size = params.get("font_size", 20)
             try:
@@ -440,6 +453,7 @@ def draw_on_image(
     except Exception as e:
         return {"error": str(e)}
 
+
 @tool
 def generate_simple_image(
     image_type: str,
@@ -447,7 +461,15 @@ def generate_simple_image(
     height: int = 500,
     params: Optional[Dict[str, Any]] = None,
 ) -> Dict[str, Any]:
-    """Generate a simple image (gradient, noise, pattern, chart)."""
+    """
+    Generate a simple image (gradient, noise, pattern, chart).
+    Args:
+        image_type (str): Type of image
+        width (int), height (int)
+        params (Dict[str, Any], optional): Specific parameters
+    Returns:
+        Dictionary with generated image (base64)
+    """
     try:
         params = params or {}
 
@@ -461,20 +483,33 @@ def generate_simple_image(
 
             if direction == "horizontal":
                 for x in range(width):
-                    r = int(start_color[0] + (end_color[0] - start_color[0]) * x / width)
-                    g = int(start_color[1] + (end_color[1] - start_color[1]) * x / width)
-                    b = int(start_color[2] + (end_color[2] - start_color[2]) * x / width)
+                    r = int(
+                        start_color[0] + (end_color[0] - start_color[0]) * x / width
+                    )
+                    g = int(
+                        start_color[1] + (end_color[1] - start_color[1]) * x / width
+                    )
+                    b = int(
+                        start_color[2] + (end_color[2] - start_color[2]) * x / width
+                    )
                     draw.line([(x, 0), (x, height)], fill=(r, g, b))
             else:
                 for y in range(height):
-                    r = int(start_color[0] + (end_color[0] - start_color[0]) * y / height)
-                    g = int(start_color[1] + (end_color[1] - start_color[1]) * y / height)
-                    b = int(start_color[2] + (end_color[2] - start_color[2]) * y / height)
+                    r = int(
+                        start_color[0] + (end_color[0] - start_color[0]) * y / height
+                    )
+                    g = int(
+                        start_color[1] + (end_color[1] - start_color[1]) * y / height
+                    )
+                    b = int(
+                        start_color[2] + (end_color[2] - start_color[2]) * y / height
+                    )
                     draw.line([(0, y), (width, y)], fill=(r, g, b))
 
         elif image_type == "noise":
             noise_array = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
             img = Image.fromarray(noise_array, "RGB")
+
         else:
             return {"error": f"Unsupported image_type {image_type}"}
 
@@ -485,11 +520,20 @@ def generate_simple_image(
     except Exception as e:
         return {"error": str(e)}
 
+
 @tool
 def combine_images(
     images_base64: List[str], operation: str, params: Optional[Dict[str, Any]] = None
 ) -> Dict[str, Any]:
-    """Combine multiple images (collage, stack, blend)."""
+    """
+    Combine multiple images (collage, stack, blend).
+    Args:
+        images_base64 (List[str]): List of base64 images
+        operation (str): Combination type
+        params (Dict[str, Any], optional)
+    Returns:
+        Dictionary with combined image (base64)
+    """
     try:
         images = [decode_image(b64) for b64 in images_base64]
         params = params or {}
@@ -522,157 +566,87 @@ def combine_images(
     except Exception as e:
         return {"error": str(e)}
 
-# ============== SPECIALIZED AGENT CLASSES ============== #
-
-class SpecializedAgent:
-    """Base class for specialized agents"""
-    def __init__(self, name: str, system_prompt: str, tools: List):
-        self.name = name
-        self.system_prompt = system_prompt
-        self.tools = tools
-        self.llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", api_key=os.getenv("GOOGLE_API_KEY"))
-        self.llm_with_tools = self.llm.bind_tools(tools)
-        self.graph = self._build_graph()
+
+@tool
+def download_task_file(task_id: str, api_url: str = "https://agents-course-unit4-scoring.hf.space") -> str:
+    """
+    Download a file associated with a task from the evaluation API.
+    Args:
+        task_id (str): The task ID to download the file for
+        api_url (str): The base API URL (defaults to the evaluation server)
+    """
+    try:
+        # Construct the file download URL
+        file_url = f"{api_url}/files/{task_id}"
         
-    def _build_graph(self):
-        def agent_node(state: MessagesState) -> MessagesState:
-            messages = state["messages"]
-            if not messages or not isinstance(messages[0], SystemMessage):
-                messages = [SystemMessage(content=self.system_prompt)] + messages
-            return {"messages": [self.llm_with_tools.invoke(messages)]}
+        # Create temporary file
+        temp_dir = tempfile.gettempdir()
+        filename = f"task_{task_id}.png"  # Most files are images
+        filepath = os.path.join(temp_dir, filename)
         
-        builder = StateGraph(MessagesState)
-        builder.add_node("agent", agent_node)
-        builder.add_node("tools", ToolNode(self.tools))
+        # Download the file
+        response = requests.get(file_url, stream=True)
+        response.raise_for_status()
         
-        builder.add_edge(START, "agent")
-        builder.add_conditional_edges("agent", tools_condition)
-        builder.add_edge("tools", "agent")
+        # Save the file
+        with open(filepath, "wb") as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
         
-        return builder.compile()
-    
-    def __call__(self, question: str) -> str:
-        try:
-            messages = [HumanMessage(content=question)]
-            result = self.graph.invoke({"messages": messages})
-            return result["messages"][-1].content
-        except Exception as e:
-            return f"Error in {self.name}: {str(e)}"
+        return f"Task file downloaded to {filepath}. You can now analyze this file."
+    except Exception as e:
+        return f"Error downloading task file: {str(e)}"
 
-# Agent tool groupings
-RESEARCH_TOOLS = [wikidata_search, tavily_search_tool, youtube_search_tool]
-MATH_TOOLS = [multiply, add, subtract, divide]
-DATA_ANALYSIS_TOOLS = [analyze_csv_file, analyze_excel_file, extract_text_from_image]
-IMAGE_PROCESSING_TOOLS = [analyze_image, transform_image, draw_on_image, generate_simple_image, combine_images]
-FILE_MANAGEMENT_TOOLS = [save_and_read_file, download_file_from_url, download_task_file]
 
-# ============== MULTI-AGENT SYSTEM ============== #
+tools = [multiply, add, subtract, divide, wikidata_search, tavily_search_tool, load_youtube_transcript, combine_images, analyze_image, transform_image, draw_on_image, generate_simple_image, analyze_csv_file, analyze_excel_file, save_and_read_file, download_file_from_url, extract_text_from_image, download_task_file]
 
-class MultiAgentSystem:
-    def __init__(self):
-        # Initialize specialized agents
-        self.research_agent = SpecializedAgent("Research Agent", RESEARCH_AGENT_PROMPT, RESEARCH_TOOLS)
-        self.math_agent = SpecializedAgent("Math Agent", MATH_AGENT_PROMPT, MATH_TOOLS)
-        self.data_agent = SpecializedAgent("Data Analysis Agent", DATA_ANALYSIS_AGENT_PROMPT, DATA_ANALYSIS_TOOLS)
-        self.image_agent = SpecializedAgent("Image Processing Agent", IMAGE_PROCESSING_AGENT_PROMPT, IMAGE_PROCESSING_TOOLS)
-        self.file_agent = SpecializedAgent("File Management Agent", FILE_MANAGEMENT_AGENT_PROMPT, FILE_MANAGEMENT_TOOLS)
-        
-        # Coordinator LLM
-        self.coordinator_llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", api_key=os.getenv("GOOGLE_API_KEY"))
+def build_graph():
+    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", api_key=os.getenv("GOOGLE_API_KEY"))
+    llm_with_tools = llm.bind_tools(tools)
+
+    def agent_node(state: MessagesState) -> MessagesState:
+        """This is the agent node with ReAct methodology"""
+        messages = state["messages"]
         
-        print("Multi-Agent System initialized with 5 specialized agents.")
+        # Add system prompt if not already present
+        if not messages or not isinstance(messages[0], SystemMessage):
+            messages = [SystemMessage(content=REACT_SYSTEM_PROMPT)] + messages
+            
+        return {"messages": [llm_with_tools.invoke(messages)]}
     
-    def _classify_task(self, question: str) -> Dict[str, Any]:
-        """Use the coordinator to classify the task and determine which agents to use"""
-        classification_prompt = f"""
-        As a task classifier, analyze this question and determine which specialized agents are needed:
-        
-        Question: {question}
-        
-        Available agents:
-        - research: For Wikipedia, web search, YouTube search
-        - math: For mathematical calculations
-        - data_analysis: For CSV/Excel analysis, OCR
-        - image_processing: For image analysis, manipulation, generation
-        - file_management: For file operations, downloads
-        
-        Respond with a JSON object containing:
-        {{
-            "primary_agent": "agent_name",
-            "supporting_agents": ["agent1", "agent2"],
-            "task_breakdown": "explanation of how to approach this task",
-            "requires_coordination": true/false
-        }}
-        """
-        
-        response = self.coordinator_llm.invoke([HumanMessage(content=classification_prompt)])
-        
-        # Simple classification logic as fallback
-        question_lower = question.lower()
-        
-        classification = {
-            "primary_agent": "research",
-            "supporting_agents": [],
-            "task_breakdown": "Research-based question",
-            "requires_coordination": False
-        }
-        
-        # Determine primary agent based on keywords
-        if any(word in question_lower for word in ['calculate', 'multiply', 'add', 'subtract', 'divide', 'math']):
-            classification["primary_agent"] = "math"
-        elif any(word in question_lower for word in ['csv', 'excel', 'data', 'analyze data', 'spreadsheet']):
-            classification["primary_agent"] = "data_analysis"
-        elif any(word in question_lower for word in ['image', 'photo', 'picture', 'draw', 'generate image']):
-            classification["primary_agent"] = "image_processing"
-        elif any(word in question_lower for word in ['download', 'file', 'save']):
-            classification["primary_agent"] = "file_management"
-        
-        return classification
     
-    def __call__(self, question: str) -> str:
-        """Route the question to appropriate agents and coordinate the response"""
-        try:
-            # Classify the task
-            classification = self._classify_task(question)
-            primary_agent = classification["primary_agent"]
-            
-            # Route to primary agent
-            if primary_agent == "research":
-                response = self.research_agent(question)
-            elif primary_agent == "math":
-                response = self.math_agent(question)
-            elif primary_agent == "data_analysis":
-                response = self.data_agent(question)
-            elif primary_agent == "image_processing":
-                response = self.image_agent(question)
-            elif primary_agent == "file_management":
-                response = self.file_agent(question)
-            else:
-                response = self.research_agent(question)  # Default fallback
-            
-            # For now, return the primary agent's response
-            # In a more sophisticated system, we would coordinate between multiple agents
-            return response
-            
-        except Exception as e:
-            return f"Error in Multi-Agent System: {str(e)}"
 
-# ============== MAIN AGENT CLASS (for backward compatibility) ============== #
+    builder = StateGraph(MessagesState)
+    builder.add_node("agent", agent_node)
+    builder.add_node("tools", ToolNode(tools))
+    
+    
+    builder.add_edge(START, "agent")
+    builder.add_conditional_edges("agent", tools_condition)
+    builder.add_edge("tools", "agent")
+
+    return builder.compile()
 
 class LangGraphAgent:
     def __init__(self):
-        self.multi_agent_system = MultiAgentSystem()
-        print("LangGraphAgent initialized with Multi-Agent System.")
+        self.graph = build_graph()
+        print("LangGraphAgent initialized with tools.")
         
     def __call__(self, question: str) -> str:
-        """Run the multi-agent system on a question and return the answer"""
-        return self.multi_agent_system(question)
+        """Run the agent on a question and return the answer"""
+        try:
+            messages = [HumanMessage(content=question)]
+            result = self.graph.invoke({"messages": messages})
+            for m in result["messages"]:
+                m.pretty_print()
+            return result["messages"][-1].content
+        except Exception as e:
+            return f"Error: {str(e)}"
 
 if __name__ == "__main__":
     agent = LangGraphAgent()
     question = "How many studio albums were published by Mercedes Sosa between 2000 and 2009 (included)? You can use the latest 2022 version of english wikipedia."
     answer = agent(question)
-    print(f"\nFinal Answer: {answer}")