feat: add outputs field to support mid-circuit measurements

src/braket/task_result/gate_model_task_result_v1.py

@@ -12,13 +12,25 @@
 # language governing permissions and limitations under the License
 
 
-from pydantic.v1 import BaseModel, Field, confloat, conint, conlist, constr
+from typing import TypeAlias
+
+from pydantic.v1 import BaseModel, Field, StrictBool, confloat, conint, conlist, constr, validator
 
 from braket.ir.jaqcd.program_v1 import Results
 from braket.schema_common import BraketSchemaBase, BraketSchemaHeader
 from braket.task_result.additional_metadata import AdditionalMetadata
 from braket.task_result.task_metadata_v1 import TaskMetadata
 
+StrictInt = conint(strict=True)
+StrictFloat = confloat(allow_inf_nan=False, strict=True)
+
+ScalarValue: TypeAlias = StrictBool | StrictInt | StrictFloat
+ScalarOrNull: TypeAlias = ScalarValue | None
+
+OutputValue: TypeAlias = (
+    ScalarOrNull | list[StrictBool | None] | list[StrictInt | None] | list[StrictFloat | None]
+)
+
 
 class ResultTypeValue(BaseModel):
     """
@@ -51,6 +63,11 @@ class GateModelTaskResult(BraketSchemaBase):
             Indicates which qubits are in `measurements`. Default is `None`.
         resultTypes (list[ResultTypeValue]): Requested result types and their values.
             Default is `None`.
+        outputs (list[dict[str, OutputValue]]): Per-shot values of OpenQASM 3
+            ``output``-declared variables. Each element is one shot, mapping each
+            output variable name to its value (a scalar, or a homogeneous list for
+            registers). An undefined value is represented by ``None``. Default is
+            `None`.
         taskMetadata (TaskMetadata): The task metadata
         additionalMetadata (AdditionalMetadata): Additional metadata of the task
     """
@@ -70,5 +87,25 @@ class GateModelTaskResult(BraketSchemaBase):
     )
     resultTypes: list[ResultTypeValue] | None
     measuredQubits: conlist(conint(ge=0), min_items=1) | None
+    outputs: conlist(dict[constr(min_length=1), OutputValue], min_items=1) | None
     taskMetadata: TaskMetadata
     additionalMetadata: AdditionalMetadata
+
+    @validator("outputs", each_item=True)
+    def validate_non_empty_shot(cls, shot):
+        """
+        Rejects empty per-shot dicts. Every declared output variable appears in
+        each shot (its value or None), so an empty shot signals malformed data.
+
+        Args:
+            shot (dict): One per-shot output dict from the outputs list.
+
+        Returns:
+            dict: The same shot, validated to be non-empty.
+
+        Raises:
+            ValueError: If the shot dict is empty.
+        """
+        if not shot:
+            raise ValueError("each output shot must contain at least one variable")
+        return shot

test/unit_tests/braket/task_result/test_gate_model_task_result_v1.py

@@ -182,3 +182,150 @@ def test_incorrect_result_type_attribute_type():
 @pytest.mark.xfail(raises=ValidationError)
 def test_incorrect_result_type_attribute_value():
     ResultTypeValue(type={"type": "unknown"}, value=1)
