mx-bluesky 1.4.6__py3-none-any.whl → 1.4.8__py3-none-any.whl

mx_bluesky/hyperion/experiment_plans/robot_load_then_centre_plan.py

@@ -3,7 +3,6 @@ from __future__ import annotations
 from math import isclose
 from typing import cast
 
-import bluesky.preprocessors as bpp
 import pydantic
 from blueapi.core import BlueskyContext
 from bluesky import plan_stubs as bps
@@ -11,13 +10,14 @@ from bluesky.utils import MsgGenerator
 from dodal.devices.aperturescatterguard import ApertureScatterguard
 from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
 from dodal.devices.backlight import Backlight
-from dodal.devices.dcm import DCM
 from dodal.devices.detector.detector_motion import DetectorMotion
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.fast_grid_scan import PandAFastGridScan, ZebraFastGridScan
 from dodal.devices.flux import Flux
 from dodal.devices.focusing_mirror import FocusingMirrorWithStripes, MirrorVoltages
-from dodal.devices.i03.beamstop import Beamstop
+from dodal.devices.i03 import Beamstop
+from dodal.devices.i03.dcm import DCM
+from dodal.devices.i03.undulator_dcm import UndulatorDCM
 from dodal.devices.motors import XYZPositioner
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.pin_image_recognition import PinTipDetection
@@ -27,7 +27,6 @@ from dodal.devices.smargon import Smargon
 from dodal.devices.synchrotron import Synchrotron
 from dodal.devices.thawer import Thawer
 from dodal.devices.undulator import Undulator
-from dodal.devices.undulator_dcm import UndulatorDCM
 from dodal.devices.webcam import Webcam
 from dodal.devices.xbpm_feedback import XBPMFeedback
 from dodal.devices.zebra.zebra import Zebra
@@ -37,14 +36,10 @@ from dodal.log import LOGGER
 from ophyd_async.fastcs.panda import HDFPanda
 
 from mx_bluesky.common.parameters.constants import OavConstants
-from mx_bluesky.common.xrc_result import XRayCentreEventHandler
 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.grid_detect_then_xray_centre_plan import (
     GridDetectThenXRayCentreComposite,
 )
