From d94bdaf6ef2b126dde05d60e26eab23bd913e28a Mon Sep 17 00:00:00 2001
From: Sankalp Sanand <sankalp@agnostiq.ai>
Date: Mon, 11 Dec 2023 16:05:52 -0500
Subject: [PATCH] Minor fixes to pydantic models (#1884)

* Minor fixes encountered recently

* changelog updated

* removed import error suppression from covalent init for qelectrons
---
 CHANGELOG.md                               | 2 ++
 covalent/_shared_files/schemas/electron.py | 4 ++--
 covalent/_shared_files/schemas/lattice.py  | 4 ++--
 covalent/_workflow/transport.py            | 2 +-
 requirements-client.txt                    | 1 +
 5 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ec165d88b..0dd0d6b66 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -18,6 +18,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   important, whether the size is reported to be zero or positive does
   have consequences.
 - Pack deps, call_before, and call_after assets into one file.
+- `qelectron_db`, `qelectron_data_exists`, `python_version`, and `covalent_version`
+   are now optional in the pydantic model definitions.
 
 ### Fixed
 
diff --git a/covalent/_shared_files/schemas/electron.py b/covalent/_shared_files/schemas/electron.py
index 4a6541383..b245cc93d 100644
--- a/covalent/_shared_files/schemas/electron.py
+++ b/covalent/_shared_files/schemas/electron.py
@@ -86,7 +86,7 @@ class ElectronAssets(BaseModel):
     error: Optional[AssetSchema] = None
     stdout: Optional[AssetSchema] = None
     stderr: Optional[AssetSchema] = None
-    qelectron_db: AssetSchema
+    qelectron_db: Optional[AssetSchema] = None
 
     # user dependent assets
     hooks: AssetSchema
@@ -97,7 +97,7 @@ class ElectronMetadata(BaseModel):
     name: str
     executor: str
     executor_data: dict
-    qelectron_data_exists: bool
+    qelectron_data_exists: Optional[bool] = None
     sub_dispatch_id: Optional[str] = None
     status: StatusEnum
     start_time: Optional[datetime] = None
diff --git a/covalent/_shared_files/schemas/lattice.py b/covalent/_shared_files/schemas/lattice.py
index 0fdf35f16..6a3e2bbf9 100644
--- a/covalent/_shared_files/schemas/lattice.py
+++ b/covalent/_shared_files/schemas/lattice.py
@@ -98,8 +98,8 @@ class LatticeMetadata(BaseModel):
     executor_data: dict
     workflow_executor: str
     workflow_executor_data: dict
-    python_version: str
-    covalent_version: str
+    python_version: Optional[str] = None
+    covalent_version: Optional[str] = None
 
 
 class LatticeSchema(BaseModel):
diff --git a/covalent/_workflow/transport.py b/covalent/_workflow/transport.py
index 2ab1db739..77ed5a14d 100644
--- a/covalent/_workflow/transport.py
+++ b/covalent/_workflow/transport.py
@@ -212,7 +212,7 @@ def get_node_value(self, node_key: int, value_key: str) -> Any:
         """
         return self._graph.nodes[node_key][value_key]
 
-    def set_node_value(self, node_key: int, value_key: int, value: Any) -> None:
+    def set_node_value(self, node_key: str, value_key: int, value: Any) -> None:
         """
         Set a certain value of a node. This allows for saving custom data
         in the graph nodes.
diff --git a/requirements-client.txt b/requirements-client.txt
index 9f7e806aa..ede6a20e3 100644
--- a/requirements-client.txt
+++ b/requirements-client.txt
@@ -5,6 +5,7 @@ dask[distributed]>=2022.6.0
 filelock>=3.12.2
 furl>=2.1.3
 networkx>=2.8.6
+pydantic>=2.1.1
 requests>=2.24.0
 simplejson>=3.17.6
 toml>=0.10.2