mirror of
https://github.com/langgenius/dify.git
synced 2024-11-16 19:59:50 +08:00
230 lines
8.7 KiB
Python
230 lines
8.7 KiB
Python
from copy import deepcopy
|
|
from typing import Any
|
|
|
|
from pydantic import BaseModel
|
|
|
|
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.entities.tool_entities import (
|
|
ToolParameter,
|
|
ToolProviderCredentials,
|
|
)
|
|
from core.tools.provider.tool_provider import ToolProviderController
|
|
from core.tools.tool.tool import Tool
|
|
|
|
|
|
class ToolConfigurationManager(BaseModel):
|
|
tenant_id: str
|
|
provider_controller: ToolProviderController
|
|
|
|
def _deep_copy(self, credentials: dict[str, str]) -> dict[str, str]:
|
|
"""
|
|
deep copy credentials
|
|
"""
|
|
return deepcopy(credentials)
|
|
|
|
def encrypt_tool_credentials(self, credentials: dict[str, str]) -> dict[str, str]:
|
|
"""
|
|
encrypt tool credentials with tenant id
|
|
|
|
return a deep copy of credentials with encrypted values
|
|
"""
|
|
credentials = self._deep_copy(credentials)
|
|
|
|
# get fields need to be decrypted
|
|
fields = self.provider_controller.get_credentials_schema()
|
|
for field_name, field in fields.items():
|
|
if field.type == ToolProviderCredentials.CredentialsType.SECRET_INPUT:
|
|
if field_name in credentials:
|
|
encrypted = encrypter.encrypt_token(self.tenant_id, credentials[field_name])
|
|
credentials[field_name] = encrypted
|
|
|
|
return credentials
|
|
|
|
def mask_tool_credentials(self, credentials: dict[str, Any]) -> dict[str, Any]:
|
|
"""
|
|
mask tool credentials
|
|
|
|
return a deep copy of credentials with masked values
|
|
"""
|
|
credentials = self._deep_copy(credentials)
|
|
|
|
# get fields need to be decrypted
|
|
fields = self.provider_controller.get_credentials_schema()
|
|
for field_name, field in fields.items():
|
|
if field.type == ToolProviderCredentials.CredentialsType.SECRET_INPUT:
|
|
if field_name in credentials:
|
|
if len(credentials[field_name]) > 6:
|
|
credentials[field_name] = \
|
|
credentials[field_name][:2] + \
|
|
'*' * (len(credentials[field_name]) - 4) + \
|
|
credentials[field_name][-2:]
|
|
else:
|
|
credentials[field_name] = '*' * len(credentials[field_name])
|
|
|
|
return credentials
|
|
|
|
def decrypt_tool_credentials(self, credentials: 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_controller.provider_type.value}.{self.provider_controller.identity.name}',
|
|
cache_type=ToolProviderCredentialsCacheType.PROVIDER
|
|
)
|
|
cached_credentials = cache.get()
|
|
if cached_credentials:
|
|
return cached_credentials
|
|
credentials = self._deep_copy(credentials)
|
|
# get fields need to be decrypted
|
|
fields = self.provider_controller.get_credentials_schema()
|
|
for field_name, field in fields.items():
|
|
if field.type == ToolProviderCredentials.CredentialsType.SECRET_INPUT:
|
|
if field_name in credentials:
|
|
try:
|
|
credentials[field_name] = encrypter.decrypt_token(self.tenant_id, credentials[field_name])
|
|
except:
|
|
pass
|
|
|
|
cache.set(credentials)
|
|
return credentials
|
|
|
|
def delete_tool_credentials_cache(self):
|
|
cache = ToolProviderCredentialsCache(
|
|
tenant_id=self.tenant_id,
|
|
identity_id=f'{self.provider_controller.provider_type.value}.{self.provider_controller.identity.name}',
|
|
cache_type=ToolProviderCredentialsCacheType.PROVIDER
|
|
)
|
|
cache.delete()
|
|
|
|
class ToolParameterConfigurationManager(BaseModel):
|
|
"""
|
|
Tool parameter configuration manager
|
|
"""
|
|
tenant_id: str
|
|
tool_runtime: Tool
|
|
provider_name: str
|
|
provider_type: str
|
|
identity_id: str
|
|
|
|
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.parameters or []
|
|
# get tool runtime parameters
|
|
runtime_parameters = self.tool_runtime.get_runtime_parameters() or []
|
|
# 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}.{self.provider_name}',
|
|
tool_name=self.tool_runtime.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:
|
|
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}.{self.provider_name}',
|
|
tool_name=self.tool_runtime.identity.name,
|
|
cache_type=ToolParameterCacheType.PARAMETER,
|
|
identity_id=self.identity_id
|
|
)
|
|
cache.delete()
|