Skip to content

Commit fed346a

Browse files
committed
fix(sdk/python): add model_dump compatibility
1 parent 698dd9b commit fed346a

File tree

1 file changed

+9
-4
lines changed

1 file changed

+9
-4
lines changed

assemblyai/streaming/v3/client.py

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,12 @@ def _user_agent() -> str:
4141
)
4242

4343

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+
4450
class StreamingClient:
4551
def __init__(self, options: StreamingClientOptions):
4652
self._options = options
@@ -56,7 +62,7 @@ def __init__(self, options: StreamingClientOptions):
5662
self._stop_event = threading.Event()
5763

5864
def connect(self, params: StreamingParameters) -> None:
59-
params_dict = params.model_dump(exclude_none=True)
65+
params_dict = _dump_model(params)
6066
params_encoded = urlencode(params_dict)
6167

6268
uri = f"wss://{self._options.api_host}/v3/ws?{params_encoded}"
@@ -105,7 +111,7 @@ def stream(
105111
self._write_queue.put(chunk)
106112

107113
def set_params(self, params: StreamingSessionParameters):
108-
message = UpdateConfiguration(**params.model_dump())
114+
message = UpdateConfiguration(**_dump_model(params))
109115
self._write_queue.put(message)
110116

111117
def on(self, event: StreamingEvents, handler: Callable) -> None:
@@ -126,8 +132,7 @@ def _write_message(self) -> None:
126132
if isinstance(data, bytes):
127133
self._websocket.send(data)
128134
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))
131136
else:
132137
raise ValueError(f"Attempted to send invalid message: {type(data)}")
133138
except websockets.exceptions.ConnectionClosed as exc:

0 commit comments

Comments
 (0)