File size: 1,578 Bytes
f61048a |
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 |
# Libs
import os
import requests
import pandas as pd
from smolagents import Tool
# Local
from consts import DEFAULT_API_URL
class GetTaskFileTool(Tool):
name = "get_task_file_tool"
description = """This tool downloads the file content associated with the given task_id if exists. Returns absolute file path"""
inputs = {
"task_id": {"type": "string", "description": "Task id"},
"file_name": {"type": "string", "description": "File name"},
}
output_type = "string"
def forward(self, task_id: str, file_name: str) -> str:
response = requests.get(f"{DEFAULT_API_URL}/files/{task_id}", timeout=15)
response.raise_for_status()
with open(file_name, 'wb') as file:
file.write(response.content)
return os.path.abspath(file_name)
class LoadXlsxFileTool(Tool):
name = "load_xlsx_file_tool"
description = """This tool loads xlsx file into pandas and returns it"""
inputs = {
"file_path": {"type": "string", "description": "File path"}
}
output_type = "object"
def forward(self, file_path: str) -> object:
return pd.read_excel(file_path)
class LoadTextFileTool(Tool):
name = "load_text_file_tool"
description = """This tool loads any text file"""
inputs = {
"file_path": {"type": "string", "description": "File path"}
}
output_type = "string"
def forward(self, file_path: str) -> object:
with open(file_path, 'r', encoding='utf-8') as file:
return file.read()
|