@@ -117,6 +112,7 @@ def _flyscan_plan_from_robot_load_params(
     yield from pin_centre_then_flyscan_plan(
         cast(GridDetectThenXRayCentreComposite, composite),
         params.pin_centre_then_xray_centre_params,
+        oav_config_file,
     )
 
 
@@ -133,30 +129,10 @@ def _robot_load_then_flyscan_plan(
     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,
+    oav_config_file: str = OavConstants.OAV_CONFIG_JSON,
 ) -> MsgGenerator:
     """Perform pin-tip detection followed by a flyscan to determine centres of interest.
     Performs a robot load if necessary."""
@@ -180,10 +156,7 @@ def robot_load_then_xray_centre(
 
     if doing_sample_load:
         LOGGER.info("Pin not loaded, loading and centring")
-        plan = _robot_load_then_flyscan_plan(
-            composite,
-            parameters,
-        )
+        plan = _robot_load_then_flyscan_plan(composite, parameters, oav_config_file)
     else:
         # Robot load normally sets the energy so we should do this explicitly if no load is
         # being done
@@ -194,7 +167,9 @@ def robot_load_then_xray_centre(
         )
 
         if doing_chi_change:
-            plan = _flyscan_plan_from_robot_load_params(composite, parameters)
+            plan = _flyscan_plan_from_robot_load_params(
+                composite, parameters, oav_config_file
+            )
             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

@@ -10,11 +10,11 @@ from bluesky.utils import MsgGenerator
 from dodal.devices.aperturescatterguard import ApertureScatterguard
 from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
 from dodal.devices.backlight import Backlight
-from dodal.devices.dcm import DCM
 from dodal.devices.detector.detector_motion import DetectorMotion
 from dodal.devices.eiger import EigerDetector
 from dodal.devices.flux import Flux
-from dodal.devices.i03.beamstop import Beamstop
+from dodal.devices.i03 import Beamstop
+from dodal.devices.i03.dcm import DCM
 from dodal.devices.oav.oav_detector import OAV
 from dodal.devices.oav.oav_parameters import OAVParameters
 from dodal.devices.robot import BartRobot
@@ -26,20 +26,27 @@ from dodal.devices.xbpm_feedback import XBPMFeedback
 from dodal.devices.zebra.zebra import RotationDirection, Zebra
 from dodal.devices.zebra.zebra_controlled_shutter import ZebraShutter
 from dodal.plan_stubs.check_topup import check_topup_and_wait_if_necessary
+from dodal.plans.preprocessors.verify_undulator_gap import (
+    verify_undulator_gap_before_run_decorator,
+)
 
+from mx_bluesky.common.device_setup_plans.manipulate_sample import (
+    cleanup_sample_environment,
+    move_phi_chi_omega,
+    move_x_y_z,
+    setup_sample_environment,
+)
+from mx_bluesky.common.parameters.components import WithSnapshot
 from mx_bluesky.common.plans.read_hardware import (
     read_hardware_for_zocalo,
     standard_read_hardware_during_collection,
     standard_read_hardware_pre_collection,
 )
+from mx_bluesky.common.preprocessors.preprocessors import (
+    transmission_and_xbpm_feedback_for_collection_decorator,
+)
 from mx_bluesky.common.utils.context import device_composite_from_context
 from mx_bluesky.common.utils.log import LOGGER
-from mx_bluesky.hyperion.device_setup_plans.manipulate_sample import (
-    cleanup_sample_environment,
-    move_phi_chi_omega,
-    move_x_y_z,
-    setup_sample_environment,
-)
 from mx_bluesky.hyperion.device_setup_plans.setup_zebra import (
     arm_zebra,
     setup_zebra_for_rotation,
@@ -48,9 +55,6 @@ from mx_bluesky.hyperion.device_setup_plans.setup_zebra import (
 from mx_bluesky.hyperion.device_setup_plans.utils import (
     start_preparing_data_collection_then_do_plan,
 )
-from mx_bluesky.hyperion.device_setup_plans.xbpm_feedback import (
-    transmission_and_xbpm_feedback_for_collection_decorator,
-)
 from mx_bluesky.hyperion.experiment_plans.oav_snapshot_plan import (
     OavSnapshotComposite,
     oav_snapshot_plan,
@@ -352,66 +356,30 @@ def _move_and_rotation(
                 group=CONST.WAIT.ROTATION_READY_FOR_DC,
             )
         yield from oav_snapshot_plan(composite, params, oav_params)
-    yield from rotation_scan_plan(
-        composite,
-        params,
-        motion_values,
-    )
+    yield from rotation_scan_plan(composite, params, motion_values)
 
 
-def rotation_scan(
+def multi_rotation_scan(
     composite: RotationScanComposite,
-    parameters: RotationScan,
+    parameters: MultiRotationScan,
     oav_params: OAVParameters | None = None,
 ) -> MsgGenerator:
-    parameters.features.update_self_from_server()
-
-    if not oav_params:
-        oav_params = OAVParameters(context="xrayCentring")
-
-    @bpp.set_run_key_decorator("rotation_scan")
-    @bpp.run_decorator(  # attach experiment metadata to the start document
+    @bpp.set_run_key_decorator(CONST.PLAN.ROTATION_MULTI_OUTER)
+    @bpp.run_decorator(
         md={
-            "subplan_name": CONST.PLAN.ROTATION_OUTER,
-            "mx_bluesky_parameters": parameters.model_dump_json(),
-            "activate_callbacks": [
-                "RotationISPyBCallback",
-                "RotationNexusFileCallback",
-            ],
+            "activate_callbacks": ["BeamDrawingCallback"],
+            "with_snapshot": parameters.model_dump_json(
+                include=WithSnapshot.model_fields.keys()  # type: ignore
+            ),
         }
     )
-    @transmission_and_xbpm_feedback_for_collection_decorator(
-        composite.undulator,
-        composite.xbpm_feedback,
-        composite.attenuator,
-        composite.dcm,
-        parameters.transmission_frac,
-    )
-    def rotation_scan_plan_with_stage_and_cleanup(
-        params: RotationScan,
-    ):
-        eiger: EigerDetector = composite.eiger
-        eiger.set_detector_parameters(params.detector_params)
-
-        @bpp.finalize_decorator(lambda: _cleanup_plan(composite))
-        def rotation_with_cleanup_and_stage(params: RotationScan):
-            yield from _move_and_rotation(composite, params, oav_params)
-
-        LOGGER.info("setting up and staging eiger...")
-        yield from start_preparing_data_collection_then_do_plan(
-            composite.beamstop,
-            eiger,
-            composite.detector_motion,
-            params.detector_distance_mm,
-            rotation_with_cleanup_and_stage(params),
-            group=CONST.WAIT.ROTATION_READY_FOR_DC,
-        )
-        yield from bps.unstage(eiger)
+    def _wrapped_multi_rotation_scan():
+        yield from multi_rotation_scan_internal(composite, parameters, oav_params)
 
-    yield from rotation_scan_plan_with_stage_and_cleanup(parameters)
+    yield from _wrapped_multi_rotation_scan()
 
 
-def multi_rotation_scan(
+def multi_rotation_scan_internal(
     composite: RotationScanComposite,
     parameters: MultiRotationScan,
     oav_params: OAVParameters | None = None,
@@ -422,6 +390,10 @@ def multi_rotation_scan(
     eiger: EigerDetector = composite.eiger
     eiger.set_detector_parameters(parameters.detector_params)
 
+    @transmission_and_xbpm_feedback_for_collection_decorator(
+        composite,
+        parameters.transmission_frac,
+    )
     @bpp.set_run_key_decorator("multi_rotation_scan")
     @bpp.run_decorator(
         md={
@@ -434,17 +406,11 @@ def multi_rotation_scan(
             ],
         }
     )
-    @transmission_and_xbpm_feedback_for_collection_decorator(
-        composite.undulator,
-        composite.xbpm_feedback,
-        composite.attenuator,
-        composite.dcm,
-        parameters.transmission_frac,
-    )
     @bpp.finalize_decorator(lambda: _cleanup_plan(composite))
     def _multi_rotation_scan():
         for single_scan in parameters.single_rotation_scans:
 
+            @verify_undulator_gap_before_run_decorator(composite)
             @bpp.set_run_key_decorator("rotation_scan")
             @bpp.run_decorator(  # attach experiment metadata to the start document
                 md={
@@ -459,6 +425,8 @@ def multi_rotation_scan(
 
             yield from rotation_scan_core(single_scan)
 
+        yield from bps.unstage(eiger)
+
     LOGGER.info("setting up and staging eiger...")
     yield from start_preparing_data_collection_then_do_plan(
         composite.beamstop,
@@ -468,4 +436,3 @@ def multi_rotation_scan(
         _multi_rotation_scan(),
         group=CONST.WAIT.ROTATION_READY_FOR_DC,
     )
-    yield from bps.unstage(eiger)

mx_bluesky/hyperion/experiment_plans/set_energy_plan.py

@@ -5,18 +5,21 @@
 * reenable feedback
 """
 
+import bluesky.preprocessors as bpp
 import pydantic
 from bluesky import plan_stubs as bps
 from dodal.devices.attenuator.attenuator import BinaryFilterAttenuator
-from dodal.devices.dcm import DCM
 from dodal.devices.focusing_mirror import FocusingMirrorWithStripes, MirrorVoltages
-from dodal.devices.undulator_dcm import UndulatorDCM
+from dodal.devices.i03.dcm import DCM
+from dodal.devices.i03.undulator_dcm import UndulatorDCM
+from dodal.devices.undulator import Undulator
 from dodal.devices.xbpm_feedback import XBPMFeedback
 
-from mx_bluesky.hyperion.device_setup_plans import dcm_pitch_roll_mirror_adjuster
-from mx_bluesky.hyperion.device_setup_plans.xbpm_feedback import (
+from mx_bluesky.common.parameters.constants import PlanNameConstants
+from mx_bluesky.common.preprocessors.preprocessors import (
     transmission_and_xbpm_feedback_for_collection_wrapper,
 )
+from mx_bluesky.hyperion.device_setup_plans import dcm_pitch_roll_mirror_adjuster
 
 DESIRED_TRANSMISSION_FRACTION = 0.1
 
@@ -33,6 +36,17 @@ class SetEnergyComposite:
     attenuator: BinaryFilterAttenuator
 
 
+# Remove composite after https://github.com/DiamondLightSource/dodal/issues/1092
+@pydantic.dataclasses.dataclass(config={"arbitrary_types_allowed": True})
+class XBPMWrapperComposite:
+    undulator: Undulator
+    xbpm_feedback: XBPMFeedback
+    attenuator: BinaryFilterAttenuator
+    dcm: DCM
+
+
+@bpp.set_run_key_decorator(PlanNameConstants.SET_ENERGY)
+@bpp.run_decorator()
 def _set_energy_plan(
     energy_kev,
     composite: SetEnergyComposite,
@@ -51,12 +65,17 @@ def set_energy_plan(
     energy_ev: float | None,
     composite: SetEnergyComposite,
 ):
+    # Remove conversion after https://github.com/DiamondLightSource/dodal/issues/1092
+    composite_for_wrapper = XBPMWrapperComposite(
+        composite.undulator_dcm.undulator_ref._obj,  # noqa: SLF001
+        composite.xbpm_feedback,
+        composite.attenuator,
+        composite.dcm,
+    )
+
     if energy_ev:
         yield from transmission_and_xbpm_feedback_for_collection_wrapper(
             _set_energy_plan(energy_ev / 1000, composite),
-            composite.undulator_dcm.undulator_ref(),
-            composite.xbpm_feedback,
-            composite.attenuator,
-            composite.dcm,
+            composite_for_wrapper,
             DESIRED_TRANSMISSION_FRACTION,
         )

mx_bluesky/hyperion/external_interaction/agamemnon.py

@@ -1,14 +1,25 @@
 import dataclasses
 import json
 import re
-from typing import TypeVar
+import traceback
+from os import path
+from typing import Any, TypeVar
 
 import requests
+from deepdiff.diff import DeepDiff
 from dodal.utils import get_beamline_name
-
-from mx_bluesky.common.parameters.components import WithVisit
-from mx_bluesky.common.parameters.constants import GridscanParamConstants
+from jsonschema import ValidationError
+from pydantic_extra_types.semantic_version import SemanticVersion
+
+from mx_bluesky.common.parameters.components import (
+    PARAMETER_VERSION,
+    WithVisit,
+)
+from mx_bluesky.common.parameters.constants import (
+    GridscanParamConstants,
+)
 from mx_bluesky.common.utils.log import LOGGER
+from mx_bluesky.common.utils.utils import convert_angstrom_to_eV
 from mx_bluesky.hyperion.parameters.load_centre_collect import LoadCentreCollect
 
 T = TypeVar("T", bound=WithVisit)
@@ -16,6 +27,7 @@ AGAMEMNON_URL = "http://agamemnon.diamond.ac.uk/"
 MULTIPIN_PREFIX = "multipin"
 MULTIPIN_FORMAT_DESC = "Expected multipin format is multipin_{number_of_wells}x{well_size}+{distance_between_tip_and_first_well}"
 MULTIPIN_REGEX = rf"^{MULTIPIN_PREFIX}_(\d+)x(\d+(?:\.\d+)?)\+(\d+(?:\.\d+)?)$"
+MX_GENERAL_ROOT_REGEX = r"^/dls/(?P<beamline>[^/]+)/data/[^/]*/(?P<visit>[^/]+)(?:/|$)"
 
 
 @dataclasses.dataclass
@@ -60,7 +72,7 @@ def _get_parameters_from_url(url: str) -> dict:
         raise KeyError(f"Unexpected json from agamemnon: {response_json}") from e
 
 
-def _get_pin_type_from_agamemnon_parameters(parameters: dict) -> PinType:
+def get_pin_type_from_agamemnon_parameters(parameters: dict) -> PinType:
     loop_type_name: str | None = parameters["sample"]["loopType"]
     if loop_type_name:
         regex_search = re.search(MULTIPIN_REGEX, loop_type_name)
@@ -81,15 +93,130 @@ def get_next_instruction(beamline: str) -> dict:
     return _get_parameters_from_url(AGAMEMNON_URL + f"getnextcollect/{beamline}")
 
 
-def get_pin_type_from_agamemnon(beamline: str) -> PinType:
-    params = get_next_instruction(beamline)
-    return _get_pin_type_from_agamemnon_parameters(params)
+def get_withvisit_parameters_from_agamemnon(parameters: dict) -> tuple:
+    try:
+        prefix = parameters["prefix"]
+        collection = parameters["collection"]
+        # Assuming distance is identical for multiple collections. Remove after https://github.com/DiamondLightSource/mx-bluesky/issues/773
+        detector_distance = collection[0]["distance"]
+    except KeyError as e:
+        raise KeyError("Unexpected json from agamemnon") from e
+
+    match = re.match(MX_GENERAL_ROOT_REGEX, prefix) if prefix else None
+
+    if match:
+        return (match.group("visit"), detector_distance)
+
+    raise ValueError(
+        f"Agamemnon prefix '{prefix}' does not match MX-General root structure"
+    )
+
+
+def get_withenergy_parameters_from_agamemnon(parameters: dict) -> dict[str, Any]:
+    try:
+        first_collection: dict = parameters["collection"][0]
+        wavelength = first_collection.get("wavelength")
+        assert isinstance(wavelength, float)
+        demand_energy_ev = convert_angstrom_to_eV(wavelength)
+        return {"demand_energy_ev": demand_energy_ev}
+    except (KeyError, IndexError, AttributeError, TypeError):
+        return {"demand_energy_ev": None}
+
+
+def get_param_version() -> SemanticVersion:
+    return SemanticVersion.validate_from_str(str(PARAMETER_VERSION))
+
+
+def populate_parameters_from_agamemnon(agamemnon_params):
+    visit, detector_distance = get_withvisit_parameters_from_agamemnon(agamemnon_params)
+    with_energy_params = get_withenergy_parameters_from_agamemnon(agamemnon_params)
+    pin_type = get_pin_type_from_agamemnon_parameters(agamemnon_params)
+    first_collection = agamemnon_params["collection"][0]
+    visit_directory, file_name = path.split(agamemnon_params["prefix"])
+    return LoadCentreCollect.model_validate(
+        {
+            "parameter_model_version": get_param_version(),
+            "visit": visit,
+            "detector_distance_mm": detector_distance,
+            "sample_id": agamemnon_params["sample"]["id"],
+            "sample_puck": agamemnon_params["sample"]["container"],
+            "sample_pin": agamemnon_params["sample"]["position"],
+            "select_centres": {
+                "name": "TopNByMaxCount",
+                "n": pin_type.expected_number_of_crystals,
+            },
+            "robot_load_then_centre": {
+                "storage_directory": str(visit_directory) + "/xraycentring",
+                "file_name": file_name,
+                "tip_offset_um": pin_type.full_width / 2,
+                "grid_width_um": pin_type.full_width,
+                "omega_start_deg": 0.0,
+                "chi_start_deg": first_collection["chi"],
+                "transmission_frac": 1.0,
+                "features": {"use_gpu_results": True},
+                **with_energy_params,
+            },
+            "multi_rotation_scan": {
+                "comment": first_collection["comment"],
+                "storage_directory": str(visit_directory),
+                "exposure_time_s": first_collection["exposure_time"],
+                "file_name": file_name,
+                "transmission_frac": first_collection["transmission"],
+                "rotation_increment_deg": first_collection["omega_increment"],
+                "ispyb_experiment_type": first_collection["experiment_type"],
+                "snapshot_omegas_deg": [0.0, 90.0, 180.0, 270.0],
+                "rotation_scans": [
+                    {
+                        "scan_width_deg": (
+                            first_collection["number_of_images"]
+                            * first_collection["omega_increment"]
+                        ),
+                        "omega_start_deg": first_collection["omega_start"],
+                        "phi_start_deg": first_collection["phi_start"],
+                        "chi_start_deg": first_collection["chi"],
+                        "rotation_direction": "Positive",
+                    }
+                ],
+                **with_energy_params,
+            },
+        }
+    )
+
+
+def create_parameters_from_agamemnon() -> LoadCentreCollect | None:
+    beamline_name = get_beamline_name("i03")
+    agamemnon_params = get_next_instruction(beamline_name)
+    return (
+        populate_parameters_from_agamemnon(agamemnon_params)
+        if agamemnon_params
+        else None
+    )
+
+
+def compare_params(load_centre_collect_params):
+    try:
+        parameters = create_parameters_from_agamemnon()
+        # Log differences against GDA populated parameters
+        differences = DeepDiff(
+            parameters, load_centre_collect_params, math_epsilon=1e-5
+        )
+        if differences:
+            LOGGER.info(
+                f"Different parameters found when directly reading from Hyperion: {differences}"
+            )
+    except (ValueError, KeyError):
+        LOGGER.warning(f"Failed to compare parameters: {traceback.format_exc()}")
+    except Exception:
+        LOGGER.warning(
+            f"Unexpected error occurred. Failed to compare parameters: {traceback.format_exc()}"
+        )
 
 
 def update_params_from_agamemnon(parameters: T) -> T:
     try:
         beamline_name = get_beamline_name("i03")
-        pin_type = get_pin_type_from_agamemnon(beamline_name)
+        agamemnon_params = get_next_instruction(beamline_name)
+        pin_type = get_pin_type_from_agamemnon_parameters(agamemnon_params)
         if isinstance(parameters, LoadCentreCollect):
             parameters.robot_load_then_centre.tip_offset_um = pin_type.full_width / 2
             parameters.robot_load_then_centre.grid_width_um = pin_type.full_width
@@ -99,6 +226,9 @@ def update_params_from_agamemnon(parameters: T) -> T:
                 # Before we do https://github.com/DiamondLightSource/mx-bluesky/issues/226
                 # this will give no snapshots but that's preferable
                 parameters.multi_rotation_scan.snapshot_omegas_deg = []
+    except (ValueError, ValidationError) as e:
+        LOGGER.warning(f"Failed to update parameters: {e}")
     except Exception as e:
-        LOGGER.warning(f"Failed to get pin type from agamemnon, using single pin {e}")
+        LOGGER.warning(f"Unexpected error occurred. Failed to update parameters: {e}")
+
     return parameters

mx_bluesky/hyperion/external_interaction/callbacks/__main__.py

@@ -1,7 +1,7 @@
 import logging
 from collections.abc import Callable, Sequence
 from threading import Thread
-from time import sleep
+from time import sleep  # noqa
 
 from bluesky.callbacks import CallbackBase
 from bluesky.callbacks.zmq import Proxy, RemoteDispatcher
@@ -14,6 +14,9 @@ from mx_bluesky.common.external_interaction.callbacks.common.log_uid_tag_callbac
 from mx_bluesky.common.external_interaction.callbacks.common.zocalo_callback import (
     ZocaloCallback,
 )
+from mx_bluesky.common.external_interaction.callbacks.sample_handling.sample_handling_callback import (
+    SampleHandlingCallback,
+)
 from mx_bluesky.common.external_interaction.callbacks.xray_centre.ispyb_callback import (
     GridscanISPyBCallback,
 )
@@ -23,7 +26,7 @@ from mx_bluesky.common.external_interaction.callbacks.xray_centre.nexus_callback
 from mx_bluesky.common.utils.log import (
     ISPYB_ZOCALO_CALLBACK_LOGGER,
     NEXUS_LOGGER,
-    _get_logging_dir,
+    _get_logging_dirs,
     tag_filter,
 )
 from mx_bluesky.hyperion.external_interaction.callbacks.robot_load.ispyb_callback import (
@@ -35,8 +38,8 @@ from mx_bluesky.hyperion.external_interaction.callbacks.rotation.ispyb_callback
 from mx_bluesky.hyperion.external_interaction.callbacks.rotation.nexus_callback import (
     RotationNexusFileCallback,
 )
-from mx_bluesky.hyperion.external_interaction.callbacks.sample_handling.sample_handling_callback import (
-    SampleHandlingCallback,
+from mx_bluesky.hyperion.external_interaction.callbacks.snapshot_callback import (
+    BeamDrawingCallback,
 )
 from mx_bluesky.hyperion.parameters.cli import parse_callback_dev_mode_arg
 from mx_bluesky.hyperion.parameters.constants import CONST
@@ -67,15 +70,18 @@ def create_rotation_callbacks() -> tuple[
     return (
         RotationNexusFileCallback(),
         RotationISPyBCallback(
-            emit=ZocaloCallback(CONST.PLAN.ROTATION_MAIN, CONST.ZOCALO_ENV)
+            emit=ZocaloCallback(CONST.PLAN.ROTATION_MULTI, CONST.ZOCALO_ENV)
         ),
     )
 
 
 def setup_callbacks() -> list[CallbackBase]:
+    rot_nexus_cb, rot_ispyb_cb = create_rotation_callbacks()
+    snapshot_cb = BeamDrawingCallback(emit=rot_ispyb_cb)
     return [
         *create_gridscan_callbacks(),
-        *create_rotation_callbacks(),
+        rot_nexus_cb,
+        snapshot_cb,
         LogUidTaggingCallback(),
         RobotLoadISPyBCallback(),
         SampleHandlingCallback(),
@@ -87,14 +93,16 @@ def setup_logging(dev_mode: bool):
         (ISPYB_ZOCALO_CALLBACK_LOGGER, "hyperion_ispyb_callback.log"),
         (NEXUS_LOGGER, "hyperion_nexus_callback.log"),
     ]:
+        logging_path, debug_logging_path = _get_logging_dirs()
         if logger.handlers == []:
             handlers = set_up_all_logging_handlers(
                 logger,
-                _get_logging_dir(),
+                logging_path,
                 filename,
                 dev_mode,
-                error_log_buffer_lines=ERROR_LOG_BUFFER_LINES,
-                graylog_port=CONST.GRAYLOG_PORT,
+                ERROR_LOG_BUFFER_LINES,
+                CONST.GRAYLOG_PORT,
+                debug_logging_path,
             )
             handlers["graylog_handler"].addFilter(tag_filter)
     log_info(f"Loggers initialised with dev_mode={dev_mode}")