feat(api): build out well state for liquid height tracking (#15681)

Overview

re EXEC-603

This PR adds a new WellState class. This class currently has only one attribute, measured_liquid_heights, which is used in liquid height tracking. We created this new class(rather than just adding this attribute to labware state or something related), to allow for more detailed well state tracking in the future.
Details

measured_liquid_heights contains a Dict[labware_id, Dict[well_name, LiquidHeightInfo]]. This dictionary structure is meant to give time-efficient access to all wells of a particular labware, since the dictionary is primarily indexed by labware_id.

The presence/non-presence of a well in the WellState class indicates the following:

    If a well is not present: that well has never been probed.
    If a well is present but has a height value of 0: the most recent probe into that well found no liquid.
    If a well is present and has a height value of > 0: the most recent probe into that well found liquid.

A quick word on the implementation of the new LiquidHeightInfo type. It has two attributes, height and last_measured. The last_measured attribute currently stores as a datetime, but it could be useful in the future to consider other alternatives. Maybe protocol steps?(or something similar). The reason we have a time-related attribute at all is because right now we don't track things like aspirates and dispenses, or even user inputted liquid loading(if that's the correct term). All liquid height state is determined only by doing a liquid_probe, so it's useful to know what the last time was we did a liquid_probe on this well.

For the purposes of the StateSummary, a new type was defined called LiquidHeightSummary. It has identifying information about the labware/well and the contents of the LiquidHeightInfo type.
Test Plan

Unit tests written and passed for well_view and well_store.
Changelog

    Created WellState class, along with complementary WellView and WellStore
    Created LiquidHeightInfo type, which is used to store info about recent liquid height measurements.
    Created LiquidHeightSummary type, which is used in StateSummary.

Review requests
Risk assessment

Low.

api/src/opentrons/cli/analyze.py

@@ -275,6 +275,7 @@ async def _do_analyze(protocol_source: ProtocolSource) -> RunResult:
                 modules=[],
                 labwareOffsets=[],
                 liquids=[],
+                wells=[],
                 hasEverEnteredErrorRecovery=False,
             ),
             parameters=[],

api/src/opentrons/protocol_engine/state/state.py

@@ -26,6 +26,7 @@
 from .modules import ModuleState, ModuleStore, ModuleView
 from .liquids import LiquidState, LiquidView, LiquidStore
 from .tips import TipState, TipView, TipStore
+from .wells import WellState, WellView, WellStore
 from .geometry import GeometryView
 from .motion import MotionView
 from .config import Config
@@ -48,6 +49,7 @@ class State:
     modules: ModuleState
     liquids: LiquidState
     tips: TipState
+    wells: WellState
 
 
 class StateView(HasState[State]):
@@ -61,6 +63,7 @@ class StateView(HasState[State]):
     _modules: ModuleView
     _liquid: LiquidView
     _tips: TipView
+    _wells: WellView
     _geometry: GeometryView
     _motion: MotionView
     _config: Config
@@ -100,6 +103,11 @@ def tips(self) -> TipView:
         """Get state view selectors for tip state."""
         return self._tips
 
+    @property
+    def wells(self) -> WellView:
+        """Get state view selectors for well state."""
+        return self._wells
+
     @property
     def geometry(self) -> GeometryView:
         """Get state view selectors for derived geometry state."""
@@ -129,6 +137,7 @@ def get_summary(self) -> StateSummary:
             completedAt=self._state.commands.run_completed_at,
             startedAt=self._state.commands.run_started_at,
             liquids=self._liquid.get_all(),
+            wells=self._wells.get_all(),
             hasEverEnteredErrorRecovery=self._commands.get_has_entered_recovery_mode(),
         )
 
@@ -196,6 +205,7 @@ def __init__(
         )
         self._liquid_store = LiquidStore()
         self._tip_store = TipStore()
