fix: resolve 36288 mypy errors (#37850)

Co-authored-by: WH-2099 <wh2099@pm.me>

api/core/helper/code_executor/code_executor.py

@@ -13,7 +13,7 @@ from core.helper.code_executor.jinja2.jinja2_transformer import Jinja2TemplateTr
 from core.helper.code_executor.python3.python3_transformer import Python3TemplateTransformer
 from core.helper.code_executor.template_transformer import TemplateTransformer
 from core.helper.http_client_pooling import get_pooled_http_client
-from graphon.nodes.code.entities import CodeLanguage
+from graphon.nodes.code.entities import CodeLanguage as CodeLanguage  # noqa: PLC0414
 
 logger = logging.getLogger(__name__)
 code_execution_endpoint_url = URL(str(dify_config.CODE_EXECUTION_ENDPOINT))
@@ -133,7 +133,9 @@ class CodeExecutor:
         return response_code.data.stdout or ""
 
     @classmethod
-    def execute_workflow_code_template(cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]):
+    def execute_workflow_code_template(
+        cls, language: CodeLanguage, code: str, inputs: Mapping[str, Any]
+    ) -> dict[str, Any]:
         """
         Execute code
         :param language: code language

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

@@ -11,7 +11,7 @@ class Jinja2TemplateTransformer(TemplateTransformer):
 
     @classmethod
     @override
-    def transform_response(cls, response: str):
+    def transform_response(cls, response: str) -> dict[str, Any]:
         """
         Transform response to dict
         :param response: response

api/core/helper/code_executor/template_transformer.py

@@ -36,14 +36,14 @@ class TemplateTransformer(ABC):
         return runner_script, preload_script
 
     @classmethod
-    def extract_result_str_from_response(cls, response: str):
+    def extract_result_str_from_response(cls, response: str) -> str:
         result = re.search(rf"{cls._result_tag}(.*){cls._result_tag}", response, re.DOTALL)
         if not result:
             raise ValueError(f"Failed to parse result: no result tag found in response. Response: {response[:200]}...")
         return result.group(1)
 
     @classmethod
