tjdify/api/core/external_data_tool/api/api.py

from typing import Optional

from core.extension.api_based_extension_requestor import APIBasedExtensionRequestor
from core.external_data_tool.base import ExternalDataTool
from core.helper import encrypter
from extensions.ext_database import db
from models.api_based_extension import APIBasedExtension, APIBasedExtensionPoint


class ApiExternalDataTool(ExternalDataTool):
    """
    The api external data tool.
    """

    name: str = "api"
    """the unique name of external data tool"""

    @classmethod
    def validate_config(cls, tenant_id: str, config: dict) -> None:
        """
        Validate the incoming form config data.

        :param tenant_id: the id of workspace
        :param config: the form config data
        :return:
        """
        # own validation logic
        api_based_extension_id = config.get("api_based_extension_id")
        if not api_based_extension_id:
            raise ValueError("api_based_extension_id is required")

        # get api_based_extension
        api_based_extension = db.session.query(APIBasedExtension).filter(
            APIBasedExtension.tenant_id == tenant_id,
            APIBasedExtension.id == api_based_extension_id
        ).first()

        if not api_based_extension:
            raise ValueError("api_based_extension_id is invalid")

    def query(self, inputs: dict, query: Optional[str] = None) -> str:
        """
        Query the external data tool.

        :param inputs: user inputs
        :param query: the query of chat app
        :return: the tool query result
        """
        # get params from config
        api_based_extension_id = self.config.get("api_based_extension_id")

        # get api_based_extension
        api_based_extension = db.session.query(APIBasedExtension).filter(
            APIBasedExtension.tenant_id == self.tenant_id,
            APIBasedExtension.id == api_based_extension_id
        ).first()

        if not api_based_extension:
            raise ValueError("[External data tool] API query failed, variable: {}, "
                             "error: api_based_extension_id is invalid"
                             .format(self.config.get('variable')))

        # decrypt api_key
        api_key = encrypter.decrypt_token(
            tenant_id=self.tenant_id,
            token=api_based_extension.api_key
        )

        try:
            # request api
            requestor = APIBasedExtensionRequestor(
                api_endpoint=api_based_extension.api_endpoint,
                api_key=api_key
            )
        except Exception as e:
            raise ValueError("[External data tool] API query failed, variable: {}, error: {}".format(
                self.config.get('variable'),
                e
            ))

        response_json = requestor.request(point=APIBasedExtensionPoint.APP_EXTERNAL_DATA_TOOL_QUERY, params={
            'app_id': self.app_id,
            'tool_variable': self.variable,
            'inputs': inputs,
            'query': query
        })

        if 'result' not in response_json:
            raise ValueError("[External data tool] API query failed, variable: {}, error: result not found in response"
                             .format(self.config.get('variable')))

        return response_json['result']