Merge remote-tracking branch 'origin/main' into feat/trigger

api/core/tools/__base/tool.py

@@ -1,7 +1,7 @@
 from abc import ABC, abstractmethod
 from collections.abc import Generator
 from copy import deepcopy
-from typing import TYPE_CHECKING, Any, Optional
+from typing import TYPE_CHECKING, Any
 
 if TYPE_CHECKING:
     from models.model import File
@@ -20,7 +20,7 @@ class Tool(ABC):
     The base class of a tool
     """
 
-    def __init__(self, entity: ToolEntity, runtime: ToolRuntime) -> None:
+    def __init__(self, entity: ToolEntity, runtime: ToolRuntime):
         self.entity = entity
         self.runtime = runtime
 
@@ -46,9 +46,9 @@ class Tool(ABC):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage]:
         if self.runtime and self.runtime.runtime_parameters:
             tool_parameters.update(self.runtime.runtime_parameters)
@@ -96,17 +96,17 @@ class Tool(ABC):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> ToolInvokeMessage | list[ToolInvokeMessage] | Generator[ToolInvokeMessage, None, None]:
         pass
 
     def get_runtime_parameters(
         self,
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> list[ToolParameter]:
         """
         get the runtime parameters
@@ -119,9 +119,9 @@ class Tool(ABC):
 
     def get_merged_runtime_parameters(
         self,
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> list[ToolParameter]:
         """
         get merged runtime parameters
@@ -196,7 +196,7 @@ class Tool(ABC):
             message=ToolInvokeMessage.TextMessage(text=text),
         )
 
-    def create_blob_message(self, blob: bytes, meta: Optional[dict] = None) -> ToolInvokeMessage:
+    def create_blob_message(self, blob: bytes, meta: dict | None = None) -> ToolInvokeMessage:
         """
         create a blob message
 

api/core/tools/__base/tool_provider.py

@@ -12,7 +12,7 @@ from core.tools.errors import ToolProviderCredentialValidationError
 
 
 class ToolProviderController(ABC):
-    def __init__(self, entity: ToolProviderEntity) -> None:
+    def __init__(self, entity: ToolProviderEntity):
         self.entity = entity
 
     def get_credentials_schema(self) -> list[ProviderConfig]:
@@ -41,7 +41,7 @@ class ToolProviderController(ABC):
         """
         return ToolProviderType.BUILT_IN
 
-    def validate_credentials_format(self, credentials: dict[str, Any]) -> None:
+    def validate_credentials_format(self, credentials: dict[str, Any]):
         """
         validate the format of the credentials of the provider and set the default value if needed
 

api/core/tools/__base/tool_runtime.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any
 
 from openai import BaseModel
 from pydantic import Field
@@ -14,9 +14,9 @@ class ToolRuntime(BaseModel):
     """
 
     tenant_id: str
-    tool_id: Optional[str] = None
-    invoke_from: Optional[InvokeFrom] = None
-    tool_invoke_from: Optional[ToolInvokeFrom] = None
+    tool_id: str | None = None
+    invoke_from: InvokeFrom | None = None
+    tool_invoke_from: ToolInvokeFrom | None = None
     credentials: dict[str, Any] = Field(default_factory=dict)
     credential_type: CredentialType = Field(default=CredentialType.API_KEY)
     runtime_parameters: dict[str, Any] = Field(default_factory=dict)

api/core/tools/builtin_tool/provider.py

@@ -18,20 +18,20 @@ from core.tools.entities.values import ToolLabelEnum, default_tool_label_dict
 from core.tools.errors import (
     ToolProviderNotFoundError,
 )
-from core.tools.utils.yaml_utils import load_yaml_file
+from core.tools.utils.yaml_utils import load_yaml_file_cached
 
 
 class BuiltinToolProviderController(ToolProviderController):
     tools: list[BuiltinTool]
 
-    def __init__(self, **data: Any) -> None:
+    def __init__(self, **data: Any):
         self.tools = []
 
         # load provider yaml
         provider = self.__class__.__module__.split(".")[-1]
         yaml_path = path.join(path.dirname(path.realpath(__file__)), "providers", provider, f"{provider}.yaml")
         try:
-            provider_yaml = load_yaml_file(yaml_path, ignore_error=False)
+            provider_yaml = load_yaml_file_cached(yaml_path)
         except Exception as e:
             raise ToolProviderNotFoundError(f"can not load provider yaml for {provider}: {e}")
 
@@ -71,10 +71,10 @@ class BuiltinToolProviderController(ToolProviderController):
         for tool_file in tool_files:
             # get tool name
             tool_name = tool_file.split(".")[0]
-            tool = load_yaml_file(path.join(tool_path, tool_file), ignore_error=False)
+            tool = load_yaml_file_cached(path.join(tool_path, tool_file))
 
             # get tool class, import the module
-            assistant_tool_class: type[BuiltinTool] = load_single_subclass_from_source(
+            assistant_tool_class: type = load_single_subclass_from_source(
                 module_name=f"core.tools.builtin_tool.providers.{provider}.tools.{tool_name}",
                 script_path=path.join(
                     path.dirname(path.realpath(__file__)),
@@ -197,7 +197,7 @@ class BuiltinToolProviderController(ToolProviderController):
         """
         return self.entity.identity.tags or []
 
-    def validate_credentials(self, user_id: str, credentials: dict[str, Any]) -> None:
+    def validate_credentials(self, user_id: str, credentials: dict[str, Any]):
         """
         validate the credentials of the provider
 
@@ -211,7 +211,7 @@ class BuiltinToolProviderController(ToolProviderController):
         self._validate_credentials(user_id, credentials)
 
     @abstractmethod
-    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]) -> None:
+    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]):
         """
         validate the credentials of the provider
 

api/core/tools/builtin_tool/providers/audio/audio.py

@@ -4,5 +4,5 @@ from core.tools.builtin_tool.provider import BuiltinToolProviderController
 
 
 class AudioToolProvider(BuiltinToolProviderController):
-    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]) -> None:
+    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]):
         pass

api/core/tools/builtin_tool/providers/audio/tools/asr.py

@@ -1,6 +1,6 @@
 import io
 from collections.abc import Generator
-from typing import Any, Optional
+from typing import Any
 
 from core.file.enums import FileType
 from core.file.file_manager import download
@@ -18,9 +18,9 @@ class ASRTool(BuiltinTool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         file = tool_parameters.get("audio_file")
         if file.type != FileType.AUDIO:  # type: ignore
@@ -56,9 +56,9 @@ class ASRTool(BuiltinTool):
 
     def get_runtime_parameters(
         self,
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> list[ToolParameter]:
         parameters = []
 

api/core/tools/builtin_tool/providers/audio/tools/tts.py

@@ -1,6 +1,6 @@
 import io
 from collections.abc import Generator
-from typing import Any, Optional
+from typing import Any
 
 from core.model_manager import ModelManager
 from core.model_runtime.entities.model_entities import ModelPropertyKey, ModelType
@@ -16,9 +16,9 @@ class TTSTool(BuiltinTool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         provider, model = tool_parameters.get("model").split("#")  # type: ignore
         voice = tool_parameters.get(f"voice#{provider}#{model}")
@@ -72,9 +72,9 @@ class TTSTool(BuiltinTool):
 
     def get_runtime_parameters(
         self,
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> list[ToolParameter]:
         parameters = []
 

api/core/tools/builtin_tool/providers/code/code.py

@@ -4,5 +4,5 @@ from core.tools.builtin_tool.provider import BuiltinToolProviderController
 
 
 class CodeToolProvider(BuiltinToolProviderController):
-    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]) -> None:
+    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]):
         pass

api/core/tools/builtin_tool/providers/code/tools/simple_code.py

@@ -1,5 +1,5 @@
 from collections.abc import Generator
-from typing import Any, Optional
+from typing import Any
 
 from core.helper.code_executor.code_executor import CodeExecutor, CodeLanguage
 from core.tools.builtin_tool.tool import BuiltinTool
@@ -12,9 +12,9 @@ class SimpleCode(BuiltinTool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         invoke simple code

api/core/tools/builtin_tool/providers/time/time.py

@@ -4,5 +4,5 @@ from core.tools.builtin_tool.provider import BuiltinToolProviderController
 
 
 class WikiPediaProvider(BuiltinToolProviderController):
-    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]) -> None:
+    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]):
         pass

api/core/tools/builtin_tool/providers/time/tools/current_time.py

@@ -1,6 +1,6 @@
 from collections.abc import Generator
 from datetime import UTC, datetime
-from typing import Any, Optional
+from typing import Any
 
 from pytz import timezone as pytz_timezone
 
