diff --git a/api/core/provider_manager.py b/api/core/provider_manager.py index 5d536e0e324..552de66f8bc 100644 --- a/api/core/provider_manager.py +++ b/api/core/provider_manager.py @@ -1,11 +1,10 @@ from __future__ import annotations import contextlib -import json from collections import defaultdict from collections.abc import Sequence from json import JSONDecodeError -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any from graphon.model_runtime.entities.model_entities import ModelType from graphon.model_runtime.entities.provider_entities import ( @@ -15,6 +14,7 @@ from graphon.model_runtime.entities.provider_entities import ( ProviderEntity, ) from graphon.model_runtime.model_providers.model_provider_factory import ModelProviderFactory +from pydantic import TypeAdapter from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session @@ -58,6 +58,8 @@ from services.feature_service import FeatureService if TYPE_CHECKING: from graphon.model_runtime.runtime import ModelRuntime +_credentials_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any]) + class ProviderManager: """ @@ -875,8 +877,8 @@ class ProviderManager: return {"openai_api_key": encrypted_config} try: - credentials = cast(dict, json.loads(encrypted_config)) - except JSONDecodeError: + credentials = _credentials_adapter.validate_json(encrypted_config) + except (ValueError, JSONDecodeError): return {} # Decrypt secret variables @@ -1015,7 +1017,7 @@ class ProviderManager: if not cached_provider_credentials: provider_credentials: dict[str, Any] = {} if provider_records and provider_records[0].encrypted_config: - provider_credentials = json.loads(provider_records[0].encrypted_config) + provider_credentials = _credentials_adapter.validate_json(provider_records[0].encrypted_config) # Get provider credential secret variables provider_credential_secret_variables = self._extract_secret_variables( @@ -1162,8 +1164,10 @@ class ProviderManager: if not cached_provider_model_credentials: try: - provider_model_credentials = json.loads(load_balancing_model_config.encrypted_config) - except JSONDecodeError: + provider_model_credentials = _credentials_adapter.validate_json( + load_balancing_model_config.encrypted_config + ) + except (ValueError, JSONDecodeError): continue # Get decoding rsa key and cipher for decrypting credentials @@ -1176,7 +1180,7 @@ class ProviderManager: if variable in provider_model_credentials: try: provider_model_credentials[variable] = encrypter.decrypt_token_with_decoding( - provider_model_credentials.get(variable), + provider_model_credentials.get(variable) or "", self.decoding_rsa_key, self.decoding_cipher_rsa, ) diff --git a/api/core/tools/entities/tool_entities.py b/api/core/tools/entities/tool_entities.py index 96268d029ed..46f079937bf 100644 --- a/api/core/tools/entities/tool_entities.py +++ b/api/core/tools/entities/tool_entities.py @@ -6,7 +6,17 @@ from collections.abc import Mapping from enum import StrEnum, auto from typing import Any, Union -from pydantic import BaseModel, ConfigDict, Field, ValidationInfo, field_serializer, field_validator, model_validator +from pydantic import ( + BaseModel, + ConfigDict, + Field, + TypeAdapter, + ValidationInfo, + field_serializer, + field_validator, + model_validator, +) +from typing_extensions import TypedDict from core.entities.provider_entities import ProviderConfig from core.plugin.entities.parameters import ( @@ -23,6 +33,14 @@ from core.tools.entities.common_entities import I18nObject from core.tools.entities.constants import TOOL_SELECTOR_MODEL_IDENTITY +class EmojiIconDict(TypedDict): + background: str + content: str + + +emoji_icon_adapter: TypeAdapter[EmojiIconDict] = TypeAdapter(EmojiIconDict) + + class ToolLabelEnum(StrEnum): SEARCH = "search" IMAGE = "image" diff --git a/api/core/tools/tool_manager.py b/api/core/tools/tool_manager.py index a58d3103137..34fbec0c8b6 100644 --- a/api/core/tools/tool_manager.py +++ b/api/core/tools/tool_manager.py @@ -5,12 +5,14 @@ import time from collections.abc import Generator, Mapping from os import listdir, path from threading import Lock -from typing import TYPE_CHECKING, Any, Literal, Optional, Protocol, TypedDict, Union, cast +from typing import TYPE_CHECKING, Any, Literal, Optional, Protocol, Union, cast import sqlalchemy as sa from graphon.runtime import VariablePool +from pydantic import TypeAdapter from sqlalchemy import select from sqlalchemy.orm import Session +from typing_extensions import TypedDict from yarl import URL import contexts @@ -49,9 +51,11 @@ from core.tools.entities.api_entities import ToolProviderApiEntity, ToolProvider from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_entities import ( ApiProviderAuthType, + EmojiIconDict, ToolInvokeFrom, ToolParameter, ToolProviderType, + emoji_icon_adapter, ) from core.tools.errors import ToolProviderNotFoundError from core.tools.tool_label_manager import ToolLabelManager @@ -72,9 +76,7 @@ class ApiProviderControllerItem(TypedDict): controller: ApiToolProviderController -class EmojiIconDict(TypedDict): - background: str - content: str +_credentials_adapter: TypeAdapter[dict[str, Any]] = TypeAdapter(dict[str, Any]) class WorkflowToolRuntimeSpec(Protocol): @@ -885,7 +887,7 @@ class ToolManager: raise ValueError(f"you have not added provider {provider_name}") try: - credentials = json.loads(provider_obj.credentials_str) or {} + credentials = _credentials_adapter.validate_json(provider_obj.credentials_str) or {} except Exception: credentials = {} @@ -910,7 +912,7 @@ class ToolManager: masked_credentials = encrypter.mask_plugin_credentials(encrypter.decrypt(credentials)) try: - icon = json.loads(provider_obj.icon) + icon = emoji_icon_adapter.validate_json(provider_obj.icon) except Exception: icon = {"background": "#252525", "content": "\ud83d\ude01"} @@ -973,7 +975,7 @@ class ToolManager: if workflow_provider is None: raise ToolProviderNotFoundError(f"workflow provider {provider_id} not found") - icon = json.loads(workflow_provider.icon) + icon = emoji_icon_adapter.validate_json(workflow_provider.icon) return icon except Exception: return {"background": "#252525", "content": "\ud83d\ude01"} @@ -990,7 +992,7 @@ class ToolManager: if api_provider is None: raise ToolProviderNotFoundError(f"api provider {provider_id} not found") - icon = json.loads(api_provider.icon) + icon = emoji_icon_adapter.validate_json(api_provider.icon) return icon except Exception: return {"background": "#252525", "content": "\ud83d\ude01"} diff --git a/api/libs/helper.py b/api/libs/helper.py index a7b3da77ff8..ece53e88062 100644 --- a/api/libs/helper.py +++ b/api/libs/helper.py @@ -18,8 +18,9 @@ from flask import Response, stream_with_context from flask_restx import fields from graphon.file import helpers as file_helpers from graphon.model_runtime.utils.encoders import jsonable_encoder -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from pydantic.functional_validators import AfterValidator +from typing_extensions import TypedDict from configs import dify_config from core.app.features.rate_limiting.rate_limit import RateLimitGenerator @@ -32,6 +33,17 @@ if TYPE_CHECKING: logger = logging.getLogger(__name__) +class _TokenData(TypedDict, total=False): + account_id: str | None + email: str + token_type: str + code: str + old_email: str + + +_token_data_adapter: TypeAdapter[_TokenData] = TypeAdapter(_TokenData) + + def _stream_with_request_context(response: object) -> Any: """Bridge Flask's loosely-typed streaming helper without leaking casts into callers.""" return cast(Any, stream_with_context)(response) @@ -443,7 +455,7 @@ class TokenManager: if token_data_json is None: logger.warning("%s token %s not found with key %s", token_type, token, key) return None - token_data: dict[str, Any] | None = json.loads(token_data_json) + token_data = dict(_token_data_adapter.validate_json(token_data_json)) return token_data @classmethod diff --git a/api/services/tools/tools_transform_service.py b/api/services/tools/tools_transform_service.py index 7cd61e3162a..b24f0011338 100644 --- a/api/services/tools/tools_transform_service.py +++ b/api/services/tools/tools_transform_service.py @@ -1,4 +1,3 @@ -import json import logging from collections.abc import Mapping from typing import Any, Union @@ -21,6 +20,7 @@ from core.tools.entities.tool_entities import ( ApiProviderAuthType, ToolParameter, ToolProviderType, + emoji_icon_adapter, ) from core.tools.plugin_tool.provider import PluginToolProviderController from core.tools.utils.encryption import create_provider_encrypter, create_tool_provider_encrypter @@ -53,11 +53,14 @@ class ToolTransformService: elif provider_type in {ToolProviderType.API, ToolProviderType.WORKFLOW}: try: if isinstance(icon, str): - return json.loads(icon) - return icon - except (json.JSONDecodeError, ValueError): + parsed = emoji_icon_adapter.validate_json(icon) + return {"background": parsed["background"], "content": parsed["content"]} + return {"background": icon["background"], "content": icon["content"]} + except (ValueError, ValidationError, KeyError): return {"background": "#252525", "content": "\ud83d\ude01"} elif provider_type == ToolProviderType.MCP: + if isinstance(icon, Mapping): + return {"background": icon.get("background", ""), "content": icon.get("content", "")} return icon return "" diff --git a/api/services/tools/workflow_tools_manage_service.py b/api/services/tools/workflow_tools_manage_service.py index dc0b281e15e..8f5144c8667 100644 --- a/api/services/tools/workflow_tools_manage_service.py +++ b/api/services/tools/workflow_tools_manage_service.py @@ -8,7 +8,7 @@ from sqlalchemy.orm import Session from core.tools.__base.tool_provider import ToolProviderController from core.tools.entities.api_entities import ToolApiEntity, ToolProviderApiEntity -from core.tools.entities.tool_entities import WorkflowToolParameterConfiguration +from core.tools.entities.tool_entities import WorkflowToolParameterConfiguration, emoji_icon_adapter from core.tools.tool_label_manager import ToolLabelManager from core.tools.utils.workflow_configuration_sync import WorkflowToolConfigurationUtils from core.tools.workflow_as_tool.provider import WorkflowToolProviderController @@ -313,7 +313,7 @@ class WorkflowToolManageService: "label": db_tool.label, "workflow_tool_id": db_tool.id, "workflow_app_id": db_tool.app_id, - "icon": json.loads(db_tool.icon), + "icon": emoji_icon_adapter.validate_json(db_tool.icon), "description": db_tool.description, "parameters": jsonable_encoder(db_tool.parameter_configurations), "output_schema": output_schema,