DAQ: updates to tellupdater and spreadsheetupdater from master branch.

This commit is contained in:
2026-01-22 11:25:54 +01:00
parent a8ccc88636
commit 7f1c495d40
2 changed files with 129 additions and 63 deletions
+59 -40
View File
@@ -3,11 +3,11 @@ import json
import websocket
import time
from aareDBclient.models import PuckWithTellPosition
from aare.common.models import SampleShortInfoList, SampleShortInfo, DewarAddress
from aaredaqlib.models import SampleShortInfoList, SampleShortInfo, DewarAddress
from config import BeamlineConfig
from aare.common.beamline import MXBeamline
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
@@ -16,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
@@ -39,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.
+70 -23
View File
@@ -3,29 +3,30 @@ import json
import websocket
import sseclient
import requests
import threading
import time
from aareDBclient.models import PuckWithTellPosition
from aare.devices.tell_client import TellClient
from aaredb import AareWrapper # Make sure the import path fits your project
from aare.common.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: