Merge branch 'main' into fix/chore-fix

api/core/helper/code_executor/code_executor.py

@@ -38,7 +38,7 @@ class CodeLanguage(StrEnum):
 
 
 class CodeExecutor:
-    dependencies_cache = {}
+    dependencies_cache: dict[str, str] = {}
     dependencies_cache_lock = Lock()
 
     code_template_transformers: dict[CodeLanguage, type[TemplateTransformer]] = {
@@ -103,22 +103,22 @@ class CodeExecutor:
             )
 
         try:
-            response = response.json()
+            response_data = response.json()
         except:
             raise CodeExecutionError("Failed to parse response")
 
-        if (code := response.get("code")) != 0:
-            raise CodeExecutionError(f"Got error code: {code}. Got error msg: {response.get('message')}")
+        if (code := response_data.get("code")) != 0:
+            raise CodeExecutionError(f"Got error code: {code}. Got error msg: {response_data.get('message')}")
 
-        response = CodeExecutionResponse(**response)
+        response_code = CodeExecutionResponse(**response_data)
 
-        if response.data.error:
-            raise CodeExecutionError(response.data.error)
+        if response_code.data.error:
+            raise CodeExecutionError(response_code.data.error)
 
-        return response.data.stdout or ""
+        return response_code.data.stdout or ""
 
     @classmethod
-    def execute_workflow_code_template(cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]) -> dict:
+    def execute_workflow_code_template(cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]):
         """
         Execute code
         :param language: code language

api/core/helper/code_executor/jinja2/jinja2_formatter.py

@@ -1,9 +1,11 @@
+from collections.abc import Mapping
+
 from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
 
 
 class Jinja2Formatter:
     @classmethod
-    def format(cls, template: str, inputs: dict) -> str:
+    def format(cls, template: str, inputs: Mapping[str, str]) -> str:
         """
         Format template
         :param template: template
