|
|
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_dotenv() |
|
|
|
|
|
|
|
|
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 |
|
|
|