-    def transform_response(cls, response: str) -> Mapping[str, Any]:
+    def transform_response(cls, response: str) -> dict[str, Any]:
         """
         Transform response to dict
         :param response: response
@@ -71,7 +71,7 @@ class TemplateTransformer(ABC):
         return result
 
     @classmethod
-    def _post_process_result(cls, result: dict[Any, Any]) -> dict[Any, Any]:
+    def _post_process_result(cls, result: dict[str, Any]) -> dict[str, Any]:
         """
         Post-process the result to convert scientific notation strings back to numbers
         """
@@ -89,7 +89,7 @@ class TemplateTransformer(ABC):
                     return [convert_scientific_notation(v) for v in value]
             return value
 
-        return convert_scientific_notation(result)
+        return {key: convert_scientific_notation(value) for key, value in result.items()}
 
     @classmethod
     @abstractmethod

api/core/helper/creators.py

@@ -24,7 +24,7 @@ def upload_dsl(dsl_file_bytes: bytes, filename: str = "template.yaml") -> str:
     response.raise_for_status()
     data = response.json()
     claim_code = data.get("data", {}).get("claim_code")
-    if not claim_code:
+    if not isinstance(claim_code, str) or not claim_code:
         raise ValueError("Creators Platform did not return a valid claim_code")
     return claim_code
 

api/core/helper/credential_utils.py

@@ -45,7 +45,7 @@ def is_credential_exists(credential_id: str, credential_type: "PluginCredentialT
 
 def runtime_check_credential_policy_compliance(
     credential_id: str, provider: str, credential_type: "PluginCredentialType", check_existence: bool = True
-):
+) -> None:
     if dify_config.ENTERPRISE_DISABLE_RUNTIME_CREDENTIAL_CHECK:
         return
     check_credential_policy_compliance(

api/core/helper/download.py

@@ -1,4 +1,7 @@
-def download_with_size_limit(url, max_download_size: int, **kwargs):
+from typing import Any
+
+
+def download_with_size_limit(url: str, max_download_size: int, **kwargs: Any) -> bytes:
     from core.file import remote_fetcher
 
     response = remote_fetcher.make_request("GET", url, follow_redirects=True, **kwargs)

api/core/helper/encrypter.py

@@ -1,5 +1,7 @@
 import base64
 
+from Crypto.PublicKey import RSA
+
 from libs import rsa
 
 
@@ -11,13 +13,13 @@ def obfuscated_token(token: str) -> str:
     return token[:6] + "*" * 12 + token[-2:]
 
 
-def full_mask_token(token_length=20):
+def full_mask_token(token_length: int = 20) -> str:
     return "*" * token_length
 
 
-def encrypt_token(tenant_id: str, token: str):
-    from extensions.ext_database import db
+def encrypt_token(tenant_id: str, token: str) -> str:
     from models.account import Tenant
+    from models.engine import db
 
     if not (tenant := db.session.get(Tenant, tenant_id)):
         raise ValueError(f"Tenant with id {tenant_id} not found")
@@ -30,15 +32,15 @@ def decrypt_token(tenant_id: str, token: str) -> str:
     return rsa.decrypt(base64.b64decode(token), tenant_id)
 
 
-def batch_decrypt_token(tenant_id: str, tokens: list[str]):
+def batch_decrypt_token(tenant_id: str, tokens: list[str]) -> list[str]:
     rsa_key, cipher_rsa = rsa.get_decrypt_decoding(tenant_id)
 
     return [rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa) for token in tokens]
 
 
-def get_decrypt_decoding(tenant_id: str):
+def get_decrypt_decoding(tenant_id: str) -> tuple[RSA.RsaKey, object]:
     return rsa.get_decrypt_decoding(tenant_id)
 
 
-def decrypt_token_with_decoding(token: str, rsa_key, cipher_rsa):
+def decrypt_token_with_decoding(token: str, rsa_key: RSA.RsaKey, cipher_rsa: object) -> str:
     return rsa.decrypt_token_with_decoding(base64.b64decode(token), rsa_key, cipher_rsa)

api/core/helper/marketplace.py

@@ -1,5 +1,6 @@
 import logging
 from collections.abc import Sequence
+from typing import Any
 from urllib.parse import urlencode
 
 import httpx
@@ -21,7 +22,7 @@ def get_plugin_pkg_url(plugin_unique_identifier: str) -> str:
     return f"{marketplace_api_url / 'api/v1/plugins/download'}?{query}"
 
 
-def download_plugin_pkg(plugin_unique_identifier: str):
+def download_plugin_pkg(plugin_unique_identifier: str) -> bytes:
     return download_with_size_limit(get_plugin_pkg_url(plugin_unique_identifier), dify_config.PLUGIN_MAX_PACKAGE_SIZE)
 
 
@@ -41,7 +42,7 @@ def batch_fetch_plugin_manifests(plugin_ids: list[str]) -> Sequence[MarketplaceP
     return [MarketplacePluginDeclaration.model_validate(plugin) for plugin in response.json()["data"]["plugins"]]
 
 
-def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict]:
+def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict[str, Any]]:
     if not plugin_ids:
         return []
 
@@ -55,10 +56,19 @@ def batch_fetch_plugin_by_ids(plugin_ids: list[str]) -> list[dict]:
     response.raise_for_status()
 
     data = response.json()
-    return data.get("data", {}).get("plugins", [])
+    plugins = data.get("data", {}).get("plugins", [])
+    if not isinstance(plugins, list):
+        raise ValueError("Marketplace did not return a valid plugins list")
+
+    result: list[dict[str, Any]] = []
+    for plugin in plugins:
+        if not isinstance(plugin, dict) or not all(isinstance(key, str) for key in plugin):
+            raise ValueError("Marketplace did not return a valid plugins list")
+        result.append(plugin)
+    return result
 
 
-def record_install_plugin_event(plugin_unique_identifier: str):
+def record_install_plugin_event(plugin_unique_identifier: str) -> None:
     url = str(marketplace_api_url / "api/v1/stats/plugins/install_count")
     response = httpx.post(url, json={"unique_identifier": plugin_unique_identifier}, timeout=MARKETPLACE_TIMEOUT)
     response.raise_for_status()

api/core/helper/model_provider_cache.py

@@ -34,7 +34,7 @@ class ProviderCredentialsCache:
         else:
             return None
 
-    def set(self, credentials: dict[str, Any]):
+    def set(self, credentials: dict[str, Any]) -> None:
         """
         Cache model provider credentials.
 
@@ -43,7 +43,7 @@ class ProviderCredentialsCache:
         """
         redis_client.setex(self.cache_key, 86400, json.dumps(credentials))
 
-    def delete(self):
+    def delete(self) -> None:
         """
         Delete cached model provider credentials.
 

api/core/helper/module_import_helper.py

