File size: 3,594 Bytes
fd5668b
87394ed
 
fd5668b
87394ed
ff5ced3
dbb7d44
 
 
87394ed
 
fd5668b
87394ed
ff5ced3
fd5668b
87394ed
ff5ced3
dbb7d44
87394ed
 
 
 
 
 
 
 
 
 
ff5ced3
fd5668b
87394ed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
dbb7d44
 
 
 
 
 
 
 
 
87394ed
 
 
 
 
 
 
dbb7d44
 
 
87394ed
 
 
 
 
 
 
 
 
 
 
 
fd5668b
 
87394ed
 
fd5668b
dbb7d44
 
fd5668b
 
 
 
 
87394ed
fd5668b
 
 
 
87394ed
 
 
 
 
 
fd5668b
 
dbb7d44
 
 
 
 
 
 
 
 
87394ed
fd5668b
 
ff5ced3
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
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