瀏覽代碼

fix(tool_file_manager): raise ValueError when get timeout (#11928)

Signed-off-by: -LAN- <laipz8200@outlook.com>
-LAN- 4 月之前
父節點
當前提交
90323cd355
共有 1 個文件被更改,包括 5 次插入5 次删除
  1. 5 5
      api/core/tools/tool_file_manager.py

+ 5 - 5
api/core/tools/tool_file_manager.py

@@ -8,9 +8,10 @@ from mimetypes import guess_extension, guess_type
 from typing import Optional, Union
 from uuid import uuid4
 
-from httpx import get
+import httpx
 
 from configs import dify_config
+from core.helper import ssrf_proxy
 from extensions.ext_database import db
 from extensions.ext_storage import storage
 from models.model import MessageFile
@@ -94,12 +95,11 @@ class ToolFileManager:
     ) -> ToolFile:
         # try to download image
         try:
-            response = get(file_url)
+            response = ssrf_proxy.get(file_url)
             response.raise_for_status()
             blob = response.content
-        except Exception as e:
-            logger.exception(f"Failed to download file from {file_url}")
-            raise
+        except httpx.TimeoutException as e:
+            raise ValueError(f"timeout when downloading file from {file_url}")
 
         mimetype = guess_type(file_url)[0] or "octet/stream"
         extension = guess_extension(mimetype) or ".bin"