@@ -13,9 +13,9 @@ class CurrentTimeTool(BuiltinTool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         invoke tools

api/core/tools/builtin_tool/providers/time/tools/localtime_to_timestamp.py

@@ -1,6 +1,6 @@
 from collections.abc import Generator
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any
 
 import pytz
 
@@ -14,9 +14,9 @@ class LocaltimeToTimestampTool(BuiltinTool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         Convert localtime to timestamp

api/core/tools/builtin_tool/providers/time/tools/timestamp_to_localtime.py

@@ -1,6 +1,6 @@
 from collections.abc import Generator
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any
 
 import pytz
 
@@ -14,9 +14,9 @@ class TimestampToLocaltimeTool(BuiltinTool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         Convert timestamp to localtime

api/core/tools/builtin_tool/providers/time/tools/timezone_conversion.py

@@ -1,6 +1,6 @@
 from collections.abc import Generator
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any
 
 import pytz
 
@@ -14,9 +14,9 @@ class TimezoneConversionTool(BuiltinTool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         Convert time to equivalent time zone

api/core/tools/builtin_tool/providers/time/tools/weekday.py

@@ -1,7 +1,7 @@
 import calendar
 from collections.abc import Generator
 from datetime import datetime
-from typing import Any, Optional
+from typing import Any
 
 from core.tools.builtin_tool.tool import BuiltinTool
 from core.tools.entities.tool_entities import ToolInvokeMessage
@@ -12,9 +12,9 @@ class WeekdayTool(BuiltinTool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         Calculate the day of the week for a given date

api/core/tools/builtin_tool/providers/webscraper/tools/webscraper.py

@@ -1,5 +1,5 @@
 from collections.abc import Generator
-from typing import Any, Optional
+from typing import Any
 
 from core.tools.builtin_tool.tool import BuiltinTool
 from core.tools.entities.tool_entities import ToolInvokeMessage
@@ -12,9 +12,9 @@ class WebscraperTool(BuiltinTool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         invoke tools

api/core/tools/builtin_tool/providers/webscraper/webscraper.py

@@ -4,7 +4,7 @@ from core.tools.builtin_tool.provider import BuiltinToolProviderController
 
 
 class WebscraperProvider(BuiltinToolProviderController):
-    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]) -> None:
+    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]):
         """
         Validate credentials
         """

api/core/tools/custom_tool/provider.py

@@ -1,4 +1,5 @@
 from pydantic import Field
+from sqlalchemy import select
 
 from core.entities.provider_entities import ProviderConfig
 from core.tools.__base.tool_provider import ToolProviderController
@@ -24,7 +25,7 @@ class ApiToolProviderController(ToolProviderController):
     tenant_id: str
     tools: list[ApiTool] = Field(default_factory=list)
 
-    def __init__(self, entity: ToolProviderEntity, provider_id: str, tenant_id: str) -> None:
+    def __init__(self, entity: ToolProviderEntity, provider_id: str, tenant_id: str):
         super().__init__(entity)
         self.provider_id = provider_id
         self.tenant_id = tenant_id
@@ -176,11 +177,11 @@ class ApiToolProviderController(ToolProviderController):
         tools: list[ApiTool] = []
 
         # get tenant api providers
-        db_providers: list[ApiToolProvider] = (
-            db.session.query(ApiToolProvider)
-            .where(ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.name == self.entity.identity.name)
-            .all()
-        )
+        db_providers = db.session.scalars(
+            select(ApiToolProvider).where(
+                ApiToolProvider.tenant_id == tenant_id, ApiToolProvider.name == self.entity.identity.name
+            )
+        ).all()
 
         if db_providers and len(db_providers) != 0:
             for db_provider in db_providers:
@@ -191,7 +192,7 @@ class ApiToolProviderController(ToolProviderController):
         self.tools = tools
         return tools
 
-    def get_tool(self, tool_name: str):
+    def get_tool(self, tool_name: str) -> ApiTool:
         """
         get tool by name
 

api/core/tools/custom_tool/tool.py

@@ -2,7 +2,7 @@ import json
 from collections.abc import Generator
 from dataclasses import dataclass
 from os import getenv
-from typing import Any, Optional, Union
+from typing import Any, Union
 from urllib.parse import urlencode
 
 import httpx
@@ -302,7 +302,7 @@ class ApiTool(Tool):
 
     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 []:
@@ -337,7 +337,7 @@ class ApiTool(Tool):
         # If no option succeeded, you might want to return the value as is or raise an error
         return value  # or raise ValueError(f"Cannot convert value '{value}' to any specified type in anyOf")
 
-    def _convert_body_property_type(self, property: dict[str, Any], value: Any) -> Any:
+    def _convert_body_property_type(self, property: dict[str, Any], value: Any):
         try:
             if "type" in property:
                 if property["type"] == "integer" or property["type"] == "int":
@@ -376,9 +376,9 @@ class ApiTool(Tool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         invoke http request
@@ -396,6 +396,10 @@ class ApiTool(Tool):
         # assemble invoke message based on response type
         if parsed_response.is_json and isinstance(parsed_response.content, dict):
             yield self.create_json_message(parsed_response.content)
+
+            # FIXES: https://github.com/langgenius/dify/pull/23456#issuecomment-3182413088
+            # We need never break the original flows
+            yield self.create_text_message(response.text)
         else:
             # Convert to string if needed and create text message
             text_response = (

api/core/tools/entities/api_entities.py

@@ -1,5 +1,6 @@
+from collections.abc import Mapping
 from datetime import datetime
-from typing import Any, Literal, Optional
+from typing import Any, Literal
 
 from pydantic import BaseModel, Field, field_validator
 
@@ -15,12 +16,12 @@ class ToolApiEntity(BaseModel):
     name: str  # identifier
     label: I18nObject  # label
     description: I18nObject
-    parameters: Optional[list[ToolParameter]] = None
+    parameters: list[ToolParameter] | None = None
     labels: list[str] = Field(default_factory=list)
-    output_schema: Optional[dict] = None
+    output_schema: Mapping[str, object] = Field(default_factory=dict)
 
 
-ToolProviderTypeApiLiteral = Optional[Literal["builtin", "api", "workflow", "mcp"]]
+ToolProviderTypeApiLiteral = Literal["builtin", "api", "workflow", "mcp"] | None
 
 
 class ToolProviderApiEntity(BaseModel):
@@ -28,29 +29,33 @@ class ToolProviderApiEntity(BaseModel):
     author: str
     name: str  # identifier
     description: I18nObject
-    icon: str | dict
-    icon_dark: Optional[str | dict] = Field(default=None, description="The dark icon of the tool")
+    icon: str | Mapping[str, str]
+    icon_dark: str | Mapping[str, str] = ""
     label: I18nObject  # label
     type: ToolProviderType
-    masked_credentials: Optional[dict] = None
-    original_credentials: Optional[dict] = None
+    masked_credentials: Mapping[str, object] = Field(default_factory=dict)
+    original_credentials: Mapping[str, object] = Field(default_factory=dict)
     is_team_authorization: bool = False
     allow_delete: bool = True
-    plugin_id: Optional[str] = Field(default="", description="The plugin id of the tool")
-    plugin_unique_identifier: Optional[str] = Field(default="", description="The unique identifier of the tool")
-    tools: list[ToolApiEntity] = Field(default_factory=list)
+    plugin_id: str | None = Field(default="", description="The plugin id of the tool")
+    plugin_unique_identifier: str | None = Field(default="", description="The unique identifier of the tool")
+    tools: list[ToolApiEntity] = Field(default_factory=list[ToolApiEntity])
     labels: list[str] = Field(default_factory=list)
     # MCP
-    server_url: Optional[str] = Field(default="", description="The server url of the tool")
+    server_url: str | None = Field(default="", description="The server url of the tool")
     updated_at: int = Field(default_factory=lambda: int(datetime.now().timestamp()))
-    server_identifier: Optional[str] = Field(default="", description="The server identifier of the MCP tool")
+    server_identifier: str | None = Field(default="", description="The server identifier of the MCP tool")
+    timeout: float | None = Field(default=30.0, description="The timeout of the MCP tool")
+    sse_read_timeout: float | None = Field(default=300.0, description="The SSE read timeout of the MCP tool")
+    masked_headers: dict[str, str] | None = Field(default=None, description="The masked headers of the MCP tool")
+    original_headers: dict[str, str] | None = Field(default=None, description="The original headers of the MCP tool")
 
     @field_validator("tools", mode="before")
     @classmethod
     def convert_none_to_empty_list(cls, v):
         return v if v is not None else []
 
-    def to_dict(self) -> dict:
+    def to_dict(self):
         # -------------
         # overwrite tool parameter types for temp fix
         tools = jsonable_encoder(self.tools)
@@ -66,6 +71,10 @@ class ToolProviderApiEntity(BaseModel):
         if self.type == ToolProviderType.MCP:
             optional_fields.update(self.optional_field("updated_at", self.updated_at))
             optional_fields.update(self.optional_field("server_identifier", self.server_identifier))
+            optional_fields.update(self.optional_field("timeout", self.timeout))
+            optional_fields.update(self.optional_field("sse_read_timeout", self.sse_read_timeout))
+            optional_fields.update(self.optional_field("masked_headers", self.masked_headers))
+            optional_fields.update(self.optional_field("original_headers", self.original_headers))
         return {
             "id": self.id,
             "author": self.author,
@@ -85,7 +94,7 @@ class ToolProviderApiEntity(BaseModel):
             **optional_fields,
         }
 
-    def optional_field(self, key: str, value: Any) -> dict:
+    def optional_field(self, key: str, value: Any):
         """Return dict with key-value if value is truthy, empty dict otherwise."""
         return {key: value} if value else {}
 
@@ -98,7 +107,7 @@ class ToolProviderCredentialApiEntity(BaseModel):
     is_default: bool = Field(
         default=False, description="Whether the credential is the default credential for the provider in the workspace"
     )
-    credentials: dict = Field(description="The credentials of the provider")
+    credentials: Mapping[str, object] = Field(description="The credentials of the provider", default_factory=dict)
 
 
 class ToolProviderCredentialInfoApiEntity(BaseModel):

api/core/tools/entities/common_entities.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from pydantic import BaseModel, Field
 
 
@@ -9,9 +7,9 @@ class I18nObject(BaseModel):
     """
 
     en_US: str
-    zh_Hans: Optional[str] = Field(default=None)
-    pt_BR: Optional[str] = Field(default=None)
-    ja_JP: Optional[str] = Field(default=None)
+    zh_Hans: str | None = Field(default=None)
+    pt_BR: str | None = Field(default=None)
+    ja_JP: str | None = Field(default=None)
 
     def __init__(self, **data):
         super().__init__(**data)
@@ -19,5 +17,5 @@ class I18nObject(BaseModel):
         self.pt_BR = self.pt_BR or self.en_US
         self.ja_JP = self.ja_JP or self.en_US
 
-    def to_dict(self) -> dict:
+    def to_dict(self):
         return {"zh_Hans": self.zh_Hans, "en_US": self.en_US, "pt_BR": self.pt_BR, "ja_JP": self.ja_JP}

api/core/tools/entities/tool_bundle.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from pydantic import BaseModel
 
 from core.tools.entities.tool_entities import ToolParameter
@@ -16,14 +14,14 @@ class ApiToolBundle(BaseModel):
     # method
     method: str
     # summary
-    summary: Optional[str] = None
+    summary: str | None = None
     # operation_id
-    operation_id: Optional[str] = None
+    operation_id: str | None = None
     # parameters
-    parameters: Optional[list[ToolParameter]] = None
+    parameters: list[ToolParameter] | None = None
     # author
     author: str
     # icon
-    icon: Optional[str] = None
+    icon: str | None = None
     # openapi operation
     openapi: dict

api/core/tools/entities/tool_entities.py

@@ -1,9 +1,8 @@
 import base64
 import contextlib
-import enum
 from collections.abc import Mapping
-from enum import Enum
-from typing import Any, Optional, Union
+from enum import StrEnum, auto
+from typing import Any, Union
 
 from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_serializer, field_validator, model_validator
 
@@ -22,7 +21,7 @@ from core.tools.entities.common_entities import I18nObject
 from core.tools.entities.constants import TOOL_SELECTOR_MODEL_IDENTITY
 
 
-class ToolLabelEnum(Enum):
+class ToolLabelEnum(StrEnum):
     SEARCH = "search"
     IMAGE = "image"
     VIDEOS = "videos"
@@ -38,21 +37,22 @@ class ToolLabelEnum(Enum):
     BUSINESS = "business"
     ENTERTAINMENT = "entertainment"
     UTILITIES = "utilities"
+    RAG = "rag"
     OTHER = "other"
 
 
-class ToolProviderType(enum.StrEnum):
+class ToolProviderType(StrEnum):
     """
     Enum class for tool provider
     """
 
-    PLUGIN = "plugin"
+    PLUGIN = auto()
     BUILT_IN = "builtin"
-    WORKFLOW = "workflow"
-    API = "api"
-    APP = "app"
+    WORKFLOW = auto()
+    API = auto()
+    APP = auto()
     DATASET_RETRIEVAL = "dataset-retrieval"
-    MCP = "mcp"
+    MCP = auto()
 
     @classmethod
     def value_of(cls, value: str) -> "ToolProviderType":
@@ -68,15 +68,15 @@ class ToolProviderType(enum.StrEnum):
         raise ValueError(f"invalid mode value {value}")
 
 
-class ApiProviderSchemaType(Enum):
+class ApiProviderSchemaType(StrEnum):
     """
     Enum class for api provider schema type.
     """
 
-    OPENAPI = "openapi"
-    SWAGGER = "swagger"
-    OPENAI_PLUGIN = "openai_plugin"
-    OPENAI_ACTIONS = "openai_actions"
+    OPENAPI = auto()
+    SWAGGER = auto()
+    OPENAI_PLUGIN = auto()
+    OPENAI_ACTIONS = auto()
 
     @classmethod
     def value_of(cls, value: str) -> "ApiProviderSchemaType":
@@ -92,14 +92,14 @@ class ApiProviderSchemaType(Enum):
         raise ValueError(f"invalid mode value {value}")
 
 
-class ApiProviderAuthType(Enum):
+class ApiProviderAuthType(StrEnum):
     """
     Enum class for api provider auth type.
     """
 
-    NONE = "none"
-    API_KEY_HEADER = "api_key_header"
-    API_KEY_QUERY = "api_key_query"
+    NONE = auto()
+    API_KEY_HEADER = auto()
+    API_KEY_QUERY = auto()
 
     @classmethod
     def value_of(cls, value: str) -> "ApiProviderAuthType":
@@ -150,7 +150,7 @@ class ToolInvokeMessage(BaseModel):
 
         @model_validator(mode="before")
         @classmethod
-        def transform_variable_value(cls, values) -> Any:
+        def transform_variable_value(cls, values):
             """
             Only basic types and lists are allowed.
             """
@@ -176,36 +176,36 @@ class ToolInvokeMessage(BaseModel):
             return value
 
     class LogMessage(BaseModel):
-        class LogStatus(Enum):
-            START = "start"
-            ERROR = "error"
-            SUCCESS = "success"
+        class LogStatus(StrEnum):
+            START = auto()
+            ERROR = auto()
+            SUCCESS = auto()
 
         id: str
         label: str = Field(..., description="The label of the log")
-        parent_id: Optional[str] = Field(default=None, description="Leave empty for root log")
-        error: Optional[str] = Field(default=None, description="The error message")
+        parent_id: str | None = Field(default=None, description="Leave empty for root log")
+        error: str | None = Field(default=None, description="The error message")
         status: LogStatus = Field(..., description="The status of the log")
         data: Mapping[str, Any] = Field(..., description="Detailed log data")
-        metadata: Optional[Mapping[str, Any]] = Field(default=None, description="The metadata of the log")
+        metadata: Mapping[str, Any] = Field(default_factory=dict, description="The metadata of the log")
 
     class RetrieverResourceMessage(BaseModel):
         retriever_resources: list[RetrievalSourceMetadata] = Field(..., description="retriever resources")
         context: str = Field(..., description="context")
 
-    class MessageType(Enum):
-        TEXT = "text"
-        IMAGE = "image"
-        LINK = "link"
-        BLOB = "blob"
-        JSON = "json"
-        IMAGE_LINK = "image_link"
-        BINARY_LINK = "binary_link"
-        VARIABLE = "variable"
-        FILE = "file"
-        LOG = "log"
-        BLOB_CHUNK = "blob_chunk"
-        RETRIEVER_RESOURCES = "retriever_resources"
+    class MessageType(StrEnum):
+        TEXT = auto()
+        IMAGE = auto()
+        LINK = auto()
+        BLOB = auto()
+        JSON = auto()
+        IMAGE_LINK = auto()
+        BINARY_LINK = auto()
+        VARIABLE = auto()
+        FILE = auto()
+        LOG = auto()
+        BLOB_CHUNK = auto()
+        RETRIEVER_RESOURCES = auto()
 
     type: MessageType = MessageType.TEXT
     """
@@ -242,7 +242,7 @@ class ToolInvokeMessage(BaseModel):
 class ToolInvokeMessageBinary(BaseModel):
     mimetype: str = Field(..., description="The mimetype of the binary")
     url: str = Field(..., description="The url of the binary")
-    file_var: Optional[dict[str, Any]] = None
+    file_var: dict[str, Any] | None = None
 
 
 class ToolParameter(PluginParameter):
@@ -250,29 +250,29 @@ class ToolParameter(PluginParameter):
     Overrides type
     """
 
-    class ToolParameterType(enum.StrEnum):
+    class ToolParameterType(StrEnum):
         """
         removes TOOLS_SELECTOR from PluginParameterType
         """
 
-        STRING = PluginParameterType.STRING.value
-        NUMBER = PluginParameterType.NUMBER.value
-        BOOLEAN = PluginParameterType.BOOLEAN.value
-        SELECT = PluginParameterType.SELECT.value
-        SECRET_INPUT = PluginParameterType.SECRET_INPUT.value
-        FILE = PluginParameterType.FILE.value
-        FILES = PluginParameterType.FILES.value
-        APP_SELECTOR = PluginParameterType.APP_SELECTOR.value
-        MODEL_SELECTOR = PluginParameterType.MODEL_SELECTOR.value
-        ANY = PluginParameterType.ANY.value
-        DYNAMIC_SELECT = PluginParameterType.DYNAMIC_SELECT.value
+        STRING = PluginParameterType.STRING
+        NUMBER = PluginParameterType.NUMBER
+        BOOLEAN = PluginParameterType.BOOLEAN
+        SELECT = PluginParameterType.SELECT
+        SECRET_INPUT = PluginParameterType.SECRET_INPUT
+        FILE = PluginParameterType.FILE
+        FILES = PluginParameterType.FILES
+        APP_SELECTOR = PluginParameterType.APP_SELECTOR
+        MODEL_SELECTOR = PluginParameterType.MODEL_SELECTOR
+        ANY = PluginParameterType.ANY
+        DYNAMIC_SELECT = PluginParameterType.DYNAMIC_SELECT
 
         # MCP object and array type parameters
-        ARRAY = MCPServerParameterType.ARRAY.value
-        OBJECT = MCPServerParameterType.OBJECT.value
+        ARRAY = MCPServerParameterType.ARRAY
+        OBJECT = MCPServerParameterType.OBJECT
 
         # deprecated, should not use.
-        SYSTEM_FILES = PluginParameterType.SYSTEM_FILES.value
+        SYSTEM_FILES = PluginParameterType.SYSTEM_FILES
 
         def as_normal_type(self):
             return as_normal_type(self)
@@ -280,17 +280,17 @@ class ToolParameter(PluginParameter):
         def cast_value(self, value: Any):
             return cast_parameter_value(self, value)
 
-    class ToolParameterForm(Enum):
-        SCHEMA = "schema"  # should be set while adding tool
-        FORM = "form"  # should be set before invoking tool
-        LLM = "llm"  # will be set by LLM
+    class ToolParameterForm(StrEnum):
+        SCHEMA = auto()  # should be set while adding tool
+        FORM = auto()  # should be set before invoking tool
+        LLM = auto()  # will be set by LLM
 
     type: ToolParameterType = Field(..., description="The type of the parameter")
-    human_description: Optional[I18nObject] = Field(default=None, description="The description presented to the user")
+    human_description: I18nObject | None = Field(default=None, description="The description presented to the user")
     form: ToolParameterForm = Field(..., description="The form of the parameter, schema/form/llm")
-    llm_description: Optional[str] = None
+    llm_description: str | None = None
     # MCP object and array type parameters use this field to store the schema
-    input_schema: Optional[dict] = None
+    input_schema: dict | None = None
 
     @classmethod
     def get_simple_instance(
@@ -299,7 +299,7 @@ class ToolParameter(PluginParameter):
         llm_description: str,
         typ: ToolParameterType,
         required: bool,
-        options: Optional[list[str]] = None,
+        options: list[str] | None = None,
     ) -> "ToolParameter":
         """
         get a simple tool parameter
@@ -340,9 +340,9 @@ class ToolProviderIdentity(BaseModel):
     name: str = Field(..., description="The name of the tool")
     description: I18nObject = Field(..., description="The description of the tool")
     icon: str = Field(..., description="The icon of the tool")
-    icon_dark: Optional[str] = Field(default=None, description="The dark icon of the tool")
+    icon_dark: str | None = Field(default=None, description="The dark icon of the tool")
     label: I18nObject = Field(..., description="The label of the tool")
-    tags: Optional[list[ToolLabelEnum]] = Field(
+    tags: list[ToolLabelEnum] | None = Field(
         default=[],
         description="The tags of the tool",
     )
@@ -353,7 +353,7 @@ class ToolIdentity(BaseModel):
     name: str = Field(..., description="The name of the tool")
     label: I18nObject = Field(..., description="The label of the tool")
     provider: str = Field(..., description="The provider of the tool")
-    icon: Optional[str] = None
+    icon: str | None = None
 
 
 class ToolDescription(BaseModel):
@@ -363,9 +363,9 @@ class ToolDescription(BaseModel):
 
 class ToolEntity(BaseModel):
     identity: ToolIdentity
-    parameters: list[ToolParameter] = Field(default_factory=list)
-    description: Optional[ToolDescription] = None
-    output_schema: Optional[dict] = None
+    parameters: list[ToolParameter] = Field(default_factory=list[ToolParameter])
+    description: ToolDescription | None = None
+    output_schema: Mapping[str, object] = Field(default_factory=dict)
     has_runtime_parameters: bool = Field(default=False, description="Whether the tool has runtime parameters")
 
     # pydantic configs
@@ -378,21 +378,23 @@ class ToolEntity(BaseModel):
 
 
 class OAuthSchema(BaseModel):
-    client_schema: list[ProviderConfig] = Field(default_factory=list, description="The schema of the OAuth client")
+    client_schema: list[ProviderConfig] = Field(
+        default_factory=list[ProviderConfig], description="The schema of the OAuth client"
+    )
     credentials_schema: list[ProviderConfig] = Field(
-        default_factory=list, description="The schema of the OAuth credentials"
+        default_factory=list[ProviderConfig], description="The schema of the OAuth credentials"
     )
 
 
 class ToolProviderEntity(BaseModel):
     identity: ToolProviderIdentity
-    plugin_id: Optional[str] = None
-    credentials_schema: list[ProviderConfig] = Field(default_factory=list)
-    oauth_schema: Optional[OAuthSchema] = None
+    plugin_id: str | None = None
+    credentials_schema: list[ProviderConfig] = Field(default_factory=list[ProviderConfig])
+    oauth_schema: OAuthSchema | None = None
 
 
 class ToolProviderEntityWithPlugin(ToolProviderEntity):
-    tools: list[ToolEntity] = Field(default_factory=list)
+    tools: list[ToolEntity] = Field(default_factory=list[ToolEntity])
 
 
 class WorkflowToolParameterConfiguration(BaseModel):
@@ -411,8 +413,8 @@ class ToolInvokeMeta(BaseModel):
     """
 
     time_cost: float = Field(..., description="The time cost of the tool invoke")
-    error: Optional[str] = None
-    tool_config: Optional[dict] = None
+    error: str | None = None
+    tool_config: dict | None = None
 
     @classmethod
     def empty(cls) -> "ToolInvokeMeta":
@@ -428,7 +430,7 @@ class ToolInvokeMeta(BaseModel):
         """
         return cls(time_cost=0.0, error=error, tool_config={})
 
-    def to_dict(self) -> dict:
+    def to_dict(self):
         return {
             "time_cost": self.time_cost,
             "error": self.error,
@@ -446,14 +448,14 @@ class ToolLabel(BaseModel):
     icon: str = Field(..., description="The icon of the tool")
 
 
-class ToolInvokeFrom(Enum):
+class ToolInvokeFrom(StrEnum):
     """
     Enum class for tool invoke
     """
 
-    WORKFLOW = "workflow"
-    AGENT = "agent"
-    PLUGIN = "plugin"
+    WORKFLOW = auto()
+    AGENT = auto()
+    PLUGIN = auto()
 
 
 class ToolSelector(BaseModel):
@@ -464,11 +466,11 @@ class ToolSelector(BaseModel):
         type: ToolParameter.ToolParameterType = Field(..., description="The type of the parameter")
         required: bool = Field(..., description="Whether the parameter is required")
         description: str = Field(..., description="The description of the parameter")
-        default: Optional[Union[int, float, str]] = None
-        options: Optional[list[PluginParameterOption]] = None
+        default: Union[int, float, str] | None = None
+        options: list[PluginParameterOption] | None = None
 
     provider_id: str = Field(..., description="The id of the provider")
-    credential_id: Optional[str] = Field(default=None, description="The id of the credential")
+    credential_id: str | None = Field(default=None, description="The id of the credential")
     tool_name: str = Field(..., description="The name of the tool")
     tool_description: str = Field(..., description="The description of the tool")
     tool_configuration: Mapping[str, Any] = Field(..., description="Configuration, type form")

api/core/tools/entities/values.py

@@ -49,6 +49,9 @@ ICONS = {
 </svg>""",  # noqa: E501
     ToolLabelEnum.OTHER: """<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 16 16" fill="none">
   <path d="M8.00052 0.666748L4.00065 7.33342H12.0007L8.00052 0.666748ZM8.00052 3.25828L9.64572 6.00008H6.35553L8.00052 3.25828ZM4.50065 13.3334C3.48813 13.3334 2.66732 12.5126 2.66732 11.5001C2.66732 10.4875 3.48813 9.66675 4.50065 9.66675C5.51317 9.66675 6.33398 10.4875 6.33398 11.5001C6.33398 12.5126 5.51317 13.3334 4.50065 13.3334ZM4.50065 14.6667C6.24955 14.6667 7.66732 13.249 7.66732 11.5001C7.66732 9.75115 6.24955 8.33342 4.50065 8.33342C2.75175 8.33342 1.33398 9.75115 1.33398 11.5001C1.33398 13.249 2.75175 14.6667 4.50065 14.6667ZM10.0007 10.3334V13.0001H12.6673V10.3334H10.0007ZM8.66732 14.3334V9.00008H14.0007V14.3334H8.66732Z" fill="#344054"/>
+</svg>""",  # noqa: E501
+    ToolLabelEnum.RAG: """<svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 16 16" fill="none">
+  <path d="M8.00065 1.3335H9.33398V2.66683H8.00065V1.3335ZM5.33398 1.3335H6.66732V2.66683H5.33398V1.3335ZM3.99935 2.66683C3.99935 2.29864 4.29783 2.00016 4.66602 2.00016H12.3327C12.7009 2.00016 13.0007 2.29864 13.0007 2.66683V13.3335C13.0007 13.7017 12.7009 14.0002 12.3327 14.0002H4.66602C4.29783 14.0002 3.99935 13.7017 3.99935 13.3335V2.66683ZM4.66602 12.6668C4.29783 12.6668 3.99935 12.3683 3.99935 12.0002V10.6668H5.33398V12.0002C5.33398 12.3683 5.0355 12.6668 4.66602 12.6668ZM5.33398 8.66683H6.66732V10.0002H5.33398V8.66683ZM5.33398 6.66683H6.66732V8.00016H5.33398V6.66683ZM3.99935 4.66683H6.66602V6.00016H3.99935V4.66683ZM6.66602 1.3335H12.3327V2.66683H6.66602V1.3335Z" fill="#344054"/>
 </svg>""",  # noqa: E501
 }
 
@@ -105,7 +108,10 @@ default_tool_label_dict = {
     ToolLabelEnum.OTHER: ToolLabel(
         name="other", label=I18nObject(en_US="Other", zh_Hans="其他"), icon=ICONS[ToolLabelEnum.OTHER]
     ),
+    ToolLabelEnum.RAG: ToolLabel(
+        name="rag", label=I18nObject(en_US="RAG", zh_Hans="RAG"), icon=ICONS[ToolLabelEnum.RAG]
+    ),
 }
 
-default_tool_labels = [v for k, v in default_tool_label_dict.items()]
+default_tool_labels = list(default_tool_label_dict.values())
 default_tool_label_name_list = [label.name for label in default_tool_labels]

api/core/tools/errors.py

@@ -29,6 +29,10 @@ class ToolApiSchemaError(ValueError):
     pass
 
 
+class ToolCredentialPolicyViolationError(ValueError):
+    pass
+
+
 class ToolEngineInvokeError(Exception):
     meta: ToolInvokeMeta
 

api/core/tools/mcp_tool/provider.py

@@ -1,5 +1,5 @@
 import json
-from typing import Any, Optional
+from typing import Any, Self
 
 from core.mcp.types import Tool as RemoteMCPTool
 from core.tools.__base.tool_provider import ToolProviderController
@@ -25,10 +25,10 @@ class MCPToolProviderController(ToolProviderController):
         provider_id: str,
         tenant_id: str,
         server_url: str,
-        headers: Optional[dict[str, str]] = None,
-        timeout: Optional[float] = None,
-        sse_read_timeout: Optional[float] = None,
-    ) -> None:
+        headers: dict[str, str] | None = None,
+        timeout: float | None = None,
+        sse_read_timeout: float | None = None,
+    ):
         super().__init__(entity)
         self.entity: ToolProviderEntityWithPlugin = entity
         self.tenant_id = tenant_id