@@ -20,17 +20,18 @@ def import_module_from_source[T: (str, bytes)](
                 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
-            # FIXME: mypy does not support the type of spec.loader
-            spec = importlib.util.spec_from_file_location(module_name, py_file_path)  # type: ignore[assignment]
-            if not spec or not spec.loader:
+            new_spec = importlib.util.spec_from_file_location(module_name, py_file_path)
+            if not new_spec or not new_spec.loader:
                 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)
+                new_spec.loader = importlib.util.LazyLoader(new_spec.loader)
+            spec = new_spec
         module = importlib.util.module_from_spec(spec)
         if not existed_spec:
             sys.modules[module_name] = module
-        spec.loader.exec_module(module)
+        if spec.loader is not None:
+            spec.loader.exec_module(module)
         return module
     except Exception as e:
         logger.exception("Failed to load module %s from script file '%s'", module_name, repr(py_file_path))

api/core/helper/provider_cache.py

@@ -9,11 +9,11 @@ from extensions.ext_redis import redis_client
 class ProviderCredentialsCache(ABC):
     """Base class for provider credentials cache"""
 
-    def __init__(self, **kwargs):
+    def __init__(self, **kwargs: Any) -> None:
         self.cache_key = self._generate_cache_key(**kwargs)
 
     @abstractmethod
-    def _generate_cache_key(self, **kwargs) -> str:
+    def _generate_cache_key(self, **kwargs: Any) -> str:
         """Generate cache key based on subclass implementation"""
         pass
 
@@ -28,11 +28,11 @@ class ProviderCredentialsCache(ABC):
                 return None
         return None
 
-    def set(self, config: dict[str, Any]):
+    def set(self, config: dict[str, Any]) -> None:
         """Cache provider credentials"""
         redis_client.setex(self.cache_key, 86400, json.dumps(config))
 
-    def delete(self):
+    def delete(self) -> None:
         """Delete cached provider credentials"""
         redis_client.delete(self.cache_key)
 
@@ -48,7 +48,7 @@ class SingletonProviderCredentialsCache(ProviderCredentialsCache):
         )
 
     @override
-    def _generate_cache_key(self, **kwargs) -> str:
+    def _generate_cache_key(self, **kwargs: Any) -> str:
         tenant_id = kwargs["tenant_id"]
         provider_type = kwargs["provider_type"]
         identity_name = kwargs["provider_identity"]
@@ -63,7 +63,7 @@ class ToolProviderCredentialsCache(ProviderCredentialsCache):
         super().__init__(tenant_id=tenant_id, provider=provider, credential_id=credential_id)
 
     @override
-    def _generate_cache_key(self, **kwargs) -> str:
+    def _generate_cache_key(self, **kwargs: Any) -> str:
         tenant_id = kwargs["tenant_id"]
         provider = kwargs["provider"]
         credential_id = kwargs["credential_id"]
@@ -77,10 +77,10 @@ class NoOpProviderCredentialCache:
         """Get cached provider credentials"""
         return None
 
-    def set(self, config: dict[str, Any]):
+    def set(self, config: dict[str, Any]) -> None:
         """Cache provider credentials"""
         pass
 
-    def delete(self):
+    def delete(self) -> None:
         """Delete cached provider credentials"""
         pass

api/core/helper/provider_encryption.py

@@ -125,5 +125,7 @@ class ProviderConfigEncrypter:
         return data
 
 
-def create_provider_encrypter(tenant_id: str, config: list[BasicProviderConfig], cache: ProviderConfigCache):
+def create_provider_encrypter(
+    tenant_id: str, config: list[BasicProviderConfig], cache: ProviderConfigCache
+) -> tuple[ProviderConfigEncrypter, ProviderConfigCache]:
     return ProviderConfigEncrypter(tenant_id=tenant_id, config=config, provider_config_cache=cache), cache

api/core/helper/tool_parameter_cache.py

@@ -37,11 +37,11 @@ class ToolParameterCache:
         else:
             return None
 
-    def set(self, parameters: dict[str, Any]):
+    def set(self, parameters: dict[str, Any]) -> None:
         """Cache model provider credentials."""
         redis_client.setex(self.cache_key, 86400, json.dumps(parameters))
 
-    def delete(self):
+    def delete(self) -> None:
         """
         Delete cached model provider credentials.
 

api/core/helper/trace_id_helper.py

@@ -61,7 +61,7 @@ def get_external_trace_id(request: Any) -> str | None:
     return None
 
 
-def extract_external_trace_id_from_args(args: Mapping[str, Any]):
+def extract_external_trace_id_from_args(args: Mapping[str, Any]) -> dict[str, Any]:
     """
     Extract 'external_trace_id' from args.
 

