From 45fdee79c6097734e30fc772c3a0c48eda24313a Mon Sep 17 00:00:00 2001
From: LuoShui <luoshuijs@outlook.com>
Date: Wed, 27 Mar 2024 11:37:48 +0800
Subject: [PATCH] :bug: FIx chat validation about new chatgpt

---
 crates/openai/src/serve/proxy/req.rs | 30 ++++++++++++++++++++++++++++
 1 file changed, 30 insertions(+)

diff --git a/crates/openai/src/serve/proxy/req.rs b/crates/openai/src/serve/proxy/req.rs
index 79f1f13fd..df61d3469 100644
--- a/crates/openai/src/serve/proxy/req.rs
+++ b/crates/openai/src/serve/proxy/req.rs
@@ -19,6 +19,7 @@ use super::header_convert;
 use super::toapi;
 use crate::serve::error::{ProxyError, ResponseError};
 use crate::serve::puid::{get_or_init, reduce_key};
+use crate::URL_CHATGPT_API;
 
 #[async_trait]
 impl SendRequestExt for reqwest::Client {
@@ -119,6 +120,15 @@ async fn handle_conv_request(req: &mut RequestExt) -> Result<(), ResponseError>
         }
     }
 
+    let chat_requirements_token = create_chat_requirements_token(&token).await?;
+    if let Some(chat_requirements_token) = chat_requirements_token {
+        req.headers.insert(
+            header::HeaderName::from_static("openai-sentinel-chat-requirements-token"),
+            header::HeaderValue::from_str(chat_requirements_token.as_str())
+                .map_err(ResponseError::BadRequest)?,
+        );
+    }
+
     // Parse model
     let model = GPTModel::from_str(model).map_err(ResponseError::BadRequest)?;
 
@@ -194,3 +204,23 @@ async fn handle_dashboard_request(req: &mut RequestExt) -> Result<(), ResponseEr
 
     Ok(())
 }
+
+async fn create_chat_requirements_token(token: &str) -> Result<Option<String>, ResponseError> {
+    let token = token.trim_start_matches("Bearer ");
+    let resp = with_context!(api_client)
+        .get(format!("{URL_CHATGPT_API}/backend-api/sentinel/chat-requirements"))
+        .bearer_auth(token)
+        .send()
+        .await
+        .map_err(ResponseError::InternalServerError)?
+        .error_for_status()
+        .map_err(ResponseError::BadRequest)?;
+    let body = resp.bytes().await?;
+    let json = serde_json::from_slice::<Value>(&body).map_err(ResponseError::BadRequest)?;
+    if let Some(token_value) = json.get("token") {
+        if let Some(token_str) = token_value.as_str() {
+            return Ok(Some(token_str.to_owned()));
+        }
+    }
+    Ok(None)
+}