@@ -48,7 +48,7 @@ class MCPToolProviderController(ToolProviderController):
         return ToolProviderType.MCP
 
     @classmethod
-    def _from_db(cls, db_provider: MCPToolProvider) -> "MCPToolProviderController":
+    def from_db(cls, db_provider: MCPToolProvider) -> Self:
         """
         from db provider
         """
@@ -72,7 +72,6 @@ class MCPToolProviderController(ToolProviderController):
                     ),
                     llm=remote_mcp_tool.description or "",
                 ),
-                output_schema=None,
                 has_runtime_parameters=len(remote_mcp_tool.inputSchema) > 0,
             )
             for remote_mcp_tool in remote_mcp_tools
@@ -94,12 +93,12 @@ class MCPToolProviderController(ToolProviderController):
             provider_id=db_provider.server_identifier or "",
             tenant_id=db_provider.tenant_id or "",
             server_url=db_provider.decrypted_server_url,
-            headers={},  # TODO: get headers from db provider
+            headers=db_provider.decrypted_headers or {},
             timeout=db_provider.timeout,
             sse_read_timeout=db_provider.sse_read_timeout,
         )
 
-    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]) -> None:
+    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]):
         """
         validate the credentials of the provider
         """

api/core/tools/mcp_tool/tool.py

