Update agent.py

agent.py

@@ -1,6 +1,6 @@
 import os
 from dotenv import load_dotenv
-from typing import List, Dict, Any, Optional
+from typing import List, Dict, Any, Optional, Literal
 from langgraph.graph import START, StateGraph, MessagesState
 from langgraph.graph.message import add_messages
 from langchain_core.messages import AnyMessage, HumanMessage, AIMessage, SystemMessage
@@ -15,77 +15,152 @@ 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()
 
-# ReAct System Prompt
-REACT_SYSTEM_PROMPT = """You are a research assistant that uses ReAct (Reasoning + Acting) methodology. For each question, follow this systematic approach:
+# ============== 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
+
+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
+
+Be thorough in your research and cross-reference sources when possible.
+"""
+
+MATH_AGENT_PROMPT = """You are a Math Agent specialized in mathematical calculations and operations.
 
-**THINK**: First, analyze the question carefully. What type of information do you need? What tools might help?
+Your expertise includes:
+- Basic arithmetic operations (add, subtract, multiply, divide)
+- Mathematical reasoning and problem-solving
 
-**ACT**: Use available tools to gather information. Search thoroughly and verify facts from multiple sources when possible.
+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
 
-**OBSERVE**: Analyze the results from your tools. Are they complete and reliable? Do you need more information?
+Always show your work and verify calculations.
+"""
+
+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
 
-**REASON**: Synthesize all information gathered. Check for consistency and identify any gaps or uncertainties.
+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.
 
-**VERIFY**: Before providing your final answer, double-check your reasoning and ensure you have sufficient evidence.
+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
 
-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
+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.
+"""
 
-When you cannot access certain content (videos, audio, images without tools), clearly state this limitation.
+FILE_MANAGEMENT_AGENT_PROMPT = """You are a File Management Agent specialized in file operations and data handling.
 
-Always finish with: FINAL ANSWER: [YOUR FINAL ANSWER]
+Your expertise includes:
+- Saving and reading files
+- Downloading files from URLs
+- Downloading task files from APIs
+- File format handling
 
-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
+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
 
-Be thorough in your research but honest about uncertainty. Quality and accuracy are more important than speed.
+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) -> int:
-    """
-    Divide two numbers
-    """
+def divide(a: int, b: int) -> float:
+    """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.
-    
-    Args:
-        query: The search query.
-    """
+    """Search for information on Wikipedia and return maximum 2 results."""
     loader = WikipediaLoader(query=query, load_max_docs=2)
     docs = loader.load()
     formatted_search_docs = "\n\n---\n\n".join(
@@ -95,23 +170,14 @@ def wikidata_search(query: str) -> str:
         ])
     return {"wiki_results": formatted_search_docs}
 
-# Initialize Tavily Search Tool
-tavily_search_tool = TavilySearch(
-    max_results=3,
-    topic="general",
-)
-
-# Initialize YouTube Search Tool
+# Initialize search tools
+tavily_search_tool = TavilySearch(max_results=3, topic="general")
 youtube_search_tool = YouTubeSearchTool()
 
+# 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.
-    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.
-    """
+    """Save content to a file and return the path."""
     temp_dir = tempfile.gettempdir()
     if filename is None:
         temp_file = tempfile.NamedTemporaryFile(delete=False, dir=temp_dir)
@@ -124,32 +190,22 @@ 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.
-    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.
-    """
+    """Download a file from a URL and save it to a temporary location."""
     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)
@@ -158,100 +214,75 @@ 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 library pytesseract (if available).
-    Args:
-        image_path (str): the path to the image file.
-    """
+    """Extract text from an image using OCR."""
     try:
-        # Open the image
+        import pytesseract
         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.
-    Args:
-        file_path (str): the path to the CSV file.
-        query (str): Question about the data
-    """
+    """Analyze a CSV file using pandas and answer a question about it."""
     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.
-    Args:
-        file_path (str): the path to the Excel file.
-        query (str): Question about the data
-    """
+    """Analyze an Excel file using pandas and answer a question about it."""
     try:
-        # Read the Excel file
         df = pd.read_excel(file_path)
-
-        # Run various analyses based on the query
-        result = (
-            f"Excel file loaded with {len(df)} rows and {len(df.columns)} columns.\n"
-        )
+        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 AND GENERATION TOOLS =============== ###
-import os
-import io
-import base64
-import uuid
-from PIL import Image, ImageDraw, ImageFont, ImageEnhance, ImageFilter
-
-# Helper functions for image processing
+# Image Processing Agent Tools - Helper functions
 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)
@@ -262,13 +293,7 @@ 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 (size, mode, color analysis, thumbnail preview).
-    Args:
-        image_base64 (str): Base64 encoded image string
-    Returns:
-        Dictionary with analysis result
-    """
+    """Analyze basic properties of an image."""
     try:
         img = decode_image(image_base64)
         width, height = img.size
