Update agent.py

agent.py

@@ -2,6 +2,7 @@ from smolagents import CodeAgent,  LiteLLMModel, tool, load_tool, DuckDuckGoSear
 import asyncio
 import os
 import re
+import pandas as pd
 from typing import Optional
 from token_bucket import Limiter
 import yaml
@@ -92,6 +93,64 @@ def SpeechToTextTool(audio_path: str) -> str:
     result = 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)}"
 
 #@tool
 #class LocalFileAudioTool:
@@ -130,10 +189,12 @@ class MagAgent:
                 WikipediaSearchTool(),
 #                ImageAnalysisTool,
                 SpeechToTextTool
+                ExcelReaderTool(),
 #                LocalFileAudioTool()
             ],
             verbosity_level=2,
-            max_steps=15
+            add_base_tools=True,
+            max_steps=20
         )
         print("MagAgent initialized.")