Skip to content

Commit

Permalink
fix: Updated getGlobalConfig to not use pydantic json() call like get…
Browse files Browse the repository at this point in the history
…ServerInfo - problematic call no longer occurs
  • Loading branch information
Christopher-R-Perkins committed Jul 17, 2024
1 parent 09bf036 commit 3265772
Show file tree
Hide file tree
Showing 2 changed files with 31 additions and 10 deletions.
11 changes: 10 additions & 1 deletion bundled/tool/type_hints.py
Original file line number Diff line number Diff line change
Expand Up @@ -75,4 +75,13 @@ class ZenmlStoreConfig(TypedDict):

class ZenmlServerInfoResp(TypedDict):
store_info: ZenmlStoreInfo
store_config: ZenmlStoreConfig
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
30 changes: 21 additions & 9 deletions bundled/tool/zenml_wrappers.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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
}
}

Expand Down

0 comments on commit 3265772

Please sign in to comment.