+
+
+def test_outputs_only_two_shots_multiple_registers(
+    task_metadata,
+    additional_metadata_gate_model,
+    measured_qubits,
+):
+    outputs = [
+        {"c1": 1, "c2": [0, 1], "float_variable": 1.57},
+        {"c1": 0, "c2": [0, 0], "float_variable": 1.57},
+    ]
+    result = GateModelTaskResult(
+        outputs=outputs,
+        measuredQubits=measured_qubits,
+        taskMetadata=task_metadata,
+        additionalMetadata=additional_metadata_gate_model,
+    )
+    assert result.outputs == outputs
+    assert result.measurements is None
+    assert result.measurementProbabilities is None
+    assert GateModelTaskResult.parse_raw(result.json()) == result
+
+
+def test_outputs_loop_same_qubit_measured_repeatedly(
+    task_metadata,
+    additional_metadata_gate_model,
+    measured_qubits,
+):
+    outputs = [{"mcm_abc": [1, 0, 0]}]
+    result = GateModelTaskResult(
+        outputs=outputs,
+        measuredQubits=measured_qubits,
+        taskMetadata=task_metadata,
+        additionalMetadata=additional_metadata_gate_model,
+    )
+    assert result.outputs == outputs
+
+
+def test_outputs_negative_integer_value(
+    task_metadata,
+    additional_metadata_gate_model,
+    measured_qubits,
+):
+    outputs = [{"mcm": [-1, 1, 0]}]
+    result = GateModelTaskResult(
+        outputs=outputs,
+        measuredQubits=measured_qubits,
+        taskMetadata=task_metadata,
+        additionalMetadata=additional_metadata_gate_model,
+    )
+    assert result.outputs == outputs
+    assert GateModelTaskResult.parse_raw(result.json()).outputs == outputs
+
+
+def test_outputs_undefined_value_represented_as_none(
+    task_metadata,
+    additional_metadata_gate_model,
+    measured_qubits,
+):
+    outputs = [
+        {"c1": 1, "c2": [0, 1]},
+        {"c1": 0, "c2": [None, 1]},
+    ]
+    result = GateModelTaskResult(
+        outputs=outputs,
+        measuredQubits=measured_qubits,
+        taskMetadata=task_metadata,
+        additionalMetadata=additional_metadata_gate_model,
+    )
+    assert result.outputs == outputs
+
+
+def test_outputs_bool_value(
+    task_metadata,
+    additional_metadata_gate_model,
+    measured_qubits,
+):
+    outputs = [{"flag": True}, {"flag": False}]
+    result = GateModelTaskResult(
+        outputs=outputs,
+        measuredQubits=measured_qubits,
+        taskMetadata=task_metadata,
+        additionalMetadata=additional_metadata_gate_model,
+    )
+    assert result.outputs == outputs
+    assert GateModelTaskResult.parse_raw(result.json()).outputs == outputs
+
+
+def test_outputs_alongside_measurements(
+    task_metadata,
+    additional_metadata_gate_model,
+    measured_qubits,
+    measurements,
+):
+    outputs = [{"c": [m_q0, m_q1]} for m_q0, m_q1 in measurements]
+    result = GateModelTaskResult(
+        measurements=measurements,
+        outputs=outputs,
+        measuredQubits=measured_qubits,
+        taskMetadata=task_metadata,
+        additionalMetadata=additional_metadata_gate_model,
+    )
+    assert result.measurements == measurements
+    assert result.outputs == outputs
+
+
+@pytest.mark.parametrize(
+    "outputs",
+    [
+        [],  # must contain at least one shot
+        [{}],  # shot dict must not be empty
+        [{"c1": "not-a-scalar"}],  # strings not allowed
+        [{"c1": [1, "0"]}],  # strings inside a list not allowed
+        [{"c1": [1, 1.5]}],  # mixed types within a list not allowed
+        [{"c1": [[0, 1]]}],  # nested lists not allowed
+        [{"c1": {"nested": "dict"}}],  # dict values not allowed
+        [{"": 1}],  # empty variable name not allowed
+    ],
+)
+@pytest.mark.xfail(raises=ValidationError)
+def test_incorrect_outputs(
+    outputs,
+    measured_qubits,
+    task_metadata,
+    additional_metadata_gate_model,
+):
+    GateModelTaskResult(
+        outputs=outputs,
+        measuredQubits=measured_qubits,
+        taskMetadata=task_metadata,
+        additionalMetadata=additional_metadata_gate_model,
+    )
+
+
+def test_outputs_default_is_none(
+    task_metadata,
+    additional_metadata_gate_model,
+    measured_qubits,
+    measurements,
+):
+    result = GateModelTaskResult(
+        measurements=measurements,
+        measuredQubits=measured_qubits,
+        taskMetadata=task_metadata,
+        additionalMetadata=additional_metadata_gate_model,
+    )
+    assert result.outputs is None