This commit is contained in:
Yeuoly
2024-09-23 13:09:46 +08:00
parent 91cb80f795
commit 435e71eb60
10 changed files with 164 additions and 67 deletions

View File

@ -12,7 +12,7 @@ from core.plugin.entities.plugin_daemon import PluginDaemonBasicResponse
plugin_daemon_inner_api_baseurl = dify_config.PLUGIN_API_URL
plugin_daemon_inner_api_key = dify_config.PLUGIN_API_KEY
T = TypeVar("T", bound=(BaseModel | dict | bool))
T = TypeVar("T", bound=(BaseModel | dict | list | bool))
class BasePluginManager:
@ -22,6 +22,7 @@ class BasePluginManager:
path: str,
headers: dict | None = None,
data: bytes | dict | None = None,
params: dict | None = None,
stream: bool = False,
) -> requests.Response:
"""
@ -30,16 +31,23 @@ class BasePluginManager:
url = URL(str(plugin_daemon_inner_api_baseurl)) / path
headers = headers or {}
headers["X-Api-Key"] = plugin_daemon_inner_api_key
response = requests.request(method=method, url=str(url), headers=headers, data=data, stream=stream)
response = requests.request(
method=method, url=str(url), headers=headers, data=data, params=params, stream=stream
)
return response
def _stream_request(
self, method: str, path: str, headers: dict | None = None, data: bytes | dict | None = None
self,
method: str,
path: str,
params: dict | None = None,
headers: dict | None = None,
data: bytes | dict | None = None,
) -> Generator[bytes, None, None]:
"""
Make a stream request to the plugin daemon inner API
"""
response = self._request(method, path, headers, data, stream=True)
response = self._request(method, path, headers, data, params, stream=True)
yield from response.iter_lines()
def _stream_request_with_model(
@ -49,29 +57,42 @@ class BasePluginManager:
type: type[T],
headers: dict | None = None,
data: bytes | dict | None = None,
params: dict | None = None,
) -> Generator[T, None, None]:
"""
Make a stream request to the plugin daemon inner API and yield the response as a model.
"""
for line in self._stream_request(method, path, headers, data):
for line in self._stream_request(method, path, params, headers, data):
yield type(**json.loads(line))
def _request_with_model(
self, method: str, path: str, type: type[T], headers: dict | None = None, data: bytes | None = None
self,
method: str,
path: str,
type: type[T],
headers: dict | None = None,
data: bytes | None = None,
params: dict | None = None,
) -> T:
"""
Make a request to the plugin daemon inner API and return the response as a model.
"""
response = self._request(method, path, headers, data)
response = self._request(method, path, headers, data, params)
return type(**response.json())
def _request_with_plugin_daemon_response(
self, method: str, path: str, type: type[T], headers: dict | None = None, data: bytes | dict | None = None
self,
method: str,
path: str,
type: type[T],
headers: dict | None = None,
data: bytes | dict | None = None,
params: dict | None = None,
) -> T:
"""
Make a request to the plugin daemon inner API and return the response as a model.
"""
response = self._request(method, path, headers, data)
response = self._request(method, path, headers, data, params)
rep = PluginDaemonBasicResponse[type](**response.json())
if rep.code != 0:
raise ValueError(f"got error from plugin daemon: {rep.message}, code: {rep.code}")
@ -81,12 +102,18 @@ class BasePluginManager:
return rep.data
def _request_with_plugin_daemon_response_stream(
self, method: str, path: str, type: type[T], headers: dict | None = None, data: bytes | dict | None = None
self,
method: str,
path: str,
type: type[T],
headers: dict | None = None,
data: bytes | dict | None = None,
params: dict | None = None,
) -> Generator[T, None, None]:
"""
Make a stream request to the plugin daemon inner API and yield the response as a model.
"""
for line in self._stream_request(method, path, headers, data):
for line in self._stream_request(method, path, params, headers, data):
line_data = json.loads(line)
rep = PluginDaemonBasicResponse[type](**line_data)
if rep.code != 0:

View File

@ -1,5 +1,13 @@
from core.model_runtime.entities.provider_entities import ProviderEntity
from core.plugin.manager.base import BasePluginManager
class PluginModelManager(BasePluginManager):
pass
def fetch_model_providers(self, tenant_id: str) -> list[ProviderEntity]:
"""
Fetch model providers for the given tenant.
"""
response = self._request_with_plugin_daemon_response(
"GET", f"plugin/{tenant_id}/models", list[ProviderEntity], params={"page": 1, "page_size": 256}
)
return response

View File

@ -1,5 +1,4 @@
from collections.abc import Generator
from urllib.parse import quote
from core.plugin.entities.plugin_daemon import InstallPluginMessage
from core.plugin.manager.base import BasePluginManager
@ -9,9 +8,8 @@ class PluginInstallationManager(BasePluginManager):
def fetch_plugin_by_identifier(self, tenant_id: str, identifier: str) -> bool:
# urlencode the identifier
identifier = quote(identifier)
return self._request_with_plugin_daemon_response(
"GET", f"/plugin/{tenant_id}/fetch/identifier?plugin_unique_identifier={identifier}", bool
"GET", f"plugin/{tenant_id}/fetch/identifier", bool, params={"plugin_unique_identifier": identifier}
)
def install_from_pkg(self, tenant_id: str, pkg: bytes) -> Generator[InstallPluginMessage, None, None]:
@ -22,21 +20,20 @@ class PluginInstallationManager(BasePluginManager):
body = {"dify_pkg": ("dify_pkg", pkg, "application/octet-stream")}
return self._request_with_plugin_daemon_response_stream(
"POST", f"/plugin/{tenant_id}/install/pkg", InstallPluginMessage, data=body
"POST", f"plugin/{tenant_id}/install/pkg", InstallPluginMessage, data=body
)
def install_from_identifier(self, tenant_id: str, identifier: str) -> bool:
"""
Install a plugin from an identifier.
"""
identifier = quote(identifier)
# exception will be raised if the request failed
return self._request_with_plugin_daemon_response(
"POST",
f"/plugin/{tenant_id}/install/identifier",
f"plugin/{tenant_id}/install/identifier",
bool,
headers={
"Content-Type": "application/json",
params={
"plugin_unique_identifier": identifier,
},
data={
"plugin_unique_identifier": identifier,
@ -48,5 +45,5 @@ class PluginInstallationManager(BasePluginManager):
Uninstall a plugin.
"""
return self._request_with_plugin_daemon_response(
"DELETE", f"/plugin/{tenant_id}/uninstall?plugin_unique_identifier={identifier}", bool
"DELETE", f"plugin/{tenant_id}/uninstall", bool, params={"plugin_unique_identifier": identifier}
)

View File

@ -1,9 +1,13 @@
from core.plugin.manager.base import BasePluginManager
from core.tools.entities.tool_entities import ToolProviderEntity
class PluginToolManager(BasePluginManager):
def fetch_tool_providers(self, asset_id: str) -> list[str]:
def fetch_tool_providers(self, tenant_id: str) -> list[ToolProviderEntity]:
"""
Fetch tool providers for the given asset.
"""
response = self._request('GET', f'/plugin/asset/{asset_id}')
response = self._request_with_plugin_daemon_response(
"GET", f"plugin/{tenant_id}/tools", list[ToolProviderEntity], params={"page": 1, "page_size": 256}
)
return response