feat: support custom tool upload file (langgenius#10796)

AlwaysBluer · Nov 20, 2024 · bf4b6e5 · bf4b6e5
1 parent 25fda7a
commit bf4b6e5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 3 deletions.
diff --git a/api/core/tools/tool/api_tool.py b/api/core/tools/tool/api_tool.py
@@ -5,6 +5,7 @@
 
 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 @@ def do_http_request(
         path_params = {}
         body = {}
         cookies = {}
+        files = []
 
         # check parameters
         for parameter in self.api_bundle.openapi.get("parameters", []):
@@ -166,8 +168,12 @@ def do_http_request(
                     properties = body_schema.get("properties", {})
                     for name, property in properties.items():
                         if name in parameters:
-                            # convert type
-                            body[name] = self._convert_body_property_type(property, parameters[name])
+                            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:
                             raise ToolParameterValidationError(
                                 f"Missing required parameter {name} in operation {self.api_bundle.operation_id}"
@@ -182,7 +188,7 @@ def do_http_request(
         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 @@ def do_http_request(
                 headers=headers,
                 cookies=cookies,
                 data=body,
+                files=files,
                 timeout=API_TOOL_DEFAULT_TIMEOUT,
                 follow_redirects=True,
             )

diff --git a/api/core/tools/utils/parser.py b/api/core/tools/utils/parser.py
@@ -161,6 +161,9 @@ def parse_openapi_to_tool_bundle(
     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"]: