tjdify/api/core/tools/provider/api_tool_provider.py

from pydantic import Field

from core.entities.provider_entities import ProviderConfig
from core.tools.entities.common_entities import I18nObject
from core.tools.entities.tool_bundle import ApiToolBundle
from core.tools.entities.tool_entities import (
    ApiProviderAuthType,
    ToolProviderType,
)
from core.tools.provider.tool_provider import ToolProviderController
from core.tools.tool.api_tool import ApiTool
from extensions.ext_database import db
from models.tools import ApiToolProvider


class ApiToolProviderController(ToolProviderController):
    provider_id: str
    tenant_id: str
    tools: list[ApiTool] = Field(default_factory=list)

    @staticmethod
    def from_db(db_provider: ApiToolProvider, auth_type: ApiProviderAuthType) -> 'ApiToolProviderController':
        credentials_schema = {
            'auth_type': ProviderConfig(
                name='auth_type',
                required=True,
                type=ProviderConfig.Type.SELECT,
                options=[
                    ProviderConfig.Option(value='none', label=I18nObject(en_US='None', zh_Hans='无')),
                    ProviderConfig.Option(value='api_key', label=I18nObject(en_US='api_key', zh_Hans='api_key'))
                ],
                default='none',
                help=I18nObject(
                    en_US='The auth type of the api provider',
                    zh_Hans='api provider 的认证类型'
                )
            )
        }
        if auth_type == ApiProviderAuthType.API_KEY:
            credentials_schema = {
                **credentials_schema,
                'api_key_header': ProviderConfig(
                    name='api_key_header',
                    required=False,
                    default='api_key',
                    type=ProviderConfig.Type.TEXT_INPUT,
                    help=I18nObject(
                        en_US='The header name of the api key',
                        zh_Hans='携带 api key 的 header 名称'
                    )
                ),
                'api_key_value': ProviderConfig(
                    name='api_key_value',
                    required=True,
                    type=ProviderConfig.Type.SECRET_INPUT,
                    help=I18nObject(
                        en_US='The api key',
                        zh_Hans='api key的值'
                    )
                ),
                'api_key_header_prefix': ProviderConfig(
                    name='api_key_header_prefix',
                    required=False,
                    default='basic',
                    type=ProviderConfig.Type.SELECT,
                    help=I18nObject(
                        en_US='The prefix of the api key header',
                        zh_Hans='api key header 的前缀'
                    ),
                    options=[
                        ProviderConfig.Option(value='basic', label=I18nObject(en_US='Basic', zh_Hans='Basic')),
                        ProviderConfig.Option(value='bearer', label=I18nObject(en_US='Bearer', zh_Hans='Bearer')),
                        ProviderConfig.Option(value='custom', label=I18nObject(en_US='Custom', zh_Hans='Custom'))
                    ]
                )
            }
        elif auth_type == ApiProviderAuthType.NONE:
            pass
        else:
            raise ValueError(f'invalid auth type {auth_type}')

        user_name = db_provider.user.name if db_provider.user_id else ''

        return ApiToolProviderController(**{
            'identity': {
                'author': user_name,
                'name': db_provider.name,
                'label': {
                    'en_US': db_provider.name,
                    'zh_Hans': db_provider.name
                },
                'description': {
                    'en_US': db_provider.description,
                    'zh_Hans': db_provider.description
                },
                'icon': db_provider.icon,
            },
            'credentials_schema': credentials_schema,
            'provider_id': db_provider.id or '',
            'tenant_id': db_provider.tenant_id or '',
        })

    @property
    def provider_type(self) -> ToolProviderType:
        return ToolProviderType.API

    def _parse_tool_bundle(self, tool_bundle: ApiToolBundle) -> ApiTool:
        """
            parse tool bundle to tool

            :param tool_bundle: the tool bundle
            :return: the tool
        """
        return ApiTool(**{
            'api_bundle': tool_bundle,
            'identity' : {
                'author': tool_bundle.author,
                'name': tool_bundle.operation_id,
                'label': {
                    'en_US': tool_bundle.operation_id,
                    'zh_Hans': tool_bundle.operation_id
                },
                'icon': self.identity.icon,
                'provider': self.provider_id,
            },
            'description': {
                'human': {
                    'en_US': tool_bundle.summary or '',
                    'zh_Hans': tool_bundle.summary or ''
                },
                'llm': tool_bundle.summary or ''
            },
            'parameters' : tool_bundle.parameters if tool_bundle.parameters else [],
        })

    def load_bundled_tools(self, tools: list[ApiToolBundle]) -> list[ApiTool]:
        """
            load bundled tools

            :param tools: the bundled tools
            :return: the tools
        """
        self.tools = [self._parse_tool_bundle(tool) for tool in tools]

        return self.tools

    def get_tools(self, tenant_id: str) -> list[ApiTool]:
        """
            fetch tools from database

            :param user_id: the user id
            :param tenant_id: the tenant id
            :return: the tools
        """
        if self.tools is not None:
            return self.tools
        
        tools: list[ApiTool] = []

        # get tenant api providers
        db_providers: list[ApiToolProvider] = db.session.query(ApiToolProvider).filter(
            ApiToolProvider.tenant_id == tenant_id,
            ApiToolProvider.name == self.identity.name
        ).all()

        if db_providers and len(db_providers) != 0:
            for db_provider in db_providers:
                for tool in db_provider.tools:
                    assistant_tool = self._parse_tool_bundle(tool)
                    assistant_tool.is_team_authorization = True
                    tools.append(assistant_tool)
        
        self.tools = tools
        return tools
    
    def get_tool(self, tool_name: str) -> ApiTool:
        """
            get tool by name

            :param tool_name: the name of the tool
            :return: the tool
        """
        if self.tools is None:
            self.get_tools(self.tenant_id)

        for tool in self.tools:
            if tool.identity.name == tool_name:
                return tool

        raise ValueError(f'tool {tool_name} not found')