Add processing feedback handling for smartem

src/murfey/util/config.py

@@ -1,12 +1,14 @@
 from __future__ import annotations
 
+import configparser
 import copy
 import os
 import socket
 from functools import lru_cache
 from importlib.metadata import entry_points
 from pathlib import Path
 from typing import Any, Literal, Optional
+from urllib.parse import quote
 
 import yaml
 from pydantic import BaseModel, ConfigDict, RootModel, ValidationInfo, field_validator
@@ -373,3 +375,27 @@ def get_smartem_keycloak_client():
                 load_keycloak_config(Path(keycloak_config))
             )
     return keycloak_client
+
+
+@lru_cache(maxsize=1)
+def get_rabbitmq_url() -> str:
+    rabbitmq_defaults = {
+        "host": "localhost",
+        "port": "5672",
+        "username": "guest",
+        "password": "guest",
+        "vhost": "/",
+    }
+    rabbitmq_credentials_file = get_security_config().rabbitmq_credentials
+    cfgparser = configparser.ConfigParser(allow_no_value=True)
+    if rabbitmq_credentials_file:
+        cfgparser.read(rabbitmq_credentials_file)
+        rabbitmq_creds = (
+            {**rabbitmq_defaults, **cfgparser["rabbit"]}
+            if cfgparser.has_section("rabbit")
+            else rabbitmq_defaults
+        )
+    user = quote(rabbitmq_creds["username"], safe="")
+    password = quote(rabbitmq_creds["password"], safe="")
+    vhost = quote(rabbitmq_creds["vhost"], safe="")
+    return f"amqp://{user}:{password}@{rabbitmq_creds['host']}:{rabbitmq_creds['port']}/{vhost}"

src/murfey/workflows/spa/ctf_estimation.py

@@ -1,8 +1,9 @@
+import asyncio
 from logging import getLogger
 
 from sqlmodel import Session, select
 
-from murfey.util.config import get_machine_config
+from murfey.util.config import get_machine_config, get_rabbitmq_url
 from murfey.util.db import (
     Movie,
     Session as MurfeySession,
@@ -20,6 +21,11 @@
         MicrographResponse,
         ProcessingFeedbackPublishResponse,
     )
+    from smartem_backend.model.mq_event import (
+        CtfCompleteBody,
+        MessageQueueEventType,
+    )
+    from smartem_backend.rmq.publisher import AioPikaPublisher
     from smartem_common.entity_status import MicrographStatus
 
     from murfey.util.config import get_smartem_keycloak_client
@@ -37,7 +43,7 @@ def ctf_estimated(message: dict, murfey_db: Session) -> dict[str, bool]:
     if not SMARTEM_ACTIVE:
         return {"success": True}
     movie = murfey_db.exec(
-        select(Movie).where(Movie.murfey_id == message["motion_correction_id"])
+        select(Movie).where(Movie.murfey_id == message["mc_uuid"])
     ).one()
     if movie.smartem_uuid:
         try:
@@ -69,6 +75,40 @@ def ctf_estimated(message: dict, murfey_db: Session) -> dict[str, bool]:
                     registered_request,
                     ProcessingFeedbackPublishResponse,
                 )
+
+                async def _publish_ctf_completed(
+                    micrograph_uuid: str, ctf_max_resolution: float
+                ) -> None:
+                    publisher = AioPikaPublisher(
+                        url=get_rabbitmq_url(),
+                        exchange_name="smartem",
+                        routing_key="smartem",
+                        exchange_type="fanout",
+                    )
+                    await publisher.connect()
+                    try:
+                        await publisher.publish_event(
+                            MessageQueueEventType.CTF_COMPLETE,
+                            CtfCompleteBody(
+                                event_type=MessageQueueEventType.CTF_COMPLETE,
+                                micrograph_uuid=micrograph_uuid,
+                                ctf_max_resolution_estimate=ctf_max_resolution,
+                            ),
+                        )
+                    except Exception:
+                        logger.warning(
+                            f"smartem failed to ctf estimation completion {micrograph_uuid}",
+                            exc_info=True,
+                        )
+                    finally:
+                        await publisher.close()
+
+                asyncio.run(
+                    _publish_ctf_completed(
+                        movie.smartem_uuid, message.get("ctf_max_resolution", 1000)
+                    )
+                )
+
         except Exception:
             logger.warning(
                 "Failed to emit CTF estimation complete event to smartem",

src/murfey/workflows/spa/motion_correction.py

@@ -1,8 +1,9 @@
+import asyncio
 from logging import getLogger
 
 from sqlmodel import Session, select
 
-from murfey.util.config import get_machine_config
+from murfey.util.config import get_machine_config, get_rabbitmq_url
 from murfey.util.db import (
     Movie,
     Session as MurfeySession,
@@ -20,6 +21,11 @@
         MicrographResponse,
         ProcessingFeedbackPublishResponse,
     )
+    from smartem_backend.model.mq_event import (
+        MessageQueueEventType,
+        MotionCorrectionCompleteBody,
+    )
+    from smartem_backend.rmq.publisher import AioPikaPublisher
     from smartem_common.entity_status import MicrographStatus
 
     from murfey.util.config import get_smartem_keycloak_client
@@ -37,7 +43,7 @@ def motion_corrected(message: dict, murfey_db: Session) -> dict[str, bool]:
     if not SMARTEM_ACTIVE:
         return {"success": True}
     movie = murfey_db.exec(
-        select(Movie).where(Movie.murfey_id == message["motion_correction_id"])
+        select(Movie).where(Movie.murfey_id == message["mc_uuid"])
     ).one()
     if movie.smartem_uuid:
         try:
@@ -71,6 +77,43 @@ def motion_corrected(message: dict, murfey_db: Session) -> dict[str, bool]:
                     registered_request,
                     ProcessingFeedbackPublishResponse,
                 )
