Update agent.py

agent.py

@@ -1,117 +1,118 @@
-import os
-from pathlib import Path
-from typing import Optional, Union
-
-import pandas as pd
-from dotenv import load_dotenv
-from smolagents import (CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool,
-                        LiteLLMModel, PythonInterpreterTool,
-                        WikipediaSearchTool)
-from smolagents.tools import Tool
-from tabulate import tabulate
-
-# Load environment variables
-load_dotenv()
-
-# Initialize the model
-model = LiteLLMModel(
-    model_id=os.getenv("GEMINI_MODEL"), api_key=os.getenv("GEMINI_API_KEY")
-)
-
-
-class ExcelToTextTool(Tool):
-    """Render an Excel worksheet as a Markdown table."""
-
-    name = "excel_to_text"
-    description = (
-        "Read an Excel file and return a Markdown table of the requested sheet. "
-        "Accepts either the sheet name or a zero-based index (as a string)."
-    )
-
-    inputs = {
-        "excel_path": {
-            "type": "string",
-            "description": "Path to the Excel file (.xlsx or .xls).",
-        },
-        "sheet_name": {
-            "type": "string",
-            "description": (
-                "Worksheet name or zero-based index (as a string). "
-                "Optional; defaults to the first sheet."
-            ),
-            "nullable": True,
-        },
-    }
-
-    output_type = "string"
-
-    def forward(self, excel_path: str, sheet_name: Optional[str] = None) -> str:
-        """Load the Excel file and return the sheet as a Markdown table.
-
-        Args:
-            excel_path: Path to the Excel file.
-            sheet_name: Optional name or index of the sheet to read. If None, reads the first sheet.
-
-        Returns:
-            A Markdown table representing the Excel sheet, or an error message if the file is not found or cannot be read.
-        """
-
-        file_path = Path(excel_path).expanduser().resolve()
-        if not file_path.is_file():
-            return f"Error: Excel file not found at {file_path}"
-
-        try:
-            sheet: Union[str, int] = (
-                int(sheet_name)
-                if sheet_name and sheet_name.isdigit()
-                else sheet_name or 0
-            )
-
-            df = pd.read_excel(file_path, sheet_name=sheet)
-
-            if hasattr(df, "to_markdown"):
-                return df.to_markdown(index=False)
-
-            return tabulate(df, headers="keys", tablefmt="github", showindex=False)
-
-        except Exception as e:
-            return f"Error reading Excel file: {e}"
-
-
-class GaiaAgent:
-    """An agent capable of using tools to answer general questions."""
-
-    def __init__(self):
-        """Initializes the GaiaAgent with a set of tools."""
-
-        print("GaiaAgent initialized with tools.")
-
-        tools = [
-            DuckDuckGoSearchTool(),
-            WikipediaSearchTool(),
-            ExcelToTextTool(),
-            PythonInterpreterTool(),
-            FinalAnswerTool(),
-        ]
-
-        self.agent = CodeAgent(
-            model=model,
-            tools=tools,
-            add_base_tools=True,
-            additional_authorized_imports=["pandas", "numpy", "csv", "subprocess"],
-        )
-
-    def __call__(self, task_id: str, question: str) -> str:
-        """Processes a question using the agent and its tools.
-
-        Args:
-            task_id: A unique identifier for the task.
-            question: The question to be answered.
-
-        Returns:
-            The answer generated by the agent.
-        """
-        print(f"Agent received task_id='{task_id}' | question='{question[:50]}...'")
-        answer = self.agent.run(question)
-        print(f"Agent returning answer: {answer}")
-        return answer
+import os
+from pathlib import Path
+from typing import Optional, Union
+import smolagent
+
+import pandas as pd
+from dotenv import load_dotenv
+from smolagents import (CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool,
+                        LiteLLMModel, PythonInterpreterTool,
+                        WikipediaSearchTool)
+from smolagents.tools import Tool
+from tabulate import tabulate
+
+# Load environment variables
+load_dotenv()
+
+# Initialize the model
+model = LiteLLMModel(
+    model_id=os.getenv("GEMINI_MODEL"), api_key=os.getenv("GEMINI_API_KEY")
+)
+
+
+class ExcelToTextTool(Tool):
+    """Render an Excel worksheet as a Markdown table."""
+
+    name = "excel_to_text"
+    description = (
+        "Read an Excel file and return a Markdown table of the requested sheet. "
+        "Accepts either the sheet name or a zero-based index (as a string)."
+    )
+
+    inputs = {
+        "excel_path": {
+            "type": "string",
+            "description": "Path to the Excel file (.xlsx or .xls).",
+        },
+        "sheet_name": {
+            "type": "string",
+            "description": (
+                "Worksheet name or zero-based index (as a string). "
+                "Optional; defaults to the first sheet."
+            ),
+            "nullable": True,
+        },
+    }
+
+    output_type = "string"
+
+    def forward(self, excel_path: str, sheet_name: Optional[str] = None) -> str:
+        """Load the Excel file and return the sheet as a Markdown table.
+
+        Args:
+            excel_path: Path to the Excel file.
+            sheet_name: Optional name or index of the sheet to read. If None, reads the first sheet.
+
+        Returns:
+            A Markdown table representing the Excel sheet, or an error message if the file is not found or cannot be read.
+        """
+
+        file_path = Path(excel_path).expanduser().resolve()
+        if not file_path.is_file():
+            return f"Error: Excel file not found at {file_path}"
+
+        try:
+            sheet: Union[str, int] = (
+                int(sheet_name)
+                if sheet_name and sheet_name.isdigit()
+                else sheet_name or 0
+            )
+
+            df = pd.read_excel(file_path, sheet_name=sheet)
+
+            if hasattr(df, "to_markdown"):
+                return df.to_markdown(index=False)
+
+            return tabulate(df, headers="keys", tablefmt="github", showindex=False)
+
+        except Exception as e:
+            return f"Error reading Excel file: {e}"
+
+
+class GaiaAgent:
+    """An agent capable of using tools to answer general questions."""
+
+    def __init__(self):
+        """Initializes the GaiaAgent with a set of tools."""
+
+        print("GaiaAgent initialized with tools.")
+
+        tools = [
+            DuckDuckGoSearchTool(),
+            WikipediaSearchTool(),
+            ExcelToTextTool(),
+            PythonInterpreterTool(),
+            FinalAnswerTool(),
+        ]
+
+        self.agent = CodeAgent(
+            model=model,
+            tools=tools,
+            add_base_tools=True,
+            additional_authorized_imports=["pandas", "numpy", "csv", "subprocess"],
+        )
+
+    def __call__(self, task_id: str, question: str) -> str:
+        """Processes a question using the agent and its tools.
+
+        Args:
+            task_id: A unique identifier for the task.
+            question: The question to be answered.
+
+        Returns:
+            The answer generated by the agent.
+        """
+        print(f"Agent received task_id='{task_id}' | question='{question[:50]}...'")
+        answer = self.agent.run(question)
+        print(f"Agent returning answer: {answer}")
+        return answer