Update websocket_storage.py

websocket_storage.py

@@ -7,6 +7,20 @@ from queue import Queue
 import time
 import asyncio
 import hashlib
+import dataclasses
+
+def custom_json_serializer(obj):
+    if hasattr(obj, '__dict__'):
+        return obj.__dict__
+    elif hasattr(obj, '_asdict'):  # For namedtuples
+        return obj._asdict()
+    elif dataclasses.is_dataclass(obj):
+        return dataclasses.asdict(obj)
+    elif isinstance(obj, (np.ndarray, np.number)):
+        return obj.tolist()
+    elif isinstance(obj, set):
+        return list(obj)
+    raise TypeError(f'Object of type {type(obj)} is not JSON serializable')
 
 class WebSocketGPUStorage:
     # Singleton instance
@@ -72,7 +86,7 @@ class WebSocketGPUStorage:
                         try:
                             while not self.message_queue.empty():
                                 msg_id, operation = self.message_queue.get()
-                                await websocket.send(json.dumps(operation))
+                                await websocket.send(json.dumps(operation, default=custom_json_serializer))
                                 
                                 # Wait for response with timeout
                                 try: