feat: timeout

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() {