@@ -11,5 +13,4 @@ class Jinja2Formatter:
         :return:
         """
         result = CodeExecutor.execute_workflow_code_template(language=CodeLanguage.JINJA2, code=template, inputs=inputs)
-
-        return result["result"]
+        return str(result.get("result", ""))

api/core/helper/code_executor/template_transformer.py

@@ -25,21 +25,28 @@ class TemplateTransformer(ABC):
         return runner_script, preload_script
 
     @classmethod
-    def extract_result_str_from_response(cls, response: str) -> str:
+    def extract_result_str_from_response(cls, response: str):
         result = re.search(rf"{cls._result_tag}(.*){cls._result_tag}", response, re.DOTALL)
         if not result:
             raise ValueError("Failed to parse result")
-        result = result.group(1)
-        return result
+        return result.group(1)
 
     @classmethod
-    def transform_response(cls, response: str) -> dict:
+    def transform_response(cls, response: str) -> Mapping[str, Any]:
         """
         Transform response to dict
         :param response: response
         :return:
         """
-        return json.loads(cls.extract_result_str_from_response(response))
+        try:
+            result = json.loads(cls.extract_result_str_from_response(response))
+        except json.JSONDecodeError:
+            raise ValueError("failed to parse response")
+        if not isinstance(result, dict):
+            raise ValueError("result must be a dict")
+        if not all(isinstance(k, str) for k in result):
+            raise ValueError("result keys must be strings")
+        return result
 
     @classmethod
     @abstractmethod

api/core/helper/encrypter.py

@@ -1,6 +1,5 @@
 import base64
 
-from extensions.ext_database import db
 from libs import rsa
 
 
@@ -14,6 +13,7 @@ def obfuscated_token(token: str):
 
 def encrypt_token(tenant_id: str, token: str):
     from models.account import Tenant
+    from models.engine import db
 
     if not (tenant := db.session.query(Tenant).filter(Tenant.id == tenant_id).first()):
         raise ValueError(f"Tenant with id {tenant_id} not found")

api/core/helper/lru_cache.py

@@ -4,7 +4,7 @@ from typing import Any
 
 class LRUCache:
     def __init__(self, capacity: int):
-        self.cache = OrderedDict()
+        self.cache: OrderedDict[Any, Any] = OrderedDict()
         self.capacity = capacity
 
     def get(self, key: Any) -> Any:

api/core/helper/model_provider_cache.py

@@ -30,7 +30,7 @@ class ProviderCredentialsCache:
             except JSONDecodeError:
                 return None
 
-            return cached_provider_credentials
+            return dict(cached_provider_credentials)
         else:
             return None
 

api/core/helper/moderation.py

@@ -54,7 +54,7 @@ def check_moderation(tenant_id: str, model_config: ModelConfigWithCredentialsEnt
 
                 if moderation_result is True:
                     return True
-            except Exception as ex:
+            except Exception:
                 logger.exception(f"Fails to check moderation, provider_name: {provider_name}")
                 raise InvokeBadRequestError("Rate limit exceeded, please try again later.")
 

api/core/helper/module_import_helper.py

@@ -14,12 +14,13 @@ def import_module_from_source(*, module_name: str, py_file_path: AnyStr, use_laz
         if existed_spec:
             spec = existed_spec
             if not spec.loader:
-                raise Exception(f"Failed to load module {module_name} from {py_file_path}")
+                raise Exception(f"Failed to load module {module_name} from {py_file_path!r}")
         else:
             # Refer to: https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly
-            spec = importlib.util.spec_from_file_location(module_name, py_file_path)
+            # FIXME: mypy does not support the type of spec.loader
+            spec = importlib.util.spec_from_file_location(module_name, py_file_path)  # type: ignore
             if not spec or not spec.loader:
-                raise Exception(f"Failed to load module {module_name} from {py_file_path}")
+                raise Exception(f"Failed to load module {module_name} from {py_file_path!r}")
             if use_lazy_loader:
                 # Refer to: https://docs.python.org/3/library/importlib.html#implementing-lazy-imports
                 spec.loader = importlib.util.LazyLoader(spec.loader)
@@ -29,7 +30,7 @@ def import_module_from_source(*, module_name: str, py_file_path: AnyStr, use_laz
         spec.loader.exec_module(module)
         return module
     except Exception as e:
-        logging.exception(f"Failed to load module {module_name} from script file '{py_file_path}'")
+        logging.exception(f"Failed to load module {module_name} from script file '{py_file_path!r}'")
         raise e
 
 
@@ -57,6 +58,6 @@ def load_single_subclass_from_source(
         case 1:
             return subclasses[0]
         case 0:
-            raise Exception(f"Missing subclass of {parent_type.__name__} in {script_path}")
+            raise Exception(f"Missing subclass of {parent_type.__name__} in {script_path!r}")
         case _:
-            raise Exception(f"Multiple subclasses of {parent_type.__name__} in {script_path}")
+            raise Exception(f"Multiple subclasses of {parent_type.__name__} in {script_path!r}")

api/core/helper/ssrf_proxy.py

@@ -24,7 +24,7 @@ BACKOFF_FACTOR = 0.5
 STATUS_FORCELIST = [429, 500, 502, 503, 504]
 
 
-class MaxRetriesExceededError(Exception):
+class MaxRetriesExceededError(ValueError):
     """Raised when the maximum number of retries is exceeded."""
 
     pass
@@ -65,11 +65,12 @@ def make_request(method, url, max_retries=SSRF_DEFAULT_MAX_RETRIES, **kwargs):
 
         except httpx.RequestError as e:
             logging.warning(f"Request to URL {url} failed on attempt {retries + 1}: {e}")
+            if max_retries == 0:
+                raise
 
         retries += 1
         if retries <= max_retries:
             time.sleep(BACKOFF_FACTOR * (2 ** (retries - 1)))
-
     raise MaxRetriesExceededError(f"Reached maximum retries ({max_retries}) for URL {url}")
 
 

api/core/helper/tool_parameter_cache.py

@@ -33,7 +33,7 @@ class ToolParameterCache:
             except JSONDecodeError:
                 return None
 
-            return cached_tool_parameter
+            return dict(cached_tool_parameter)
         else:
             return None
 

api/core/helper/tool_provider_cache.py

@@ -29,7 +29,7 @@ class ToolProviderCredentialsCache:
             except JSONDecodeError:
                 return None
 
-            return cached_provider_credentials
+            return dict(cached_provider_credentials)
         else:
             return None