mirror of
https://github.com/langgenius/dify.git
synced 2025-10-29 18:03:13 +00:00
266 lines
9.3 KiB
Python
266 lines
9.3 KiB
Python
from copy import deepcopy
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel
|
|
|
|
from core.entities.provider_entities import BasicProviderConfig
|
|
from core.helper import encrypter
|
|
from core.helper.tool_parameter_cache import ToolParameterCache, ToolParameterCacheType
|
|
from core.helper.tool_provider_cache import ToolProviderCredentialsCache, ToolProviderCredentialsCacheType
|
|
from core.tools.__base.tool import Tool
|
|
from core.tools.entities.tool_entities import (
|
|
ToolParameter,
|
|
ToolProviderType,
|
|
)
|
|
|
|
|
|
class ProviderConfigEncrypter(BaseModel):
|
|
tenant_id: str
|
|
config: list[BasicProviderConfig]
|
|
provider_type: str
|
|
provider_identity: str
|
|
|
|
def _deep_copy(self, data: dict[str, str]) -> dict[str, str]:
|
|
"""
|
|
deep copy data
|
|
"""
|
|
return deepcopy(data)
|
|
|
|
def encrypt(self, data: dict[str, str]) -> dict[str, str]:
|
|
"""
|
|
encrypt tool credentials with tenant id
|
|
|
|
return a deep copy of credentials with encrypted values
|
|
"""
|
|
data = self._deep_copy(data)
|
|
|
|
# get fields need to be decrypted
|
|
fields = dict[str, BasicProviderConfig]()
|
|
for credential in self.config:
|
|
fields[credential.name] = credential
|
|
|
|
for field_name, field in fields.items():
|
|
if field.type == BasicProviderConfig.Type.SECRET_INPUT:
|
|
if field_name in data:
|
|
encrypted = encrypter.encrypt_token(self.tenant_id, data[field_name] or "")
|
|
data[field_name] = encrypted
|
|
|
|
return data
|
|
|
|
def mask_tool_credentials(self, data: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
mask tool credentials
|
|
|
|
return a deep copy of credentials with masked values
|
|
"""
|
|
data = self._deep_copy(data)
|
|
|
|
# get fields need to be decrypted
|
|
fields = dict[str, BasicProviderConfig]()
|
|
for credential in self.config:
|
|
fields[credential.name] = credential
|
|
|
|
for field_name, field in fields.items():
|
|
if field.type == BasicProviderConfig.Type.SECRET_INPUT:
|
|
if field_name in data:
|
|
if len(data[field_name]) > 6:
|
|
data[field_name] = (
|
|
data[field_name][:2] + "*" * (len(data[field_name]) - 4) + data[field_name][-2:]
|
|
)
|
|
else:
|
|
data[field_name] = "*" * len(data[field_name])
|
|
|
|
return data
|
|
|
|
def decrypt(self, data: dict[str, str]) -> dict[str, str]:
|
|
"""
|
|
decrypt tool credentials with tenant id
|
|
|
|
return a deep copy of credentials with decrypted values
|
|
"""
|
|
cache = ToolProviderCredentialsCache(
|
|
tenant_id=self.tenant_id,
|
|
identity_id=f"{self.provider_type}.{self.provider_identity}",
|
|
cache_type=ToolProviderCredentialsCacheType.PROVIDER,
|
|
)
|
|
cached_credentials = cache.get()
|
|
if cached_credentials:
|
|
return cached_credentials
|
|
data = self._deep_copy(data)
|
|
# get fields need to be decrypted
|
|
fields = dict[str, BasicProviderConfig]()
|
|
for credential in self.config:
|
|
fields[credential.name] = credential
|
|
|
|
for field_name, field in fields.items():
|
|
if field.type == BasicProviderConfig.Type.SECRET_INPUT:
|
|
if field_name in data:
|
|
try:
|
|
# if the value is None or empty string, skip decrypt
|
|
if not data[field_name]:
|
|
continue
|
|
|
|
data[field_name] = encrypter.decrypt_token(self.tenant_id, data[field_name])
|
|
except Exception:
|
|
pass
|
|
|
|
cache.set(data)
|
|
return data
|
|
|
|
def delete_tool_credentials_cache(self):
|
|
cache = ToolProviderCredentialsCache(
|
|
tenant_id=self.tenant_id,
|
|
identity_id=f"{self.provider_type}.{self.provider_identity}",
|
|
cache_type=ToolProviderCredentialsCacheType.PROVIDER,
|
|
)
|
|
cache.delete()
|
|
|
|
|
|
class ToolParameterConfigurationManager:
|
|
"""
|
|
Tool parameter configuration manager
|
|
"""
|
|
|
|
tenant_id: str
|
|
tool_runtime: Tool
|
|
provider_name: str
|
|
provider_type: ToolProviderType
|
|
identity_id: str
|
|
|
|
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
|
|
self.provider_type = provider_type
|
|
self.identity_id = identity_id
|
|
|
|
def _deep_copy(self, parameters: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
deep copy parameters
|
|
"""
|
|
return deepcopy(parameters)
|
|
|
|
def _merge_parameters(self) -> list[ToolParameter]:
|
|
"""
|
|
merge parameters
|
|
"""
|
|
# get tool parameters
|
|
tool_parameters = self.tool_runtime.entity.parameters or []
|
|
# get tool runtime parameters
|
|
runtime_parameters = self.tool_runtime.get_runtime_parameters()
|
|
# override parameters
|
|
current_parameters = tool_parameters.copy()
|
|
for runtime_parameter in runtime_parameters:
|
|
found = False
|
|
for index, parameter in enumerate(current_parameters):
|
|
if parameter.name == runtime_parameter.name and parameter.form == runtime_parameter.form:
|
|
current_parameters[index] = runtime_parameter
|
|
found = True
|
|
break
|
|
|
|
if not found and runtime_parameter.form == ToolParameter.ToolParameterForm.FORM:
|
|
current_parameters.append(runtime_parameter)
|
|
|
|
return current_parameters
|
|
|
|
def mask_tool_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
mask tool parameters
|
|
|
|
return a deep copy of parameters with masked values
|
|
"""
|
|
parameters = self._deep_copy(parameters)
|
|
|
|
# override parameters
|
|
current_parameters = self._merge_parameters()
|
|
|
|
for parameter in current_parameters:
|
|
if (
|
|
parameter.form == ToolParameter.ToolParameterForm.FORM
|
|
and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT
|
|
):
|
|
if parameter.name in parameters:
|
|
if len(parameters[parameter.name]) > 6:
|
|
parameters[parameter.name] = (
|
|
parameters[parameter.name][:2]
|
|
+ "*" * (len(parameters[parameter.name]) - 4)
|
|
+ parameters[parameter.name][-2:]
|
|
)
|
|
else:
|
|
parameters[parameter.name] = "*" * len(parameters[parameter.name])
|
|
|
|
return parameters
|
|
|
|
def encrypt_tool_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
encrypt tool parameters with tenant id
|
|
|
|
return a deep copy of parameters with encrypted values
|
|
"""
|
|
# override parameters
|
|
current_parameters = self._merge_parameters()
|
|
|
|
parameters = self._deep_copy(parameters)
|
|
|
|
for parameter in current_parameters:
|
|
if (
|
|
parameter.form == ToolParameter.ToolParameterForm.FORM
|
|
and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT
|
|
):
|
|
if parameter.name in parameters:
|
|
encrypted = encrypter.encrypt_token(self.tenant_id, parameters[parameter.name])
|
|
parameters[parameter.name] = encrypted
|
|
|
|
return parameters
|
|
|
|
def decrypt_tool_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
decrypt tool parameters with tenant id
|
|
|
|
return a deep copy of parameters with decrypted values
|
|
"""
|
|
|
|
cache = ToolParameterCache(
|
|
tenant_id=self.tenant_id,
|
|
provider=f"{self.provider_type.value}.{self.provider_name}",
|
|
tool_name=self.tool_runtime.entity.identity.name,
|
|
cache_type=ToolParameterCacheType.PARAMETER,
|
|
identity_id=self.identity_id,
|
|
)
|
|
cached_parameters = cache.get()
|
|
if cached_parameters:
|
|
return cached_parameters
|
|
|
|
# override parameters
|
|
current_parameters = self._merge_parameters()
|
|
has_secret_input = False
|
|
|
|
for parameter in current_parameters:
|
|
if (
|
|
parameter.form == ToolParameter.ToolParameterForm.FORM
|
|
and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT
|
|
):
|
|
if parameter.name in parameters:
|
|
try:
|
|
has_secret_input = True
|
|
parameters[parameter.name] = encrypter.decrypt_token(self.tenant_id, parameters[parameter.name])
|
|
except Exception:
|
|
pass
|
|
|
|
if has_secret_input:
|
|
cache.set(parameters)
|
|
|
|
return parameters
|
|
|
|
def delete_tool_parameters_cache(self):
|
|
cache = ToolParameterCache(
|
|
tenant_id=self.tenant_id,
|
|
provider=f"{self.provider_type.value}.{self.provider_name}",
|
|
tool_name=self.tool_runtime.entity.identity.name,
|
|
cache_type=ToolParameterCacheType.PARAMETER,
|
|
identity_id=self.identity_id,
|
|
)
|
|
cache.delete()
|