@@ -41,6 +41,12 @@ def _user_agent() -> str:
41
41
)
42
42
43
43
44
+ def _dump_model (model : BaseModel ):
45
+ if hasattr (model , "model_dump" ):
46
+ return model .model_dump (exclude_none = True )
47
+ return model .dict (exclude_none = True )
48
+
49
+
44
50
class StreamingClient :
45
51
def __init__ (self , options : StreamingClientOptions ):
46
52
self ._options = options
@@ -56,7 +62,7 @@ def __init__(self, options: StreamingClientOptions):
56
62
self ._stop_event = threading .Event ()
57
63
58
64
def connect (self , params : StreamingParameters ) -> None :
59
- params_dict = params . model_dump ( exclude_none = True )
65
+ params_dict = _dump_model ( params )
60
66
params_encoded = urlencode (params_dict )
61
67
62
68
uri = f"wss://{ self ._options .api_host } /v3/ws?{ params_encoded } "
@@ -105,7 +111,7 @@ def stream(
105
111
self ._write_queue .put (chunk )
106
112
107
113
def set_params (self , params : StreamingSessionParameters ):
108
- message = UpdateConfiguration (** params . model_dump ( ))
114
+ message = UpdateConfiguration (** _dump_model ( params ))
109
115
self ._write_queue .put (message )
110
116
111
117
def on (self , event : StreamingEvents , handler : Callable ) -> None :
@@ -126,8 +132,7 @@ def _write_message(self) -> None:
126
132
if isinstance (data , bytes ):
127
133
self ._websocket .send (data )
128
134
elif isinstance (data , BaseModel ):
129
- message = data .model_dump_json (exclude_none = True )
130
- self ._websocket .send (message )
135
+ self ._websocket .send (json .dumps (data ))
131
136
else :
132
137
raise ValueError (f"Attempted to send invalid message: { type (data )} " )
133
138
except websockets .exceptions .ConnectionClosed as exc :
0 commit comments