+        self._well_store = WellStore()
 
         self._substores: List[HandlesActions] = [
             self._command_store,
@@ -205,6 +215,7 @@ def __init__(
             self._module_store,
             self._liquid_store,
             self._tip_store,
+            self._well_store,
         ]
         self._config = config
         self._change_notifier = change_notifier or ChangeNotifier()
@@ -321,6 +332,7 @@ def _get_next_state(self) -> State:
             modules=self._module_store.state,
             liquids=self._liquid_store.state,
             tips=self._tip_store.state,
+            wells=self._well_store.state,
         )
 
     def _initialize_state(self) -> None:
@@ -336,6 +348,7 @@ def _initialize_state(self) -> None:
         self._modules = ModuleView(state.modules)
         self._liquid = LiquidView(state.liquids)
         self._tips = TipView(state.tips)
+        self._wells = WellView(state.wells)
 
         # Derived states
         self._geometry = GeometryView(
@@ -365,6 +378,7 @@ def _update_state_views(self) -> None:
         self._modules._state = next_state.modules
         self._liquid._state = next_state.liquids
         self._tips._state = next_state.tips
+        self._wells._state = next_state.wells
         self._change_notifier.notify()
         if self._notify_robot_server is not None:
             self._notify_robot_server()

api/src/opentrons/protocol_engine/state/state_summary.py

@@ -6,6 +6,7 @@
 from ..errors import ErrorOccurrence
 from ..types import (
     EngineStatus,
+    LiquidHeightSummary,
     LoadedLabware,
     LabwareOffset,
     LoadedModule,
@@ -29,3 +30,4 @@ class StateSummary(BaseModel):
     startedAt: Optional[datetime]
     completedAt: Optional[datetime]
     liquids: List[Liquid] = Field(default_factory=list)
+    wells: List[LiquidHeightSummary] = Field(default_factory=list)

api/src/opentrons/protocol_engine/state/wells.py

@@ -0,0 +1,129 @@
+"""Basic well data state and store."""
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Dict, List, Optional
+from opentrons.protocol_engine.actions.actions import (
+    FailCommandAction,
+    SucceedCommandAction,
+)
+from opentrons.protocol_engine.commands.liquid_probe import LiquidProbeResult
+from opentrons.protocol_engine.commands.pipetting_common import LiquidNotFoundError
+from opentrons.protocol_engine.types import LiquidHeightInfo, LiquidHeightSummary
+
+from ._abstract_store import HasState, HandlesActions
+from ..actions import Action
+from ..commands import Command
+
+
+@dataclass
+class WellState:
+    """State of all wells."""
+
+    measured_liquid_heights: Dict[str, Dict[str, LiquidHeightInfo]]
+
+
+class WellStore(HasState[WellState], HandlesActions):
+    """Well state container."""
+
+    _state: WellState
+
+    def __init__(self) -> None:
+        """Initialize a well store and its state."""
+        self._state = WellState(measured_liquid_heights={})
+
+    def handle_action(self, action: Action) -> None:
+        """Modify state in reaction to an action."""
+        if isinstance(action, SucceedCommandAction):
+            self._handle_succeeded_command(action.command)
+        if isinstance(action, FailCommandAction):
+            self._handle_failed_command(action)
+
+    def _handle_succeeded_command(self, command: Command) -> None:
+        if isinstance(command.result, LiquidProbeResult):
+            self._set_liquid_height(
+                labware_id=command.params.labwareId,
+                well_name=command.params.wellName,
+                height=command.result.z_position,
+                time=command.createdAt,
+            )
+
+    def _handle_failed_command(self, action: FailCommandAction) -> None:
+        if isinstance(action.error, LiquidNotFoundError):
+            self._set_liquid_height(
+                labware_id=action.error.private.labware_id,
+                well_name=action.error.private.well_name,
+                height=0,
+                time=action.failed_at,
+            )
+
+    def _set_liquid_height(
+        self, labware_id: str, well_name: str, height: float, time: datetime
+    ) -> None:
+        """Set the liquid height of the well."""
+        lhi = LiquidHeightInfo(height=height, last_measured=time)
+        if labware_id not in self._state.measured_liquid_heights:
+            self._state.measured_liquid_heights[labware_id] = {}
+        self._state.measured_liquid_heights[labware_id][well_name] = lhi
+
+
+class WellView(HasState[WellState]):
+    """Read-only well state view."""
+
+    _state: WellState
+
+    def __init__(self, state: WellState) -> None:
+        """Initialize the computed view of well state.
+
+        Arguments:
+            state: Well state dataclass used for all calculations.
+        """
+        self._state = state
+
+    def get_all(self) -> List[LiquidHeightSummary]:
+        """Get all well liquid heights."""
+        all_heights: List[LiquidHeightSummary] = []
+        for labware, wells in self._state.measured_liquid_heights.items():
+            for well, lhi in wells.items():
+                lhs = LiquidHeightSummary(
+                    labware_id=labware,
+                    well_name=well,
+                    height=lhi.height,
+                    last_measured=lhi.last_measured,
+                )
+            all_heights.append(lhs)
+        return all_heights
+
+    def get_all_in_labware(self, labware_id: str) -> List[LiquidHeightSummary]:
+        """Get all well liquid heights for a particular labware."""
+        all_heights: List[LiquidHeightSummary] = []
+        for well, lhi in self._state.measured_liquid_heights[labware_id].items():
+            lhs = LiquidHeightSummary(
+                labware_id=labware_id,
+                well_name=well,
+                height=lhi.height,
+                last_measured=lhi.last_measured,
+            )
+            all_heights.append(lhs)
+        return all_heights
+
+    def get_last_measured_liquid_height(
+        self, labware_id: str, well_name: str
+    ) -> Optional[float]:
+        """Returns the height of the liquid according to the most recent liquid level probe to this well.
+
+        Returns None if no liquid probe has been done.
+        """
+        try:
+            height = self._state.measured_liquid_heights[labware_id][well_name].height
+            return height
+        except KeyError:
+            return None
+
+    def has_measured_liquid_height(self, labware_id: str, well_name: str) -> bool:
+        """Returns True if the well has been liquid level probed previously."""
+        try:
+            return bool(
+                self._state.measured_liquid_heights[labware_id][well_name].height
+            )
+        except KeyError:
+            return False

api/src/opentrons/protocol_engine/types.py

@@ -311,6 +311,22 @@ class CurrentWell:
     well_name: str
 
 
+class LiquidHeightInfo(BaseModel):
+    """Payload required to store recent measured liquid heights."""
+
+    height: float
+    last_measured: datetime
+
+
+class LiquidHeightSummary(BaseModel):
+    """Payload for liquid state height in StateSummary."""
+
+    labware_id: str
+    well_name: str
+    height: float
+    last_measured: datetime
+
+
 @dataclass(frozen=True)
 class CurrentAddressableArea:
     """The latest addressable area the robot has accessed."""

api/tests/opentrons/protocol_engine/state/command_fixtures.py

@@ -295,6 +295,32 @@ def create_dispense_in_place_command(
     )
 
 
+def create_liquid_probe_command(
+    pipette_id: str = "pippete-id",
+    labware_id: str = "labware-id",
+    well_name: str = "well-name",
+    well_location: Optional[WellLocation] = None,
+    destination: DeckPoint = DeckPoint(x=0, y=0, z=0),
+) -> cmd.LiquidProbe:
+    """Get a completed Liquid Probe command."""
+    params = cmd.LiquidProbeParams(
+        pipetteId=pipette_id,
+        labwareId=labware_id,
+        wellName=well_name,
+        wellLocation=well_location or WellLocation(),
+    )
+    result = cmd.LiquidProbeResult(position=destination, z_position=0.5)
+
+    return cmd.LiquidProbe(
+        id="command-id",
+        key="command-key",
+        status=cmd.CommandStatus.SUCCEEDED,
+        createdAt=datetime.now(),
+        params=params,
+        result=result,
+    )
+
+
 def create_pick_up_tip_command(
     pipette_id: str,
     labware_id: str = "labware-id",

api/tests/opentrons/protocol_engine/state/test_well_store.py

@@ -0,0 +1,28 @@
+"""Well state store tests."""
+import pytest
+from opentrons.protocol_engine.state.wells import WellStore
+from opentrons.protocol_engine.actions.actions import SucceedCommandAction
+
+from .command_fixtures import create_liquid_probe_command
+
+
+@pytest.fixture
+def subject() -> WellStore:
+    """Well store test subject."""
+    return WellStore()
+
+
+def test_handles_liquid_probe_success(subject: WellStore) -> None:
+    """It should add the well to the state after a successful liquid probe."""
+    labware_id = "labware-id"
+    well_name = "well-name"
+
+    liquid_probe = create_liquid_probe_command()
+
+    subject.handle_action(
+        SucceedCommandAction(private_result=None, command=liquid_probe)
+    )
+
+    assert len(subject.state.measured_liquid_heights) == 1
+
+    assert subject.state.measured_liquid_heights[labware_id][well_name].height == 0.5

api/tests/opentrons/protocol_engine/state/test_well_view.py

@@ -0,0 +1,51 @@
+"""Well view tests."""
+from datetime import datetime
+from opentrons.protocol_engine.types import LiquidHeightInfo
+import pytest
+from opentrons.protocol_engine.state.wells import WellState, WellView
+
+
+@pytest.fixture
+def subject() -> WellView:
+    """Get a well view test subject."""
+    labware_id = "labware-id"
+    well_name = "well-name"
+    height_info = LiquidHeightInfo(height=0.5, last_measured=datetime.now())
+    state = WellState(measured_liquid_heights={labware_id: {well_name: height_info}})
+
+    return WellView(state)
+
+
+def test_get_all(subject: WellView) -> None:
+    """Should return a list of well heights."""
+    assert subject.get_all()[0].height == 0.5
+
+
+def test_get_last_measured_liquid_height(subject: WellView) -> None:
+    """Should return the height of a single well correctly for valid wells."""
+    labware_id = "labware-id"
+    well_name = "well-name"
+
+    invalid_labware_id = "invalid-labware-id"
+    invalid_well_name = "invalid-well-name"
+
+    assert (
+        subject.get_last_measured_liquid_height(invalid_labware_id, invalid_well_name)
+        is None
+    )
+    assert subject.get_last_measured_liquid_height(labware_id, well_name) == 0.5
+
+
+def test_has_measured_liquid_height(subject: WellView) -> None:
+    """Should return True for measured wells and False for ones that have no measurements."""
+    labware_id = "labware-id"
+    well_name = "well-name"
+
+    invalid_labware_id = "invalid-labware-id"
+    invalid_well_name = "invalid-well-name"
+
+    assert (
+        subject.has_measured_liquid_height(invalid_labware_id, invalid_well_name)
+        is False
+    )
+    assert subject.has_measured_liquid_height(labware_id, well_name) is True

robot-server/robot_server/maintenance_runs/maintenance_run_data_manager.py

@@ -32,6 +32,7 @@ def _build_run(
         pipettes=[],
         modules=[],
         liquids=[],
+        wells=[],
         hasEverEnteredErrorRecovery=False,
     )
     return MaintenanceRun.construct(

robot-server/robot_server/runs/run_data_manager.py

@@ -71,6 +71,7 @@ def _build_run(
             pipettes=[],
             modules=[],
             liquids=[],
+            wells=[],
             hasEverEnteredErrorRecovery=False,
         )
         errors.append(state_summary.dataError)

robot-server/tests/maintenance_runs/test_run_data_manager.py

@@ -69,6 +69,7 @@ def engine_state_summary() -> StateSummary:
         pipettes=[LoadedPipette.construct(id="some-pipette-id")],  # type: ignore[call-arg]
         modules=[LoadedModule.construct(id="some-module-id")],  # type: ignore[call-arg]
         liquids=[Liquid(id="some-liquid-id", displayName="liquid", description="desc")],
+        wells=[],
     )