api_entities.py 2.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566
  1. from typing import Literal, Optional
  2. from pydantic import BaseModel, Field
  3. from core.entities.provider_entities import ProviderConfig
  4. from core.model_runtime.utils.encoders import jsonable_encoder
  5. from core.tools.entities.common_entities import I18nObject
  6. from core.tools.entities.tool_entities import ToolProviderType
  7. from core.tools.tool.tool import ToolParameter
  8. class UserTool(BaseModel):
  9. author: str
  10. name: str # identifier
  11. label: I18nObject # label
  12. description: I18nObject
  13. parameters: Optional[list[ToolParameter]] = None
  14. labels: list[str] = Field(default_factory=list)
  15. UserToolProviderTypeLiteral = Optional[Literal[
  16. 'builtin', 'api', 'workflow'
  17. ]]
  18. class UserToolProvider(BaseModel):
  19. id: str
  20. author: str
  21. name: str # identifier
  22. description: I18nObject
  23. icon: str
  24. label: I18nObject # label
  25. type: ToolProviderType
  26. masked_credentials: Optional[dict] = None
  27. original_credentials: Optional[dict] = None
  28. is_team_authorization: bool = False
  29. allow_delete: bool = True
  30. tools: list[UserTool] = Field(default_factory=list)
  31. labels: list[str] = Field(default_factory=list)
  32. def to_dict(self) -> dict:
  33. # -------------
  34. # overwrite tool parameter types for temp fix
  35. tools = jsonable_encoder(self.tools)
  36. for tool in tools:
  37. if tool.get('parameters'):
  38. for parameter in tool.get('parameters'):
  39. if parameter.get('type') == ToolParameter.ToolParameterType.FILE.value:
  40. parameter['type'] = 'files'
  41. # -------------
  42. return {
  43. 'id': self.id,
  44. 'author': self.author,
  45. 'name': self.name,
  46. 'description': self.description.to_dict(),
  47. 'icon': self.icon,
  48. 'label': self.label.to_dict(),
  49. 'type': self.type.value,
  50. 'team_credentials': self.masked_credentials,
  51. 'is_team_authorization': self.is_team_authorization,
  52. 'allow_delete': self.allow_delete,
  53. 'tools': tools,
  54. 'labels': self.labels,
  55. }
  56. class UserToolProviderCredentials(BaseModel):
  57. credentials: dict[str, ProviderConfig]