@@ -1,7 +1,7 @@
 import base64
 import json
 from collections.abc import Generator
-from typing import Any, Optional
+from typing import Any
 
 from core.mcp.error import MCPAuthError, MCPConnectionError
 from core.mcp.mcp_client import MCPClient
@@ -20,10 +20,10 @@ class MCPTool(Tool):
         icon: str,
         server_url: str,
         provider_id: str,
-        headers: Optional[dict[str, str]] = None,
-        timeout: Optional[float] = None,
-        sse_read_timeout: Optional[float] = None,
-    ) -> None:
+        headers: dict[str, str] | None = None,
+        timeout: float | None = None,
+        sse_read_timeout: float | None = None,
+    ):
         super().__init__(entity, runtime)
         self.tenant_id = tenant_id
         self.icon = icon
@@ -40,9 +40,9 @@ class MCPTool(Tool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         from core.tools.errors import ToolInvokeError
 
@@ -67,22 +67,42 @@ class MCPTool(Tool):
 
         for content in result.content:
             if isinstance(content, TextContent):
-                try:
-                    content_json = json.loads(content.text)
-                    if isinstance(content_json, dict):
-                        yield self.create_json_message(content_json)
-                    elif isinstance(content_json, list):
-                        for item in content_json:
-                            yield self.create_json_message(item)
-                    else:
-                        yield self.create_text_message(content.text)
-                except json.JSONDecodeError:
-                    yield self.create_text_message(content.text)
-
+                yield from self._process_text_content(content)
             elif isinstance(content, ImageContent):
-                yield self.create_blob_message(
-                    blob=base64.b64decode(content.data), meta={"mime_type": content.mimeType}
-                )
+                yield self._process_image_content(content)
+
+    def _process_text_content(self, content: TextContent) -> Generator[ToolInvokeMessage, None, None]:
+        """Process text content and yield appropriate messages."""
+        try:
+            content_json = json.loads(content.text)
+            yield from self._process_json_content(content_json)
+        except json.JSONDecodeError:
+            yield self.create_text_message(content.text)
+
+    def _process_json_content(self, content_json: Any) -> Generator[ToolInvokeMessage, None, None]:
+        """Process JSON content based on its type."""
+        if isinstance(content_json, dict):
+            yield self.create_json_message(content_json)
+        elif isinstance(content_json, list):
+            yield from self._process_json_list(content_json)
+        else:
+            # For primitive types (str, int, bool, etc.), convert to string
+            yield self.create_text_message(str(content_json))
+
+    def _process_json_list(self, json_list: list) -> Generator[ToolInvokeMessage, None, None]:
+        """Process a list of JSON items."""
+        if any(not isinstance(item, dict) for item in json_list):
+            # If the list contains any non-dict item, treat the entire list as a text message.
+            yield self.create_text_message(str(json_list))
+            return
+
+        # Otherwise, process each dictionary as a separate JSON message.
+        for item in json_list:
+            yield self.create_json_message(item)
+
+    def _process_image_content(self, content: ImageContent) -> ToolInvokeMessage:
+        """Process image content and return a blob message."""
+        return self.create_blob_message(blob=base64.b64decode(content.data), meta={"mime_type": content.mimeType})
 
     def fork_tool_runtime(self, runtime: ToolRuntime) -> "MCPTool":
         return MCPTool(

api/core/tools/plugin_tool/provider.py

@@ -16,7 +16,7 @@ class PluginToolProviderController(BuiltinToolProviderController):
 
     def __init__(
         self, entity: ToolProviderEntityWithPlugin, plugin_id: str, plugin_unique_identifier: str, tenant_id: str
-    ) -> None:
+    ):
         self.entity = entity
         self.tenant_id = tenant_id
         self.plugin_id = plugin_id
@@ -31,7 +31,7 @@ class PluginToolProviderController(BuiltinToolProviderController):
         """
         return ToolProviderType.PLUGIN
 
-    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]) -> None:
+    def _validate_credentials(self, user_id: str, credentials: dict[str, Any]):
         """
         validate the credentials of the provider
         """

api/core/tools/plugin_tool/tool.py

@@ -1,5 +1,5 @@
 from collections.abc import Generator
-from typing import Any, Optional
+from typing import Any
 
 from core.plugin.impl.tool import PluginToolManager
 from core.plugin.utils.converter import convert_parameters_to_plugin_format
@@ -11,12 +11,12 @@ from core.tools.entities.tool_entities import ToolEntity, ToolInvokeMessage, Too
 class PluginTool(Tool):
     def __init__(
         self, entity: ToolEntity, runtime: ToolRuntime, tenant_id: str, icon: str, plugin_unique_identifier: str
-    ) -> None:
+    ):
         super().__init__(entity, runtime)
         self.tenant_id = tenant_id
         self.icon = icon
         self.plugin_unique_identifier = plugin_unique_identifier
-        self.runtime_parameters: Optional[list[ToolParameter]] = None
+        self.runtime_parameters: list[ToolParameter] | None = None
 
     def tool_provider_type(self) -> ToolProviderType:
         return ToolProviderType.PLUGIN
@@ -25,9 +25,9 @@ class PluginTool(Tool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         manager = PluginToolManager()
 
@@ -57,9 +57,9 @@ class PluginTool(Tool):
 
     def get_runtime_parameters(
         self,
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> list[ToolParameter]:
         """
         get the runtime parameters

api/core/tools/tool_engine.py

@@ -4,7 +4,7 @@ from collections.abc import Generator, Iterable
 from copy import deepcopy
 from datetime import UTC, datetime
 from mimetypes import guess_type
-from typing import Any, Optional, Union, cast
+from typing import Any, Union, cast
 
 from yarl import URL
 
@@ -51,10 +51,10 @@ class ToolEngine:
         message: Message,
         invoke_from: InvokeFrom,
         agent_tool_callback: DifyAgentCallbackHandler,
-        trace_manager: Optional[TraceQueueManager] = None,
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        trace_manager: TraceQueueManager | None = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> tuple[str, list[str], ToolInvokeMeta]:
         """
         Agent invokes the tool with the given arguments.
@@ -152,10 +152,9 @@ class ToolEngine:
         user_id: str,
         workflow_tool_callback: DifyWorkflowCallbackHandler,
         workflow_call_depth: int,
-        thread_pool_id: Optional[str] = None,
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         Workflow invokes the tool with the given arguments.
@@ -166,7 +165,6 @@ class ToolEngine:
 
             if isinstance(tool, WorkflowTool):
                 tool.workflow_call_depth = workflow_call_depth + 1
-                tool.thread_pool_id = thread_pool_id
 
             if tool.runtime and tool.runtime.runtime_parameters:
                 tool_parameters = {**tool.runtime.runtime_parameters, **tool_parameters}
@@ -196,9 +194,9 @@ class ToolEngine:
         tool: Tool,
         tool_parameters: dict,
         user_id: str,
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage | ToolInvokeMeta, None, None]:
         """
         Invoke the tool with the given arguments.

api/core/tools/tool_file_manager.py

@@ -6,7 +6,7 @@ import os
 import time
 from collections.abc import Generator
 from mimetypes import guess_extension, guess_type
-from typing import Optional, Union
+from typing import Union
 from uuid import uuid4
 
 import httpx
@@ -72,10 +72,10 @@ class ToolFileManager:
         *,
         user_id: str,
         tenant_id: str,
-        conversation_id: Optional[str],
+        conversation_id: str | None,
         file_binary: bytes,
         mimetype: str,
-        filename: Optional[str] = None,
+        filename: str | None = None,
     ) -> ToolFile:
         extension = guess_extension(mimetype) or ".bin"
         unique_name = uuid4().hex
