SergeyO7 commited on
Commit
ab5793d
·
verified ·
1 Parent(s): 8f70d65

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +76 -1
agent.py CHANGED
@@ -9,7 +9,7 @@ import yaml
9
  from PIL import Image
10
  import requests
11
  from io import BytesIO
12
-
13
  import whisper
14
 
15
  # Simulated additional tools (implementation depends on external APIs or setup)
@@ -80,6 +80,79 @@ import whisper
80
  #
81
  # return f"The image description: '{response}'"
82
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
83
  @tool
84
  def SpeechToTextTool(audio_path: str) -> str:
85
  """Tool for converting an audio file to text using OpenAI Whisper.
@@ -108,6 +181,7 @@ class ExcelReaderTool(Tool):
108
  },
109
  "sheet_name": {
110
  "type": "string",
 
111
  "description": "The name of the sheet to read (optional, defaults to first sheet)",
112
  "nullable": True
113
  }
@@ -187,6 +261,7 @@ class MagAgent:
187
  model= model,
188
  tools=[
189
  # GoogleSearchTool,
 
190
  DuckDuckGoSearchTool(),
191
  WikipediaSearchTool(),
192
  # ImageAnalysisTool,
 
9
  from PIL import Image
10
  import requests
11
  from io import BytesIO
12
+ from markdownify import markdownify
13
  import whisper
14
 
15
  # Simulated additional tools (implementation depends on external APIs or setup)
 
80
  #
81
  # return f"The image description: '{response}'"
82
 
83
+ class VisitWebpageTool(Tool):
84
+ name = "visit_webpage"
85
+ description = "Visits a webpage at the given url and reads its content as a markdown string. Use this to browse webpages."
86
+ inputs = {'url': {'type': 'string', 'description': 'The url of the webpage to visit.'}}
87
+ output_type = "string"
88
+
89
+ def forward(self, url: str) -> str:
90
+ try:
91
+ import requests
92
+ from markdownify import markdownify
93
+ from requests.exceptions import RequestException
94
+
95
+ from smolagents.utils import truncate_content
96
+ except ImportError as e:
97
+ raise ImportError(
98
+ "You must install packages `markdownify` and `requests` to run this tool: for instance run `pip install markdownify requests`."
99
+ ) from e
100
+ try:
101
+ # Send a GET request to the URL with a 20-second timeout
102
+ response = requests.get(url, timeout=20)
103
+ response.raise_for_status() # Raise an exception for bad status codes
104
+
105
+ # Convert the HTML content to Markdown
106
+ markdown_content = markdownify(response.text).strip()
107
+
108
+ # Remove multiple line breaks
109
+ markdown_content = re.sub(r"\n{3,}", "\n\n", markdown_content)
110
+
111
+ return truncate_content(markdown_content, 10000)
112
+
113
+ except requests.exceptions.Timeout:
114
+ return "The request timed out. Please try again later or check the URL."
115
+ except RequestException as e:
116
+ return f"Error fetching the webpage: {str(e)}"
117
+ except Exception as e:
118
+ return f"An unexpected error occurred: {str(e)}"
119
+
120
+ def __init__(self, *args, **kwargs):
121
+ self.is_initialized = False
122
+
123
+ class DownloadTaskAttachmentTool(Tool):
124
+ name = "download_file"
125
+ description = "Downloads the file attached to the task ID"
126
+ inputs = {'task_id': {'type': 'string', 'description': 'The task id to download attachment from.'}}
127
+ output_type = "string"
128
+
129
+ def forward(self, task_id: str) -> str:
130
+ """
131
+ Downloads a file associated with the given task ID.
132
+ Returns the file path where the file is saved locally.
133
+ """
134
+ file_url = f"{DEFAULT_API_URL}/files/{task_id}"
135
+ local_file_path = f"downloads/{task_id}.file"
136
+
137
+ print(f"Downloading file for task ID {task_id} from {file_url}...")
138
+ try:
139
+ response = requests.get(file_url, stream=True, timeout=15)
140
+ response.raise_for_status()
141
+
142
+ os.makedirs("downloads", exist_ok=True)
143
+ with open(local_file_path, "wb") as file:
144
+ for chunk in response.iter_content(chunk_size=8192):
145
+ file.write(chunk)
146
+
147
+ print(f"File downloaded successfully: {local_file_path}")
148
+ return local_file_path
149
+ except requests.exceptions.RequestException as e:
150
+ print(f"Error downloading file for task {task_id}: {e}")
151
+ raise
152
+
153
+ def __init__(self, *args, **kwargs):
154
+ self.is_initialized = False
155
+
156
  @tool
157
  def SpeechToTextTool(audio_path: str) -> str:
158
  """Tool for converting an audio file to text using OpenAI Whisper.
 
181
  },
182
  "sheet_name": {
183
  "type": "string",
184
+
185
  "description": "The name of the sheet to read (optional, defaults to first sheet)",
186
  "nullable": True
187
  }
 
261
  model= model,
262
  tools=[
263
  # GoogleSearchTool,
264
+ DownloadTaskAttachmentTool(),
265
  DuckDuckGoSearchTool(),
266
  WikipediaSearchTool(),
267
  # ImageAnalysisTool,