@@ -301,42 +326,29 @@ 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.
-    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)
-    """
+    """Apply transformations: resize, rotate, crop, flip, brightness, contrast, blur, sharpen, grayscale."""
     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)
@@ -362,20 +374,11 @@ 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.
-    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)
-    """
+    """Draw shapes (rectangle, circle, line) or text onto an image."""
     try:
         img = decode_image(image_base64)
         draw = ImageDraw.Draw(img)
@@ -395,16 +398,12 @@ 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:
@@ -427,7 +426,6 @@ def draw_on_image(
     except Exception as e:
         return {"error": str(e)}
 
-
 @tool
 def generate_simple_image(
     image_type: str,
@@ -435,15 +433,7 @@ def generate_simple_image(
     height: int = 500,
     params: Optional[Dict[str, Any]] = None,
 ) -> Dict[str, Any]:
-    """
-    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)
-    """
+    """Generate a simple image (gradient, noise, pattern, chart)."""
     try:
         params = params or {}
 
@@ -457,33 +447,20 @@ 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}"}
 
@@ -494,20 +471,11 @@ 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).
-    Args:
-        images_base64 (List[str]): List of base64 images
-        operation (str): Combination type
-        params (Dict[str, Any], optional)
-    Returns:
-        Dictionary with combined image (base64)
-    """
+    """Combine multiple images (collage, stack, blend)."""
     try:
         images = [decode_image(b64) for b64 in images_base64]
         params = params or {}
@@ -540,87 +508,157 @@ def combine_images(
     except Exception as e:
         return {"error": 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.
-    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}"
+# ============== 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()
         
-        # Create temporary file
-        temp_dir = tempfile.gettempdir()
-        filename = f"task_{task_id}.png"  # Most files are images
-        filepath = os.path.join(temp_dir, filename)
+    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)]}
         
-        # Download the file
-        response = requests.get(file_url, stream=True)
-        response.raise_for_status()
+        builder = StateGraph(MessagesState)
+        builder.add_node("agent", agent_node)
+        builder.add_node("tools", ToolNode(self.tools))
         
-        # Save the file
-        with open(filepath, "wb") as f:
-            for chunk in response.iter_content(chunk_size=8192):
-                f.write(chunk)
+        builder.add_edge(START, "agent")
+        builder.add_conditional_edges("agent", tools_condition)
+        builder.add_edge("tools", "agent")
         
-        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)}"
-
+        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)}"
 
-tools = [multiply, add, subtract, divide, wikidata_search, tavily_search_tool, youtube_search_tool, 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]
+# 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]
 
-def build_graph():
-    llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", api_key=os.getenv("GOOGLE_API_KEY"))
-    llm_with_tools = llm.bind_tools(tools)
+# ============== MULTI-AGENT SYSTEM ============== #
 
-    def agent_node(state: MessagesState) -> MessagesState:
-        """This is the agent node with ReAct methodology"""
-        messages = state["messages"]
+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)
         
-        # 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)]}
-    
-    
-
-    builder = StateGraph(MessagesState)
-    builder.add_node("agent", agent_node)
-    builder.add_node("tools", ToolNode(tools))
+        # Coordinator LLM
+        self.coordinator_llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", api_key=os.getenv("GOOGLE_API_KEY"))
+        
+        print("Multi-Agent System initialized with 5 specialized agents.")
     
+    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
     
-    builder.add_edge(START, "agent")
-    builder.add_conditional_edges("agent", tools_condition)
-    builder.add_edge("tools", "agent")
+    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)}"
 
-    return builder.compile()
+# ============== MAIN AGENT CLASS (for backward compatibility) ============== #
 
 class LangGraphAgent:
     def __init__(self):
-        self.graph = build_graph()
-        print("LangGraphAgent initialized with tools.")
+        self.multi_agent_system = MultiAgentSystem()
+        print("LangGraphAgent initialized with Multi-Agent System.")
         
     def __call__(self, question: str) -> str:
-        """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)}"
+        """Run the multi-agent system on a question and return the answer"""
+        return self.multi_agent_system(question)
 
 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}")