mx-bluesky 1.4.0__py3-none-any.whl → 1.4.1a0__py3-none-any.whl

mx_bluesky/hyperion/experiment_plans/pin_centre_then_xray_centre_plan.py

@@ -2,15 +2,27 @@ from __future__ import annotations
 
 import json
 
-from blueapi.core import BlueskyContext, MsgGenerator
+import bluesky.preprocessors as bpp
+from blueapi.core import BlueskyContext
+from bluesky.utils import MsgGenerator
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.oav.oav_parameters import OAVParameters
 
 from mx_bluesky.common.parameters.constants import OavConstants
+from mx_bluesky.common.parameters.gridscan import (
+    GridScanWithEdgeDetect,
+    PinTipCentreThenXrayCentre,
+)
 from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import move_phi_chi_omega
 from mx_bluesky.hyperion.device_setup_plans.utils import (
     start_preparing_data_collection_then_do_plan,
 )
+from mx_bluesky.hyperion.experiment_plans.change_aperture_then_move_plan import (
+    change_aperture_then_move_to_xtal,
+)
+from mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan import (
+    XRayCentreEventHandler,
+)
 from mx_bluesky.hyperion.experiment_plans.grid_detect_then_xray_centre_plan import (
     GridDetectThenXRayCentreComposite,
     detect_grid_and_do_gridscan,
@@ -27,10 +39,6 @@ from mx_bluesky.hyperion.external_interaction.callbacks.xray_centre.ispyb_callba
 )
 from mx_bluesky.hyperion.log import LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.gridscan import (
-    GridScanWithEdgeDetect,
-    PinTipCentreThenXrayCentre,
-)
 from mx_bluesky.hyperion.utils.context import device_composite_from_context
 
 
