dify/api/core/plugin/manager/tool.py

90 lines
2.9 KiB
Python
Raw Normal View History

2024-09-23 18:06:16 +08:00
from collections.abc import Generator
from typing import Any
2024-09-23 21:13:02 +08:00
from core.plugin.entities.plugin_daemon import PluginBasicBooleanResponse, PluginToolProviderEntity
2024-09-20 13:55:09 +08:00
from core.plugin.manager.base import BasePluginManager
2024-09-23 18:06:16 +08:00
from core.tools.entities.tool_entities import ToolInvokeMessage
2024-09-20 13:55:09 +08:00
class PluginToolManager(BasePluginManager):
2024-09-23 18:06:16 +08:00
def fetch_tool_providers(self, tenant_id: str) -> list[PluginToolProviderEntity]:
2024-09-20 15:08:39 +08:00
"""
Fetch tool providers for the given asset.
"""
2024-09-23 21:13:02 +08:00
def transformer(json_response: dict[str, Any]) -> dict:
for provider in json_response.get("data", []):
declaration = provider.get("declaration", {}) or {}
provider_name = declaration.get("identity", {}).get("name")
for tool in declaration.get("tools", []):
tool["identity"]["provider"] = provider_name
return json_response
2024-09-23 13:09:46 +08:00
response = self._request_with_plugin_daemon_response(
2024-09-23 21:13:02 +08:00
"GET",
f"plugin/{tenant_id}/management/tools",
list[PluginToolProviderEntity],
params={"page": 1, "page_size": 256},
transformer=transformer,
2024-09-23 18:06:16 +08:00
)
return response
def invoke(
self,
tenant_id: str,
user_id: str,
plugin_id: str,
2024-09-23 18:06:16 +08:00
tool_provider: str,
tool_name: str,
credentials: dict[str, Any],
tool_parameters: dict[str, Any],
) -> Generator[ToolInvokeMessage, None, None]:
response = self._request_with_plugin_daemon_response_stream(
"POST",
2024-09-23 21:13:02 +08:00
f"plugin/{tenant_id}/dispatch/tool/invoke",
2024-09-23 18:06:16 +08:00
ToolInvokeMessage,
data={
"user_id": user_id,
"data": {
"provider": tool_provider,
"tool": tool_name,
"credentials": credentials,
"tool_parameters": tool_parameters,
},
},
2024-09-23 21:13:02 +08:00
headers={
"X-Plugin-ID": plugin_id,
2024-09-23 21:13:02 +08:00
"Content-Type": "application/json",
},
2024-09-23 18:06:16 +08:00
)
return response
def validate_provider_credentials(
self, tenant_id: str, user_id: str, plugin_id: str, provider: str, credentials: dict[str, Any]
2024-09-23 18:06:16 +08:00
) -> bool:
"""
validate the credentials of the provider
"""
2024-09-23 21:13:02 +08:00
response = self._request_with_plugin_daemon_response_stream(
2024-09-23 18:06:16 +08:00
"POST",
2024-09-23 21:13:02 +08:00
f"plugin/{tenant_id}/dispatch/tool/validate_credentials",
PluginBasicBooleanResponse,
2024-09-23 18:06:16 +08:00
data={
"user_id": user_id,
"data": {
"provider": provider,
"credentials": credentials,
},
},
2024-09-23 21:13:02 +08:00
headers={
"X-Plugin-ID": plugin_id,
2024-09-23 21:13:02 +08:00
"Content-Type": "application/json",
},
2024-09-23 13:09:46 +08:00
)
2024-09-23 21:13:02 +08:00
for resp in response:
return resp.result
2024-09-23 21:13:02 +08:00
return False