tool.py 6.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188
  1. from collections.abc import Generator
  2. from typing import Any
  3. from pydantic import BaseModel
  4. from core.plugin.entities.plugin_daemon import PluginBasicBooleanResponse, PluginToolProviderEntity
  5. from core.plugin.manager.base import BasePluginManager
  6. from core.tools.entities.tool_entities import ToolInvokeMessage, ToolParameter
  7. class PluginToolManager(BasePluginManager):
  8. def _split_provider(self, provider: str) -> tuple[str, str]:
  9. """
  10. split the provider to plugin_id and provider_name
  11. provider follows format: plugin_id/provider_name
  12. """
  13. if "/" in provider:
  14. parts = provider.split("/", -1)
  15. if len(parts) >= 2:
  16. return "/".join(parts[:-1]), parts[-1]
  17. raise ValueError(f"invalid provider format: {provider}")
  18. raise ValueError(f"invalid provider format: {provider}")
  19. def fetch_tool_providers(self, tenant_id: str) -> list[PluginToolProviderEntity]:
  20. """
  21. Fetch tool providers for the given tenant.
  22. """
  23. def transformer(json_response: dict[str, Any]) -> dict:
  24. for provider in json_response.get("data", []):
  25. declaration = provider.get("declaration", {}) or {}
  26. provider_name = declaration.get("identity", {}).get("name")
  27. for tool in declaration.get("tools", []):
  28. tool["identity"]["provider"] = provider_name
  29. return json_response
  30. response = self._request_with_plugin_daemon_response(
  31. "GET",
  32. f"plugin/{tenant_id}/management/tools",
  33. list[PluginToolProviderEntity],
  34. params={"page": 1, "page_size": 256},
  35. transformer=transformer,
  36. )
  37. for provider in response:
  38. provider.declaration.identity.name = f"{provider.plugin_id}/{provider.declaration.identity.name}"
  39. # override the provider name for each tool to plugin_id/provider_name
  40. for tool in provider.declaration.tools:
  41. tool.identity.provider = provider.declaration.identity.name
  42. return response
  43. def fetch_tool_provider(self, tenant_id: str, provider: str) -> PluginToolProviderEntity:
  44. """
  45. Fetch tool provider for the given tenant and plugin.
  46. """
  47. plugin_id, provider_name = self._split_provider(provider)
  48. def transformer(json_response: dict[str, Any]) -> dict:
  49. for tool in json_response.get("data", {}).get("declaration", {}).get("tools", []):
  50. tool["identity"]["provider"] = provider_name
  51. return json_response
  52. response = self._request_with_plugin_daemon_response(
  53. "GET",
  54. f"plugin/{tenant_id}/management/tool",
  55. PluginToolProviderEntity,
  56. params={"provider": provider_name, "plugin_id": plugin_id},
  57. transformer=transformer,
  58. )
  59. response.declaration.identity.name = f"{response.plugin_id}/{response.declaration.identity.name}"
  60. # override the provider name for each tool to plugin_id/provider_name
  61. for tool in response.declaration.tools:
  62. tool.identity.provider = response.declaration.identity.name
  63. return response
  64. def invoke(
  65. self,
  66. tenant_id: str,
  67. user_id: str,
  68. tool_provider: str,
  69. tool_name: str,
  70. credentials: dict[str, Any],
  71. tool_parameters: dict[str, Any],
  72. ) -> Generator[ToolInvokeMessage, None, None]:
  73. """
  74. Invoke the tool with the given tenant, user, plugin, provider, name, credentials and parameters.
  75. """
  76. plugin_id, provider_name = self._split_provider(tool_provider)
  77. response = self._request_with_plugin_daemon_response_stream(
  78. "POST",
  79. f"plugin/{tenant_id}/dispatch/tool/invoke",
  80. ToolInvokeMessage,
  81. data={
  82. "user_id": user_id,
  83. "data": {
  84. "provider": provider_name,
  85. "tool": tool_name,
  86. "credentials": credentials,
  87. "tool_parameters": tool_parameters,
  88. },
  89. },
  90. headers={
  91. "X-Plugin-ID": plugin_id,
  92. "Content-Type": "application/json",
  93. },
  94. )
  95. return response
  96. def validate_provider_credentials(
  97. self, tenant_id: str, user_id: str, provider: str, credentials: dict[str, Any]
  98. ) -> bool:
  99. """
  100. validate the credentials of the provider
  101. """
  102. plugin_id, provider_name = self._split_provider(provider)
  103. response = self._request_with_plugin_daemon_response_stream(
  104. "POST",
  105. f"plugin/{tenant_id}/dispatch/tool/validate_credentials",
  106. PluginBasicBooleanResponse,
  107. data={
  108. "user_id": user_id,
  109. "data": {
  110. "provider": provider_name,
  111. "credentials": credentials,
  112. },
  113. },
  114. headers={
  115. "X-Plugin-ID": plugin_id,
  116. "Content-Type": "application/json",
  117. },
  118. )
  119. for resp in response:
  120. return resp.result
  121. return False
  122. def get_runtime_parameters(
  123. self,
  124. tenant_id: str,
  125. user_id: str,
  126. provider: str,
  127. credentials: dict[str, Any],
  128. tool: str,
  129. ) -> list[ToolParameter]:
  130. """
  131. get the runtime parameters of the tool
  132. """
  133. plugin_id, provider_name = self._split_provider(provider)
  134. class RuntimeParametersResponse(BaseModel):
  135. parameters: list[ToolParameter]
  136. response = self._request_with_plugin_daemon_response_stream(
  137. "POST",
  138. f"plugin/{tenant_id}/dispatch/tool/get_runtime_parameters",
  139. RuntimeParametersResponse,
  140. data={
  141. "user_id": "user_id",
  142. "data": {
  143. "provider": provider_name,
  144. "tool": tool,
  145. "credentials": credentials,
  146. },
  147. },
  148. headers={
  149. "X-Plugin-ID": plugin_id,
  150. "Content-Type": "application/json",
  151. },
  152. )
  153. for resp in response:
  154. return resp.parameters
  155. return []