Create agent.py

agent.py

@@ -0,0 +1,422 @@
+from smolagents import CodeAgent,  LiteLLMModel, tool, Tool, load_tool, DuckDuckGoSearchTool, WikipediaSearchTool #, HfApiModel, OpenAIServerModel
+import asyncio
+import os
+import re
+import pandas as pd
+from typing import Optional
+from token_bucket import Limiter, MemoryStorage
+import yaml
+from PIL import Image, ImageOps
+import requests
+from io import BytesIO
+from markdownify import markdownify
+import whisper
+
+import time
+import shutil
+import traceback
+
+
+@tool
+def GoogleSearchTool(query: str) -> str:
+    """Tool for performing Google searches using Custom Search JSON API
+    Args:
+        query (str): Search query string
+    Returns:
+        str: Formatted search results
+    """
+    cse_id = os.environ.get("GOOGLE_CSE_ID")
+    if not api_key or not cse_id:
+        raise ValueError("GOOGLE_API_KEY and GOOGLE_CSE_ID must be set in environment variables.")
+    url = "https://www.googleapis.com/customsearch/v1"
+    params = {
+        "key": api_key,
+        "cx": cse_id,
+        "q": query,
+        "num": 5  # Number of results to return
+    }
+    try:
+        response = requests.get(url, params=params)
+        response.raise_for_status()
+        results = response.json().get("items", [])
+        return "\n".join([f"{item['title']}: {item['link']}" for item in results]) or "No results found."
+    except Exception as e:
+        return f"Error performing Google search: {str(e)}"
+
+#@tool
+
+#def ImageAnalysisTool(question: str, model: LiteLLMModel) -> str:
+#    """Tool for analyzing images mentioned in the question.
+#    Args:
+#        question (str): The question text which may contain an image URL.
+#    Returns:
+#        str: Image description or error message.
+#    """
+#    # Extract URL from question using regex
+#    url_pattern = r'https?://\S+'
+#
+#    match = re.search(url_pattern, question)
+#    if not match:
+#        return "No image URL found in the question."
+#    image_url = match.group(0)
+#
+#    headers = {
+#        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/114.0.0.0 Safari/537.36"
+#    }
+#    try:
+#        response = requests.get(image_url, headers=headers)
+
+#        response.raise_for_status()
+#        image = Image.open(BytesIO(response.content)).convert("RGB")
+#    except Exception as e:
+#        return f"Error fetching image: {e}"
+#
+#    agent = CodeAgent(
+#        tools=[],
+#        model=model,
+#        max_steps=10,
+#        verbosity_level=2
+#    )
+#
+#    response = agent.run(
+#        "Describe in details the chess position you see in the image.",
+#        images=[image]
+#    )
+#    
+#    return f"The image description: '{response}'"
+
+class VisitWebpageTool(Tool):
+    name = "visit_webpage"
+    description = "Visits a webpage at the given url and reads its content as a markdown string. Use this to browse webpages."
+    inputs = {'url': {'type': 'string', 'description': 'The url of the webpage to visit.'}}
+    output_type = "string"
+
+    def forward(self, url: str) -> str:
+        try:
+            response = requests.get(url, timeout=20)
+            response.raise_for_status()
+            markdown_content = markdownify(response.text).strip()
+            markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
+            from smolagents.utils import truncate_content
+            return truncate_content(markdown_content, 10000)
+        except requests.exceptions.Timeout:
+            return "The request timed out. Please try again later or check the URL."
+        except requests.exceptions.RequestException as e:
+            return f"Error fetching the webpage: {str(e)}"
+        except Exception as e:
+            return f"An unexpected error occurred: {str(e)}"
+
+    def __init__(self, *args, **kwargs):
+        self.is_initialized = False
+
+class DownloadTaskAttachmentTool(Tool):
+    name = "download_file"
+    description = "Downloads the file attached to the task ID and returns the local file path. Supports Excel (.xlsx), image (.png, .jpg), audio (.mp3), PDF (.pdf), and Python (.py) files."
+    inputs = {'task_id': {'type': 'string', 'description': 'The task id to download attachment from.'}}
+    output_type = "string"
+    DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
+
+    def __init__(self, rate_limiter: Optional[Limiter] = None, default_api_url: str = DEFAULT_API_URL, *args, **kwargs):
+        self.is_initialized = False
+        self.rate_limiter = rate_limiter
+        self.default_api_url = default_api_url
+
+    def forward(self, task_id: str) -> str:
+        file_url = f"{self.default_api_url}/files/{task_id}"
+        print(f"Downloading file for task ID {task_id} from {file_url}...")
+        try:
+            if self.rate_limiter:
+                while not self.rate_limiter.consume(1):
+                    print(f"Rate limit reached for downloading file for task {task_id}. Waiting...")
+                    time.sleep(60 / 15)  # Assuming 15 RPM
+            response = requests.get(file_url, stream=True, timeout=15)
+            response.raise_for_status()
+            
+            # Determine file extension based on Content-Type
+            content_type = response.headers.get('Content-Type', '').lower()
+            if 'image/png' in content_type:
+                extension = '.png'
+            elif 'image/jpeg' in content_type:
+                extension = '.jpg'
+            elif 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet' in content_type:
+                extension = '.xlsx'
+            elif 'audio/mpeg' in content_type:
+                extension = '.mp3'
+            elif 'application/pdf' in content_type:
+                extension = '.pdf'
+            elif 'text/x-python' in content_type:
+                extension = '.py'
+            else:
+                return f"Error: Unsupported file type {content_type} for task {task_id}. Try using visit_webpage or web_search if the content is online."
+            
+            local_file_path = f"downloads/{task_id}{extension}"
+            os.makedirs("downloads", exist_ok=True)
+            with open(local_file_path, "wb") as file:
+                for chunk in response.iter_content(chunk_size=8192):
+                    file.write(chunk)
+            print(f"File downloaded successfully: {local_file_path}")
+            return local_file_path
+        except requests.exceptions.HTTPError as e:
+            if e.response.status_code == 429:
+                return f"Error: Rate limit exceeded for task {task_id}. Try again later."
+            return f"Error downloading file for task {task_id}: {str(e)}"
+        except requests.exceptions.RequestException as e:
+            return f"Error downloading file for task {task_id}: {str(e)}"
+
+class SpeechToTextTool(Tool):
+    name = "speech_to_text"
+    description = (
+        "Converts an audio file to text using OpenAI Whisper."
+    )
+    inputs = {
+        "audio_path": {"type": "string", "description": "Path to audio file (.mp3, .wav)"},
+    }
+    output_type = "string"
+
+    def __init__(self):
+        super().__init__()
+        self.model = whisper.load_model("base")
+
+    def forward(self, audio_path: str) -> str:
+        if not os.path.exists(audio_path):
+            return f"Error: File not found at {audio_path}"
+        result = self.model.transcribe(audio_path)
+        return result.get("text", "")
+
+class ExcelReaderTool(Tool):
+    name = "excel_reader"
+
+    description = """
+    This tool reads and processes Excel files (.xlsx, .xls).
+    It can extract data, calculate statistics, and perform data analysis on spreadsheets.
+    """
+    inputs = {
+        "excel_path": {
+            "type": "string"
+,
+            "description": "The path to the Excel file to read",
+        },
+        "sheet_name": {
+            "type": "string",
+
+            "description": "The name of the sheet to read (optional, defaults to first sheet)",
+            "nullable": True
+        }
+    }
+    output_type = "string"
+    
+    def forward(self, excel_path: str, sheet_name: str = None) -> str:
+        """
+        Reads and processes the given Excel file.
+        """
+        try:
+            # Check if the file exists
+            if not os.path.exists(excel_path):
+                return f"Error: Excel file not found at {excel_path}"
+                
+            import pandas as pd
+            
+            # Read the Excel file
+            if sheet_name:
+                df = pd.read_excel(excel_path, sheet_name=sheet_name)
+            else:
+                df = pd.read_excel(excel_path)
+                
+            # Get basic info about the data
+            info = {
+                "shape": df.shape,
+                "columns": list(df.columns),
+                "dtypes": df.dtypes.to_dict(),
+                "head": df.head(5).to_dict()
+            }
+            
+            # Return formatted info
+            result = f"Excel file: {excel_path}\n"
+            result += f"Shape: {info['shape'][0]} rows × {info['shape'][1]} columns\n\n"
+            result += "Columns:\n"
+            for col in info['columns']:
+                result += f"- {col} ({info['dtypes'].get(col)})\n"
+            
+            result += "\nPreview (first 5 rows):\n"
+            result += df.head(5).to_string()
+            
+            return result
+            
+        except Exception as e:
+            return f"Error reading Excel file: {str(e)}"
+
+
+
+
+    class DownloadImageTool(Tool):
+        name = "download_chess_image"
+        description = "Downloads chess position image from task ID"
+        inputs = {'task_id': {'type': 'string'}}
+        output_type = "string"
+
+        def forward(self, task_id: str) -> str:
+            try:
+                response = requests.get(
+                    f"https://agents-course-unit4-scoring.hf.space/files/{task_id}",
+                    stream=True
+                )
+                response.raise_for_status()
+                
+                img_path = f"chess_{task_id}.png"
+                with open(img_path, "wb") as f:
+                    for chunk in response.iter_content(8192):
+                        f.write(chunk)
+                return img_path
+            except Exception as e:
+                raise RuntimeError(f"Image download failed: {str(e)}")
+
+
+    
+    class ChessEngineTool(Tool):
+        import chess
+        import chess.engine
+        name = "stockfish_analysis"
+        description = "Analyzes chess position using Stockfish"
+        inputs = {'fen': {'type': 'string'}}
+        output_type = "string"
+
+        def forward(self, fen: str) -> str:
+            try:
+                board = chess.Board(fen)
+                engine = chess.engine.SimpleEngine.popen_uci("stockfish")
+                result = engine.play(board, chess.engine.Limit(time=2.0))
+                engine.quit()
+                return board.san(result.move)
+            except Exception as e:
+                return f"Engine error: {str(e)}"
+
+    async def analyze_position(self, task_id: str):
+        try:
+            # Step 1: Download image
+            img_path = await self.tools[0](task_id)
+            
+            # Step 2: Get multimodal analysis
+            response = await self.model.acreate(
+                messages=[{
+                    "role": "user",
+                    "content": [
+                        {"type": "text", "text": """Analyze this chess position. 
+                        It's black's turn. Provide the winning move in algebraic notation.
+                        Respond ONLY with the move, nothing else."""},
+                        {"type": "image_url", "image_url": {"url": f"file://{img_path}"}}
+                    ]
+                }],
+                temperature=0.1
+            )
+            
+            return response.choices[0].message.content
+            
+        except Exception as e:
+            return f"Analysis failed: {str(e)}"
+
+
+
+        
+class PythonCodeReaderTool(Tool):
+    name = "read_python_code"
+    description = "Reads a Python (.py) file and returns its content as a string."
+    inputs = {
+        "file_path": {"type": "string", "description": "The path to the Python file to read"}
+    }
+    output_type = "string"
+
+    def forward(self, file_path: str) -> str:
+        try:
+            if not os.path.exists(file_path):
+                return f"Error: Python file not found at {file_path}"
+            with open(file_path, "r", encoding="utf-8") as file:
+                content = file.read()
+            return content
+        except Exception as e:
+            return f"Error reading Python file: {str(e)}"
+            
+class MagAgent:
+    def __init__(self, rate_limiter: Optional[Limiter] = None):
+        """Initialize the MagAgent with search tools."""
+        self.rate_limiter = rate_limiter
+
+        print("Initializing MagAgent with search tools...")
+#        model = LiteLLMModel(
+#            model_id="gemini/gemini-2.0-flash-preview-image-generation",
+#            api_key= os.environ.get("GEMINI_KEY"),
+#            max_tokens=8192
+#        )
+
+        self.model = LiteLLMModel(
+            model_id="gemini/gemini-1.5-flash",
+            api_key=os.environ.get("GEMINI_KEY"),
+            api_base="https://generativelanguage.googleapis.com/v1beta",
+            max_tokens=2048
+        )
+        
+        self.tools = [
+            self.DownloadImageTool(),
+            self.ChessEngineTool()
+        ]
+
+                
+        # Load prompt templates
+        with open("prompts.yaml", 'r') as stream:
+            prompt_templates = yaml.safe_load(stream)
+
+        # Initialize rate limiter for DuckDuckGoSearchTool
+        search_rate_limiter = Limiter(rate=30/60, capacity=30, storage=MemoryStorage()) if not rate_limiter else rate_limiter
+        
+        self.agent = CodeAgent(
+            model= model,
+            tools=[
+                DownloadTaskAttachmentTool(rate_limiter=rate_limiter),
+#                DuckDuckGoSearchTool(),
+#                WikipediaSearchTool(),
+                SpeechToTextTool(),
+                ExcelReaderTool(),
+                VisitWebpageTool(),
+                PythonCodeReaderTool(),
+                PNG2FENTool,
+                ChessEngineTool(),
+#                GoogleSearchTool,
+#                ImageAnalysisTool,
+            ],
+            verbosity_level=2,
+            prompt_templates=prompt_templates,
+            add_base_tools=True,
+            max_steps=15
+        )
+        print("MagAgent initialized.")
+
+    async def __call__(self, question: str, task_id: str) -> str:
+        """Process a question asynchronously using the MagAgent."""
+        print(f"MagAgent received question (first 50 chars): {question[:50]}... Task ID: {task_id}")
+        try:
+            if self.rate_limiter:
+                while not self.rate_limiter.consume(1):
+                    print(f"Rate limit reached for task {task_id}. Waiting...")
+                    await asyncio.sleep(60 / 15)  # Assuming 15 RPM
+            # Include task_id in the task prompt to guide the agent
+            task = (
+#                f"Answer the following question accurately and concisely: \n"
+                f"{question} \n"
+                f"If the question references an attachment, use tool to download it with task_id: {task_id}\n"
+#                f"Return the answer as a string."
+            )
+            print(f"Calling agent.run for task {task_id}...")
+            response = await asyncio.to_thread(
+                self.agent.run,
+                task=task
+            )
+            print(f"Agent.run completed for task {task_id}.")
+            response = str(response)
+            if not response:
+                print(f"No answer found for task {task_id}.")
+                response = "No answer found."
+            print(f"MagAgent response: {response[:50]}...")
+            return response
+        except Exception as e:
+            error_msg = f"Error processing question for task {task_id}: {str(e)}. Check API key or network connectivity."
+            print(error_msg)
+            return error_msg