Underground-Digital/Workflow-Engine

1from typing import Optional2 3from core.extension.api_based_extension_requestor import APIBasedExtensionRequestor4from core.external_data_tool.base import ExternalDataTool5from core.helper import encrypter6from extensions.ext_database import db7from models.api_based_extension import APIBasedExtension, APIBasedExtensionPoint8 9 10class ApiExternalDataTool(ExternalDataTool):11    """12    The api external data tool.13    """14 15    name: str = "api"16    """the unique name of external data tool"""17 18    @classmethod19    def validate_config(cls, tenant_id: str, config: dict) -> None:20        """21        Validate the incoming form config data.22 23        :param tenant_id: the id of workspace24        :param config: the form config data25        :return:26        """27        # own validation logic28        api_based_extension_id = config.get("api_based_extension_id")29        if not api_based_extension_id:30            raise ValueError("api_based_extension_id is required")31 32        # get api_based_extension33        api_based_extension = (34            db.session.query(APIBasedExtension)35            .filter(APIBasedExtension.tenant_id == tenant_id, APIBasedExtension.id == api_based_extension_id)36            .first()37        )38 39        if not api_based_extension:40            raise ValueError("api_based_extension_id is invalid")41 42    def query(self, inputs: dict, query: Optional[str] = None) -> str:43        """44        Query the external data tool.45 46        :param inputs: user inputs47        :param query: the query of chat app48        :return: the tool query result49        """50        # get params from config51        api_based_extension_id = self.config.get("api_based_extension_id")52 53        # get api_based_extension54        api_based_extension = (55            db.session.query(APIBasedExtension)56            .filter(APIBasedExtension.tenant_id == self.tenant_id, APIBasedExtension.id == api_based_extension_id)57            .first()58        )59 60        if not api_based_extension:61            raise ValueError(62                "[External data tool] API query failed, variable: {}, "63                "error: api_based_extension_id is invalid".format(self.variable)64            )65 66        # decrypt api_key67        api_key = encrypter.decrypt_token(tenant_id=self.tenant_id, token=api_based_extension.api_key)68 69        try:70            # request api71            requestor = APIBasedExtensionRequestor(api_endpoint=api_based_extension.api_endpoint, api_key=api_key)72        except Exception as e:73            raise ValueError("[External data tool] API query failed, variable: {}, error: {}".format(self.variable, e))74 75        response_json = requestor.request(76            point=APIBasedExtensionPoint.APP_EXTERNAL_DATA_TOOL_QUERY,77            params={"app_id": self.app_id, "tool_variable": self.variable, "inputs": inputs, "query": query},78        )79 80        if "result" not in response_json:81            raise ValueError(82                "[External data tool] API query failed, variable: {}, error: result not found in response".format(83                    self.variable84                )85            )86 87        if not isinstance(response_json["result"], str):88            raise ValueError(89                "[External data tool] API query failed, variable: {}, error: result is not string".format(self.variable)90            )91 92        return response_json["result"]93