@@ -98,6 +98,7 @@ class ToolFileManager:
                 mimetype=mimetype,
                 name=present_filename,
                 size=len(file_binary),
+                original_url=None,
             )
 
             session.add(tool_file)
@@ -111,7 +112,7 @@ class ToolFileManager:
         user_id: str,
         tenant_id: str,
         file_url: str,
-        conversation_id: Optional[str] = None,
+        conversation_id: str | None = None,
     ) -> ToolFile:
         # try to download image
         try:
@@ -131,7 +132,6 @@ class ToolFileManager:
         filename = f"{unique_name}{extension}"
         filepath = f"tools/{tenant_id}/{filename}"
         storage.save(filepath, blob)
-
         with Session(self._engine, expire_on_commit=False) as session:
             tool_file = ToolFile(
                 user_id=user_id,
@@ -217,7 +217,7 @@ class ToolFileManager:
 
         return blob, tool_file.mimetype
 
-    def get_file_generator_by_tool_file_id(self, tool_file_id: str) -> tuple[Optional[Generator], Optional[ToolFile]]:
+    def get_file_generator_by_tool_file_id(self, tool_file_id: str) -> tuple[Generator | None, ToolFile | None]:
         """
         get file binary
 

api/core/tools/tool_label_manager.py

@@ -1,3 +1,5 @@
+from sqlalchemy import select
+
 from core.tools.__base.tool_provider import ToolProviderController
 from core.tools.builtin_tool.provider import BuiltinToolProviderController
 from core.tools.custom_tool.provider import ApiToolProviderController
@@ -24,7 +26,7 @@ class ToolLabelManager:
         labels = cls.filter_tool_labels(labels)
 
         if isinstance(controller, ApiToolProviderController | WorkflowToolProviderController):
-            provider_id = controller.provider_id
+            provider_id = controller.provider_id  # ty: ignore [unresolved-attribute]
         else:
             raise ValueError("Unsupported tool type")
 
@@ -49,22 +51,18 @@ class ToolLabelManager:
         Get tool labels
         """
         if isinstance(controller, ApiToolProviderController | WorkflowToolProviderController):
-            provider_id = controller.provider_id
+            provider_id = controller.provider_id  # ty: ignore [unresolved-attribute]
         elif isinstance(controller, BuiltinToolProviderController):
             return controller.tool_labels
         else:
             raise ValueError("Unsupported tool type")
-
-        labels = (
-            db.session.query(ToolLabelBinding.label_name)
-            .where(
-                ToolLabelBinding.tool_id == provider_id,
-                ToolLabelBinding.tool_type == controller.provider_type.value,
-            )
-            .all()
+        stmt = select(ToolLabelBinding.label_name).where(
+            ToolLabelBinding.tool_id == provider_id,
+            ToolLabelBinding.tool_type == controller.provider_type.value,
         )
+        labels = db.session.scalars(stmt).all()
 
-        return [label.label_name for label in labels]
+        return list(labels)
 
     @classmethod
     def get_tools_labels(cls, tool_providers: list[ToolProviderController]) -> dict[str, list[str]]:
@@ -87,11 +85,9 @@ class ToolLabelManager:
         provider_ids = []
         for controller in tool_providers:
             assert isinstance(controller, ApiToolProviderController | WorkflowToolProviderController)
-            provider_ids.append(controller.provider_id)
+            provider_ids.append(controller.provider_id)  # ty: ignore [unresolved-attribute]
 
-        labels: list[ToolLabelBinding] = (
-            db.session.query(ToolLabelBinding).where(ToolLabelBinding.tool_id.in_(provider_ids)).all()
-        )
+        labels = db.session.scalars(select(ToolLabelBinding).where(ToolLabelBinding.tool_id.in_(provider_ids))).all()
 
         tool_labels: dict[str, list[str]] = {label.tool_id: [] for label in labels}
 

api/core/tools/tool_manager.py

@@ -9,36 +9,23 @@ from typing import TYPE_CHECKING, Any, Literal, Optional, Union, cast
 
 import sqlalchemy as sa
 from pydantic import TypeAdapter
+from sqlalchemy import select
 from sqlalchemy.orm import Session
 from yarl import URL
 
 import contexts
-from core.helper.provider_cache import ToolProviderCredentialsCache
-from core.plugin.entities.plugin import ToolProviderID
-from core.plugin.impl.oauth import OAuthHandler
-from core.plugin.impl.tool import PluginToolManager
-from core.tools.__base.tool_provider import ToolProviderController
-from core.tools.__base.tool_runtime import ToolRuntime
-from core.tools.mcp_tool.provider import MCPToolProviderController
-from core.tools.mcp_tool.tool import MCPTool
-from core.tools.plugin_tool.provider import PluginToolProviderController
-from core.tools.plugin_tool.tool import PluginTool
-from core.tools.utils.uuid_utils import is_valid_uuid
-from core.tools.workflow_as_tool.provider import WorkflowToolProviderController
-from core.workflow.entities.variable_pool import VariablePool
-from services.tools.mcp_tools_manage_service import MCPToolManageService
-
-if TYPE_CHECKING:
-    from core.workflow.nodes.tool.entities import ToolEntity
-
 from configs import dify_config
 from core.agent.entities import AgentToolEntity
 from core.app.entities.app_invoke_entities import InvokeFrom
 from core.helper.module_import_helper import load_single_subclass_from_source
 from core.helper.position_helper import is_filtered
+from core.helper.provider_cache import ToolProviderCredentialsCache
 from core.model_runtime.utils.encoders import jsonable_encoder
 from core.plugin.entities.plugin_daemon import CredentialType
+from core.plugin.impl.tool import PluginToolManager
 from core.tools.__base.tool import Tool
+from core.tools.__base.tool_provider import ToolProviderController
+from core.tools.__base.tool_runtime import ToolRuntime
 from core.tools.builtin_tool.provider import BuiltinToolProviderController
 from core.tools.builtin_tool.providers._positions import BuiltinToolProviderSort
 from core.tools.builtin_tool.tool import BuiltinTool
@@ -53,16 +40,27 @@ from core.tools.entities.tool_entities import (
     ToolProviderType,
 )
 from core.tools.errors import ToolProviderNotFoundError
+from core.tools.mcp_tool.provider import MCPToolProviderController
+from core.tools.mcp_tool.tool import MCPTool
+from core.tools.plugin_tool.provider import PluginToolProviderController
+from core.tools.plugin_tool.tool import PluginTool
 from core.tools.tool_label_manager import ToolLabelManager
-from core.tools.utils.configuration import (
-    ToolParameterConfigurationManager,
-)
+from core.tools.utils.configuration import ToolParameterConfigurationManager
 from core.tools.utils.encryption import create_provider_encrypter, create_tool_provider_encrypter
+from core.tools.utils.uuid_utils import is_valid_uuid
+from core.tools.workflow_as_tool.provider import WorkflowToolProviderController
 from core.tools.workflow_as_tool.tool import WorkflowTool
 from extensions.ext_database import db
+from models.provider_ids import ToolProviderID
 from models.tools import ApiToolProvider, BuiltinToolProvider, MCPToolProvider, WorkflowToolProvider
+from services.enterprise.plugin_manager_service import PluginCredentialType
+from services.tools.mcp_tools_manage_service import MCPToolManageService
 from services.tools.tools_transform_service import ToolTransformService
 
+if TYPE_CHECKING:
+    from core.workflow.entities import VariablePool
+    from core.workflow.nodes.tool.entities import ToolEntity
+
 logger = logging.getLogger(__name__)
 
 
@@ -117,6 +115,7 @@ class ToolManager:
         get the plugin provider
         """
         # check if context is set
+
         try:
             contexts.plugin_tool_providers.get()
         except LookupError:
@@ -157,7 +156,7 @@ class ToolManager:
         tenant_id: str,
         invoke_from: InvokeFrom = InvokeFrom.DEBUGGER,
         tool_invoke_from: ToolInvokeFrom = ToolInvokeFrom.AGENT,
-        credential_id: Optional[str] = None,
+        credential_id: str | None = None,
     ) -> Union[BuiltinTool, PluginTool, ApiTool, WorkflowTool, MCPTool]:
         """
         get the tool runtime
@@ -172,6 +171,7 @@ class ToolManager:
 
         :return: the tool
         """
+
         if provider_type == ToolProviderType.BUILT_IN:
             # check if the builtin tool need credentials
             provider_controller = cls.get_builtin_provider(provider_id, tenant_id)
@@ -198,14 +198,11 @@ class ToolManager:
                 # get specific credentials
                 if is_valid_uuid(credential_id):
                     try:
-                        builtin_provider = (
-                            db.session.query(BuiltinToolProvider)
-                            .where(
-                                BuiltinToolProvider.tenant_id == tenant_id,
-                                BuiltinToolProvider.id == credential_id,
-                            )
-                            .first()
+                        builtin_provider_stmt = select(BuiltinToolProvider).where(
+                            BuiltinToolProvider.tenant_id == tenant_id,
+                            BuiltinToolProvider.id == credential_id,
                         )
+                        builtin_provider = db.session.scalar(builtin_provider_stmt)
                     except Exception as e:
                         builtin_provider = None
                         logger.info("Error getting builtin provider %s:%s", credential_id, e, exc_info=True)
@@ -216,16 +213,16 @@ class ToolManager:
                 # fallback to the default provider
                 if builtin_provider is None:
                     # use the default provider
-                    builtin_provider = (
-                        db.session.query(BuiltinToolProvider)
-                        .where(
-                            BuiltinToolProvider.tenant_id == tenant_id,
-                            (BuiltinToolProvider.provider == str(provider_id_entity))
-                            | (BuiltinToolProvider.provider == provider_id_entity.provider_name),
+                    with Session(db.engine) as session:
+                        builtin_provider = session.scalar(
+                            sa.select(BuiltinToolProvider)
+                            .where(
+                                BuiltinToolProvider.tenant_id == tenant_id,
+                                (BuiltinToolProvider.provider == str(provider_id_entity))
+                                | (BuiltinToolProvider.provider == provider_id_entity.provider_name),
+                            )
+                            .order_by(BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc())
                         )
-                        .order_by(BuiltinToolProvider.is_default.desc(), BuiltinToolProvider.created_at.asc())
-                        .first()
-                    )
                     if builtin_provider is None:
                         raise ToolProviderNotFoundError(f"no default provider for {provider_id}")
             else:
@@ -239,6 +236,16 @@ class ToolManager:
                 if builtin_provider is None:
                     raise ToolProviderNotFoundError(f"builtin provider {provider_id} not found")
 
+            # check if the credential is allowed to be used
+            from core.helper.credential_utils import check_credential_policy_compliance
+
+            check_credential_policy_compliance(
+                credential_id=builtin_provider.id,
+                provider=provider_id,
+                credential_type=PluginCredentialType.TOOL,
+                check_existence=False,
+            )
+
             encrypter, cache = create_provider_encrypter(
                 tenant_id=tenant_id,
                 config=[
@@ -256,6 +263,7 @@ class ToolManager:
             # check if the credentials is expired
             if builtin_provider.expires_at != -1 and (builtin_provider.expires_at - 60) < int(time.time()):
                 # TODO: circular import
+                from core.plugin.impl.oauth import OAuthHandler
                 from services.tools.builtin_tools_manage_service import BuiltinToolManageService
 
                 # refresh the credentials
@@ -263,6 +271,7 @@ class ToolManager:
                 provider_name = tool_provider.provider_name
                 redirect_uri = f"{dify_config.CONSOLE_API_URL}/console/api/oauth/plugin/{provider_id}/tool/callback"
                 system_credentials = BuiltinToolManageService.get_oauth_client(tenant_id, provider_id)
+
                 oauth_handler = OAuthHandler()
                 # refresh the credentials
                 refreshed_credentials = oauth_handler.refresh_credentials(
@@ -305,23 +314,19 @@ class ToolManager:
                 tenant_id=tenant_id,
                 controller=api_provider,
             )
-            return cast(
-                ApiTool,
-                api_provider.get_tool(tool_name).fork_tool_runtime(
-                    runtime=ToolRuntime(
-                        tenant_id=tenant_id,
-                        credentials=encrypter.decrypt(credentials),
-                        invoke_from=invoke_from,
-                        tool_invoke_from=tool_invoke_from,
-                    )
-                ),
+            return api_provider.get_tool(tool_name).fork_tool_runtime(
+                runtime=ToolRuntime(
+                    tenant_id=tenant_id,
+                    credentials=encrypter.decrypt(credentials),
+                    invoke_from=invoke_from,
+                    tool_invoke_from=tool_invoke_from,
+                )
             )
         elif provider_type == ToolProviderType.WORKFLOW:
-            workflow_provider = (
-                db.session.query(WorkflowToolProvider)
-                .where(WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id)
-                .first()
+            workflow_provider_stmt = select(WorkflowToolProvider).where(
+                WorkflowToolProvider.tenant_id == tenant_id, WorkflowToolProvider.id == provider_id
             )
+            workflow_provider = db.session.scalar(workflow_provider_stmt)
 
             if workflow_provider is None:
                 raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found")
@@ -355,7 +360,7 @@ class ToolManager:
         app_id: str,
         agent_tool: AgentToolEntity,
         invoke_from: InvokeFrom = InvokeFrom.DEBUGGER,
-        variable_pool: Optional[VariablePool] = None,
+        variable_pool: Optional["VariablePool"] = None,
     ) -> Tool:
         """
         get the agent tool runtime
@@ -397,7 +402,7 @@ class ToolManager:
         node_id: str,
         workflow_tool: "ToolEntity",
         invoke_from: InvokeFrom = InvokeFrom.DEBUGGER,
-        variable_pool: Optional[VariablePool] = None,
+        variable_pool: Optional["VariablePool"] = None,
     ) -> Tool:
         """
         get the workflow tool runtime
@@ -440,7 +445,7 @@ class ToolManager:
         provider: str,
         tool_name: str,
         tool_parameters: dict[str, Any],
-        credential_id: Optional[str] = None,
+        credential_id: str | None = None,
     ) -> Tool:
         """
         get tool runtime from plugin
@@ -513,6 +518,7 @@ class ToolManager:
         """
         list all the plugin providers
         """
+
         manager = PluginToolManager()
         provider_entities = manager.fetch_tool_providers(tenant_id)
         return [
@@ -648,7 +654,7 @@ class ToolManager:
                         include_set=dify_config.POSITION_TOOL_INCLUDES_SET,
                         exclude_set=dify_config.POSITION_TOOL_EXCLUDES_SET,
                         data=provider,
-                        name_func=lambda x: x.identity.name,
+                        name_func=lambda x: x.entity.identity.name,
                     ):
                         continue
                     user_provider = ToolTransformService.builtin_provider_to_user_provider(
@@ -664,9 +670,9 @@ class ToolManager:
 
             # get db api providers
             if "api" in filters:
-                db_api_providers: list[ApiToolProvider] = (
-                    db.session.query(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant_id).all()
-                )
+                db_api_providers = db.session.scalars(
+                    select(ApiToolProvider).where(ApiToolProvider.tenant_id == tenant_id)
+                ).all()
 
                 api_provider_controllers: list[dict[str, Any]] = [
                     {"provider": provider, "controller": ToolTransformService.api_provider_to_controller(provider)}
@@ -687,9 +693,9 @@ class ToolManager:
 
             if "workflow" in filters:
                 # get workflow providers
-                workflow_providers: list[WorkflowToolProvider] = (
-                    db.session.query(WorkflowToolProvider).where(WorkflowToolProvider.tenant_id == tenant_id).all()
-                )
+                workflow_providers = db.session.scalars(
+                    select(WorkflowToolProvider).where(WorkflowToolProvider.tenant_id == tenant_id)
+                ).all()
 
                 workflow_provider_controllers: list[WorkflowToolProviderController] = []
                 for workflow_provider in workflow_providers:
@@ -779,12 +785,12 @@ class ToolManager:
         if provider is None:
             raise ToolProviderNotFoundError(f"mcp provider {provider_id} not found")
 
-        controller = MCPToolProviderController._from_db(provider)
+        controller = MCPToolProviderController.from_db(provider)
 
         return controller
 
     @classmethod
-    def user_get_api_provider(cls, provider: str, tenant_id: str) -> dict:
+    def user_get_api_provider(cls, provider: str, tenant_id: str):
         """
         get api provider
         """
@@ -879,7 +885,7 @@ class ToolManager:
         )
 
     @classmethod
-    def generate_workflow_tool_icon_url(cls, tenant_id: str, provider_id: str) -> dict:
+    def generate_workflow_tool_icon_url(cls, tenant_id: str, provider_id: str) -> Mapping[str, str]:
         try:
             workflow_provider: WorkflowToolProvider | None = (
                 db.session.query(WorkflowToolProvider)
@@ -890,13 +896,13 @@ class ToolManager:
             if workflow_provider is None:
                 raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found")
 
-            icon: dict = json.loads(workflow_provider.icon)
+            icon = json.loads(workflow_provider.icon)
             return icon
         except Exception:
             return {"background": "#252525", "content": "\ud83d\ude01"}
 
     @classmethod
-    def generate_api_tool_icon_url(cls, tenant_id: str, provider_id: str) -> dict:
+    def generate_api_tool_icon_url(cls, tenant_id: str, provider_id: str) -> Mapping[str, str]:
         try:
             api_provider: ApiToolProvider | None = (
                 db.session.query(ApiToolProvider)
@@ -907,13 +913,13 @@ class ToolManager:
             if api_provider is None:
                 raise ToolProviderNotFoundError(f"api provider {provider_id} not found")
 
-            icon: dict = json.loads(api_provider.icon)
+            icon = json.loads(api_provider.icon)
             return icon
         except Exception:
             return {"background": "#252525", "content": "\ud83d\ude01"}
 
     @classmethod
-    def generate_mcp_tool_icon_url(cls, tenant_id: str, provider_id: str) -> dict[str, str] | str:
+    def generate_mcp_tool_icon_url(cls, tenant_id: str, provider_id: str) -> Mapping[str, str] | str:
         try:
             mcp_provider: MCPToolProvider | None = (
                 db.session.query(MCPToolProvider)
@@ -934,7 +940,7 @@ class ToolManager:
         tenant_id: str,
         provider_type: ToolProviderType,
         provider_id: str,
-    ) -> Union[str, dict]:
+    ) -> str | Mapping[str, str]:
         """
         get the tool icon
 
@@ -959,11 +965,10 @@ class ToolManager:
             return cls.generate_workflow_tool_icon_url(tenant_id, provider_id)
         elif provider_type == ToolProviderType.PLUGIN:
             provider = ToolManager.get_plugin_provider(provider_id, tenant_id)
-            if isinstance(provider, PluginToolProviderController):
-                try:
-                    return cls.generate_plugin_tool_icon_url(tenant_id, provider.entity.identity.icon)
-                except Exception:
-                    return {"background": "#252525", "content": "\ud83d\ude01"}
+            try:
+                return cls.generate_plugin_tool_icon_url(tenant_id, provider.entity.identity.icon)
+            except Exception:
+                return {"background": "#252525", "content": "\ud83d\ude01"}
             raise ValueError(f"plugin provider {provider_id} not found")
         elif provider_type == ToolProviderType.MCP:
             return cls.generate_mcp_tool_icon_url(tenant_id, provider_id)
@@ -974,7 +979,7 @@ class ToolManager:
     def _convert_tool_parameters_type(
         cls,
         parameters: list[ToolParameter],
-        variable_pool: Optional[VariablePool],
+        variable_pool: Optional["VariablePool"],
         tool_configurations: dict[str, Any],
         typ: Literal["agent", "workflow", "tool"] = "workflow",
     ) -> dict[str, Any]:

api/core/tools/utils/configuration.py

@@ -24,7 +24,7 @@ class ToolParameterConfigurationManager:
 
     def __init__(
         self, tenant_id: str, tool_runtime: Tool, provider_name: str, provider_type: ToolProviderType, identity_id: str
-    ) -> None:
+    ):
         self.tenant_id = tenant_id
         self.tool_runtime = tool_runtime
         self.provider_name = provider_name

api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py

@@ -3,6 +3,7 @@ from typing import Any
 
 from flask import Flask, current_app
 from pydantic import BaseModel, Field
+from sqlalchemy import select
 
 from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
 from core.model_manager import ModelManager
@@ -85,17 +86,14 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
 
         document_context_list = []
         index_node_ids = [document.metadata["doc_id"] for document in all_documents if document.metadata]
-        segments = (
-            db.session.query(DocumentSegment)
-            .where(
-                DocumentSegment.dataset_id.in_(self.dataset_ids),
-                DocumentSegment.completed_at.isnot(None),
-                DocumentSegment.status == "completed",
-                DocumentSegment.enabled == True,
-                DocumentSegment.index_node_id.in_(index_node_ids),
-            )
-            .all()
+        document_segment_stmt = select(DocumentSegment).where(
+            DocumentSegment.dataset_id.in_(self.dataset_ids),
+            DocumentSegment.completed_at.isnot(None),
+            DocumentSegment.status == "completed",
+            DocumentSegment.enabled == True,
+            DocumentSegment.index_node_id.in_(index_node_ids),
         )
+        segments = db.session.scalars(document_segment_stmt).all()
 
         if segments:
             index_node_id_to_position = {id: position for position, id in enumerate(index_node_ids)}
@@ -112,15 +110,12 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
                 resource_number = 1
                 for segment in sorted_segments:
                     dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first()
-                    document = (
-                        db.session.query(Document)
-                        .where(
-                            Document.id == segment.document_id,
-                            Document.enabled == True,
-                            Document.archived == False,
-                        )
-                        .first()
+                    document_stmt = select(Document).where(
+                        Document.id == segment.document_id,
+                        Document.enabled == True,
+                        Document.archived == False,
                     )
+                    document = db.session.scalar(document_stmt)
                     if dataset and document:
                         source = RetrievalSourceMetadata(
                             position=resource_number,
@@ -162,9 +157,8 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
         hit_callbacks: list[DatasetIndexToolCallbackHandler],
     ):
         with flask_app.app_context():
-            dataset = (
-                db.session.query(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == dataset_id).first()
-            )
+            stmt = select(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == dataset_id)
+            dataset = db.session.scalar(stmt)
 
             if not dataset:
                 return []

api/core/tools/utils/dataset_retriever/dataset_retriever_base_tool.py

@@ -1,7 +1,5 @@
-from abc import abstractmethod
-from typing import Optional
+from abc import ABC, abstractmethod
 
-from msal_extensions.persistence import ABC  # type: ignore
 from pydantic import BaseModel, ConfigDict
 
 from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
@@ -14,7 +12,7 @@ class DatasetRetrieverBaseTool(BaseModel, ABC):
     description: str = "use this to retrieve a dataset. "
     tenant_id: str
     top_k: int = 4
-    score_threshold: Optional[float] = None
+    score_threshold: float | None = None
     hit_callbacks: list[DatasetIndexToolCallbackHandler] = []
     return_resource: bool
     retriever_from: str

api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py

@@ -1,6 +1,7 @@
-from typing import Any, Optional, cast
+from typing import Any, cast
 
 from pydantic import BaseModel, Field
+from sqlalchemy import select
 
 from core.app.app_config.entities import DatasetRetrieveConfigEntity, ModelConfig
 from core.rag.datasource.retrieval_service import RetrievalService
@@ -36,7 +37,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
     args_schema: type[BaseModel] = DatasetRetrieverToolInput
     description: str = "use this to retrieve a dataset. "
     dataset_id: str
-    user_id: Optional[str] = None
+    user_id: str | None = None
     retrieve_config: DatasetRetrieveConfigEntity
     inputs: dict
 
@@ -56,9 +57,8 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
         )
 
     def _run(self, query: str) -> str:
-        dataset = (
-            db.session.query(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == self.dataset_id).first()
-        )
+        dataset_stmt = select(Dataset).where(Dataset.tenant_id == self.tenant_id, Dataset.id == self.dataset_id)
+        dataset = db.session.scalar(dataset_stmt)
 
         if not dataset:
             return ""
@@ -188,15 +188,12 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
                         for record in records:
                             segment = record.segment
                             dataset = db.session.query(Dataset).filter_by(id=segment.dataset_id).first()
-                            document = (
-                                db.session.query(DatasetDocument)  # type: ignore
-                                .where(
-                                    DatasetDocument.id == segment.document_id,
-                                    DatasetDocument.enabled == True,
-                                    DatasetDocument.archived == False,
-                                )
-                                .first()
+                            dataset_document_stmt = select(DatasetDocument).where(
+                                DatasetDocument.id == segment.document_id,
+                                DatasetDocument.enabled == True,
+                                DatasetDocument.archived == False,
                             )
+                            document = db.session.scalar(dataset_document_stmt)  # type: ignore
                             if dataset and document:
                                 source = RetrievalSourceMetadata(
                                     dataset_id=dataset.id,

api/core/tools/utils/dataset_retriever_tool.py

@@ -1,5 +1,5 @@
 from collections.abc import Generator
-from typing import Any, Optional
+from typing import Any
 
 from core.app.app_config.entities import DatasetRetrieveConfigEntity
 from core.app.entities.app_invoke_entities import InvokeFrom
@@ -20,7 +20,7 @@ from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import Datas
 
 
 class DatasetRetrieverTool(Tool):
-    def __init__(self, entity: ToolEntity, runtime: ToolRuntime, retrieval_tool: DatasetRetrieverBaseTool) -> None:
+    def __init__(self, entity: ToolEntity, runtime: ToolRuntime, retrieval_tool: DatasetRetrieverBaseTool):
         super().__init__(entity, runtime)
         self.retrieval_tool = retrieval_tool
 
@@ -87,9 +87,9 @@ class DatasetRetrieverTool(Tool):
 
     def get_runtime_parameters(
         self,
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> list[ToolParameter]:
         return [
             ToolParameter(
@@ -112,9 +112,9 @@ class DatasetRetrieverTool(Tool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         invoke dataset retriever tool

api/core/tools/utils/message_transformer.py

@@ -3,7 +3,6 @@ from collections.abc import Generator
 from datetime import date, datetime
 from decimal import Decimal
 from mimetypes import guess_extension
-from typing import Optional
 from uuid import UUID
 
 import numpy as np
@@ -60,7 +59,7 @@ class ToolFileMessageTransformer:
         messages: Generator[ToolInvokeMessage, None, None],
         user_id: str,
         tenant_id: str,
-        conversation_id: Optional[str] = None,
+        conversation_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         Transform tool message and handle file download
@@ -165,5 +164,5 @@ class ToolFileMessageTransformer:
                 yield message
 
     @classmethod
-    def get_tool_file_url(cls, tool_file_id: str, extension: Optional[str]) -> str:
+    def get_tool_file_url(cls, tool_file_id: str, extension: str | None) -> str:
         return f"/files/tools/{tool_file_id}{extension or '.bin'}"

api/core/tools/utils/model_invocation_utils.py

@@ -5,7 +5,7 @@ Therefore, a model manager is needed to list/invoke/validate models.
 """
 
 import json
-from typing import Optional, cast
+from typing import cast
 
 from core.model_manager import ModelManager
 from core.model_runtime.entities.llm_entities import LLMResult
@@ -51,7 +51,7 @@ class ModelInvocationUtils:
         if not schema:
             raise InvokeModelError("No model schema found")
 
-        max_tokens: Optional[int] = schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE, None)
+        max_tokens: int | None = schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE, None)
         if max_tokens is None:
             return 2048
 

api/core/tools/utils/parser.py

@@ -2,7 +2,6 @@ import re
 from json import dumps as json_dumps
 from json import loads as json_loads
 from json.decoder import JSONDecodeError
-from typing import Optional
 
 from flask import request
 from requests import get
@@ -198,9 +197,9 @@ class ApiBasedToolSchemaParser:
         return bundles
 
     @staticmethod
-    def _get_tool_parameter_type(parameter: dict) -> Optional[ToolParameter.ToolParameterType]:
+    def _get_tool_parameter_type(parameter: dict) -> ToolParameter.ToolParameterType | None:
         parameter = parameter or {}
-        typ: Optional[str] = None
+        typ: str | None = None
         if parameter.get("format") == "binary":
             return ToolParameter.ToolParameterType.FILE
 
@@ -242,7 +241,7 @@ class ApiBasedToolSchemaParser:
         return ApiBasedToolSchemaParser.parse_openapi_to_tool_bundle(openapi, extra_info=extra_info, warning=warning)
 
     @staticmethod
-    def parse_swagger_to_openapi(swagger: dict, extra_info: dict | None = None, warning: dict | None = None) -> dict:
+    def parse_swagger_to_openapi(swagger: dict, extra_info: dict | None = None, warning: dict | None = None):
         warning = warning or {}
         """
         parse swagger to openapi

api/core/tools/utils/system_oauth_encryption.py

@@ -2,7 +2,7 @@ import base64
 import hashlib
 import logging
 from collections.abc import Mapping
-from typing import Any, Optional
+from typing import Any
 
 from Crypto.Cipher import AES
 from Crypto.Random import get_random_bytes
@@ -28,7 +28,7 @@ class SystemOAuthEncrypter:
     using AES-CBC mode with a key derived from the application's SECRET_KEY.
     """
 
-    def __init__(self, secret_key: Optional[str] = None):
+    def __init__(self, secret_key: str | None = None):
         """
         Initialize the OAuth encrypter.
 
@@ -130,7 +130,7 @@ class SystemOAuthEncrypter:
 
 
 # Factory function for creating encrypter instances
-def create_system_oauth_encrypter(secret_key: Optional[str] = None) -> SystemOAuthEncrypter:
+def create_system_oauth_encrypter(secret_key: str | None = None) -> SystemOAuthEncrypter:
     """
     Create an OAuth encrypter instance.
 
@@ -144,7 +144,7 @@ def create_system_oauth_encrypter(secret_key: Optional[str] = None) -> SystemOAu
 
 
 # Global encrypter instance (for backward compatibility)
-_oauth_encrypter: Optional[SystemOAuthEncrypter] = None
+_oauth_encrypter: SystemOAuthEncrypter | None = None
 
 
 def get_system_oauth_encrypter() -> SystemOAuthEncrypter:

api/core/tools/utils/web_reader_tool.py

@@ -2,7 +2,7 @@ import mimetypes
 import re
 from collections.abc import Sequence
 from dataclasses import dataclass
-from typing import Any, Optional, cast
+from typing import Any, cast
 from urllib.parse import unquote
 
 import chardet
@@ -27,7 +27,7 @@ def page_result(text: str, cursor: int, max_length: int) -> str:
     return text[cursor : cursor + max_length]
 
 
-def get_url(url: str, user_agent: Optional[str] = None) -> str:
+def get_url(url: str, user_agent: str | None = None) -> str:
     """Fetch URL and return the contents as a string."""
     headers = {
         "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko)"

api/core/tools/utils/yaml_utils.py

@@ -1,4 +1,5 @@
 import logging
+from functools import lru_cache
 from pathlib import Path
 from typing import Any
 
@@ -8,28 +9,25 @@ from yaml import YAMLError
 logger = logging.getLogger(__name__)
 
 
-def load_yaml_file(file_path: str, ignore_error: bool = True, default_value: Any = {}) -> Any:
-    """
-    Safe loading a YAML file
-    :param file_path: the path of the YAML file
-    :param ignore_error:
-        if True, return default_value if error occurs and the error will be logged in debug level
-        if False, raise error if error occurs
-    :param default_value: the value returned when errors ignored
-    :return: an object of the YAML content
-    """
+def _load_yaml_file(*, file_path: str):
     if not file_path or not Path(file_path).exists():
-        if ignore_error:
-            return default_value
-        else:
-            raise FileNotFoundError(f"File not found: {file_path}")
+        raise FileNotFoundError(f"File not found: {file_path}")
 
     with open(file_path, encoding="utf-8") as yaml_file:
         try:
             yaml_content = yaml.safe_load(yaml_file)
-            return yaml_content or default_value
+            return yaml_content
         except Exception as e:
-            if ignore_error:
-                return default_value
-            else:
-                raise YAMLError(f"Failed to load YAML file {file_path}: {e}") from e
+            raise YAMLError(f"Failed to load YAML file {file_path}: {e}") from e
+
+
+@lru_cache(maxsize=128)
+def load_yaml_file_cached(file_path: str) -> Any:
+    """
+    Cached version of load_yaml_file for static configuration files.
+    Only use for files that don't change during runtime (e.g., position files)
+
+    :param file_path: the path of the YAML file
+    :return: an object of the YAML content
+    """
+    return _load_yaml_file(file_path=file_path)

api/core/tools/workflow_as_tool/provider.py

@@ -1,5 +1,4 @@
 from collections.abc import Mapping
-from typing import Optional
 
 from pydantic import Field
 
@@ -207,7 +206,7 @@ class WorkflowToolProviderController(ToolProviderController):
 
         return self.tools
 
-    def get_tool(self, tool_name: str) -> Optional[WorkflowTool]:  # type: ignore
+    def get_tool(self, tool_name: str) -> WorkflowTool | None:  # type: ignore
         """
         get tool by name
 

api/core/tools/workflow_as_tool/tool.py

@@ -1,7 +1,9 @@
 import json
 import logging
 from collections.abc import Generator
-from typing import Any, Optional
+from typing import Any
+
+from sqlalchemy import select
 
 from core.file import FILE_MODEL_IDENTITY, File, FileTransferMethod
 from core.tools.__base.tool import Tool
@@ -37,14 +39,12 @@ class WorkflowTool(Tool):
         entity: ToolEntity,
         runtime: ToolRuntime,
         label: str = "Workflow",
-        thread_pool_id: Optional[str] = None,
     ):
         self.workflow_app_id = workflow_app_id
         self.workflow_as_tool_id = workflow_as_tool_id
         self.version = version
         self.workflow_entities = workflow_entities
         self.workflow_call_depth = workflow_call_depth
-        self.thread_pool_id = thread_pool_id
         self.label = label
 
         super().__init__(entity=entity, runtime=runtime)
@@ -61,9 +61,9 @@ class WorkflowTool(Tool):
         self,
         user_id: str,
         tool_parameters: dict[str, Any],
-        conversation_id: Optional[str] = None,
-        app_id: Optional[str] = None,
-        message_id: Optional[str] = None,
+        conversation_id: str | None = None,
+        app_id: str | None = None,
+        message_id: str | None = None,
     ) -> Generator[ToolInvokeMessage, None, None]:
         """
         invoke the tool
@@ -88,7 +88,6 @@ class WorkflowTool(Tool):
             invoke_from=self.runtime.invoke_from,
             streaming=False,
             call_depth=self.workflow_call_depth + 1,
-            workflow_thread_pool_id=self.thread_pool_id,
         )
         assert isinstance(result, dict)
         data = result.get("data", {})
@@ -136,7 +135,8 @@ class WorkflowTool(Tool):
                 .first()
             )
         else:
-            workflow = db.session.query(Workflow).where(Workflow.app_id == app_id, Workflow.version == version).first()
+            stmt = select(Workflow).where(Workflow.app_id == app_id, Workflow.version == version)
+            workflow = db.session.scalar(stmt)
 
         if not workflow:
             raise ValueError("workflow not found or not published")
@@ -147,7 +147,8 @@ class WorkflowTool(Tool):
         """
         get the app by app id
         """
-        app = db.session.query(App).where(App.id == app_id).first()
+        stmt = select(App).where(App.id == app_id)
+        app = db.session.scalar(stmt)
         if not app:
             raise ValueError("app not found")
 
@@ -219,7 +220,7 @@ class WorkflowTool(Tool):
 
         return result, files
 
-    def _update_file_mapping(self, file_dict: dict) -> dict:
+    def _update_file_mapping(self, file_dict: dict):
         transfer_method = FileTransferMethod.value_of(file_dict.get("transfer_method"))
         if transfer_method == FileTransferMethod.TOOL_FILE:
             file_dict["tool_file_id"] = file_dict.get("related_id")