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 or [],
            }
        )

    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")