From 8a7a4d7fc1d81a219dad5a3cead0030463e61a41 Mon Sep 17 00:00:00 2001 From: GotthardG <51994228+GotthardG@users.noreply.github.com> Date: Thu, 22 Jan 2026 10:19:07 +0100 Subject: [PATCH] DAQ: tellupdater script and spreadsheetupdater updated for X10SA --- daq/src/aaredaq/spreadsheetupdater.py | 98 ++++++++++++++++----------- daq/src/aaredaq/tellupdater.py | 93 ++++++++++++++++++------- 2 files changed, 128 insertions(+), 63 deletions(-) diff --git a/daq/src/aaredaq/spreadsheetupdater.py b/daq/src/aaredaq/spreadsheetupdater.py index 1303c2d9..110e0002 100644 --- a/daq/src/aaredaq/spreadsheetupdater.py +++ b/daq/src/aaredaq/spreadsheetupdater.py @@ -2,13 +2,12 @@ import os import json import websocket import time -import signal -from aareDBclient.models import PuckWithTellPosition, DataCollectionParameters +from aareDBclient.models import PuckWithTellPosition from aaredaqlib.models import SampleShortInfoList, SampleShortInfo, DewarAddress from config import BeamlineConfig from aaredaqlib.beamline import MXBeamline -SLOT_IDENTIFIER = "X06DA" +SLOT_IDENTIFIER = "X10SA" WS_URL = f"wss://mx-db-01.psi.ch/dispatcher/protected_router/tell_runner/ws/samples-spreadsheet/{SLOT_IDENTIFIER}" # Ensure the environment variable for the shared password is set @@ -17,7 +16,7 @@ if not password: raise ValueError("The AAREDB_SHARED_PASSWORD environment variable is not set.") WS_HEADERS = [f"X-Shared-Password: {password}"] -beamline = MXBeamline.X06DA +beamline = MXBeamline.X10SA config = BeamlineConfig(bl=beamline) # Cache the current spreadsheet for change detection @@ -40,55 +39,74 @@ def on_message(ws, message): try: data = json.loads(message) - # Process the pucks data from the "samples" key in the message - if isinstance(data, dict) and "samples" in data: - pucks_data = data["samples"] - else: - pucks_data = data # fallback: might be a list directly - - # Convert dictionary data into PuckWithTellPosition Pydantic model instances + pucks_data = data["samples"] if isinstance(data, dict) and "samples" in data else data pucks = [PuckWithTellPosition(**item) for item in pucks_data] - sample_short_infos = [] + print(f"INFO: Received pucks: {[p.puck_name for p in pucks]}") + + normal_short_infos = [] + reference_short_infos = [] for p in pucks: - # Iterate through samples in the puck; `s` is a Pydantic `Sample` object for s in p.samples or []: - # Handle dewar_address dewar_address = None - if p.tell_position is not None and len(p.tell_position) == 2: - dewar_address = DewarAddress( - segment=p.tell_position[0], pos=p.tell_position[1] - ) + target_list = None - aaredb_params = s.data_collection_parameters + # Reference tool + if isinstance(p.tell_position, str) and p.tell_position.startswith("X"): + dewar_address = DewarAddress(segment="X", pos=int(p.tell_position[1:]) if len(p.tell_position) > 1 else 1) + target_list = reference_short_infos - # Build SampleShortInfo using attribute access. - sample_short_infos.append( - SampleShortInfo( - db_id=s.id, - puck_name=p.puck_name, - dewar_name="Unknown" if p.dewar_name is None else p.dewar_name, - sample_name=s.sample_name, - user=s.pgroup, - pin=s.position, # NOTE: This might need conversion if `pin` is an int and `position` is a string. - location=dewar_address, - priority=s.priority if hasattr(s, "priority") and s.priority is not None else 1.0, - comment=s.comments if hasattr(s, "comments") else None, - mount_count=s.mount_count or 0, - aaredb_params=aaredb_params - ) + # Normal puck like "A1", "B5", "C3" + elif isinstance(p.tell_position, str): + segment = p.tell_position[0] + pos = int(p.tell_position[1:]) + dewar_address = DewarAddress(segment=segment, pos=pos) + target_list = normal_short_infos + + else: + print(f"[WARN] Skipping unknown tell_position format: {p.tell_position}") + continue + + info = SampleShortInfo( + db_id=s.id, + puck_name=p.puck_name, + dewar_name=p.dewar_name or "Unknown", + sample_name=s.sample_name, + run_number=s.run_number, + user=s.pgroup, + pin=s.position, + location=dewar_address, + priority=s.priority,#getattr(s, "priority", 1.0), + comment=getattr(s, "comments", None), + mount_count=s.mount_count or 0, + aaredb_params=s.data_collection_parameters ) - # Update Redis with the new spreadsheet - ret = SampleShortInfoList(s=sample_short_infos) - config.spreadsheet = ret - print("DEBUG: Writing to redis key:", config._BeamlineConfig__bl + ":sample_spreadsheet") - print("[REDIS][INFO] Written spreadsheet to redis via .spreadsheet property") + target_list.append(info) + + # Write normal pucks to sample_spreadsheet + normal_key = f"{config._BeamlineConfig__bl}:sample_spreadsheet" + normal_list = SampleShortInfoList(s=normal_short_infos+reference_short_infos) + config._BeamlineConfig__client.set(normal_key, normal_list.json()) + print("[REDIS][INFO] Written normal spreadsheet to:", normal_key) + + # Write reference tools to reference-tools + ref_key = f"{config._BeamlineConfig__bl}:reference-tools" + if reference_short_infos: + ref_list = SampleShortInfoList(s=reference_short_infos) + config._BeamlineConfig__client.set(ref_key, ref_list.json()) + print("[REDIS][INFO] Written reference tools to:", ref_key) + else: + # Clear key if empty + try: + config._BeamlineConfig__client.delete(ref_key) + print("[REDIS][INFO] Cleared reference tools key:", ref_key) + except Exception: + pass except Exception as exc: print("[WS][ERROR] Failed to parse or convert message:", exc) - def on_error(ws, error): """ Handle WebSocket errors. diff --git a/daq/src/aaredaq/tellupdater.py b/daq/src/aaredaq/tellupdater.py index d901320d..b38b5231 100644 --- a/daq/src/aaredaq/tellupdater.py +++ b/daq/src/aaredaq/tellupdater.py @@ -3,29 +3,30 @@ import json import websocket import sseclient import requests -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 +#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}" +SLOT_IDENTIFIER = "X10SA" +#WS_URL = f"wss://mx-db-01.psi.ch/dispatcher/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}" +WS_URL = f"wss://localhost:8001/protected_router/wstell/ws/slot/{SLOT_IDENTIFIER}" WS_HEADERS = [f"X-Shared-Password: {os.getenv('AAREDB_SHARED_PASSWORD')}"] +print(WS_HEADERS) # Initialize TELL client and DB wrapper -beamline = MXBeamline.X06DA # Use your beamline enum/value -tell_client = TellClient(bl=beamline) -aare_db = AareWrapper(bl=beamline) +beamline = None #MXBeamline.X06DA # Use your beamline enum/value +tell_client = None #TellClient(bl=beamline) +aare_db = None #AareWrapper(bl=beamline) # Track current state current_pucks = [] last_pucks_sent = [] def listen_to_sse(): - print("[DEBUG] Entered listen_to_sse()") + print("[SSE][listen_to_sse] Entered listen_to_sse()") if not tell_client.url: print(f"[SSE][WARN] No TELL URL configured – SSE listener not started. (tell_client.url={tell_client.url})") return @@ -33,23 +34,43 @@ def listen_to_sse(): try: response = requests.get(sse_url, stream=True) client = sseclient.SSEClient(response) + + # Immediately fetch current detected pucks once on connect so the system starts with + # up-to-date state (instead of waiting for the first DewarContentUpdate event). + try: + print("[SSE][listen_to_sse] Initial detected pucks fetch on connect") + handle_tell_change_event() + except Exception as exc: + print(f"[SSE][listen_to_sse][WARN] Initial fetch failed: {exc}") + for event in client.events(): if event.event == "DewarContentUpdate": - print(f"event: {event.event}, data: {event.data}") + print(f"[SSE][listen_to_sse] event: {event.event}, data: {event.data}") on_sse_event(event) except requests.exceptions.RequestException as exc: - print(f"[SSE][ERROR] Failed to connect to {sse_url}: {exc}") + print(f"[SSE][listen_to_sse][ERROR] Failed to connect to {sse_url}: {exc}") 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 + """ + Compare two lists of pucks and return (joined, left). + Ignore reference-tools puck(s) with tell_position 'X1'. + """ + # Filter out reference-tools + old_filtered = [p for p in old if getattr(p, "tell_position", None) != "X1"] + new_filtered = [p for p in new if getattr(p, "tell_position", None) != "X1"] + + old_ids = set(key_func(p) for p in old_filtered) + new_ids = set(key_func(p) for p in new_filtered) + + joined = new_ids - old_ids # newly added + left = old_ids - new_ids # removed + return joined, left def ws_update_samples_info(pucks): """Send sample info to TELL robot.""" - tell_client.set_samples_info(pucks) + #tell_client.set_samples_info(pucks) + print(pucks) def handle_tell_change_event(): """Fetch the latest detected pucks from TELL and update the database.""" @@ -68,16 +89,42 @@ def on_sse_event(event): handle_tell_change_event() def on_message(ws, message): + #print("[WS][MESSAGE]", message) try: data = json.loads(message) print("[WS][MESSAGE]", data) - if isinstance(data, list): - new_pucks = [PuckWithTellPosition(**item) for item in data] - # Always forward the current state to TELL or your API + if not isinstance(data, list): + print("[WS][ERROR] Unknown message format") + return + + new_pucks = [PuckWithTellPosition(**item) for item in data] + + # Keep module-level current_pucks in sync + global current_pucks + + # First, try to refresh DB/TELL state to avoid overwriting a recent SSE update. + try: + # This will fetch detected pucks from TELL and update DB via your existing handler + handle_tell_change_event() + # Small wait gives SSE/DB handler a moment to settle (optional) + time.sleep(0.2) + except Exception as exc: + + print(f"[WS][WARN] handle_tell_change_event() failed: {exc}") + + # Compare lists and ignore reference puck at X1 if needed + joined, left = compare_and_report_change(current_pucks, new_pucks, lambda p: getattr(p, "puck_name", None)) + + # Update local cache + current_pucks = new_pucks + + if joined or left: + print(f"[WS][INFO] Changes detected joined={joined} left={left}") + # Send only when there is a real change ws_update_samples_info(new_pucks) print("[WS][INFO] Sent current pucks state to TELL.") else: - print("[WS][ERROR] Unknown message format") + print("[WS][INFO] No change, skipping send to TELL.") except Exception as exc: print("[WS][ERROR] Failed to parse websocket message:", exc) @@ -92,8 +139,8 @@ def on_open(ws): def main(): # Start SSE listener in a separate background thread - sse_thread = threading.Thread(target=listen_to_sse, daemon=True) - sse_thread.start() + #sse_thread = threading.Thread(target=listen_to_sse, daemon=True) + #sse_thread.start() # Main thread runs websocket client loop while True: