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 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 ) # Initialize core tools self.download_tool = self.DownloadTaskAttachmentTool(rate_limiter=rate_limiter) self.chess_engine = 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 # Configure agent self.agent = CodeAgent( model=self.model, tools=[ self.download_tool, self.chess_engine, self.SpeechToTextTool(), self.ExcelReaderTool(), self.VisitWebpageTool(), self.PythonCodeReaderTool() ], verbosity_level=2, prompt_templates=prompt_templates, add_base_tools=True, max_steps=15 ) print("MagAgent initialized.") 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): name = "stockfish_analysis" description = "Analyzes chess position using Stockfish and returns best move" inputs = { "fen": { "type": "string", "description": "FEN string of the current chess position" }, "time_limit": { "type": "number", "description": "Analysis time in seconds", "nullable": True } } output_type = "string" def forward(self, fen: str, time_limit: float = 0.1) -> str: # Add time_limit parameter """Analyzes chess position using Stockfish engine""" try: import chess import chess.engine board = chess.Board(fen) engine = chess.engine.SimpleEngine.popen_uci("stockfish") result = engine.play(board, chess.engine.Limit(time=time_limit)) # Use parameter engine.quit() return board.san(result.move) except Exception as e: return f"Engine error: {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)}" 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: # Unified processing flow img_path = self.download_tool(task_id) response = await asyncio.to_thread( self.model, messages=[{ "role": "user", "content": [ {"type": "text", "text": f"{question}\nProvide answer in algebraic notation."}, {"type": "image_url", "image_url": {"url": f"file://{img_path}"}} ] }], temperature=0.1 ) return response.choices[0].message.content # 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