+
+                async def _publish_motion_correction_completed(
+                    micrograph_uuid: str, total_motion: float, average_motion: float
+                ) -> None:
+                    publisher = AioPikaPublisher(
+                        url=get_rabbitmq_url(),
+                        exchange_name="smartem",
+                        routing_key="smartem",
+                        exchange_type="fanout",
+                    )
+                    await publisher.connect()
+                    try:
+                        await publisher.publish_event(
+                            MessageQueueEventType.MOTION_CORRECTION_COMPLETE,
+                            MotionCorrectionCompleteBody(
+                                event_type=MessageQueueEventType.MOTION_CORRECTION_COMPLETE,
+                                micrograph_uuid=micrograph_uuid,
+                                total_motion=total_motion,
+                                average_motion=average_motion,
+                            ),
+                        )
+                    except Exception:
+                        logger.warning(
+                            f"smartem failed to motion correction completion {micrograph_uuid}",
+                            exc_info=True,
+                        )
+                    finally:
+                        await publisher.close()
+
+                asyncio.run(
+                    _publish_motion_correction_completed(
+                        movie.smartem_uuid,
+                        message.get("total_motion", 1000),
+                        message.get("average_motion", 1000),
+                    )
+                )
+
         except Exception:
             logger.warning(
                 "Failed to emit motion correction complete event to smartem",

src/murfey/workflows/spa/picking.py

@@ -1,3 +1,4 @@
+import asyncio
 from logging import getLogger
 from typing import List
 
@@ -11,7 +12,7 @@
     _app_id,
     _pj_id,
 )
-from murfey.util.config import get_machine_config
+from murfey.util.config import get_machine_config, get_rabbitmq_url
 from murfey.util.db import (
     AutoProcProgram,
     ClassificationFeedbackParameters,
@@ -33,6 +34,11 @@
     from smartem_backend.api_client import SmartEMAPIClient
     from smartem_backend.model.http_request import MicrographUpdateRequest
     from smartem_backend.model.http_response import MicrographResponse
+    from smartem_backend.model.mq_event import (
+        MessageQueueEventType,
+        ParticlePickingCompleteBody,
+    )
+    from smartem_backend.rmq.publisher import AioPikaPublisher
     from smartem_common.entity_status import MicrographStatus
 
     from murfey.util.config import get_smartem_keycloak_client
@@ -411,6 +417,45 @@ def particles_picked(message: dict, murfey_db: Session) -> dict[str, bool]:
                     update,
                     MicrographResponse,
                 )
+
+                async def _publish_particle_picking_completed(
+                    micrograph_uuid: str, number_of_particles_picked: int
+                ) -> None:
+                    publisher = AioPikaPublisher(
+                        url=get_rabbitmq_url(),
+                        exchange_name="smartem",
+                        routing_key="smartem",
+                        exchange_type="fanout",
+                    )
+                    await publisher.connect()
+                    try:
+                        await publisher.publish_event(
+                            MessageQueueEventType.PARTICLE_PICKING_COMPLETE,
+                            ParticlePickingCompleteBody(
+                                event_type=MessageQueueEventType.PARTICLE_PICKING_COMPLETE,
+                                micrograph_uuid=micrograph_uuid,
+                                number_of_particles_picked=number_of_particles_picked,
+                            ),
+                        )
+                    except Exception:
+                        logger.warning(
+                            f"smartem failed to publish picking completion {micrograph_uuid}",
+                            exc_info=True,
+                        )
+                    finally:
+                        await publisher.close()
+
+                number_of_particles_picked = message.get("particle_count")
+                if number_of_particles_picked is None:
+                    number_of_particles_picked = len(
+                        message.get("particle_diameters") or []
+                    )
+                asyncio.run(
+                    _publish_particle_picking_completed(
+                        movie.smartem_uuid, number_of_particles_picked
+                    )
+                )
+
         except Exception:
             logger.warning(
                 "Failed to emit particle picking complete event to smartem",