from pydantic import Field from core.entities.provider_entities import ProviderConfig from core.tools.__base.tool_provider import ToolProviderController from core.tools.custom_tool.tool import ApiTool from core.tools.entities.common_entities import I18nObject from core.tools.entities.tool_bundle import ApiToolBundle from core.tools.entities.tool_entities import ( ApiProviderAuthType, ToolProviderEntity, ToolProviderIdentity, ToolProviderType, ) 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) def __init__(self, entity: ToolProviderEntity, provider_id: str, tenant_id: str) -> None: super().__init__(entity) self.provider_id = provider_id self.tenant_id = tenant_id self.tools = [] @classmethod def from_db(cls, db_provider: ApiToolProvider, auth_type: ApiProviderAuthType): 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 user = db_provider.user user_name = user.name if user else "" return ApiToolProviderController( entity=ToolProviderEntity( identity=ToolProviderIdentity( author=user_name, name=db_provider.name, label=I18nObject(en_US=db_provider.name, zh_Hans=db_provider.name), description=I18nObject(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.entity.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.entity.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) 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.entity.identity.name == tool_name: return tool raise ValueError(f"tool {tool_name} not found")