diff --git a/examples/chat_stream_cli.rs b/examples/chat_stream_cli.rs index 4cd854c..6ca4e90 100644 --- a/examples/chat_stream_cli.rs +++ b/examples/chat_stream_cli.rs @@ -5,7 +5,7 @@ use openai::{ Credentials, }; use std::io::{stdin, stdout, Write}; -use tokio::sync::mpsc::Receiver; +use tokio::sync::mpsc::{error::TryRecvError, Receiver}; #[tokio::main] async fn main() { @@ -47,25 +47,35 @@ async fn main() { async fn listen_for_tokens(mut chat_stream: Receiver) -> ChatCompletion { let mut merged: Option = None; - while let Some(delta) = chat_stream.recv().await { - let choice = &delta.choices[0]; - if let Some(role) = &choice.delta.role { - print!("{:#?}: ", role); - } - if let Some(content) = &choice.delta.content { - print!("{}", content); - } - if let Some(_) = &choice.finish_reason { - // The message being streamed has been fully received. - print!("\n"); - } - stdout().flush().unwrap(); - // Merge completion into accrued. - match merged.as_mut() { - Some(c) => { - c.merge(delta).unwrap(); + + let mut d = true; + while d { + match chat_stream.try_recv() { + Ok(delta) => { + let choice = &delta.choices[0]; + if let Some(role) = &choice.delta.role { + print!("{:#?}: ", role); + } + if let Some(content) = &choice.delta.content { + print!("{}", content); + } + stdout().flush().unwrap(); + + // Merge completion into accrued. + match merged.as_mut() { + Some(c) => { + c.merge(delta).unwrap(); + } + None => merged = Some(delta), + }; + } + Err(TryRecvError::Empty) => { + let d = std::time::Duration::from_millis(100); + std::thread::sleep(d); + } + Err(TryRecvError::Disconnected) => { + d = false; } - None => merged = Some(delta), }; } merged.unwrap().into()