api/core/plugin/entities/plugin_daemon.py

@@ -228,7 +228,7 @@ class CredentialType(enum.StrEnum):
     OAUTH2 = "oauth2"
     UNAUTHORIZED = "unauthorized"
 
-    def get_name(self):
+    def get_name(self) -> str:
         if self == CredentialType.API_KEY:
             return "API KEY"
         elif self == CredentialType.OAUTH2:

api/providers/vdb/vdb-qdrant/src/dify_vdb_qdrant/qdrant_vector.py

@@ -33,7 +33,6 @@ from models.dataset import Dataset, DatasetCollectionBinding
 
 if TYPE_CHECKING:
     from qdrant_client.conversions import common_types
-    from qdrant_client.http import models as rest
 
     type DictFilter = dict[str, str | int | bool | dict | list]
     type MetadataFilter = DictFilter | common_types.Filter

api/providers/vdb/vdb-tidb-on-qdrant/src/dify_vdb_tidb_on_qdrant/tidb_on_qdrant_vector.py

@@ -41,7 +41,6 @@ from models.enums import TidbAuthBindingStatus
 if TYPE_CHECKING:
     from qdrant_client import grpc  # noqa
     from qdrant_client.conversions import common_types
-    from qdrant_client.http import models as rest
 
     type DictFilter = dict[str, str | int | bool | dict | list]
     type MetadataFilter = DictFilter | common_types.Filter

api/tests/unit_tests/core/helper/code_executor/test_template_transformer.py

@@ -1,6 +1,5 @@
 import json
 from base64 import b64decode
-from collections.abc import Mapping
 from typing import Any
 
 import pytest
@@ -44,7 +43,7 @@ def test_serialize_inputs_encodes_payload() -> None:
 def test_transform_response_parses_json_result_and_converts_scientific_notation() -> None:
     response = '<<RESULT>>{"value": "1e+3", "nested": {"x": "2E-2"}, "arr": ["3e+1"]}<<RESULT>>'
 
-    result: Mapping[str, Any] = _DummyTransformer.transform_response(response)
+    result: dict[str, Any] = _DummyTransformer.transform_response(response)
 
     assert result == {"value": 1000.0, "nested": {"x": 0.02}, "arr": [30.0]}
 

api/tests/unit_tests/core/helper/test_creators.py

@@ -46,6 +46,18 @@ class TestUploadDSL:
         with pytest.raises(ValueError, match="claim_code"):
             upload_dsl(b"app: demo")
 
+    @patch("core.helper.creators.httpx.post")
+    def test_raises_on_non_string_claim_code(self, mock_post):
+        mock_response = MagicMock(spec=httpx.Response)
+        mock_response.json.return_value = {"data": {"claim_code": 123}}
+        mock_response.raise_for_status = MagicMock()
+        mock_post.return_value = mock_response
+
+        from core.helper.creators import upload_dsl
+
+        with pytest.raises(ValueError, match="claim_code"):
+            upload_dsl(b"app: demo")
+
     @patch("core.helper.creators.httpx.post")
     def test_raises_on_http_error(self, mock_post):
         mock_response = MagicMock(spec=httpx.Response)

api/tests/unit_tests/core/helper/test_marketplace.py

@@ -1,6 +1,7 @@
 from types import SimpleNamespace
 from unittest.mock import MagicMock
 
+import pytest
 from pytest_mock import MockerFixture
 
 from core.helper.marketplace import (
@@ -53,6 +54,16 @@ def test_batch_fetch_plugin_by_ids_returns_plugins_from_response(mocker: MockerF
     response.raise_for_status.assert_called_once()
 
 
+def test_batch_fetch_plugin_by_ids_rejects_invalid_plugins_response(mocker: MockerFixture) -> None:
+    response = MagicMock()
+    response.json.return_value = {"data": {"plugins": ["p1"]}}
+    response.raise_for_status.return_value = None
+    mocker.patch("core.helper.marketplace.httpx.post", return_value=response)
+
+    with pytest.raises(ValueError, match="plugins list"):
+        batch_fetch_plugin_by_ids(["p1"])
+
+
 def test_batch_fetch_plugin_manifests_returns_empty_for_empty_input(mocker: MockerFixture) -> None:
     post_mock = mocker.patch("core.helper.marketplace.httpx.post")