diff --git a/common/pyproject.toml b/common/pyproject.toml index 78094179..d5cc3433 100644 --- a/common/pyproject.toml +++ b/common/pyproject.toml @@ -7,12 +7,20 @@ requires-python = ">=3.11" dependencies = [ "pydantic==2.11.4", "numpy==2.2.5", - "jfjoch_client==1.0.0rc45" + "jfjoch_client==1.0.0rc45", + "aareDB==0.1.1a6" ] [lint] ignore = ["F401", "F541", "W503", "W504"] +[tool.uv.sources] +aaredb = { index = "psi"} + +[[tool.uv.index]] +name = "psi" +url = "https://gitea.psi.ch/api/packages/mx/pypi/simple" + [build-system] requires = ["setuptools>=75.6.0"] build-backend = "setuptools.build_meta" diff --git a/daq/pyproject.toml b/daq/pyproject.toml index ff6c7e10..72764df5 100644 --- a/daq/pyproject.toml +++ b/daq/pyproject.toml @@ -13,14 +13,22 @@ dependencies = [ "fastapi==0.115.12", "uvicorn==0.34.2", "ultralytics==8.3.133", - "aareDB==0.1.1a4", + "aareDB==0.1.1a6", "opencv-python-headless==4.11.0.86", - "python_multipart==0.0.20" + "python_multipart==0.0.20", + "websocket-client" ] [lint] ignore = ["F401", "F541", "W503", "W504"] +[tool.uv.sources] +aaredb = { index = "psi"} + +[[tool.uv.index]] +name = "psi" +url = "https://gitea.psi.ch/api/packages/mx/pypi/simple" + [build-system] requires = ["setuptools>=75.6.0"] build-backend = "setuptools.build_meta" diff --git a/daq/src/aaredaq/aaredb.py b/daq/src/aaredaq/aaredb.py index e8b7ff58..de7a57c4 100644 --- a/daq/src/aaredaq/aaredb.py +++ b/daq/src/aaredaq/aaredb.py @@ -1,15 +1,18 @@ import io +import os from typing import List, Optional import aareDBclient import cv2 import numpy as np import requests -from aareDBclient import SetTellPosition, SampleEventCreate, SetTellPositionRequest +from aareDBclient import ( + SetTellPosition, + SampleEventCreate, + SetTellPositionRequest ) from aaredaqlib.models import ( SampleShortInfo, PuckLoadedInfo, - PuckInfo, DewarAddress, SampleShortInfoList, ) @@ -21,38 +24,18 @@ class AareWrapper: def __init__( self, bl: MXBeamline, - host: str = "https://mx-aare-test.psi.ch:1492", + host: str = "https://mx-db-01.psi.ch/dispatcher", ): configuration = aareDBclient.Configuration(host=host) configuration.verify_ssl = False # Disable SSL verification self.client = aareDBclient.ApiClient(configuration) + self.client.default_headers["X-Shared-Password"] = os.getenv("AAREDB_SHARED_PASSWORD") self.__host = host - self.__puck_api = aareDBclient.PucksApi(self.client) - self.__sample_api = aareDBclient.SamplesApi(self.client) + self.__tell_api = aareDBclient.TellsRunnerApi(self.client) + self.__sample_api = aareDBclient.SamplesRunnerApi(self.client) self.__bl = bl - def get_pucks_beamline(self) -> List[PuckInfo]: - puck_list = self.__puck_api.get_pucks_by_slot_pucks_slot_slot_identifier_get( - slot_identifier=self.__bl.value.upper() - ) - ret = [] - for i in puck_list: - if i.tell_position is not None and len(i.tell_position) == 2: - d = DewarAddress(segment=i.tell_position[0], pos=i.tell_position[1]) - else: - d = None - ret.append( - PuckInfo( - db_id=i.id, - puck_name=i.puck_name, - dewar_name=i.dewar_name, - user=i.pgroup, - location=d, - ) - ) - return ret - def set_pucks_beamline(self, input_list: List[PuckLoadedInfo]): o = [] @@ -64,14 +47,14 @@ class AareWrapper: ) o.append(t) payload = SetTellPositionRequest(pucks = o, tell=self.__bl.value.upper()) - ret = self.__puck_api.set_tell_positions_pucks_set_tell_positions_put( + ret = self.__tell_api.set_tell_positions( set_tell_position_request=payload, ) print(ret) def get_sample_info(self) -> SampleShortInfoList: sample_list = ( - self.__puck_api.get_pucks_with_tell_position_pucks_with_tell_position_get(tell=self.__bl.value.upper()) + self.__tell_api.get_pucks_with_tell_position(tell=self.__bl.value.upper()) ) ret = [] @@ -102,7 +85,7 @@ class AareWrapper: def sample_mounted(self, s: Optional[SampleShortInfo]): if s is not None: try: - self.__sample_api.create_sample_event_samples_samples_sample_id_events_post( + self.__sample_api.create_sample_event( sample_id=s.db_id, sample_event_create=SampleEventCreate(event_type="Mounted"), ) @@ -112,7 +95,7 @@ class AareWrapper: def sample_unmounted(self, s: Optional[SampleShortInfo]): if s is not None: try: - self.__sample_api.create_sample_event_samples_samples_sample_id_events_post( + self.__sample_api.create_sample_event( sample_id=s.db_id, sample_event_create=SampleEventCreate(event_type="Unmounted"), ) diff --git a/daq/src/aaredaq/tellupdater.py b/daq/src/aaredaq/tellupdater.py new file mode 100644 index 00000000..5d712ce5 --- /dev/null +++ b/daq/src/aaredaq/tellupdater.py @@ -0,0 +1,116 @@ +import os +import json +import websocket +import threading +import time +from aareDBclient.models import PuckWithTellPosition +from mxlibs3.tell_client import TellClient +from aaredb import AareWrapper # Make sure the import path fits your project +from aaredaqlib.beamline import MXBeamline + +# Configuration +SLOT_IDENTIFIER = "X06DA" +WS_URL = f"wss://mx-db-01.psi.ch/dispatcher/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}" +WS_HEADERS = [f"X-Shared-Password: {os.getenv('AAREDB_SHARED_PASSWORD')}"] +POLL_INTERVAL = 30 # seconds + +# Initialize TELL client and DB wrapper +beamline = MXBeamline.X06DA # Use your beamline enum/value +tell_client = TellClient(bl=beamline) +aare_db = AareWrapper(bl=beamline) + +# Track current state +current_pucks = [] +last_pucks_sent = [] + +def compare_and_report_change(old, new, key_func): + old_ids = set(key_func(p) for p in old) + new_ids = set(key_func(p) for p in new) + joined = new_ids - old_ids + left = old_ids - new_ids + return joined, left + +def ws_update_samples_info(pucks): + """Send sample info to TELL robot.""" + tell_client.set_samples_info(pucks) + +def convert_pucks(puck_list): + return [PuckWithTellPosition( + id=p.puck_id, + puck_name=p.puck_name, + puck_type=getattr(p, 'puck_type', None), + puck_location_in_dewar=getattr(p, 'puck_location_in_dewar', None), + dewar_id=getattr(p, 'dewar_id', None), + dewar_name=getattr(p, 'dewar_name', None), + pgroup=getattr(p, 'pgroup', None), + samples=getattr(p, 'samples', None), + tell_position=getattr(p, 'tell_position', None), + ) for p in puck_list] + +def pucks_equal(pucks1, pucks2): + # Implement a simple equality check, e.g., comparing sorted ids, or full data + if len(pucks1) != len(pucks2): return False + return all(p1.id == p2.id for p1, p2 in zip(sorted(pucks1, key=lambda p: p.id), + sorted(pucks2, key=lambda p: p.id))) + +def timer_update_pucks_beamline(): + global last_pucks_sent + try: + detected_pucks = tell_client.get_detected_pucks() + aare_db.set_pucks_beamline(detected_pucks) + + except Exception as exc: + print(f"[ERROR][Timer] {exc}") + +def on_message(ws, message): + global current_pucks + try: + data = json.loads(message) + print("[WS][MESSAGE]", data) + if isinstance(data, list): + new_pucks = [PuckWithTellPosition(**item) for item in data] + joined, left = compare_and_report_change( + current_pucks, new_pucks, key_func=lambda p: p.id + ) + if joined or left: + print(f"[WS][UPDATE] joined={joined}, left={left}") + current_pucks = new_pucks + ws_update_samples_info(new_pucks) + else: + print("[WS][NO CHANGE]") + else: + print("[WS][ERROR] Unknown message format") + except Exception as exc: + print("[WS][ERROR] Failed to parse websocket message:", exc) + +def on_error(ws, error): + print("[WS][ERROR]", error) + +def on_close(ws, close_status_code, close_msg): + print("[WS][CLOSE]", close_status_code, close_msg) + +def on_open(ws): + print("[WS][OPEN] WebSocket opened.") + +def periodic_polling(): + while True: + timer_update_pucks_beamline() + time.sleep(POLL_INTERVAL) + +def main(): + ws = websocket.WebSocketApp( + WS_URL, + header=WS_HEADERS, + on_message=on_message, + on_error=on_error, + on_close=on_close, + on_open=on_open, + ) + + poll_thread = threading.Thread(target=periodic_polling, daemon=True) + poll_thread.start() + + ws.run_forever(sslopt={"cert_reqs": 0}) + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/daq/src/aaredaq/workflows.py b/daq/src/aaredaq/workflows.py index f08601cd..3a82394d 100644 --- a/daq/src/aaredaq/workflows.py +++ b/daq/src/aaredaq/workflows.py @@ -287,7 +287,7 @@ def sa2dh(devs: BeamlineDevices, cfg: BeamlineConfig): print(f"Error for unmounting: {e}") devs.tell.dry(wait_cold=-1, wait=False) # devs.tell.move_park(wait=True) - devs.tell.set_samples_info(devs.aare.get_pucks_beamline()) + #devs.tell.set_samples_info(devs.aare.get_pucks_beamline())# replaced by script stream def dh2sa(devs: BeamlineDevices, cfg: BeamlineConfig): diff --git a/daq/src/mxlibs3/tell_client.py b/daq/src/mxlibs3/tell_client.py index 97c29ce1..d7d05f66 100755 --- a/daq/src/mxlibs3/tell_client.py +++ b/daq/src/mxlibs3/tell_client.py @@ -12,9 +12,11 @@ from aaredaqlib.models import ( SampleShortInfo, DewarAddress, SampleDewarAddress, - PuckInfo, + #PuckInfo, +) +from aareDBclient import ( + PuckWithTellPosition, ) - from aaredaqlib.beamline import MXBeamline # noqa: F401 from mxlibs3.pshell_client import PShellClient @@ -191,24 +193,19 @@ class TellClient: ) return output - def set_samples_info(self, info: List[PuckInfo]): + def set_samples_info(self, info: List[PuckWithTellPosition]): if self.__simulation: return j = [] for x in info: - if x.location is None: - puck_address = "" - else: - puck_address = "{:1s}{:1d}".format(x.location.segment, x.location.pos) - j.append( { - "userName": x.user, - "dewarName": x.dewar_name, + "userName": x.pgroup, + "dewarName": x.dewar_name or "", "puckName": x.puck_name, - "puckType": "Unipuck", - "puckAddress": puck_address, + "puckType": "Unipuck", # could use x.puck_type + "puckAddress": x.tell_position or "", "puckBarcode": x.puck_name, "sampleBarcode": "", "sampleMountCount": 0,