Skip to content

Commit

Permalink
enhance: custom tool timeout (#2754)
Browse files Browse the repository at this point in the history
  • Loading branch information
Yeuoly authored Mar 8, 2024
1 parent 40c646c commit b6b58da
Showing 1 changed file with 8 additions and 7 deletions.
15 changes: 8 additions & 7 deletions api/core/tools/tool/api_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,7 @@
from core.tools.errors import ToolProviderCredentialValidationError
from core.tools.tool.tool import Tool

API_TOOL_DEFAULT_TIMEOUT = (10, 60)

class ApiTool(Tool):
api_bundle: ApiBasedToolBundle
Expand Down Expand Up @@ -211,19 +212,19 @@ def do_http_request(self, url: str, method: str, headers: dict[str, Any], parame

# do http request
if method == 'get':
response = ssrf_proxy.get(url, params=params, headers=headers, cookies=cookies, timeout=10, follow_redirects=True)
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=10, follow_redirects=True)
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=10, follow_redirects=True)
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=10, allow_redirects=True)
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=10, follow_redirects=True)
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=10, follow_redirects=True)
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=10, follow_redirects=True)
response = ssrf_proxy.options(url, params=params, headers=headers, cookies=cookies, timeout=API_TOOL_DEFAULT_TIMEOUT, follow_redirects=True)
else:
raise ValueError(f'Invalid http method {method}')

Expand Down

0 comments on commit b6b58da

Please sign in to comment.