@@ -53,13 +61,12 @@ def create_parameters_for_grid_detection(
     return grid_detect_and_xray_centre
 
 
-def pin_centre_then_xray_centre_plan(
+def pin_centre_then_flyscan_plan(
     composite: GridDetectThenXRayCentreComposite,
     parameters: PinTipCentreThenXrayCentre,
     oav_config_file: str = OavConstants.OAV_CONFIG_JSON,
 ):
-    """Plan that perfoms a pin tip centre followed by an xray centre to completely
-    centre the sample"""
+    """Plan that performs a pin tip centre followed by a flyscan to determine the centres of interest"""
 
     pin_tip_centring_composite = PinTipCentringComposite(
         oav=composite.oav,
@@ -68,7 +75,7 @@ def pin_centre_then_xray_centre_plan(
         pin_tip_detection=composite.pin_tip_detection,
     )
 
-    def _pin_centre_then_xray_centre_plan():
+    def _pin_centre_then_flyscan_plan():
         yield from setup_beamline_for_OAV(
             composite.smargon, composite.backlight, composite.aperture_scatterguard
         )
@@ -96,7 +103,7 @@ def pin_centre_then_xray_centre_plan(
             oav_params,
         )
 
-    yield from ispyb_activation_wrapper(_pin_centre_then_xray_centre_plan(), parameters)
+    yield from ispyb_activation_wrapper(_pin_centre_then_flyscan_plan(), parameters)
 
 
 def pin_tip_centre_then_xray_centre(
@@ -105,15 +112,27 @@ def pin_tip_centre_then_xray_centre(
     oav_config_file: str = OavConstants.OAV_CONFIG_JSON,
 ) -> MsgGenerator:
     """Starts preparing for collection then performs the pin tip centre and xray centre"""
-
     eiger: EigerDetector = composite.eiger
 
     eiger.set_detector_parameters(parameters.detector_params)
 
-    return start_preparing_data_collection_then_do_plan(
-        eiger,
-        composite.detector_motion,
-        parameters.detector_params.detector_distance,
-        pin_centre_then_xray_centre_plan(composite, parameters, oav_config_file),
-        group=CONST.WAIT.GRID_READY_FOR_DC,
+    flyscan_event_handler = XRayCentreEventHandler()
+
+    @bpp.subs_decorator(flyscan_event_handler)
+    def pin_centre_flyscan_then_fetch_results() -> MsgGenerator:
+        yield from start_preparing_data_collection_then_do_plan(
+            eiger,
+            composite.detector_motion,
+            parameters.detector_params.detector_distance,
+            pin_centre_then_flyscan_plan(composite, parameters, oav_config_file),
+            group=CONST.WAIT.GRID_READY_FOR_DC,
+        )
+
+    yield from pin_centre_flyscan_then_fetch_results()
+    flyscan_results = flyscan_event_handler.xray_centre_results
+    assert (
+        flyscan_results
+    ), "Flyscan result event not received or no crystal found and exception not raised"
+    yield from change_aperture_then_move_to_xtal(
+        flyscan_results[0], composite.smargon, composite.aperture_scatterguard
     )

mx_bluesky/hyperion/experiment_plans/pin_tip_centring_plan.py

@@ -1,7 +1,7 @@
-import dataclasses
 from collections.abc import Generator
 
 import bluesky.plan_stubs as bps
+import pydantic
 from blueapi.core import BlueskyContext
 from bluesky.utils import Msg
 from dodal.devices.backlight import Backlight
@@ -27,7 +27,7 @@ from mx_bluesky.hyperion.utils.context import device_composite_from_context
 DEFAULT_STEP_SIZE = 0.5
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class PinTipCentringComposite:
     """All devices which are directly or indirectly required by this plan"""
 

mx_bluesky/hyperion/experiment_plans/robot_load_and_change_energy.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import dataclasses
 from collections.abc import Generator
 from datetime import datetime
 from pathlib import Path
@@ -8,9 +7,10 @@ from typing import cast
 
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
+import pydantic
 from blueapi.core import BlueskyContext
 from bluesky.utils import Msg
-from dodal.devices.aperturescatterguard import ApertureScatterguard, ApertureValue
+from dodal.devices.aperturescatterguard import ApertureScatterguard
 from dodal.devices.attenuator import Attenuator
 from dodal.devices.dcm import DCM
 from dodal.devices.focusing_mirror import FocusingMirrorWithStripes, MirrorVoltages
@@ -22,18 +22,18 @@ from dodal.devices.thawer import Thawer
 from dodal.devices.undulator_dcm import UndulatorDCM
 from dodal.devices.webcam import Webcam
 from dodal.devices.xbpm_feedback import XBPMFeedback
-from dodal.plans.motor_util_plans import MoveTooLarge, home_and_reset_wrapper
+from dodal.plan_stubs.motor_utils import MoveTooLarge, home_and_reset_wrapper
 
+from mx_bluesky.common.parameters.robot_load import RobotLoadAndEnergyChange
 from mx_bluesky.hyperion.experiment_plans.set_energy_plan import (
     SetEnergyComposite,
     set_energy_plan,
 )
 from mx_bluesky.hyperion.log import LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.robot_load import RobotLoadAndEnergyChange
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class RobotLoadAndEnergyChangeComposite:
     # SetEnergyComposite fields
     vfm: FocusingMirrorWithStripes
@@ -94,11 +94,7 @@ def take_robot_snapshots(oav: OAV, webcam: Webcam, directory: Path):
 def prepare_for_robot_load(
     aperture_scatterguard: ApertureScatterguard, smargon: Smargon
 ):
-    yield from bps.abs_set(
-        aperture_scatterguard,
-        ApertureValue.ROBOT_LOAD,
-        group="prepare_robot_load",
-    )
+    yield from bps.trigger(aperture_scatterguard.move_out, group="prepare_robot_load")
 
     yield from bps.mv(smargon.stub_offsets, StubPosition.RESET_TO_ROBOT_LOAD)  # type: ignore # See: https://github.com/bluesky/bluesky/issues/1809
 

mx_bluesky/hyperion/experiment_plans/robot_load_then_centre_plan.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
-import dataclasses
 from typing import cast
 
-from blueapi.core import BlueskyContext, MsgGenerator
+import bluesky.preprocessors as bpp
+import pydantic
+from blueapi.core import BlueskyContext
+from bluesky.utils import MsgGenerator
 from dodal.devices.aperturescatterguard import ApertureScatterguard
 from dodal.devices.attenuator import Attenuator
 from dodal.devices.backlight import Backlight
@@ -32,15 +34,22 @@ from dodal.log import LOGGER
 from ophyd_async.fastcs.panda import HDFPanda
 
 from mx_bluesky.common.parameters.constants import OavConstants
+from mx_bluesky.common.parameters.gridscan import RobotLoadThenCentre
 from mx_bluesky.hyperion.device_setup_plans.utils import (
     fill_in_energy_if_not_supplied,
     start_preparing_data_collection_then_do_plan,
 )
+from mx_bluesky.hyperion.experiment_plans.change_aperture_then_move_plan import (
+    change_aperture_then_move_to_xtal,
+)
+from mx_bluesky.hyperion.experiment_plans.flyscan_xray_centre_plan import (
+    XRayCentreEventHandler,
+)
 from mx_bluesky.hyperion.experiment_plans.grid_detect_then_xray_centre_plan import (
     GridDetectThenXRayCentreComposite,
 )
 from mx_bluesky.hyperion.experiment_plans.pin_centre_then_xray_centre_plan import (
-    pin_centre_then_xray_centre_plan,
+    pin_centre_then_flyscan_plan,
 )
 from mx_bluesky.hyperion.experiment_plans.robot_load_and_change_energy import (
     RobotLoadAndEnergyChangeComposite,
@@ -48,10 +57,9 @@ from mx_bluesky.hyperion.experiment_plans.robot_load_and_change_energy import (
     robot_load_and_change_energy_plan,
 )
 from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.gridscan import RobotLoadThenCentre
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class RobotLoadThenCentreComposite:
     # common fields
     xbpm_feedback: XBPMFeedback
@@ -88,6 +96,10 @@ class RobotLoadThenCentreComposite:
     webcam: Webcam
     lower_gonio: XYZPositioner
 
+    @property
+    def sample_motors(self):
+        return self.smargon
+
 
 def create_devices(context: BlueskyContext) -> RobotLoadThenCentreComposite:
     from mx_bluesky.hyperion.utils.context import device_composite_from_context
@@ -95,18 +107,18 @@ def create_devices(context: BlueskyContext) -> RobotLoadThenCentreComposite:
     return device_composite_from_context(context, RobotLoadThenCentreComposite)
 
 
-def centring_plan_from_robot_load_params(
+def _flyscan_plan_from_robot_load_params(
     composite: RobotLoadThenCentreComposite,
     params: RobotLoadThenCentre,
     oav_config_file: str = OavConstants.OAV_CONFIG_JSON,
 ):
-    yield from pin_centre_then_xray_centre_plan(
+    yield from pin_centre_then_flyscan_plan(
         cast(GridDetectThenXRayCentreComposite, composite),
         params.pin_centre_then_xray_centre_params(),
     )
 
 
-def robot_load_then_centre_plan(
+def _robot_load_then_flyscan_plan(
     composite: RobotLoadThenCentreComposite,
     params: RobotLoadThenCentre,
     oav_config_file: str = OavConstants.OAV_CONFIG_JSON,
@@ -116,13 +128,36 @@ def robot_load_then_centre_plan(
         params.robot_load_params(),
     )
 
-    yield from centring_plan_from_robot_load_params(composite, params, oav_config_file)
+    yield from _flyscan_plan_from_robot_load_params(composite, params, oav_config_file)
 
 
 def robot_load_then_centre(
     composite: RobotLoadThenCentreComposite,
     parameters: RobotLoadThenCentre,
 ) -> MsgGenerator:
+    """Perform pin-tip detection followed by a flyscan to determine centres of interest.
+    Performs a robot load if necessary. Centre on the best diffracting centre.
+    """
+
+    xray_centre_event_handler = XRayCentreEventHandler()
+
+    yield from bpp.subs_wrapper(
+        robot_load_then_xray_centre(composite, parameters), xray_centre_event_handler
+    )
+    flyscan_results = xray_centre_event_handler.xray_centre_results
+    if flyscan_results is not None:
+        yield from change_aperture_then_move_to_xtal(
+            flyscan_results[0], composite.smargon, composite.aperture_scatterguard
+        )
+    # else no chi change, no need to recentre.
+
+
+def robot_load_then_xray_centre(
+    composite: RobotLoadThenCentreComposite,
+    parameters: RobotLoadThenCentre,
+) -> MsgGenerator:
+    """Perform pin-tip detection followed by a flyscan to determine centres of interest.
+    Performs a robot load if necessary."""
     eiger: EigerDetector = composite.eiger
 
     # TODO: get these from one source of truth #254
@@ -138,13 +173,13 @@ def robot_load_then_centre(
     doing_chi_change = parameters.chi_start_deg is not None
 
     if doing_sample_load:
-        plan = robot_load_then_centre_plan(
+        plan = _robot_load_then_flyscan_plan(
             composite,
             parameters,
         )
         LOGGER.info("Pin not loaded, loading and centring")
     elif doing_chi_change:
-        plan = centring_plan_from_robot_load_params(composite, parameters)
+        plan = _flyscan_plan_from_robot_load_params(composite, parameters)
         LOGGER.info("Pin already loaded but chi changed so centring")
     else:
         LOGGER.info("Pin already loaded and chi not changed so doing nothing")

mx_bluesky/hyperion/experiment_plans/rotation_scan_plan.py

@@ -4,7 +4,9 @@ import dataclasses
 
 import bluesky.plan_stubs as bps
 import bluesky.preprocessors as bpp
-from blueapi.core import BlueskyContext, MsgGenerator
+import pydantic
+from blueapi.core import BlueskyContext
+from bluesky.utils import MsgGenerator
 from dodal.devices.aperturescatterguard import ApertureScatterguard
 from dodal.devices.attenuator import Attenuator
 from dodal.devices.backlight import Backlight
@@ -22,7 +24,7 @@ from dodal.devices.undulator import Undulator
 from dodal.devices.xbpm_feedback import XBPMFeedback
 from dodal.devices.zebra import RotationDirection, Zebra
 from dodal.devices.zebra_controlled_shutter import ZebraShutter
-from dodal.plans.check_topup import check_topup_and_wait_if_necessary
+from dodal.plan_stubs.check_topup import check_topup_and_wait_if_necessary
 
 from mx_bluesky.common.device_setup_plans.read_hardware_for_setup import (
     read_hardware_for_zocalo,
@@ -62,7 +64,7 @@ from mx_bluesky.hyperion.parameters.rotation import (
 from mx_bluesky.hyperion.utils.context import device_composite_from_context
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class RotationScanComposite(OavSnapshotComposite):
     """All devices which are directly or indirectly required by this plan"""
 
@@ -329,6 +331,13 @@ def _move_and_rotation(
         yield from setup_beamline_for_OAV(
             composite.smargon, composite.backlight, composite.aperture_scatterguard
         )
+        yield from bps.wait(group=CONST.WAIT.READY_FOR_OAV)
+        if params.selected_aperture:
+            yield from bps.abs_set(
+                composite.aperture_scatterguard.aperture_outside_beam,
+                params.selected_aperture,
+                group=CONST.WAIT.ROTATION_READY_FOR_DC,
+            )
         yield from oav_snapshot_plan(composite, params, oav_params)
     yield from rotation_scan_plan(
         composite,
@@ -409,6 +418,11 @@ def multi_rotation_scan(
         }
     )
     @bpp.stage_decorator([eiger])
+    @transmission_and_xbpm_feedback_for_collection_decorator(
+        composite.xbpm_feedback,
+        composite.attenuator,
+        parameters.transmission_frac,
+    )
     @bpp.finalize_decorator(lambda: _cleanup_plan(composite))
     def _multi_rotation_scan():
         for single_scan in parameters.single_rotation_scans:
@@ -418,6 +432,7 @@ def multi_rotation_scan(
                 md={
                     "subplan_name": CONST.PLAN.ROTATION_OUTER,
                     CONST.TRIGGER.ZOCALO: CONST.PLAN.ROTATION_MAIN,
+                    "zocalo_environment": CONST.ZOCALO_ENV,
                     "hyperion_parameters": single_scan.model_dump_json(),
                 }
             )

mx_bluesky/hyperion/experiment_plans/set_energy_plan.py

@@ -5,8 +5,7 @@
 * reenable feedback
 """
 
-import dataclasses
-
+import pydantic
 from bluesky import plan_stubs as bps
 from dodal.devices.attenuator import Attenuator
 from dodal.devices.dcm import DCM
@@ -24,7 +23,7 @@ DESIRED_TRANSMISSION_FRACTION = 0.1
 UNDULATOR_GROUP = "UNDULATOR_GROUP"
 
 
-@dataclasses.dataclass
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
 class SetEnergyComposite:
     vfm: FocusingMirrorWithStripes
     mirror_voltages: MirrorVoltages

mx_bluesky/hyperion/external_interaction/callbacks/common/abstract_event.py

@@ -0,0 +1,66 @@
+import builtins
+import dataclasses
+import time
+from abc import ABC
+from typing import Literal
+
+from bluesky.protocols import Readable, Reading
+from event_model import DataKey
+
+
+@dataclasses.dataclass(frozen=True)
+class AbstractEvent(Readable, ABC):
+    """An abstract superclass that can be extended to provide lightweight software events
+    for bluesky plans, without having to incur the overhead of creating ophyd-async devices
+    specifically for the purpose.
+
+    The currently supported types for field annotations in the event are ``str``, ``int``, ``float``, ``bool``
+
+    In future array types may be supported.
+
+    Examples:
+        Subclasses should extend this class and decorate with::
+
+            @dataclasses.dataclass(frozen=True)
+
+        To raise an event, simply construct the event and then ``read`` it as you would a device::
+
+            yield from bps.create("MY_EVENT_NAME")
+            my_event = MyEvent(an_int=1)
+            yield from bps.read(my_event)
+            yield from bps.save()
+    """
+
+    def read(self) -> dict[str, Reading]:
+        return {
+            f.name: AbstractEvent._reading_from_value(getattr(self, f.name))
+            for f in dataclasses.fields(self)
+        }
+
+    def describe(self) -> dict[str, DataKey]:
+        return {
+            f.name: DataKey(dtype=AbstractEvent._dtype_of(f.type), shape=[], source="")
+            for f in dataclasses.fields(self)
+        }
+
+    @classmethod
+    def _reading_from_value(cls, value):
+        return Reading(timestamp=time.time(), value=value)
+
+    @classmethod
+    def _dtype_of(cls, t) -> Literal["string", "number", "boolean", "integer"]:
+        match t:
+            case builtins.str:
+                return "string"
+            case builtins.bool:
+                return "boolean"
+            case builtins.int:
+                return "integer"
+            case builtins.float:
+                return "number"
+        # TODO array support
+        raise ValueError(f"Unsupported type for AbstractEvent: {t}")
+
+    @property
+    def name(self) -> str:
+        return type(self).__name__

mx_bluesky/hyperion/external_interaction/callbacks/grid_detection_callback.py

@@ -28,7 +28,7 @@ class GridDetectionCallback(CallbackBase):
         *args,
     ) -> None:
         super().__init__(*args)
-        self.start_positions: list = []
+        self.start_positions_mm: list = []
         self.box_numbers: list = []
 
     def event(self, doc: Event):
@@ -55,16 +55,16 @@ class GridDetectionCallback(CallbackBase):
         beam_x = data["oav-beam_centre_i"]
         beam_y = data["oav-beam_centre_j"]
 
-        position_grid_start = calculate_x_y_z_of_pixel(
+        position_grid_start_mm = calculate_x_y_z_of_pixel(
             current_xyz,
             smargon_omega,
             centre_of_first_box,
             (beam_x, beam_y),
             (microns_per_pixel_x, microns_per_pixel_y),
         )
-        LOGGER.info(f"Calculated start position {position_grid_start}")
+        LOGGER.info(f"Calculated start position {position_grid_start_mm}")
 
-        self.start_positions.append(position_grid_start)
+        self.start_positions_mm.append(position_grid_start_mm)
         self.box_numbers.append(
             (
                 data["oav-grid_snapshot-num_boxes_x"],
@@ -72,22 +72,22 @@ class GridDetectionCallback(CallbackBase):
             )
         )
 
-        self.x_step_size_mm = box_width_px * microns_per_pixel_x / 1000
-        self.y_step_size_mm = box_width_px * microns_per_pixel_y / 1000
-        self.z_step_size_mm = box_width_px * microns_per_pixel_y / 1000
+        self.x_step_size_um = box_width_px * microns_per_pixel_x
+        self.y_step_size_um = box_width_px * microns_per_pixel_y
+        self.z_step_size_um = box_width_px * microns_per_pixel_y
         return doc
 
     def get_grid_parameters(self) -> GridParamUpdate:
         return {
-            "x_start_um": self.start_positions[0][0],
-            "y_start_um": self.start_positions[0][1],
-            "y2_start_um": self.start_positions[0][1],
-            "z_start_um": self.start_positions[1][2],
-            "z2_start_um": self.start_positions[1][2],
+            "x_start_um": self.start_positions_mm[0][0] * 1000,
+            "y_start_um": self.start_positions_mm[0][1] * 1000,
+            "y2_start_um": self.start_positions_mm[0][1] * 1000,
+            "z_start_um": self.start_positions_mm[1][2] * 1000,
+            "z2_start_um": self.start_positions_mm[1][2] * 1000,
             "x_steps": self.box_numbers[0][0],
             "y_steps": self.box_numbers[0][1],
             "z_steps": self.box_numbers[1][1],
-            "x_step_size_um": self.x_step_size_mm,
-            "y_step_size_um": self.y_step_size_mm,
-            "z_step_size_um": self.z_step_size_mm,
+            "x_step_size_um": self.x_step_size_um,
+            "y_step_size_um": self.y_step_size_um,
+            "z_step_size_um": self.z_step_size_um,
         }

mx_bluesky/hyperion/external_interaction/callbacks/rotation/ispyb_callback.py

@@ -61,7 +61,9 @@ class RotationISPyBCallback(BaseISPyBCallback):
             ISPYB_LOGGER.info(
                 "ISPyB callback received start document with experiment parameters."
             )
-            self.params = RotationScan.from_json(doc.get("hyperion_parameters"))
+            hyperion_params = doc.get("hyperion_parameters")
+            assert isinstance(hyperion_params, str)
+            self.params = RotationScan.model_validate_json(hyperion_params)
             dcgid = (
                 self.ispyb_ids.data_collection_group_id
                 if (self.params.sample_id == self.last_sample_id)

mx_bluesky/hyperion/external_interaction/callbacks/rotation/nexus_callback.py

@@ -78,12 +78,14 @@ class RotationNexusFileCallback(PlanReactiveCallback):
             self.meta_data_run_number = doc.get("meta_data_run_number")
         if doc.get("subplan_name") == CONST.PLAN.ROTATION_OUTER:
             self.run_uid = doc.get("uid")
-            json_params = doc.get("hyperion_parameters")
+            hyperion_params = doc.get("hyperion_parameters")
+            assert isinstance(hyperion_params, str)
             NEXUS_LOGGER.info(
-                f"Nexus writer received start document with experiment parameters {json_params}"
+                f"Nexus writer received start document with experiment parameters {hyperion_params}"
             )
-            parameters = RotationScan.from_json(json_params)
+            parameters = RotationScan.model_validate_json(hyperion_params)
             NEXUS_LOGGER.info("Setting up nexus file...")
+
             det_size = (
                 parameters.detector_params.detector_size_constants.det_size_pixels
             )

mx_bluesky/hyperion/external_interaction/callbacks/xray_centre/ispyb_callback.py

@@ -5,8 +5,8 @@ from time import time
 from typing import TYPE_CHECKING, Any
 
 import numpy as np
-from blueapi.core import MsgGenerator
 from bluesky import preprocessors as bpp
+from bluesky.utils import MsgGenerator
 from dodal.devices.zocalo.zocalo_results import (
     ZOCALO_READING_PLAN_NAME,
     get_processing_results_from_event,
@@ -98,7 +98,9 @@ class GridscanISPyBCallback(BaseISPyBCallback):
                 "ISPyB callback received start document with experiment parameters and "
                 f"uid: {self.uid_to_finalize_on}"
             )
-            self.params = GridCommon.from_json(doc.get("hyperion_parameters"))
+            hyperion_params = doc.get("hyperion_parameters")
+            assert isinstance(hyperion_params, str)
+            self.params = GridCommon.model_validate_json(hyperion_params)
             self.ispyb = StoreInIspyb(self.ispyb_config)
             data_collection_group_info = populate_data_collection_group(self.params)
 
@@ -152,6 +154,7 @@ class GridscanISPyBCallback(BaseISPyBCallback):
         ISPYB_LOGGER.info(
             f"Amending comment based on Zocalo reading doc: {format_doc_for_log(doc)}"
         )
+
         raw_results = get_processing_results_from_event("zocalo", doc)
         if len(raw_results) > 0:
             for n, res in enumerate(raw_results):

mx_bluesky/hyperion/external_interaction/callbacks/xray_centre/nexus_callback.py

@@ -12,7 +12,7 @@ from mx_bluesky.hyperion.external_interaction.nexus.nexus_utils import (
 from mx_bluesky.hyperion.external_interaction.nexus.write_nexus import NexusWriter
 from mx_bluesky.hyperion.log import NEXUS_LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
-from mx_bluesky.hyperion.parameters.gridscan import ThreeDGridScan
+from mx_bluesky.hyperion.parameters.gridscan import HyperionThreeDGridScan
 
 if TYPE_CHECKING:
     from event_model.documents import Event, EventDescriptor, RunStart
@@ -45,11 +45,12 @@ class GridscanNexusFileCallback(PlanReactiveCallback):
 
     def activity_gated_start(self, doc: RunStart):
         if doc.get("subplan_name") == CONST.PLAN.GRIDSCAN_OUTER:
-            json_params = doc.get("hyperion_parameters")
+            hyperion_params = doc.get("hyperion_parameters")
+            assert isinstance(hyperion_params, str)
             NEXUS_LOGGER.info(
-                f"Nexus writer received start document with experiment parameters {json_params}"
+                f"Nexus writer received start document with experiment parameters {hyperion_params}"
             )
-            parameters = ThreeDGridScan.from_json(json_params)
+            parameters = HyperionThreeDGridScan.model_validate_json(hyperion_params)
             d_size = parameters.detector_params.detector_size_constants.det_size_pixels
             grid_n_img_1 = parameters.scan_indices[1]
             grid_n_img_2 = parameters.num_images - grid_n_img_1

mx_bluesky/hyperion/external_interaction/config_server.py

@@ -1,47 +1,18 @@
+from functools import cache
+
 from daq_config_server.client import ConfigServer
-from pydantic import BaseModel, Field, model_validator
 
+from mx_bluesky.common.external_interaction.config_server import FeatureFlags
 from mx_bluesky.hyperion.log import LOGGER
 from mx_bluesky.hyperion.parameters.constants import CONST
 
-_CONFIG_SERVER: ConfigServer | None = None
-
 
-def config_server() -> ConfigServer:
-    global _CONFIG_SERVER
-    if _CONFIG_SERVER is None:
-        _CONFIG_SERVER = ConfigServer(CONST.CONFIG_SERVER_URL, LOGGER)
-    return _CONFIG_SERVER
+class HyperionFeatureFlags(FeatureFlags):
+    @staticmethod
+    @cache
+    def get_config_server() -> ConfigServer:
+        return ConfigServer(CONST.CONFIG_SERVER_URL, LOGGER)
 
-
-class FeatureFlags(BaseModel):
-    # The default value will be used as the fallback when doing a best-effort fetch
-    # from the service
     use_panda_for_gridscan: bool = CONST.I03.USE_PANDA_FOR_GRIDSCAN
     compare_cpu_and_gpu_zocalo: bool = CONST.I03.COMPARE_CPU_AND_GPU_ZOCALO
     set_stub_offsets: bool = CONST.I03.SET_STUB_OFFSETS
-
-    # Feature values supplied at construction will override values from the config server
-    overriden_features: dict = Field(default_factory=dict, exclude=True)
-
-    @model_validator(mode="before")
-    @classmethod
-    def mark_overridden_features(cls, values):
-        assert isinstance(values, dict)
-        values["overriden_features"] = values.copy()
-        return values
-
-    @classmethod
-    def _get_flags(cls):
-        flags = config_server().best_effort_get_all_feature_flags()
-        return {f: flags[f] for f in flags if f in cls.model_fields.keys()}
-
-    def update_self_from_server(self):
-        """Used to update the feature flags from the server during a plan. Where there are flags which were explicitly set from externally supplied parameters, these values will be used instead."""
-        for flag, value in self._get_flags().items():
-            updated_value = (
-                value
-                if flag not in self.overriden_features.keys()
-                else self.overriden_features[flag]
-            )
-            setattr(self, flag, updated_value)