From 09bf0366329c33cf6ce35801437079e16b8c300b Mon Sep 17 00:00:00 2001 From: Christopher Perkins Date: Wed, 17 Jul 2024 17:36:36 -0400 Subject: [PATCH 1/2] fix: removed use of pydantic json() call from get_server_info due to data no longer being compatible --- bundled/tool/type_hints.py | 22 +++++++++++++++++++++- bundled/tool/zenml_wrappers.py | 29 ++++++++++++++++++++++++----- 2 files changed, 45 insertions(+), 6 deletions(-) diff --git a/bundled/tool/type_hints.py b/bundled/tool/type_hints.py index bda4b20e..9b7b4404 100644 --- a/bundled/tool/type_hints.py +++ b/bundled/tool/type_hints.py @@ -55,4 +55,24 @@ class RunArtifactResponse(TypedDict): author: Dict[str, str] update: str data: Dict[str, str] - metadata: Dict[str, Any] \ No newline at end of file + metadata: Dict[str, Any] + +class ZenmlStoreInfo(TypedDict): + id: str + version: str + debug: bool + deployment_type: str + database_type: str + secrets_store_type: str + auth_scheme: str + server_url: str + dashboard_url: str + +class ZenmlStoreConfig(TypedDict): + type: str + url: str + api_token: Union[str, None] + +class ZenmlServerInfoResp(TypedDict): + store_info: ZenmlStoreInfo + store_config: ZenmlStoreConfig \ No newline at end of file diff --git a/bundled/tool/zenml_wrappers.py b/bundled/tool/zenml_wrappers.py index a119fe86..29b8d3cf 100644 --- a/bundled/tool/zenml_wrappers.py +++ b/bundled/tool/zenml_wrappers.py @@ -15,7 +15,7 @@ import json import pathlib from typing import Any, Tuple, Union -from type_hints import GraphResponse, ErrorResponse, RunStepResponse, RunArtifactResponse +from type_hints import GraphResponse, ErrorResponse, RunStepResponse, RunArtifactResponse, ZenmlServerInfoResp from zenml_grapher import Grapher @@ -167,19 +167,38 @@ def get_active_deployment(self): """Returns the function to get the active ZenML server deployment.""" return self.lazy_import("zenml.zen_server.utils", "get_active_deployment") - def get_server_info(self) -> dict: + def get_server_info(self) -> ZenmlServerInfoResp: """Fetches the ZenML server info. Returns: dict: Dictionary containing server info. """ - store_info = json.loads(self.gc.zen_store.get_store_info().json(indent=2)) + store_info = self.gc.zen_store.get_store_info() + # Handle both 'store' and 'store_configuration' depending on version store_attr_name = ( "store_configuration" if hasattr(self.gc, "store_configuration") else "store" ) - store_config = json.loads(getattr(self.gc, store_attr_name).json(indent=2)) - return {"storeInfo": store_info, "storeConfig": store_config} + store_config = getattr(self.gc, store_attr_name) + + return { + "storeInfo": { + "id": str(store_info.id), + "version": store_info.version, + "debug": store_info.debug, + "deployment_type": store_info.deployment_type, + "database_type": store_info.database_type, + "secrets_store_type": store_info.secrets_store_type, + "auth_scheme": store_info.auth_scheme, + "server_url": store_info.server_url, + "dashboard_url": store_info.dashboard_url, + }, + "storeConfig": { + "type": store_config.type, + "url": store_config.url, + "api_token": store_config.api_token if "api_token" in store_config else None + } + } def connect(self, args, **kwargs) -> dict: """Connects to a ZenML server. From 326577215beb35885040fabd334df84bad2db41a Mon Sep 17 00:00:00 2001 From: Christopher Perkins Date: Wed, 17 Jul 2024 17:57:28 -0400 Subject: [PATCH 2/2] fix: Updated getGlobalConfig to not use pydantic json() call like getServerInfo - problematic call no longer occurs --- bundled/tool/type_hints.py | 11 ++++++++++- bundled/tool/zenml_wrappers.py | 30 +++++++++++++++++++++--------- 2 files changed, 31 insertions(+), 10 deletions(-) diff --git a/bundled/tool/type_hints.py b/bundled/tool/type_hints.py index 9b7b4404..f5441944 100644 --- a/bundled/tool/type_hints.py +++ b/bundled/tool/type_hints.py @@ -75,4 +75,13 @@ class ZenmlStoreConfig(TypedDict): class ZenmlServerInfoResp(TypedDict): store_info: ZenmlStoreInfo - store_config: ZenmlStoreConfig \ No newline at end of file + store_config: ZenmlStoreConfig + +class ZenmlGlobalConfigResp(TypedDict): + user_id: str + user_email: str + analytics_opt_in: bool + version: str + active_stack_id: str + active_workspace_name: str + store: ZenmlStoreConfig \ No newline at end of file diff --git a/bundled/tool/zenml_wrappers.py b/bundled/tool/zenml_wrappers.py index 29b8d3cf..b4953ff2 100644 --- a/bundled/tool/zenml_wrappers.py +++ b/bundled/tool/zenml_wrappers.py @@ -12,10 +12,9 @@ # permissions and limitations under the License. """This module provides wrappers for ZenML configuration and operations.""" -import json import pathlib from typing import Any, Tuple, Union -from type_hints import GraphResponse, ErrorResponse, RunStepResponse, RunArtifactResponse, ZenmlServerInfoResp +from type_hints import GraphResponse, ErrorResponse, RunStepResponse, RunArtifactResponse, ZenmlServerInfoResp, ZenmlGlobalConfigResp from zenml_grapher import Grapher @@ -99,19 +98,32 @@ def set_store_configuration(self, remote_url: str, access_token: str): ) self.gc.set_store(new_store_config) - def get_global_configuration(self) -> dict: + def get_global_configuration(self) -> ZenmlGlobalConfigResp: """Get the global configuration. Returns: dict: Global configuration. """ - gc_dict = json.loads(self.gc.json(indent=2)) - user_id = gc_dict.get("user_id", "") - if user_id and user_id.startswith("UUID('") and user_id.endswith("')"): - gc_dict["user_id"] = user_id[6:-2] + store_attr_name = ( + "store_configuration" if hasattr(self.gc, "store_configuration") else "store" + ) + + store_data = getattr(self.gc, store_attr_name) - return gc_dict + return { + "user_id": str(self.gc.user_id), + "user_email": self.gc.user_email, + "analytics_opt_in": self.gc.analytics_opt_in, + "version": self.gc.version, + "active_stack_id": str(self.gc.active_stack_id), + "active_workspace_name": self.gc.active_workspace_name, + "store": { + "type": store_data.type, + "url": store_data.url, + "api_token": store_data.api_token if hasattr(store_data, "api_token") else None + } + } class ZenServerWrapper: @@ -196,7 +208,7 @@ def get_server_info(self) -> ZenmlServerInfoResp: "storeConfig": { "type": store_config.type, "url": store_config.url, - "api_token": store_config.api_token if "api_token" in store_config else None + "api_token": store_config.api_token if hasattr(store_config, "api_token") else None } }