DAQ: updates to tellupdater and spreadsheetupdater from master branch.
This commit is contained in:
@@ -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
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user