Updating task_Result device_params to device params schema (#29)

Co-authored-by: 🐅ⓚ <vundela@amazon.com>
amazon-braket · Aug 4, 2020 · 26a438d · 26a438d
1 parent 427b488
commit 26a438d
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 6 deletions.
diff --git a/src/braket/ir/_version.py b/src/braket/ir/_version.py
@@ -15,4 +15,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.2.6"
+__version__ = "0.2.7"
diff --git a/src/braket/task_result/task_metadata_v1.py b/src/braket/task_result/task_metadata_v1.py
@@ -11,10 +11,14 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License
 
-from typing import Any, Dict, Optional
+from typing import Optional, Union
 
 from pydantic import Field, conint, constr
 
+from braket.device_schema.dwave import DwaveDeviceParameters
+from braket.device_schema.ionq import IonqDeviceParameters
+from braket.device_schema.rigetti import RigettiDeviceParameters
+from braket.device_schema.simulators import GateModelSimulatorDeviceParameters
 from braket.schema_common import BraketSchemaBase, BraketSchemaHeader
 
 
@@ -29,8 +33,9 @@ class TaskMetadata(BraketSchemaBase):
         shots (str): The number of shots for the task
         deviceId (str): The ID of the device on which the task ran.
             For AWS devices, this is the device ARN.
-        deviceParameters (Dict[str, Any]): The device parameters of the task. Default is None.
-            # TODO: replace with device schema
+        deviceParameters any of (DwaveDeviceParameters, RigettiDeviceParameters,
+            IonqDeviceParameters, GateModelSimulatorDeviceParameters).
+            The device parameters of the task. Default is None.
         createdAt (str): The timestamp of creation;
             the format must be in ISO-8601/RFC3339 string format YYYY-MM-DDTHH:mm:ss.sssZ.
             Default is None.
@@ -53,7 +58,14 @@ class TaskMetadata(BraketSchemaBase):
     id: constr(min_length=1)
     shots: conint(ge=0)
     deviceId: constr(min_length=1)
-    deviceParameters: Optional[Dict[str, Any]]  # TODO: replace with device schema
+    deviceParameters: Optional[
+        Union[
+            DwaveDeviceParameters,
+            RigettiDeviceParameters,
+            IonqDeviceParameters,
+            GateModelSimulatorDeviceParameters,
+        ]
+    ]
     createdAt: Optional[constr(min_length=1, max_length=24)]
     endedAt: Optional[constr(min_length=1, max_length=24)]
     status: Optional[constr(min_length=1, max_length=20)]

diff --git a/test/braket/task_result/test_task_metadata_v1.py b/test/braket/task_result/test_task_metadata_v1.py
@@ -36,7 +36,19 @@ def test_correct_metadata_minimum(id, device_id, shots):
 
 
 def test_correct_metadata_all(id, device_id, shots):
-    device_parameters = {"test": 1}  # TODO: replace with device schema
+    device_parameters = {
+        "braketSchemaHeader": {
+            "name": "braket.device_schema.rigetti.rigetti_device_parameters",
+            "version": "1",
+        },
+        "paradigmParameters": {
+            "braketSchemaHeader": {
+                "name": "braket.device_schema.gate_model_parameters",
+                "version": "1",
+            },
+            "qubitCount": 1,
+        },
+    }
     createdAt = "2020-01-02T03:04:05.006Z"
     endedAt = "2020-01-03T03:04:05.006Z"
     status = "COMPLETED"