From 9693b5ad0c62e14de69e430c5491e541a970f796 Mon Sep 17 00:00:00 2001 From: Yeuoly Date: Fri, 20 Sep 2024 14:43:01 +0800 Subject: [PATCH] feat: debugging key --- api/core/plugin/entities/plugin_daemon.py | 15 +++++++ api/core/plugin/manager/asset.py | 9 +++- api/core/plugin/manager/base.py | 51 +++++++++++++++++++++-- api/core/plugin/manager/debugging.py | 17 ++++++++ 4 files changed, 87 insertions(+), 5 deletions(-) create mode 100644 api/core/plugin/entities/plugin_daemon.py create mode 100644 api/core/plugin/manager/debugging.py diff --git a/api/core/plugin/entities/plugin_daemon.py b/api/core/plugin/entities/plugin_daemon.py new file mode 100644 index 0000000000..ca8ea0f780 --- /dev/null +++ b/api/core/plugin/entities/plugin_daemon.py @@ -0,0 +1,15 @@ +from typing import Generic, Optional, TypeVar + +from pydantic import BaseModel + +T = TypeVar("T", bound=(BaseModel | dict)) + + +class PluginDaemonBasicResponse(BaseModel, Generic[T]): + """ + Basic response from plugin daemon. + """ + + code: int + message: str + data: Optional[T] diff --git a/api/core/plugin/manager/asset.py b/api/core/plugin/manager/asset.py index f55ff748ca..df76f56a6d 100644 --- a/api/core/plugin/manager/asset.py +++ b/api/core/plugin/manager/asset.py @@ -2,4 +2,11 @@ from core.plugin.manager.base import BasePluginManager class PluginAssetManager(BasePluginManager): - pass \ No newline at end of file + def fetch_asset(self, id: str) -> bytes: + """ + Fetch an asset by id. + """ + response = self._request(method="GET", path=f"/assets/plugin/{id}") + if response.status_code != 200: + raise ValueError(f"can not found asset {id}") + return response.content diff --git a/api/core/plugin/manager/base.py b/api/core/plugin/manager/base.py index 5f1e5cd1d2..b26ab851c7 100644 --- a/api/core/plugin/manager/base.py +++ b/api/core/plugin/manager/base.py @@ -7,6 +7,7 @@ from pydantic import BaseModel from yarl import URL from configs import dify_config +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.INNER_API_KEY_FOR_PLUGIN @@ -15,16 +16,21 @@ T = TypeVar("T", bound=(BaseModel | dict)) class BasePluginManager: - def _request(self, method: str, path: str, headers: dict, data: bytes, stream: bool = False) -> requests.Response: + def _request( + self, method: str, path: str, headers: dict | None = None, data: bytes | None = None, stream: bool = False + ) -> requests.Response: """ Make a request to the plugin daemon inner API. """ 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) return response - def _stream_request(self, method: str, path: str, headers: dict, data: bytes) -> Generator[bytes, None, None]: + def _stream_request( + self, method: str, path: str, headers: dict | None = None, data: bytes | None = None + ) -> Generator[bytes, None, None]: """ Make a stream request to the plugin daemon inner API """ @@ -32,7 +38,12 @@ class BasePluginManager: yield from response.iter_lines() def _stream_request_with_model( - self, method: str, path: str, headers: dict, data: bytes, type: type[T] + self, + method: str, + path: str, + type: type[T], + headers: dict | None = None, + data: bytes | None = None, ) -> Generator[T, None, None]: """ Make a stream request to the plugin daemon inner API and yield the response as a model. @@ -40,9 +51,41 @@ class BasePluginManager: for line in self._stream_request(method, path, headers, data): yield type(**json.loads(line)) - def _request_with_model(self, method: str, path: str, headers: dict, data: bytes, type: type[T]) -> T: + def _request_with_model( + self, method: str, path: str, type: type[T], headers: dict | None = None, data: bytes | 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) return type(**response.json()) + + def _request_with_plugin_daemon_response( + self, method: str, path: str, type: type[T], headers: dict | None = None, data: bytes | 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) + rep = PluginDaemonBasicResponse[type](**response.json()) + if rep.code != 0: + raise Exception(f"got error from plugin daemon: {rep.message}, code: {rep.code}") + if rep.data is None: + raise Exception("got empty data from plugin daemon") + + return rep.data + + def _request_with_plugin_daemon_response_stream( + self, method: str, path: str, type: type[T], headers: dict | None = None, data: bytes | 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): + line_data = json.loads(line) + rep = PluginDaemonBasicResponse[type](**line_data) + if rep.code != 0: + raise Exception(f"got error from plugin daemon: {rep.message}, code: {rep.code}") + if rep.data is None: + raise Exception("got empty data from plugin daemon") + yield rep.data \ No newline at end of file diff --git a/api/core/plugin/manager/debugging.py b/api/core/plugin/manager/debugging.py new file mode 100644 index 0000000000..6b26e3ad67 --- /dev/null +++ b/api/core/plugin/manager/debugging.py @@ -0,0 +1,17 @@ +from pydantic import BaseModel + +from core.plugin.manager.base import BasePluginManager + + +class PluginDebuggingManager(BasePluginManager): + def get_debugging_key(self, tenant_id: str) -> str: + """ + Get the debugging key for the given tenant. + """ + + class Response(BaseModel): + key: str + + response = self._request_with_plugin_daemon_response("POST", f"/plugin/{tenant_id}/debugging/key", Response) + + return response.key