agent.py 4.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113
  1. from collections.abc import Generator
  2. from typing import Any, Optional
  3. from core.agent.entities import AgentInvokeMessage
  4. from core.plugin.entities.plugin import GenericProviderID
  5. from core.plugin.entities.plugin_daemon import (
  6. PluginAgentProviderEntity,
  7. )
  8. from core.plugin.manager.base import BasePluginManager
  9. class PluginAgentManager(BasePluginManager):
  10. def fetch_agent_providers(self, tenant_id: str) -> list[PluginAgentProviderEntity]:
  11. """
  12. Fetch agent providers for the given tenant.
  13. """
  14. def transformer(json_response: dict[str, Any]) -> dict:
  15. for provider in json_response.get("data", []):
  16. declaration = provider.get("declaration", {}) or {}
  17. declaration["identity"]["name"] = (
  18. f"{provider.get('plugin_id')}/{declaration.get('identity', {}).get('name')}"
  19. )
  20. provider_name = declaration.get("identity", {}).get("name")
  21. for strategy in declaration.get("strategies", []):
  22. strategy["identity"]["provider"] = provider_name
  23. return json_response
  24. response = self._request_with_plugin_daemon_response(
  25. "GET",
  26. f"plugin/{tenant_id}/management/agents",
  27. list[PluginAgentProviderEntity],
  28. params={"page": 1, "page_size": 256},
  29. transformer=transformer,
  30. )
  31. for provider in response:
  32. provider.declaration.identity.name = f"{provider.plugin_id}/{provider.declaration.identity.name}"
  33. # override the provider name for each tool to plugin_id/provider_name
  34. for strategy in provider.declaration.strategies:
  35. strategy.identity.provider = provider.declaration.identity.name
  36. return response
  37. def fetch_agent_provider(self, tenant_id: str, provider: str) -> PluginAgentProviderEntity:
  38. """
  39. Fetch tool provider for the given tenant and plugin.
  40. """
  41. agent_provider_id = GenericProviderID(provider)
  42. def transformer(json_response: dict[str, Any]) -> dict:
  43. for strategy in json_response.get("data", {}).get("declaration", {}).get("strategies", []):
  44. strategy["identity"]["provider"] = agent_provider_id.provider_name
  45. return json_response
  46. response = self._request_with_plugin_daemon_response(
  47. "GET",
  48. f"plugin/{tenant_id}/management/agent",
  49. PluginAgentProviderEntity,
  50. params={"provider": agent_provider_id.provider_name, "plugin_id": agent_provider_id.plugin_id},
  51. transformer=transformer,
  52. )
  53. response.declaration.identity.name = f"{response.plugin_id}/{response.declaration.identity.name}"
  54. # override the provider name for each tool to plugin_id/provider_name
  55. for strategy in response.declaration.strategies:
  56. strategy.identity.provider = response.declaration.identity.name
  57. return response
  58. def invoke(
  59. self,
  60. tenant_id: str,
  61. user_id: str,
  62. agent_provider: str,
  63. agent_strategy: str,
  64. agent_params: dict[str, Any],
  65. conversation_id: Optional[str] = None,
  66. app_id: Optional[str] = None,
  67. message_id: Optional[str] = None,
  68. ) -> Generator[AgentInvokeMessage, None, None]:
  69. """
  70. Invoke the agent with the given tenant, user, plugin, provider, name and parameters.
  71. """
  72. agent_provider_id = GenericProviderID(agent_provider)
  73. response = self._request_with_plugin_daemon_response_stream(
  74. "POST",
  75. f"plugin/{tenant_id}/dispatch/agent/invoke",
  76. AgentInvokeMessage,
  77. data={
  78. "user_id": user_id,
  79. "conversation_id": conversation_id,
  80. "app_id": app_id,
  81. "message_id": message_id,
  82. "data": {
  83. "provider": agent_provider_id.provider_name,
  84. "strategy": agent_strategy,
  85. "agent_params": agent_params,
  86. },
  87. },
  88. headers={
  89. "X-Plugin-ID": agent_provider_id.plugin_id,
  90. "Content-Type": "application/json",
  91. },
  92. )
  93. return response