da03 commited on
Commit
fff0431
·
1 Parent(s): c0820d5
Files changed (1) hide show
  1. main.py +19 -6
main.py CHANGED
@@ -251,6 +251,14 @@ async def websocket_endpoint(websocket: WebSocket):
251
  async def reset_simulation():
252
  nonlocal previous_frame, hidden_states, keys_down, frame_num, is_processing, input_queue
253
 
 
 
 
 
 
 
 
 
254
  # Clear the input queue
255
  while not input_queue.empty():
256
  try:
@@ -484,7 +492,7 @@ async def websocket_endpoint(websocket: WebSocket):
484
 
485
  print(f"WebSocket connection closed: {client_id}")
486
 
487
- def log_interaction(client_id, data, generated_frame=None, is_end_of_session=False):
488
  """Log user interaction and optionally the generated frame."""
489
  timestamp = time.time()
490
 
@@ -495,11 +503,16 @@ def log_interaction(client_id, data, generated_frame=None, is_end_of_session=Fal
495
  log_entry = {
496
  "timestamp": timestamp,
497
  "client_id": client_id,
498
- "is_eos": is_end_of_session
 
499
  }
500
 
501
- # Only include input data if this isn't an EOS token or if data is provided
502
- if not is_end_of_session or data:
 
 
 
 
503
  log_entry["inputs"] = {
504
  "x": data.get("x"),
505
  "y": data.get("y"),
@@ -509,7 +522,7 @@ def log_interaction(client_id, data, generated_frame=None, is_end_of_session=Fal
509
  "keys_up": data.get("keys_up", [])
510
  }
511
  else:
512
- # For EOS records with empty data, just include minimal info
513
  log_entry["inputs"] = None
514
 
515
  # Save to a file (one file per session)
@@ -518,7 +531,7 @@ def log_interaction(client_id, data, generated_frame=None, is_end_of_session=Fal
518
  f.write(json.dumps(log_entry) + "\n")
519
 
520
  # Optionally save the frame if provided
521
- if generated_frame is not None and not is_end_of_session:
522
  frame_dir = f"interaction_logs/frames_{client_id}"
523
  os.makedirs(frame_dir, exist_ok=True)
524
  frame_file = f"{frame_dir}/{timestamp:.6f}.png"
 
251
  async def reset_simulation():
252
  nonlocal previous_frame, hidden_states, keys_down, frame_num, is_processing, input_queue
253
 
254
+ # Log the reset action
255
+ log_interaction(
256
+ client_id,
257
+ {"type": "reset"},
258
+ is_end_of_session=False,
259
+ is_reset=True # Add this parameter to the log_interaction function
260
+ )
261
+
262
  # Clear the input queue
263
  while not input_queue.empty():
264
  try:
 
492
 
493
  print(f"WebSocket connection closed: {client_id}")
494
 
495
+ def log_interaction(client_id, data, generated_frame=None, is_end_of_session=False, is_reset=False):
496
  """Log user interaction and optionally the generated frame."""
497
  timestamp = time.time()
498
 
 
503
  log_entry = {
504
  "timestamp": timestamp,
505
  "client_id": client_id,
506
+ "is_eos": is_end_of_session,
507
+ "is_reset": is_reset
508
  }
509
 
510
+ # Include type if present (for reset, etc.)
511
+ if data.get("type"):
512
+ log_entry["type"] = data.get("type")
513
+
514
+ # Only include input data if this isn't just a control message
515
+ if not is_end_of_session and not is_reset:
516
  log_entry["inputs"] = {
517
  "x": data.get("x"),
518
  "y": data.get("y"),
 
522
  "keys_up": data.get("keys_up", [])
523
  }
524
  else:
525
+ # For EOS/reset records, just include minimal info
526
  log_entry["inputs"] = None
527
 
528
  # Save to a file (one file per session)
 
531
  f.write(json.dumps(log_entry) + "\n")
532
 
533
  # Optionally save the frame if provided
534
+ if generated_frame is not None and not is_end_of_session and not is_reset:
535
  frame_dir = f"interaction_logs/frames_{client_id}"
536
  os.makedirs(frame_dir, exist_ok=True)
537
  frame_file = f"{frame_dir}/{timestamp:.6f}.png"