1use std::sync::atomic::{AtomicU64, Ordering};
2use std::time::{Duration, SystemTime, UNIX_EPOCH};
3
4use base64::Engine;
5use base64::engine::general_purpose::STANDARD;
6use futures_util::{SinkExt, StreamExt};
7use serde::{Deserialize, Serialize};
8use serde_json::{Map, Value};
9use tokio::net::TcpStream;
10use tokio::sync::mpsc;
11use tokio::task::JoinHandle;
12use tokio::time::timeout;
13use tokio_tungstenite::tungstenite::Message;
14use tokio_tungstenite::tungstenite::client::IntoClientRequest;
15use tokio_tungstenite::tungstenite::http::header::AUTHORIZATION;
16use tokio_tungstenite::{MaybeTlsStream, WebSocketStream};
17
18use crate::auth::AuthenticationProvider;
19use crate::{Result, SdkError, ZhipuAuthentication};
20
21mod events;
22pub use events::*;
23
24pub const ZHIPU_REALTIME_URL: &str = "wss://open.bigmodel.cn/api/paas/v4/realtime";
25
26type RealtimeSocket = WebSocketStream<MaybeTlsStream<TcpStream>>;
27
28static EVENT_SEQUENCE: AtomicU64 = AtomicU64::new(1);
29
30#[derive(Clone)]
31pub struct RealtimeConfig {
32 pub authentication: ZhipuAuthentication,
33 pub url: String,
34 pub connect_timeout: Duration,
35 pub channel_capacity: usize,
36}
37
38impl RealtimeConfig {
39 pub fn new(api_key: impl Into<String>) -> Self {
40 Self {
41 authentication: ZhipuAuthentication::auto(api_key),
42 url: ZHIPU_REALTIME_URL.into(),
43 connect_timeout: Duration::from_secs(15),
44 channel_capacity: 256,
45 }
46 }
47
48 pub fn authentication(mut self, value: ZhipuAuthentication) -> Self {
49 self.authentication = value;
50 self
51 }
52
53 pub fn url(mut self, value: impl Into<String>) -> Self {
54 self.url = value.into();
55 self
56 }
57
58 pub fn connect_timeout(mut self, value: Duration) -> Self {
59 self.connect_timeout = value;
60 self
61 }
62
63 pub fn channel_capacity(mut self, value: usize) -> Self {
64 self.channel_capacity = value;
65 self
66 }
67
68 pub async fn connect(self) -> Result<RealtimeConnection> {
69 RealtimeClient::from_config(self).await
70 }
71}
72
73pub struct RealtimeClient;
74
75impl RealtimeClient {
76 pub async fn connect(api_key: impl Into<String>) -> Result<RealtimeConnection> {
77 RealtimeConfig::new(api_key).connect().await
78 }
79
80 pub async fn from_config(config: RealtimeConfig) -> Result<RealtimeConnection> {
81 if !(config.url.starts_with("wss://") || config.url.starts_with("ws://")) {
82 return Err(SdkError::Configuration(
83 "realtime URL must use ws or wss".into(),
84 ));
85 }
86 if config.connect_timeout.is_zero() || config.channel_capacity == 0 {
87 return Err(SdkError::Configuration(
88 "realtime timeout and channel capacity must be greater than zero".into(),
89 ));
90 }
91 let authentication = AuthenticationProvider::zhipu(config.authentication)?;
92 let authorization = authentication.header_value()?;
93 let mut request = config.url.into_client_request()?;
94 request.headers_mut().insert(
95 AUTHORIZATION,
96 authorization
97 .to_str()
98 .map_err(|_| SdkError::Configuration("authentication header is invalid".into()))?
99 .parse()
100 .map_err(|_| SdkError::Configuration("authentication header is invalid".into()))?,
101 );
102 let (socket, _) = timeout(
103 config.connect_timeout,
104 tokio_tungstenite::connect_async(request),
105 )
106 .await
107 .map_err(|_| SdkError::Timeout("realtime connection timed out".into()))??;
108 Ok(spawn_connection(socket, config.channel_capacity))
109 }
110}
111
112#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
113pub struct RealtimeSession {
114 pub model: String,
115 pub modalities: Vec<String>,
116 #[serde(skip_serializing_if = "Option::is_none")]
117 pub instructions: Option<String>,
118 pub voice: String,
119 pub input_audio_format: String,
120 pub output_audio_format: String,
121 #[serde(skip_serializing_if = "Option::is_none")]
122 pub input_audio_noise_reduction: Option<RealtimeNoiseReduction>,
123 #[serde(skip_serializing_if = "Option::is_none")]
124 pub turn_detection: Option<RealtimeTurnDetection>,
125 #[serde(skip_serializing_if = "Option::is_none")]
126 pub temperature: Option<f32>,
127 #[serde(rename = "max_output_tokens", skip_serializing_if = "Option::is_none")]
128 pub max_response_output_tokens: Option<RealtimeMaxTokens>,
129 #[serde(default, skip_serializing_if = "Vec::is_empty")]
130 pub tools: Vec<RealtimeTool>,
131 pub beta_fields: RealtimeBetaFields,
132 #[serde(flatten, default, skip_serializing_if = "Map::is_empty")]
133 pub extra: Map<String, Value>,
134}
135
136impl Default for RealtimeSession {
137 fn default() -> Self {
138 Self {
139 model: "glm-realtime".into(),
140 modalities: vec!["text".into(), "audio".into()],
141 instructions: None,
142 voice: "tongtong".into(),
143 input_audio_format: "pcm16".into(),
144 output_audio_format: "pcm".into(),
145 input_audio_noise_reduction: None,
146 turn_detection: None,
147 temperature: None,
148 max_response_output_tokens: None,
149 tools: Vec::new(),
150 beta_fields: RealtimeBetaFields::default(),
151 extra: Map::new(),
152 }
153 }
154}
155
156impl RealtimeSession {
157 pub fn model(mut self, value: impl Into<String>) -> Self {
158 self.model = value.into();
159 self
160 }
161
162 pub fn instructions(mut self, value: impl Into<String>) -> Self {
163 self.instructions = Some(value.into());
164 self
165 }
166
167 pub fn voice(mut self, value: impl Into<String>) -> Self {
168 self.voice = value.into();
169 self
170 }
171
172 pub fn input_audio_format(mut self, value: impl Into<String>) -> Self {
173 self.input_audio_format = value.into();
174 self
175 }
176
177 pub fn server_vad(mut self, create_response: bool, interrupt_response: bool) -> Self {
178 self.turn_detection = Some(RealtimeTurnDetection {
179 kind: "server_vad".into(),
180 create_response: Some(create_response),
181 interrupt_response: Some(interrupt_response),
182 prefix_padding_ms: None,
183 silence_duration_ms: None,
184 threshold: None,
185 });
186 self
187 }
188
189 pub fn video(mut self) -> Self {
190 self.beta_fields.chat_mode = "video_passive".into();
191 self
192 }
193
194 pub fn tool(mut self, value: RealtimeTool) -> Self {
195 self.tools.push(value);
196 self
197 }
198}
199
200#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
201pub struct RealtimeNoiseReduction {
202 #[serde(rename = "type")]
203 pub kind: String,
204}
205
206#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
207pub struct RealtimeTurnDetection {
208 #[serde(rename = "type")]
209 pub kind: String,
210 #[serde(skip_serializing_if = "Option::is_none")]
211 pub create_response: Option<bool>,
212 #[serde(skip_serializing_if = "Option::is_none")]
213 pub interrupt_response: Option<bool>,
214 #[serde(skip_serializing_if = "Option::is_none")]
215 pub prefix_padding_ms: Option<u32>,
216 #[serde(skip_serializing_if = "Option::is_none")]
217 pub silence_duration_ms: Option<u32>,
218 #[serde(skip_serializing_if = "Option::is_none")]
219 pub threshold: Option<f32>,
220}
221
222#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
223pub struct RealtimeTranscriptionSession {
224 pub input_audio_format: String,
225 #[serde(skip_serializing_if = "Option::is_none")]
226 pub input_audio_noise_reduction: Option<RealtimeNoiseReduction>,
227 pub modalities: Vec<String>,
228 #[serde(skip_serializing_if = "Option::is_none")]
229 pub turn_detection: Option<RealtimeTurnDetection>,
230}
231
232impl Default for RealtimeTranscriptionSession {
233 fn default() -> Self {
234 Self {
235 input_audio_format: "pcm".into(),
236 input_audio_noise_reduction: None,
237 modalities: vec!["text".into(), "audio".into()],
238 turn_detection: None,
239 }
240 }
241}
242
243#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
244#[serde(untagged)]
245pub enum RealtimeMaxTokens {
246 Count(u16),
247 Unlimited(String),
248}
249
250#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
251pub struct RealtimeTool {
252 #[serde(rename = "type")]
253 pub kind: String,
254 pub name: String,
255 pub description: String,
256 pub parameters: Value,
257}
258
259impl RealtimeTool {
260 pub fn function(
261 name: impl Into<String>,
262 description: impl Into<String>,
263 parameters: Value,
264 ) -> Self {
265 Self {
266 kind: "function".into(),
267 name: name.into(),
268 description: description.into(),
269 parameters,
270 }
271 }
272}
273
274#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
275pub struct RealtimeBetaFields {
276 pub chat_mode: String,
277 #[serde(skip_serializing_if = "Option::is_none")]
278 pub tts_source: Option<String>,
279 #[serde(skip_serializing_if = "Option::is_none")]
280 pub auto_search: Option<bool>,
281 #[serde(skip_serializing_if = "Option::is_none")]
282 pub greeting_config: Option<RealtimeGreetingConfig>,
283}
284
285impl Default for RealtimeBetaFields {
286 fn default() -> Self {
287 Self {
288 chat_mode: "audio".into(),
289 tts_source: None,
290 auto_search: None,
291 greeting_config: None,
292 }
293 }
294}
295
296#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
297pub struct RealtimeGreetingConfig {
298 pub enable: bool,
299 #[serde(skip_serializing_if = "Option::is_none")]
300 pub content: Option<String>,
301}
302
303#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
304pub struct RealtimeContentPart {
305 #[serde(rename = "type")]
306 pub kind: String,
307 #[serde(skip_serializing_if = "Option::is_none")]
308 pub text: Option<String>,
309 #[serde(skip_serializing_if = "Option::is_none")]
310 pub audio: Option<String>,
311 #[serde(skip_serializing_if = "Option::is_none")]
312 pub transcript: Option<String>,
313}
314
315#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
316pub struct RealtimeConversationItem {
317 #[serde(skip_serializing_if = "Option::is_none")]
318 pub id: Option<String>,
319 #[serde(rename = "type")]
320 pub kind: String,
321 pub object: String,
322 #[serde(skip_serializing_if = "Option::is_none")]
323 pub status: Option<String>,
324 #[serde(skip_serializing_if = "Option::is_none")]
325 pub role: Option<String>,
326 #[serde(default, skip_serializing_if = "Vec::is_empty")]
327 pub content: Vec<RealtimeContentPart>,
328 #[serde(skip_serializing_if = "Option::is_none")]
329 pub name: Option<String>,
330 #[serde(skip_serializing_if = "Option::is_none")]
331 pub arguments: Option<String>,
332 #[serde(skip_serializing_if = "Option::is_none")]
333 pub output: Option<String>,
334}
335
336impl RealtimeConversationItem {
337 pub fn text(role: impl Into<String>, text: impl Into<String>) -> Self {
338 Self {
339 id: None,
340 kind: "message".into(),
341 object: "realtime.item".into(),
342 status: Some("completed".into()),
343 role: Some(role.into()),
344 content: vec![RealtimeContentPart {
345 kind: "input_text".into(),
346 text: Some(text.into()),
347 audio: None,
348 transcript: None,
349 }],
350 name: None,
351 arguments: None,
352 output: None,
353 }
354 }
355
356 pub fn function_output(output: impl Into<String>) -> Self {
357 Self {
358 id: None,
359 kind: "function_call_output".into(),
360 object: "realtime.item".into(),
361 status: Some("completed".into()),
362 role: None,
363 content: Vec::new(),
364 name: None,
365 arguments: None,
366 output: Some(output.into()),
367 }
368 }
369}
370
371#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
372#[serde(tag = "type")]
373pub enum RealtimeClientEvent {
374 #[serde(rename = "session.update")]
375 SessionUpdate {
376 #[serde(flatten)]
377 metadata: RealtimeEventMetadata,
378 session: Box<RealtimeSession>,
379 },
380 #[serde(rename = "transcription_session.update")]
381 TranscriptionSessionUpdate {
382 #[serde(flatten)]
383 metadata: RealtimeEventMetadata,
384 session: Box<RealtimeTranscriptionSession>,
385 },
386 #[serde(rename = "input_audio_buffer.append")]
387 InputAudioBufferAppend {
388 #[serde(flatten)]
389 metadata: RealtimeEventMetadata,
390 audio: String,
391 },
392 #[serde(rename = "input_audio_buffer.append_video_frame")]
393 InputAudioBufferAppendVideoFrame {
394 #[serde(flatten)]
395 metadata: RealtimeEventMetadata,
396 video_frame: String,
397 },
398 #[serde(rename = "input_audio_buffer.commit")]
399 InputAudioBufferCommit {
400 #[serde(flatten)]
401 metadata: RealtimeEventMetadata,
402 },
403 #[serde(rename = "input_audio_buffer.clear")]
404 InputAudioBufferClear {
405 #[serde(flatten)]
406 metadata: RealtimeEventMetadata,
407 },
408 #[serde(rename = "conversation.item.create")]
409 ConversationItemCreate {
410 #[serde(flatten)]
411 metadata: RealtimeEventMetadata,
412 item: Box<RealtimeConversationItem>,
413 },
414 #[serde(rename = "conversation.item.delete")]
415 ConversationItemDelete {
416 #[serde(flatten)]
417 metadata: RealtimeEventMetadata,
418 item_id: String,
419 },
420 #[serde(rename = "conversation.item.retrieve")]
421 ConversationItemRetrieve {
422 #[serde(flatten)]
423 metadata: RealtimeEventMetadata,
424 item_id: String,
425 },
426 #[serde(rename = "response.create")]
427 ResponseCreate {
428 #[serde(flatten)]
429 metadata: RealtimeEventMetadata,
430 },
431 #[serde(rename = "response.cancel")]
432 ResponseCancel {
433 #[serde(flatten)]
434 metadata: RealtimeEventMetadata,
435 },
436}
437
438impl RealtimeClientEvent {
439 pub fn session_update(session: RealtimeSession) -> Result<Self> {
440 Ok(Self::SessionUpdate {
441 metadata: RealtimeEventMetadata::new()?,
442 session: Box::new(session),
443 })
444 }
445
446 pub fn append_audio(bytes: &[u8]) -> Result<Self> {
447 Self::append_audio_base64(STANDARD.encode(bytes))
448 }
449
450 pub fn transcription_session_update(session: RealtimeTranscriptionSession) -> Result<Self> {
451 Ok(Self::TranscriptionSessionUpdate {
452 metadata: RealtimeEventMetadata::new()?,
453 session: Box::new(session),
454 })
455 }
456
457 pub fn append_audio_base64(value: impl Into<String>) -> Result<Self> {
458 let audio = value.into();
459 if audio.is_empty() {
460 return Err(SdkError::Validation("audio data cannot be empty".into()));
461 }
462 Ok(Self::InputAudioBufferAppend {
463 metadata: RealtimeEventMetadata::new()?,
464 audio,
465 })
466 }
467
468 pub fn append_video_frame(jpeg: &[u8]) -> Result<Self> {
469 if jpeg.is_empty() {
470 return Err(SdkError::Validation("video frame cannot be empty".into()));
471 }
472 Ok(Self::InputAudioBufferAppendVideoFrame {
473 metadata: RealtimeEventMetadata::new()?,
474 video_frame: STANDARD.encode(jpeg),
475 })
476 }
477
478 pub fn commit() -> Result<Self> {
479 Ok(Self::InputAudioBufferCommit {
480 metadata: RealtimeEventMetadata::new()?,
481 })
482 }
483
484 pub fn clear() -> Result<Self> {
485 Ok(Self::InputAudioBufferClear {
486 metadata: RealtimeEventMetadata::new()?,
487 })
488 }
489
490 pub fn create_item(item: RealtimeConversationItem) -> Result<Self> {
491 Ok(Self::ConversationItemCreate {
492 metadata: RealtimeEventMetadata::new()?,
493 item: Box::new(item),
494 })
495 }
496
497 pub fn delete_item(item_id: impl Into<String>) -> Result<Self> {
498 let item_id = require_value(item_id.into(), "item id")?;
499 Ok(Self::ConversationItemDelete {
500 metadata: RealtimeEventMetadata::new()?,
501 item_id,
502 })
503 }
504
505 pub fn retrieve_item(item_id: impl Into<String>) -> Result<Self> {
506 let item_id = require_value(item_id.into(), "item id")?;
507 Ok(Self::ConversationItemRetrieve {
508 metadata: RealtimeEventMetadata::new()?,
509 item_id,
510 })
511 }
512
513 pub fn create_response() -> Result<Self> {
514 Ok(Self::ResponseCreate {
515 metadata: RealtimeEventMetadata::new()?,
516 })
517 }
518
519 pub fn cancel_response() -> Result<Self> {
520 Ok(Self::ResponseCancel {
521 metadata: RealtimeEventMetadata::new()?,
522 })
523 }
524}
525
526#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
527pub struct RealtimeEventMetadata {
528 pub event_id: String,
529 pub client_timestamp: u64,
530}
531
532impl RealtimeEventMetadata {
533 pub fn new() -> Result<Self> {
534 let client_timestamp = unix_millis()?;
535 let sequence = EVENT_SEQUENCE.fetch_add(1, Ordering::Relaxed);
536 Ok(Self {
537 event_id: format!("rustglm-{client_timestamp}-{sequence}"),
538 client_timestamp,
539 })
540 }
541}
542
543#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
544pub struct RealtimeServerEvent {
545 #[serde(rename = "type")]
546 pub event_type: String,
547 #[serde(default)]
548 pub event_id: Option<String>,
549 #[serde(default)]
550 pub client_timestamp: Option<u64>,
551 #[serde(flatten, default)]
552 pub data: Map<String, Value>,
553}
554
555impl RealtimeServerEvent {
556 pub fn delta_text(&self) -> Option<&str> {
557 matches!(
558 self.event_type.as_str(),
559 "response.text.delta" | "response.audio_transcript.delta"
560 )
561 .then(|| self.data.get("delta").and_then(Value::as_str))
562 .flatten()
563 }
564
565 pub fn audio_base64(&self) -> Option<&str> {
566 (self.event_type == "response.audio.delta")
567 .then(|| self.data.get("delta").and_then(Value::as_str))
568 .flatten()
569 }
570
571 pub fn audio_bytes(&self) -> Result<Option<Vec<u8>>> {
572 self.audio_base64()
573 .map(|value| {
574 STANDARD
575 .decode(value)
576 .map_err(|error| SdkError::Stream(error.to_string().into()))
577 })
578 .transpose()
579 }
580
581 pub fn error(&self) -> Option<&Value> {
582 (self.event_type == "error")
583 .then(|| self.data.get("error"))
584 .flatten()
585 }
586
587 pub fn function_call(&self) -> Option<RealtimeFunctionCall<'_>> {
588 if self.event_type != "response.function_call_arguments.done" {
589 return None;
590 }
591 Some(RealtimeFunctionCall {
592 name: self.data.get("name")?.as_str()?,
593 arguments: self.data.get("arguments")?.as_str()?,
594 })
595 }
596}
597
598#[derive(Debug, Clone, Copy, PartialEq, Eq)]
599pub struct RealtimeFunctionCall<'a> {
600 pub name: &'a str,
601 pub arguments: &'a str,
602}
603
604enum RealtimeCommand {
605 Send(Message),
606 Close,
607}
608
609#[derive(Clone)]
610pub struct RealtimeSender {
611 commands: mpsc::Sender<RealtimeCommand>,
612}
613
614impl RealtimeSender {
615 pub async fn send(&self, event: &RealtimeClientEvent) -> Result<()> {
616 let value = serde_json::to_string(event)
617 .map_err(|error| SdkError::Validation(error.to_string().into()))?;
618 self.send_message(Message::Text(value.into())).await
619 }
620
621 pub async fn send_json(&self, event: &Value) -> Result<()> {
622 self.send_message(Message::Text(event.to_string().into()))
623 .await
624 }
625
626 pub async fn send_request(&self, event: &RealtimeRequest) -> Result<()> {
627 let value = serde_json::to_string(event)
628 .map_err(|error| SdkError::Validation(error.to_string().into()))?;
629 self.send_message(Message::Text(value.into())).await
630 }
631
632 pub async fn update_typed_session(&self, session: TypedRealtimeSession) -> Result<()> {
633 self.send_request(&RealtimeRequest::session_update(session)?)
634 .await
635 }
636
637 pub async fn create_typed_item(
638 &self,
639 previous_item_id: Option<String>,
640 item: TypedRealtimeItem,
641 ) -> Result<()> {
642 self.send_request(&RealtimeRequest::create_item(previous_item_id, item)?)
643 .await
644 }
645
646 pub async fn create_response_with(&self, options: RealtimeResponseOptions) -> Result<()> {
647 self.send_request(&RealtimeRequest::create_response(Some(options))?)
648 .await
649 }
650
651 pub async fn update_session(&self, session: RealtimeSession) -> Result<()> {
652 self.send(&RealtimeClientEvent::session_update(session)?)
653 .await
654 }
655
656 pub async fn append_audio(&self, bytes: &[u8]) -> Result<()> {
657 self.send(&RealtimeClientEvent::append_audio(bytes)?).await
658 }
659
660 pub async fn append_audio_base64(&self, value: impl Into<String>) -> Result<()> {
661 self.send(&RealtimeClientEvent::append_audio_base64(value)?)
662 .await
663 }
664
665 pub async fn append_video_frame(&self, jpeg: &[u8]) -> Result<()> {
666 self.send(&RealtimeClientEvent::append_video_frame(jpeg)?)
667 .await
668 }
669
670 pub async fn commit(&self) -> Result<()> {
671 self.send(&RealtimeClientEvent::commit()?).await
672 }
673
674 pub async fn clear_audio(&self) -> Result<()> {
675 self.send(&RealtimeClientEvent::clear()?).await
676 }
677
678 pub async fn update_transcription_session(
679 &self,
680 session: RealtimeTranscriptionSession,
681 ) -> Result<()> {
682 self.send(&RealtimeClientEvent::transcription_session_update(session)?)
683 .await
684 }
685
686 pub async fn create_item(&self, item: RealtimeConversationItem) -> Result<()> {
687 self.send(&RealtimeClientEvent::create_item(item)?).await
688 }
689
690 pub async fn delete_item(&self, item_id: impl Into<String>) -> Result<()> {
691 self.send(&RealtimeClientEvent::delete_item(item_id)?).await
692 }
693
694 pub async fn retrieve_item(&self, item_id: impl Into<String>) -> Result<()> {
695 self.send(&RealtimeClientEvent::retrieve_item(item_id)?)
696 .await
697 }
698
699 pub async fn create_response(&self) -> Result<()> {
700 self.send(&RealtimeClientEvent::create_response()?).await
701 }
702
703 pub async fn cancel_response(&self) -> Result<()> {
704 self.send(&RealtimeClientEvent::cancel_response()?).await
705 }
706
707 pub async fn close(&self) -> Result<()> {
708 self.commands
709 .send(RealtimeCommand::Close)
710 .await
711 .map_err(|_| SdkError::Stream("realtime connection is closed".into()))
712 }
713
714 async fn send_message(&self, message: Message) -> Result<()> {
715 self.commands
716 .send(RealtimeCommand::Send(message))
717 .await
718 .map_err(|_| SdkError::Stream("realtime connection is closed".into()))
719 }
720}
721
722pub struct RealtimeReceiver {
723 events: mpsc::Receiver<Result<RealtimeServerEvent>>,
724}
725
726impl RealtimeReceiver {
727 pub async fn next_event(&mut self) -> Option<Result<RealtimeServerEvent>> {
728 self.events.recv().await
729 }
730
731 pub async fn next_typed_event(&mut self) -> Option<Result<RealtimeServerMessage>> {
732 self.next_event()
733 .await
734 .map(|event| event.map(RealtimeServerEvent::into_typed))
735 }
736}
737
738pub struct RealtimeConnection {
739 sender: RealtimeSender,
740 receiver: RealtimeReceiver,
741 task: JoinHandle<()>,
742}
743
744impl RealtimeConnection {
745 pub fn sender(&self) -> RealtimeSender {
746 self.sender.clone()
747 }
748
749 pub async fn send(&self, event: &RealtimeClientEvent) -> Result<()> {
750 self.sender.send(event).await
751 }
752
753 pub async fn next_event(&mut self) -> Option<Result<RealtimeServerEvent>> {
754 self.receiver.next_event().await
755 }
756
757 pub async fn next_typed_event(&mut self) -> Option<Result<RealtimeServerMessage>> {
758 self.receiver.next_typed_event().await
759 }
760
761 pub async fn send_request(&self, event: &RealtimeRequest) -> Result<()> {
762 self.sender.send_request(event).await
763 }
764
765 pub fn split(self) -> (RealtimeSender, RealtimeReceiver) {
766 (self.sender, self.receiver)
767 }
768
769 pub async fn close(self) -> Result<()> {
770 self.sender.close().await?;
771 self.task
772 .await
773 .map_err(|error| SdkError::Stream(error.to_string().into()))?;
774 Ok(())
775 }
776}
777
778fn spawn_connection(socket: RealtimeSocket, capacity: usize) -> RealtimeConnection {
779 let (commands_tx, mut commands_rx) = mpsc::channel(capacity);
780 let (events_tx, events_rx) = mpsc::channel(capacity);
781 let task = tokio::spawn(async move {
782 let (mut sink, mut stream) = socket.split();
783 loop {
784 tokio::select! {
785 command = commands_rx.recv() => match command {
786 Some(RealtimeCommand::Send(message)) => {
787 if let Err(error) = sink.send(message).await {
788 let _ = events_tx.send(Err(error.into())).await;
789 break;
790 }
791 }
792 Some(RealtimeCommand::Close) | None => {
793 let _ = sink.close().await;
794 break;
795 }
796 },
797 message = stream.next() => match message {
798 Some(Ok(Message::Text(text))) => {
799 let event = serde_json::from_str(&text).map_err(|error| {
800 SdkError::Stream(format!("{error}: {text}").into())
801 });
802 if events_tx.send(event).await.is_err() {
803 break;
804 }
805 }
806 Some(Ok(Message::Binary(bytes))) => {
807 let event = serde_json::from_slice(&bytes).map_err(|error| {
808 SdkError::Stream(
809 format!("{error}: {}", String::from_utf8_lossy(&bytes)).into(),
810 )
811 });
812 if events_tx.send(event).await.is_err() {
813 break;
814 }
815 }
816 Some(Ok(Message::Ping(bytes))) => {
817 if let Err(error) = sink.send(Message::Pong(bytes)).await {
818 let _ = events_tx.send(Err(error.into())).await;
819 break;
820 }
821 }
822 Some(Ok(Message::Close(_))) | None => break,
823 Some(Ok(_)) => {}
824 Some(Err(error)) => {
825 let _ = events_tx.send(Err(error.into())).await;
826 break;
827 }
828 }
829 }
830 }
831 });
832 RealtimeConnection {
833 sender: RealtimeSender {
834 commands: commands_tx,
835 },
836 receiver: RealtimeReceiver { events: events_rx },
837 task,
838 }
839}
840
841fn unix_millis() -> Result<u64> {
842 let value = SystemTime::now()
843 .duration_since(UNIX_EPOCH)
844 .map_err(|_| SdkError::Configuration("system clock is before Unix epoch".into()))?
845 .as_millis();
846 u64::try_from(value)
847 .map_err(|_| SdkError::Configuration("timestamp exceeds supported range".into()))
848}
849
850fn require_value(value: String, name: &str) -> Result<String> {
851 if value.trim().is_empty() {
852 return Err(SdkError::Validation(
853 format!("{name} cannot be empty").into(),
854 ));
855 }
856 Ok(value)
857}
858
859#[cfg(test)]
860mod tests {
861 use std::sync::{Arc, Mutex};
862
863 use tokio::net::TcpListener;
864 use tokio_tungstenite::accept_hdr_async;
865 use tokio_tungstenite::tungstenite::handshake::server::{Request, Response};
866
867 use super::*;
868
869 #[test]
870 fn serializes_official_client_events() {
871 let mut session = RealtimeSession::default()
872 .model("glm-realtime-flash")
873 .instructions("concise")
874 .voice("xiaochen")
875 .input_audio_format("pcm24")
876 .server_vad(true, true)
877 .video()
878 .tool(RealtimeTool::function(
879 "weather",
880 "weather lookup",
881 serde_json::json!({"type":"object"}),
882 ));
883 session.max_response_output_tokens = Some(RealtimeMaxTokens::Count(1024));
884 let value =
885 serde_json::to_value(RealtimeClientEvent::session_update(session).unwrap()).unwrap();
886 assert_eq!(value["type"], "session.update");
887 assert_eq!(
888 value["session"]["beta_fields"]["chat_mode"],
889 "video_passive"
890 );
891 assert_eq!(value["session"]["turn_detection"]["type"], "server_vad");
892 assert_eq!(value["session"]["tools"][0]["name"], "weather");
893 assert_eq!(value["session"]["max_output_tokens"], 1024);
894 assert!(value["session"].get("max_response_output_tokens").is_none());
895
896 let audio =
897 serde_json::to_value(RealtimeClientEvent::append_audio(&[1, 2]).unwrap()).unwrap();
898 assert_eq!(audio["type"], "input_audio_buffer.append");
899 assert_eq!(audio["audio"], "AQI=");
900 let frame =
901 serde_json::to_value(RealtimeClientEvent::append_video_frame(&[0xff, 0xd8]).unwrap())
902 .unwrap();
903 assert_eq!(frame["type"], "input_audio_buffer.append_video_frame");
904 let transcription = serde_json::to_value(
905 RealtimeClientEvent::transcription_session_update(
906 RealtimeTranscriptionSession::default(),
907 )
908 .unwrap(),
909 )
910 .unwrap();
911 assert_eq!(transcription["type"], "transcription_session.update");
912 assert_eq!(transcription["session"]["input_audio_format"], "pcm");
913 let text_item = RealtimeConversationItem::text("user", "hello");
914 assert_eq!(text_item.content[0].kind, "input_text");
915 let output = RealtimeConversationItem::function_output("{\"ok\":true}");
916 assert_eq!(output.kind, "function_call_output");
917 for event in [
918 RealtimeClientEvent::clear().unwrap(),
919 RealtimeClientEvent::create_item(text_item).unwrap(),
920 RealtimeClientEvent::delete_item("item-1").unwrap(),
921 RealtimeClientEvent::retrieve_item("item-1").unwrap(),
922 RealtimeClientEvent::commit().unwrap(),
923 RealtimeClientEvent::create_response().unwrap(),
924 RealtimeClientEvent::cancel_response().unwrap(),
925 ] {
926 assert!(serde_json::to_value(event).unwrap()["type"].is_string());
927 }
928 assert!(RealtimeClientEvent::append_audio(&[]).is_err());
929 assert!(RealtimeClientEvent::append_video_frame(&[]).is_err());
930 assert!(RealtimeClientEvent::delete_item("").is_err());
931 assert!(RealtimeClientEvent::retrieve_item(" ").is_err());
932 }
933
934 #[test]
935 fn decodes_server_event_helpers() {
936 let text: RealtimeServerEvent = serde_json::from_value(serde_json::json!({
937 "type":"response.text.delta","delta":"hello"
938 }))
939 .unwrap();
940 assert_eq!(text.delta_text(), Some("hello"));
941 let transcript: RealtimeServerEvent = serde_json::from_value(serde_json::json!({
942 "type":"response.audio_transcript.delta","delta":"words"
943 }))
944 .unwrap();
945 assert_eq!(transcript.delta_text(), Some("words"));
946 let audio: RealtimeServerEvent = serde_json::from_value(serde_json::json!({
947 "type":"response.audio.delta","delta":"AQI="
948 }))
949 .unwrap();
950 assert_eq!(audio.audio_bytes().unwrap(), Some(vec![1, 2]));
951 let call: RealtimeServerEvent = serde_json::from_value(serde_json::json!({
952 "type":"response.function_call_arguments.done","name":"weather","arguments":"{}"
953 }))
954 .unwrap();
955 assert_eq!(call.function_call().unwrap().name, "weather");
956 let error: RealtimeServerEvent = serde_json::from_value(serde_json::json!({
957 "type":"error","error":{"code":"bad"}
958 }))
959 .unwrap();
960 assert_eq!(error.error().unwrap()["code"], "bad");
961 assert!(text.audio_base64().is_none());
962 assert!(text.audio_bytes().unwrap().is_none());
963 assert!(text.error().is_none());
964 assert!(text.function_call().is_none());
965 let invalid_audio: RealtimeServerEvent = serde_json::from_value(serde_json::json!({
966 "type":"response.audio.delta","delta":"%%%"
967 }))
968 .unwrap();
969 assert!(invalid_audio.audio_bytes().is_err());
970 }
971
972 #[tokio::test]
973 async fn rejects_invalid_realtime_configuration() {
974 assert!(
975 RealtimeConfig::new("key")
976 .url("https://example.com")
977 .connect()
978 .await
979 .is_err()
980 );
981 assert!(
982 RealtimeConfig::new("key")
983 .url("ws://127.0.0.1:1")
984 .connect_timeout(Duration::ZERO)
985 .connect()
986 .await
987 .is_err()
988 );
989 assert!(
990 RealtimeConfig::new("key")
991 .url("ws://127.0.0.1:1")
992 .channel_capacity(0)
993 .connect()
994 .await
995 .is_err()
996 );
997 }
998
999 #[tokio::test]
1000 #[allow(clippy::result_large_err)]
1001 async fn connects_sends_media_and_receives_stream_events() {
1002 let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
1003 let address = listener.local_addr().unwrap();
1004 let authorization = Arc::new(Mutex::new(String::new()));
1005 let captured = authorization.clone();
1006 let server = tokio::spawn(async move {
1007 let (stream, _) = listener.accept().await.unwrap();
1008 let mut socket =
1009 accept_hdr_async(stream, move |request: &Request, response: Response| {
1010 *captured.lock().unwrap() = request
1011 .headers()
1012 .get(AUTHORIZATION)
1013 .unwrap()
1014 .to_str()
1015 .unwrap()
1016 .to_owned();
1017 Ok(response)
1018 })
1019 .await
1020 .unwrap();
1021 socket
1022 .send(Message::Text(
1023 serde_json::json!({
1024 "type":"session.created",
1025 "session":{
1026 "id":"session-1",
1027 "model":"glm-realtime",
1028 "modalities":["text","audio"],
1029 "beta_fields":{"chat_mode":"audio"}
1030 }
1031 })
1032 .to_string()
1033 .into(),
1034 ))
1035 .await
1036 .unwrap();
1037 let mut received = Vec::new();
1038 while received.len() < 14 {
1039 if let Some(Ok(Message::Text(text))) = socket.next().await {
1040 let event: Value = serde_json::from_str(&text).unwrap();
1041 let kind = event["type"].as_str().unwrap().to_owned();
1042 received.push(event);
1043 if kind == "response.create" {
1044 socket.send(Message::Ping(vec![1].into())).await.unwrap();
1045 for event in [
1046 serde_json::json!({"type":"response.text.delta","delta":"hello"}),
1047 serde_json::json!({"type":"response.audio.delta","delta":"AQI="}),
1048 serde_json::json!({"type":"response.done","response":{"status":"completed"}}),
1049 ] {
1050 socket
1051 .send(Message::Text(event.to_string().into()))
1052 .await
1053 .unwrap();
1054 }
1055 }
1056 }
1057 }
1058 while let Some(message) = socket.next().await {
1059 if matches!(message.unwrap(), Message::Close(_)) {
1060 break;
1061 }
1062 }
1063 received
1064 });
1065
1066 let mut connection = RealtimeConfig::new("test-key")
1067 .authentication(ZhipuAuthentication::bearer("test-key"))
1068 .url(format!("ws://{address}"))
1069 .connect_timeout(Duration::from_secs(2))
1070 .channel_capacity(32)
1071 .connect()
1072 .await
1073 .unwrap();
1074 assert!(matches!(
1075 connection.next_typed_event().await.unwrap().unwrap(),
1076 RealtimeServerMessage::SessionCreated { .. }
1077 ));
1078 let sender = connection.sender();
1079 connection
1080 .send(&RealtimeClientEvent::session_update(RealtimeSession::default()).unwrap())
1081 .await
1082 .unwrap();
1083 sender
1084 .update_session(RealtimeSession::default())
1085 .await
1086 .unwrap();
1087 sender.append_audio(&[1, 2]).await.unwrap();
1088 sender.append_audio_base64("AQI=").await.unwrap();
1089 sender.append_video_frame(&[0xff, 0xd8]).await.unwrap();
1090 sender.commit().await.unwrap();
1091 sender.clear_audio().await.unwrap();
1092 sender
1093 .update_transcription_session(RealtimeTranscriptionSession::default())
1094 .await
1095 .unwrap();
1096 sender
1097 .create_item(RealtimeConversationItem::text("user", "hello"))
1098 .await
1099 .unwrap();
1100 sender.delete_item("item-1").await.unwrap();
1101 sender.retrieve_item("item-1").await.unwrap();
1102 sender.create_response().await.unwrap();
1103 sender.cancel_response().await.unwrap();
1104 sender
1105 .send_json(&serde_json::json!({"type":"custom.event"}))
1106 .await
1107 .unwrap();
1108 let mut text = String::new();
1109 let mut audio = Vec::new();
1110 while let Some(event) = connection.next_event().await {
1111 let event = event.unwrap();
1112 if let Some(delta) = event.delta_text() {
1113 text.push_str(delta);
1114 }
1115 if let Some(bytes) = event.audio_bytes().unwrap() {
1116 audio.extend(bytes);
1117 }
1118 if event.event_type == "response.done" {
1119 break;
1120 }
1121 }
1122 assert_eq!(text, "hello");
1123 assert_eq!(audio, vec![1, 2]);
1124 connection.close().await.unwrap();
1125 let received = server.await.unwrap();
1126 assert_eq!(authorization.lock().unwrap().as_str(), "Bearer test-key");
1127 assert_eq!(received[0]["type"], "session.update");
1128 assert_eq!(received[1]["type"], "session.update");
1129 assert_eq!(received[7]["type"], "transcription_session.update");
1130 assert_eq!(received[8]["type"], "conversation.item.create");
1131 assert_eq!(received[11]["type"], "response.create");
1132 assert_eq!(received[13]["type"], "custom.event");
1133 }
1134
1135 #[tokio::test]
1136 async fn typed_sender_methods_enqueue_official_requests() {
1137 let (commands, mut receiver) = mpsc::channel(8);
1138 let sender = RealtimeSender { commands };
1139
1140 sender
1141 .update_typed_session(TypedRealtimeSession::default())
1142 .await
1143 .unwrap();
1144 sender
1145 .create_typed_item(None, TypedRealtimeItem::function_output("call-1", "ok"))
1146 .await
1147 .unwrap();
1148 sender
1149 .create_response_with(RealtimeResponseOptions::default())
1150 .await
1151 .unwrap();
1152 sender
1153 .send_request(&RealtimeRequest::cancel_response().unwrap())
1154 .await
1155 .unwrap();
1156
1157 let mut kinds = Vec::new();
1158 for _ in 0..4 {
1159 let RealtimeCommand::Send(Message::Text(text)) = receiver.recv().await.unwrap() else {
1160 panic!("expected text command");
1161 };
1162 let value: Value = serde_json::from_str(&text).unwrap();
1163 kinds.push(value["type"].as_str().unwrap().to_owned());
1164 }
1165 assert_eq!(
1166 kinds,
1167 [
1168 "session.update",
1169 "conversation.item.create",
1170 "response.create",
1171 "response.cancel"
1172 ]
1173 );
1174 }
1175}