feat: support custom tool upload file (#10796)

This commit is contained in:
非法操作 2024-11-20 13:26:42 +08:00 committed by GitHub
parent 25fda7adc5
commit bf4b6e5f80
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 13 additions and 3 deletions

View File

@ -5,6 +5,7 @@ from urllib.parse import urlencode
import httpx
from core.file.file_manager import download
from core.helper import ssrf_proxy
from core.tools.entities.tool_bundle import ApiToolBundle
from core.tools.entities.tool_entities import ToolInvokeMessage, ToolProviderType
@ -138,6 +139,7 @@ class ApiTool(Tool):
path_params = {}
body = {}
cookies = {}
files = []
# check parameters
for parameter in self.api_bundle.openapi.get("parameters", []):
@ -166,6 +168,10 @@ class ApiTool(Tool):
properties = body_schema.get("properties", {})
for name, property in properties.items():
if name in parameters:
if property.get("format") == "binary":
f = parameters[name]
files.append((name, (f.filename, download(f), f.mime_type)))
else:
# convert type
body[name] = self._convert_body_property_type(property, parameters[name])
elif name in required:
@ -182,7 +188,7 @@ class ApiTool(Tool):
for name, value in path_params.items():
url = url.replace(f"{{{name}}}", f"{value}")
# parse http body data if needed, for GET/HEAD/OPTIONS/TRACE, the body is ignored
# parse http body data if needed
if "Content-Type" in headers:
if headers["Content-Type"] == "application/json":
body = json.dumps(body)
@ -198,6 +204,7 @@ class ApiTool(Tool):
headers=headers,
cookies=cookies,
data=body,
files=files,
timeout=API_TOOL_DEFAULT_TIMEOUT,
follow_redirects=True,
)

View File

@ -161,6 +161,9 @@ class ApiBasedToolSchemaParser:
def _get_tool_parameter_type(parameter: dict) -> ToolParameter.ToolParameterType:
parameter = parameter or {}
typ = None
if parameter.get("format") == "binary":
return ToolParameter.ToolParameterType.FILE
if "type" in parameter:
typ = parameter["type"]
elif "schema" in parameter and "type" in parameter["schema"]: