diff --git a/src/server/chat/completions.rs b/src/server/chat/completions.rs index aa7df9d..d5b0a69 100644 --- a/src/server/chat/completions.rs +++ b/src/server/chat/completions.rs @@ -21,10 +21,13 @@ use secrecy::ExposeSecret; use serde::Deserialize; use serde_with::DeserializeFromStr; use serde_with::SerializeDisplay; +use tokio::select; +use tokio::time::sleep; use std::collections::VecDeque; use std::convert::Infallible; use std::fmt::Display; use std::str::FromStr; +use std::time::Duration; use strum::EnumString; use strum::VariantArray; use tokio::runtime::Handle; @@ -682,12 +685,21 @@ pub async fn create_chat_completion( }); } - let upstream_response = match upstream_request.send().await { - Ok(upstream_response) => upstream_response, - Err(upstream_error) => { - tracing::error!(?upstream_error); + let upstream_response = select! { + _ = sleep(Duration::from_secs(20)) => { + tracing::error!("timed out"); continue; } + + upstream_response_result = upstream_request.send() => { + match upstream_response_result { + Ok(upstream_response) => upstream_response, + Err(upstream_error) => { + tracing::error!(?upstream_error); + continue; + } + } + } }; if let Err(upstream_error) = upstream_response.error_for_status_ref() {