Underground-Digital/Workflow-Engine

1from copy import deepcopy2from typing import Any3 4from pydantic import BaseModel5 6from core.helper import encrypter7from core.helper.tool_parameter_cache import ToolParameterCache, ToolParameterCacheType8from core.helper.tool_provider_cache import ToolProviderCredentialsCache, ToolProviderCredentialsCacheType9from core.tools.entities.tool_entities import (10    ToolParameter,11    ToolProviderCredentials,12)13from core.tools.provider.tool_provider import ToolProviderController14from core.tools.tool.tool import Tool15 16 17class ToolConfigurationManager(BaseModel):18    tenant_id: str19    provider_controller: ToolProviderController20 21    def _deep_copy(self, credentials: dict[str, str]) -> dict[str, str]:22        """23        deep copy credentials24        """25        return deepcopy(credentials)26 27    def encrypt_tool_credentials(self, credentials: dict[str, str]) -> dict[str, str]:28        """29        encrypt tool credentials with tenant id30 31        return a deep copy of credentials with encrypted values32        """33        credentials = self._deep_copy(credentials)34 35        # get fields need to be decrypted36        fields = self.provider_controller.get_credentials_schema()37        for field_name, field in fields.items():38            if field.type == ToolProviderCredentials.CredentialsType.SECRET_INPUT:39                if field_name in credentials:40                    encrypted = encrypter.encrypt_token(self.tenant_id, credentials[field_name])41                    credentials[field_name] = encrypted42 43        return credentials44 45    def mask_tool_credentials(self, credentials: dict[str, Any]) -> dict[str, Any]:46        """47        mask tool credentials48 49        return a deep copy of credentials with masked values50        """51        credentials = self._deep_copy(credentials)52 53        # get fields need to be decrypted54        fields = self.provider_controller.get_credentials_schema()55        for field_name, field in fields.items():56            if field.type == ToolProviderCredentials.CredentialsType.SECRET_INPUT:57                if field_name in credentials:58                    if len(credentials[field_name]) > 6:59                        credentials[field_name] = (60                            credentials[field_name][:2]61                            + "*" * (len(credentials[field_name]) - 4)62                            + credentials[field_name][-2:]63                        )64                    else:65                        credentials[field_name] = "*" * len(credentials[field_name])66 67        return credentials68 69    def decrypt_tool_credentials(self, credentials: dict[str, str]) -> dict[str, str]:70        """71        decrypt tool credentials with tenant id72 73        return a deep copy of credentials with decrypted values74        """75        cache = ToolProviderCredentialsCache(76            tenant_id=self.tenant_id,77            identity_id=f"{self.provider_controller.provider_type.value}.{self.provider_controller.identity.name}",78            cache_type=ToolProviderCredentialsCacheType.PROVIDER,79        )80        cached_credentials = cache.get()81        if cached_credentials:82            return cached_credentials83        credentials = self._deep_copy(credentials)84        # get fields need to be decrypted85        fields = self.provider_controller.get_credentials_schema()86        for field_name, field in fields.items():87            if field.type == ToolProviderCredentials.CredentialsType.SECRET_INPUT:88                if field_name in credentials:89                    try:90                        credentials[field_name] = encrypter.decrypt_token(self.tenant_id, credentials[field_name])91                    except:92                        pass93 94        cache.set(credentials)95        return credentials96 97    def delete_tool_credentials_cache(self):98        cache = ToolProviderCredentialsCache(99            tenant_id=self.tenant_id,100            identity_id=f"{self.provider_controller.provider_type.value}.{self.provider_controller.identity.name}",101            cache_type=ToolProviderCredentialsCacheType.PROVIDER,102        )103        cache.delete()104 105 106class ToolParameterConfigurationManager(BaseModel):107    """108    Tool parameter configuration manager109    """110 111    tenant_id: str112    tool_runtime: Tool113    provider_name: str114    provider_type: str115    identity_id: str116 117    def _deep_copy(self, parameters: dict[str, Any]) -> dict[str, Any]:118        """119        deep copy parameters120        """121        return deepcopy(parameters)122 123    def _merge_parameters(self) -> list[ToolParameter]:124        """125        merge parameters126        """127        # get tool parameters128        tool_parameters = self.tool_runtime.parameters or []129        # get tool runtime parameters130        runtime_parameters = self.tool_runtime.get_runtime_parameters() or []131        # override parameters132        current_parameters = tool_parameters.copy()133        for runtime_parameter in runtime_parameters:134            found = False135            for index, parameter in enumerate(current_parameters):136                if parameter.name == runtime_parameter.name and parameter.form == runtime_parameter.form:137                    current_parameters[index] = runtime_parameter138                    found = True139                    break140 141            if not found and runtime_parameter.form == ToolParameter.ToolParameterForm.FORM:142                current_parameters.append(runtime_parameter)143 144        return current_parameters145 146    def mask_tool_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]:147        """148        mask tool parameters149 150        return a deep copy of parameters with masked values151        """152        parameters = self._deep_copy(parameters)153 154        # override parameters155        current_parameters = self._merge_parameters()156 157        for parameter in current_parameters:158            if (159                parameter.form == ToolParameter.ToolParameterForm.FORM160                and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT161            ):162                if parameter.name in parameters:163                    if len(parameters[parameter.name]) > 6:164                        parameters[parameter.name] = (165                            parameters[parameter.name][:2]166                            + "*" * (len(parameters[parameter.name]) - 4)167                            + parameters[parameter.name][-2:]168                        )169                    else:170                        parameters[parameter.name] = "*" * len(parameters[parameter.name])171 172        return parameters173 174    def encrypt_tool_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]:175        """176        encrypt tool parameters with tenant id177 178        return a deep copy of parameters with encrypted values179        """180        # override parameters181        current_parameters = self._merge_parameters()182 183        parameters = self._deep_copy(parameters)184 185        for parameter in current_parameters:186            if (187                parameter.form == ToolParameter.ToolParameterForm.FORM188                and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT189            ):190                if parameter.name in parameters:191                    encrypted = encrypter.encrypt_token(self.tenant_id, parameters[parameter.name])192                    parameters[parameter.name] = encrypted193 194        return parameters195 196    def decrypt_tool_parameters(self, parameters: dict[str, Any]) -> dict[str, Any]:197        """198        decrypt tool parameters with tenant id199 200        return a deep copy of parameters with decrypted values201        """202        cache = ToolParameterCache(203            tenant_id=self.tenant_id,204            provider=f"{self.provider_type}.{self.provider_name}",205            tool_name=self.tool_runtime.identity.name,206            cache_type=ToolParameterCacheType.PARAMETER,207            identity_id=self.identity_id,208        )209        cached_parameters = cache.get()210        if cached_parameters:211            return cached_parameters212 213        # override parameters214        current_parameters = self._merge_parameters()215        has_secret_input = False216 217        for parameter in current_parameters:218            if (219                parameter.form == ToolParameter.ToolParameterForm.FORM220                and parameter.type == ToolParameter.ToolParameterType.SECRET_INPUT221            ):222                if parameter.name in parameters:223                    try:224                        has_secret_input = True225                        parameters[parameter.name] = encrypter.decrypt_token(self.tenant_id, parameters[parameter.name])226                    except:227                        pass228 229        if has_secret_input:230            cache.set(parameters)231 232        return parameters233 234    def delete_tool_parameters_cache(self):235        cache = ToolParameterCache(236            tenant_id=self.tenant_id,237            provider=f"{self.provider_type}.{self.provider_name}",238            tool_name=self.tool_runtime.identity.name,239            cache_type=ToolParameterCacheType.PARAMETER,240            identity_id=self.identity_id,241        )242        cache.delete()243