Tornadosky/Final_Assignment_Template
0
1from smolagents import Tool2import pandas as pd3import os4import tempfile5import requests6from urllib.parse import urlparse7import json8import re9from datetime import datetime, timedelta10 11class ReverseTextTool(Tool):12 name = "reverse_text"13 description = "Reverses the text in a string."14 inputs = {15 "text": {16 "type": "string",17 "description": "The text to reverse."18 }19 }20 output_type = "string"21 22 def forward(self, text: str) -> str:23 return text[::-1]24 25class ExtractTextFromImageTool(Tool):26 name = "extract_text_from_image"27 description = "Extracts text from an image file using OCR."28 inputs = {29 "image_path": {30 "type": "string",31 "description": "Path to the image file."32 }33 }34 output_type = "string"35 36 def forward(self, image_path: str) -> str:37 try:38 # Try to import pytesseract39 import pytesseract40 from PIL import Image41 42 # Open the image43 image = Image.open(image_path)44 45 # Try different configurations for better results46 configs = [47 '--psm 6', # Assume a single uniform block of text48 '--psm 3', # Automatic page segmentation, but no OSD49 '--psm 1', # Automatic page segmentation with OSD50 ]51 52 results = []53 for config in configs:54 try:55 text = pytesseract.image_to_string(image, config=config)56 if text.strip():57 results.append(text)58 except Exception:59 continue60 61 if results:62 # Return the longest result, which is likely the most complete63 return f"Extracted text from image:\n\n{max(results, key=len)}"64 else:65 return "No text could be extracted from the image."66 except ImportError:67 return "Error: pytesseract is not installed. Please install it with 'pip install pytesseract' and ensure Tesseract OCR is installed on your system."68 except Exception as e:69 return f"Error extracting text from image: {str(e)}"70 71class AnalyzeCSVTool(Tool):72 name = "analyze_csv_file"73 description = "Analyzes a CSV file and provides information about its contents."74 inputs = {75 "file_path": {76 "type": "string",77 "description": "Path to the CSV file."78 },79 "query": {80 "type": "string",81 "description": "Optional query about the data.",82 "default": "",83 "nullable": True84 }85 }86 output_type = "string"87 88 def forward(self, file_path: str, query: str = "") -> str:89 try:90 # Read CSV file with different encodings if needed91 for encoding in ['utf-8', 'latin1', 'iso-8859-1', 'cp1252']:92 try:93 df = pd.read_csv(file_path, encoding=encoding)94 break95 except UnicodeDecodeError:96 continue97 else:98 return "Error: Could not read the CSV file with any of the attempted encodings."99 100 # Basic information101 result = f"CSV file has {len(df)} rows and {len(df.columns)} columns.\n"102 result += f"Columns: {', '.join(df.columns)}\n\n"103 104 # If there's a specific query105 if query:106 if "count" in query.lower():107 result += f"Row count: {len(df)}\n"108 109 # Look for column-specific queries110 for col in df.columns:111 if col.lower() in query.lower():112 result += f"\nColumn '{col}' information:\n"113 if pd.api.types.is_numeric_dtype(df[col]):114 result += f"Min: {df[col].min()}\n"115 result += f"Max: {df[col].max()}\n"116 result += f"Mean: {df[col].mean()}\n"117 result += f"Median: {df[col].median()}\n"118 else:119 # For categorical data120 value_counts = df[col].value_counts().head(10)121 result += f"Unique values: {df[col].nunique()}\n"122 result += f"Top values:\n{value_counts.to_string()}\n"123 124 # General statistics for all columns125 else:126 # For numeric columns127 numeric_cols = df.select_dtypes(include=['number']).columns128 if len(numeric_cols) > 0:129 result += "Numeric columns statistics:\n"130 result += df[numeric_cols].describe().to_string()131 result += "\n\n"132 133 # For categorical columns, show counts of unique values134 cat_cols = df.select_dtypes(exclude=['number']).columns135 if len(cat_cols) > 0:136 result += "Categorical columns:\n"137 for col in cat_cols[:5]: # Limit to first 5 columns138 result += f"- {col}: {df[col].nunique()} unique values\n"139 140 return result141 except Exception as e:142 return f"Error analyzing CSV file: {str(e)}"143 144class AnalyzeExcelTool(Tool):145 name = "analyze_excel_file"146 description = "Analyzes an Excel file and provides information about its contents."147 inputs = {148 "file_path": {149 "type": "string",150 "description": "Path to the Excel file."151 },152 "query": {153 "type": "string",154 "description": "Optional query about the data.",155 "default": "",156 "nullable": True157 },158 "sheet_name": {159 "type": "string",160 "description": "Name of the sheet to analyze (defaults to first sheet).",161 "default": None,162 "nullable": True163 }164 }165 output_type = "string"166 167 def forward(self, file_path: str, query: str = "", sheet_name: str = None) -> str:168 try:169 # Read sheet names first170 excel_file = pd.ExcelFile(file_path)171 sheet_names = excel_file.sheet_names172 173 # Info about all sheets174 result = f"Excel file contains {len(sheet_names)} sheets: {', '.join(sheet_names)}\n\n"175 176 # If sheet name is specified, use it; otherwise use first sheet177 if sheet_name is None:178 sheet_name = sheet_names[0]179 elif sheet_name not in sheet_names:180 return f"Error: Sheet '{sheet_name}' not found. Available sheets: {', '.join(sheet_names)}"181 182 # Read the specified sheet183 df = pd.read_excel(file_path, sheet_name=sheet_name)184 185 # Basic information186 result += f"Sheet '{sheet_name}' has {len(df)} rows and {len(df.columns)} columns.\n"187 result += f"Columns: {', '.join(df.columns)}\n\n"188 189 # Handle query similar to CSV tool190 if query:191 if "count" in query.lower():192 result += f"Row count: {len(df)}\n"193 194 # Look for column-specific queries195 for col in df.columns:196 if col.lower() in query.lower():197 result += f"\nColumn '{col}' information:\n"198 if pd.api.types.is_numeric_dtype(df[col]):199 result += f"Min: {df[col].min()}\n"200 result += f"Max: {df[col].max()}\n"201 result += f"Mean: {df[col].mean()}\n"202 result += f"Median: {df[col].median()}\n"203 else:204 # For categorical data205 value_counts = df[col].value_counts().head(10)206 result += f"Unique values: {df[col].nunique()}\n"207 result += f"Top values:\n{value_counts.to_string()}\n"208 else:209 # For numeric columns210 numeric_cols = df.select_dtypes(include=['number']).columns211 if len(numeric_cols) > 0:212 result += "Numeric columns statistics:\n"213 result += df[numeric_cols].describe().to_string()214 result += "\n\n"215 216 # For categorical columns, show counts of unique values217 cat_cols = df.select_dtypes(exclude=['number']).columns218 if len(cat_cols) > 0:219 result += "Categorical columns:\n"220 for col in cat_cols[:5]: # Limit to first 5 columns221 result += f"- {col}: {df[col].nunique()} unique values\n"222 223 return result224 except Exception as e:225 return f"Error analyzing Excel file: {str(e)}"226 227class DateCalculatorTool(Tool):228 name = "date_calculator"229 description = "Performs date calculations like adding days, formatting dates, etc."230 inputs = {231 "query": {232 "type": "string",233 "description": "The date calculation to perform (e.g., 'What day is 10 days from today?', 'Format 2023-05-15 as MM/DD/YYYY')"234 }235 }236 output_type = "string"237 238 def forward(self, query: str) -> str:239 try:240 # Get current date/time241 if re.search(r'(today|now|current date|current time)', query, re.IGNORECASE):242 now = datetime.now()243 244 if 'time' in query.lower():245 return f"Current date and time: {now.strftime('%Y-%m-%d %H:%M:%S')}"246 else:247 return f"Today's date: {now.strftime('%Y-%m-%d')}"248 249 # Add days to a date250 add_match = re.search(r'(what|when).+?(\d+)\s+(day|days|week|weeks|month|months|year|years)\s+(from|after)\s+(.+)', query, re.IGNORECASE)251 if add_match:252 amount = int(add_match.group(2))253 unit = add_match.group(3).lower()254 date_text = add_match.group(5).strip()255 256 # Parse the date257 if date_text.lower() in ['today', 'now']:258 base_date = datetime.now()259 else:260 try:261 # Try various date formats262 for fmt in ['%Y-%m-%d', '%m/%d/%Y', '%d/%m/%Y', '%B %d, %Y']:263 try:264 base_date = datetime.strptime(date_text, fmt)265 break266 except ValueError:267 continue268 else:269 return f"Could not parse date: {date_text}"270 except Exception as e:271 return f"Error parsing date: {e}"272 273 # Calculate new date274 if 'day' in unit:275 new_date = base_date + timedelta(days=amount)276 elif 'week' in unit:277 new_date = base_date + timedelta(weeks=amount)278 elif 'month' in unit:279 # Simplified month calculation280 new_month = base_date.month + amount281 new_year = base_date.year + (new_month - 1) // 12282 new_month = ((new_month - 1) % 12) + 1283 new_date = base_date.replace(year=new_year, month=new_month)284 elif 'year' in unit:285 new_date = base_date.replace(year=base_date.year + amount)286 287 return f"Date {amount} {unit} from {base_date.strftime('%Y-%m-%d')} is {new_date.strftime('%Y-%m-%d')}"288 289 # Format a date290 format_match = re.search(r'format\s+(.+?)\s+as\s+(.+)', query, re.IGNORECASE)291 if format_match:292 date_text = format_match.group(1).strip()293 format_spec = format_match.group(2).strip()294 295 # Parse the date296 if date_text.lower() in ['today', 'now']:297 date_obj = datetime.now()298 else:299 try:300 # Try various date formats301 for fmt in ['%Y-%m-%d', '%m/%d/%Y', '%d/%m/%Y', '%B %d, %Y']:302 try:303 date_obj = datetime.strptime(date_text, fmt)304 break305 except ValueError:306 continue307 else:308 return f"Could not parse date: {date_text}"309 except Exception as e:310 return f"Error parsing date: {e}"311 312 # Convert format specification to strftime format313 format_mapping = {314 'YYYY': '%Y',315 'YY': '%y',316 'MM': '%m',317 'DD': '%d',318 'HH': '%H',319 'mm': '%M',320 'ss': '%S'321 }322 323 strftime_format = format_spec324 for key, value in format_mapping.items():325 strftime_format = strftime_format.replace(key, value)326 327 return f"Formatted date: {date_obj.strftime(strftime_format)}"328 329 return "I couldn't understand the date calculation query."330 except Exception as e:331 return f"Error performing date calculation: {str(e)}"332 333class DownloadFileTool(Tool):334 name = "download_file"335 description = "Downloads a file from a URL and saves it locally."336 inputs = {337 "url": {338 "type": "string",339 "description": "The URL to download from."340 },341 "filename": {342 "type": "string",343 "description": "Optional filename to save as (default: derived from URL).",344 "default": None,345 "nullable": True346 }347 }348 output_type = "string"349 350 def forward(self, url: str, filename: str = None) -> str:351 try:352 # Parse URL to get filename if not provided353 if not filename:354 path = urlparse(url).path355 filename = os.path.basename(path)356 if not filename:357 # Generate a random name if we couldn't extract one358 import uuid359 filename = f"downloaded_{uuid.uuid4().hex[:8]}"360 361 # Create temporary file362 temp_dir = tempfile.gettempdir()363 filepath = os.path.join(temp_dir, filename)364 365 # Download the file366 response = requests.get(url, stream=True)367 response.raise_for_status()368 369 # Save the file370 with open(filepath, 'wb') as f:371 for chunk in response.iter_content(chunk_size=8192):372 f.write(chunk)373 374 return f"File downloaded to {filepath}. You can now analyze this file."375 except Exception as e:376 return f"Error downloading file: {str(e)}"