hammaad-swe commited on
Commit
dbb7d44
·
1 Parent(s): 87394ed

feat: added download file capability

Browse files
Files changed (1) hide show
  1. agent.py +27 -14
agent.py CHANGED
@@ -4,14 +4,9 @@ from typing import Optional, Union
4
 
5
  import pandas as pd
6
  from dotenv import load_dotenv
7
- from smolagents import (
8
- CodeAgent,
9
- DuckDuckGoSearchTool,
10
- FinalAnswerTool,
11
- LiteLLMModel,
12
- PythonInterpreterTool,
13
- WikipediaSearchTool,
14
- )
15
  from smolagents.tools import Tool
16
  from tabulate import tabulate
17
 
@@ -20,8 +15,7 @@ load_dotenv()
20
 
21
  # Initialize the model
22
  model = LiteLLMModel(
23
- model_id=os.getenv("GEMINI_MODEL"),
24
- api_key=os.getenv("GEMINI_API_KEY")
25
  )
26
 
27
 
@@ -52,7 +46,15 @@ class ExcelToTextTool(Tool):
52
  output_type = "string"
53
 
54
  def forward(self, excel_path: str, sheet_name: Optional[str] = None) -> str:
55
- """Load the Excel file and return the sheet as a Markdown table."""
 
 
 
 
 
 
 
 
56
 
57
  file_path = Path(excel_path).expanduser().resolve()
58
  if not file_path.is_file():
@@ -60,9 +62,9 @@ class ExcelToTextTool(Tool):
60
 
61
  try:
62
  sheet: Union[str, int] = (
63
- int(
64
- sheet_name
65
- ) if sheet_name and sheet_name.isdigit() else sheet_name or 0
66
  )
67
 
68
  df = pd.read_excel(file_path, sheet_name=sheet)
@@ -80,6 +82,8 @@ class GaiaAgent:
80
  """An agent capable of using tools to answer general questions."""
81
 
82
  def __init__(self):
 
 
83
  print("GaiaAgent initialized with tools.")
84
 
85
  tools = [
@@ -98,6 +102,15 @@ class GaiaAgent:
98
  )
99
 
100
  def __call__(self, task_id: str, question: str) -> str:
 
 
 
 
 
 
 
 
 
101
  print(f"Agent received task_id='{task_id}' | question='{question[:50]}...'")
102
  answer = self.agent.run(question)
103
  print(f"Agent returning answer: {answer}")
 
4
 
5
  import pandas as pd
6
  from dotenv import load_dotenv
7
+ from smolagents import (CodeAgent, DuckDuckGoSearchTool, FinalAnswerTool,
8
+ LiteLLMModel, PythonInterpreterTool,
9
+ WikipediaSearchTool)
 
 
 
 
 
10
  from smolagents.tools import Tool
11
  from tabulate import tabulate
12
 
 
15
 
16
  # Initialize the model
17
  model = LiteLLMModel(
18
+ model_id=os.getenv("GEMINI_MODEL"), api_key=os.getenv("GEMINI_API_KEY")
 
19
  )
20
 
21
 
 
46
  output_type = "string"
47
 
48
  def forward(self, excel_path: str, sheet_name: Optional[str] = None) -> str:
49
+ """Load the Excel file and return the sheet as a Markdown table.
50
+
51
+ Args:
52
+ excel_path: Path to the Excel file.
53
+ sheet_name: Optional name or index of the sheet to read. If None, reads the first sheet.
54
+
55
+ Returns:
56
+ A Markdown table representing the Excel sheet, or an error message if the file is not found or cannot be read.
57
+ """
58
 
59
  file_path = Path(excel_path).expanduser().resolve()
60
  if not file_path.is_file():
 
62
 
63
  try:
64
  sheet: Union[str, int] = (
65
+ int(sheet_name)
66
+ if sheet_name and sheet_name.isdigit()
67
+ else sheet_name or 0
68
  )
69
 
70
  df = pd.read_excel(file_path, sheet_name=sheet)
 
82
  """An agent capable of using tools to answer general questions."""
83
 
84
  def __init__(self):
85
+ """Initializes the GaiaAgent with a set of tools."""
86
+
87
  print("GaiaAgent initialized with tools.")
88
 
89
  tools = [
 
102
  )
103
 
104
  def __call__(self, task_id: str, question: str) -> str:
105
+ """Processes a question using the agent and its tools.
106
+
107
+ Args:
108
+ task_id: A unique identifier for the task.
109
+ question: The question to be answered.
110
+
111
+ Returns:
112
+ The answer generated by the agent.
113
+ """
114
  print(f"Agent received task_id='{task_id}' | question='{question[:50]}...'")
115
  answer = self.agent.run(question)
116
  print(f"Agent returning answer: {answer}")