estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

chore: refactor the http executor node (#5212)

Browse Source
This commit is contained in:
非法操作
2024-06-24 16:14:59 +08:00
committed by GitHub
parent 1e28a8c033
commit f7900f298f
8 changed files with 249 additions and 230 deletions
Download Patch File Download Diff File Expand all files Collapse all files

116
api/core/tools/tool/api_tool.py
View File

@@ -1,11 +1,9 @@
import json
from json import dumps
from os import getenv
from typing import Any, Union
from typing import Any
from urllib.parse import urlencode
import httpx
import requests
import core.helper.ssrf_proxy as ssrf_proxy
from core.tools.entities.tool_bundle import ApiToolBundle
@@ -18,12 +16,14 @@ API_TOOL_DEFAULT_TIMEOUT = (
int(getenv('API_TOOL_DEFAULT_READ_TIMEOUT', '60'))
)
class ApiTool(Tool):
api_bundle: ApiToolBundle
"""
Api tool
"""
def fork_tool_runtime(self, runtime: dict[str, Any]) -> 'Tool':
"""
fork a new tool with meta data
@@ -38,8 +38,9 @@ class ApiTool(Tool):
api_bundle=self.api_bundle.model_copy() if self.api_bundle else None,
runtime=Tool.Runtime(**runtime)
)
def validate_credentials(self, credentials: dict[str, Any], parameters: dict[str, Any], format_only: bool = False) -> str:
def validate_credentials(self, credentials: dict[str, Any], parameters: dict[str, Any],
format_only: bool = False) -> str:
"""
validate the credentials for Api tool
"""
@@ -47,7 +48,7 @@ class ApiTool(Tool):
headers = self.assembling_request(parameters)
if format_only:
return
return ''
response = self.do_http_request(self.api_bundle.server_url, self.api_bundle.method, headers, parameters)
# validate response
@@ -68,12 +69,12 @@ class ApiTool(Tool):
if 'api_key_header' in credentials:
api_key_header = credentials['api_key_header']
if 'api_key_value' not in credentials:
raise ToolProviderCredentialValidationError('Missing api_key_value')
elif not isinstance(credentials['api_key_value'], str):
raise ToolProviderCredentialValidationError('api_key_value must be a string')
if 'api_key_header_prefix' in credentials:
api_key_header_prefix = credentials['api_key_header_prefix']
if api_key_header_prefix == 'basic' and credentials['api_key_value']:
@@ -82,20 +83,20 @@ class ApiTool(Tool):
credentials['api_key_value'] = f'Bearer {credentials["api_key_value"]}'
elif api_key_header_prefix == 'custom':
pass
headers[api_key_header] = credentials['api_key_value']
needed_parameters = [parameter for parameter in self.api_bundle.parameters if parameter.required]
for parameter in needed_parameters:
if parameter.required and parameter.name not in parameters:
raise ToolParameterValidationError(f"Missing required parameter {parameter.name}")
if parameter.default is not None and parameter.name not in parameters:
parameters[parameter.name] = parameter.default
return headers
def validate_and_parse_response(self, response: Union[httpx.Response, requests.Response]) -> str:
def validate_and_parse_response(self, response: httpx.Response) -> str:
"""
validate the response
"""
@@ -112,23 +113,20 @@ class ApiTool(Tool):
return json.dumps(response)
except Exception as e:
return response.text
elif isinstance(response, requests.Response):
if not response.ok:
raise ToolInvokeError(f"Request failed with status code {response.status_code} and {response.text}")
if not response.content:
return 'Empty response from the tool, please check your parameters and try again.'
try:
response = response.json()
try:
return json.dumps(response, ensure_ascii=False)
except Exception as e:
return json.dumps(response)
except Exception as e:
return response.text
else:
raise ValueError(f'Invalid response type {type(response)}')
def do_http_request(self, url: str, method: str, headers: dict[str, Any], parameters: dict[str, Any]) -> httpx.Response:
@staticmethod
def get_parameter_value(parameter, parameters):
if parameter['name'] in parameters:
return parameters[parameter['name']]
elif parameter.get('required', False):
raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}")
else:
return (parameter.get('schema', {}) or {}).get('default', '')
def do_http_request(self, url: str, method: str, headers: dict[str, Any],
parameters: dict[str, Any]) -> httpx.Response:
"""
do http request depending on api bundle
"""
@@ -141,44 +139,17 @@ class ApiTool(Tool):
# check parameters
for parameter in self.api_bundle.openapi.get('parameters', []):
value = self.get_parameter_value(parameter, parameters)
if parameter['in'] == 'path':
value = ''
if parameter['name'] in parameters:
value = parameters[parameter['name']]
elif parameter['required']:
raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}")
else:
value = (parameter.get('schema', {}) or {}).get('default', '')
path_params[parameter['name']] = value
elif parameter['in'] == 'query':
value = ''
if parameter['name'] in parameters:
value = parameters[parameter['name']]
elif parameter.get('required', False):
raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}")
else:
value = (parameter.get('schema', {}) or {}).get('default', '')
params[parameter['name']] = value
elif parameter['in'] == 'cookie':
value = ''
if parameter['name'] in parameters:
value = parameters[parameter['name']]
elif parameter.get('required', False):
raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}")
else:
value = (parameter.get('schema', {}) or {}).get('default', '')
cookies[parameter['name']] = value
elif parameter['in'] == 'header':
value = ''
if parameter['name'] in parameters:
value = parameters[parameter['name']]
elif parameter.get('required', False):
raise ToolParameterValidationError(f"Missing required parameter {parameter['name']}")
else:
value = (parameter.get('schema', {}) or {}).get('default', '')
headers[parameter['name']] = value
# check if there is a request body and handle it
@@ -203,7 +174,7 @@ class ApiTool(Tool):
else:
body[name] = None
break
# replace path parameters
for name, value in path_params.items():
url = url.replace(f'{{{name}}}', f'{value}')
@@ -211,33 +182,21 @@ class ApiTool(Tool):
# parse http body data if needed, for GET/HEAD/OPTIONS/TRACE, the body is ignored
if 'Content-Type' in headers:
if headers['Content-Type'] == 'application/json':
body = dumps(body)
body = json.dumps(body)
elif headers['Content-Type'] == 'application/x-www-form-urlencoded':
body = urlencode(body)
else:
body = body
# do http request
if method == 'get':
response = ssrf_proxy.get(url, params=params, headers=headers, cookies=cookies, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
elif method == 'post':
response = ssrf_proxy.post(url, params=params, headers=headers, cookies=cookies, data=body, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
elif method == 'put':
response = ssrf_proxy.put(url, params=params, headers=headers, cookies=cookies, data=body, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
elif method == 'delete':
response = ssrf_proxy.delete(url, params=params, headers=headers, cookies=cookies, data=body, timeout=API_TOOL_DEFAULT_TIMEOUT, allow_redirects=True)
elif method == 'patch':
response = ssrf_proxy.patch(url, params=params, headers=headers, cookies=cookies, data=body, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
elif method == 'head':
response = ssrf_proxy.head(url, params=params, headers=headers, cookies=cookies, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
elif method == 'options':
response = ssrf_proxy.options(url, params=params, headers=headers, cookies=cookies, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
if method in ('get', 'head', 'post', 'put', 'delete', 'patch'):
response = getattr(ssrf_proxy, method)(url, params=params, headers=headers, cookies=cookies, data=body,
timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
return response
else:
raise ValueError(f'Invalid http method {method}')
return response
def _convert_body_property_any_of(self, property: dict[str, Any], value: Any, any_of: list[dict[str, Any]], max_recursive=10) -> Any:
raise ValueError(f'Invalid http method {self.method}')
def _convert_body_property_any_of(self, property: dict[str, Any], value: Any, any_of: list[dict[str, Any]],
max_recursive=10) -> Any:
if max_recursive <= 0:
raise Exception("Max recursion depth reached")
for option in any_of or []:
@@ -322,4 +281,3 @@ class ApiTool(Tool):
# assemble invoke message
return self.create_text_message(response)