File size: 2,681 Bytes
b5d547f |
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 |
from smolagents import Tool
import requests
from urllib.parse import urljoin
import base64
import tempfile
class GetAttachmentTool(Tool):
name = "get_attachment"
description = """Retrieves attachment for current task in specified format."""
inputs = {
"fmt": {
"type": "string",
"description": "Format to retrieve attachment. Options are: URL (preferred), DATA_URL, LOCAL_FILE_PATH, TEXT. URL returns the URL of the file, DATA_URL returns a base64 encoded data URL, LOCAL_FILE_PATH returns a local file path to the downloaded file, and TEXT returns the content of the file as text.",
"nullable": True,
"default": "URL",
}
}
output_type = "string"
def __init__(
self,
agent_evaluation_api: str | None = None,
task_id: str | None = None,
**kwargs,
):
self.agent_evaluation_api = (
agent_evaluation_api
if agent_evaluation_api is not None
else "https://agents-course-unit4-scoring.hf.space/"
)
self.task_id = task_id
super().__init__(**kwargs)
def attachment_for(self, task_id: str | None):
self.task_id = task_id
def forward(self, fmt: str = "URL") -> str:
fmt = fmt.upper()
assert fmt in ["URL", "DATA_URL", "LOCAL_FILE_PATH", "TEXT"]
if not self.task_id:
return ""
file_url = urljoin(self.agent_evaluation_api, f"files/{self.task_id}")
if fmt == "URL":
return file_url
response = requests.get(
file_url,
headers={
"Content-Type": "application/json",
"Accept": "application/json",
},
)
if 400 <= response.status_code < 500:
return ""
response.raise_for_status()
mime = response.headers.get("content-type", "text/plain")
if fmt == "TEXT":
if mime.startswith("text/"):
return response.text
else:
raise ValueError(
f"Content of file type {mime} cannot be retrieved as TEXT."
)
elif fmt == "DATA_URL":
return f"data:{mime};base64,{base64.b64encode(response.content).decode('utf-8')}"
elif fmt == "LOCAL_FILE_PATH":
with tempfile.NamedTemporaryFile(delete=False) as tmp_file:
tmp_file.write(response.content)
return tmp_file.name
else:
raise ValueError(
f"Unsupported format: {fmt}. Supported formats are URL, DATA_URL, LOCAL_FILE_PATH, and TEXT."
)
|