chore: chat stream timeout (#1137)

* chore: increase timeout

* chore: keepalive event

libs/appflowy-ai-client/src/client.rs

@@ -263,20 +263,21 @@ impl AppFlowyAIClient {
         rag_ids,
       },
     };
-    self.stream_question_v3(model, json).await
+    self.stream_question_v3(model, json, Some(30)).await
   }
 
   pub async fn stream_question_v3(
     &self,
     model: &AIModel,
     question: ChatQuestion,
+    timeout_secs: Option<u64>,
   ) -> Result<impl Stream<Item = Result<Bytes, AIError>>, AIError> {
     let url = format!("{}/v2/chat/message/stream", self.url);
     let resp = self
       .async_http_client(Method::POST, &url)?
       .header(AI_MODEL_HEADER_KEY, model.to_str())
       .json(&question)
-      .timeout(Duration::from_secs(30))
+      .timeout(Duration::from_secs(timeout_secs.unwrap_or(30)))
       .send()
       .await?;
     AIResponse::<()>::stream_response(resp).await

libs/appflowy-ai-client/src/dto.rs

@@ -8,6 +8,7 @@ use std::str::FromStr;
 pub const STREAM_METADATA_KEY: &str = "0";
 pub const STREAM_ANSWER_KEY: &str = "1";
 pub const STREAM_IMAGE_KEY: &str = "2";
+pub const STREAM_KEEP_ALIVE_KEY: &str = "3";
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct SummarizeRowResponse {
   pub text: String,

libs/client-api-test/src/test_client.rs

@@ -1271,6 +1271,7 @@ pub async fn collect_answer(mut stream: QuestionStream) -> String {
         answer.push_str(&value);
       },
       QuestionStreamValue::Metadata { .. } => {},
+      QuestionStreamValue::KeepAlive => {},
     }
   }
   answer

libs/client-api/src/http_chat.rs

@@ -12,7 +12,7 @@ use reqwest::Method;
 use serde_json::Value;
 use shared_entity::dto::ai_dto::{
   CalculateSimilarityParams, ChatQuestionQuery, RepeatedRelatedQuestion, SimilarityResponse,
-  STREAM_ANSWER_KEY, STREAM_IMAGE_KEY, STREAM_METADATA_KEY,
+  STREAM_ANSWER_KEY, STREAM_IMAGE_KEY, STREAM_KEEP_ALIVE_KEY, STREAM_METADATA_KEY,
 };
 use shared_entity::dto::chat_dto::{ChatSettings, UpdateChatParams};
 use shared_entity::response::{AppResponse, AppResponseError};
@@ -366,6 +366,7 @@ pub enum QuestionStreamValue {
   Metadata {
     value: serde_json::Value,
   },
+  KeepAlive,
 }
 impl Stream for QuestionStream {
   type Item = Result<QuestionStreamValue, AppResponseError>;
@@ -394,6 +395,10 @@ impl Stream for QuestionStream {
             return Poll::Ready(Some(Ok(QuestionStreamValue::Answer { value: image })));
           }
 
+          if value.remove(STREAM_KEEP_ALIVE_KEY).is_some() {
+            return Poll::Ready(Some(Ok(QuestionStreamValue::KeepAlive)));
+          }
+
           error!("Invalid streaming value: {:?}", value);
           Poll::Ready(None)
         },

libs/infra/src/reqwest.rs

@@ -92,7 +92,6 @@ where
       // Poll for the next chunk of data from the underlying stream
       match ready!(this.stream.as_mut().poll_next(cx)) {
         Some(Ok(bytes)) => {
-          // Append the new bytes to the buffer
           this.buffer.extend_from_slice(&bytes);
 
           // Create a StreamDeserializer to deserialize the bytes into T
@@ -112,7 +111,6 @@ where
               return Poll::Pending;
             },
             Some(Err(err)) => {
-              // Return other deserialization errors wrapped in SE
               return Poll::Ready(Some(Err(err.into())));
             },
             None => {

src/api/chat.rs

@@ -380,7 +380,7 @@ async fn answer_stream_v3_handler(
   trace!("[Chat] stream v3 {:?}", question);
   match state
     .ai_client
-    .stream_question_v3(&ai_model, question)
+    .stream_question_v3(&ai_model, question, Some(60))
     .await
   {
     Ok(answer_stream) => {

tests/ai_test/chat_test.rs

@@ -450,6 +450,7 @@ async fn collect_answer(mut stream: QuestionStream) -> String {
         answer.push_str(&value);
       },
       QuestionStreamValue::Metadata { .. } => {},
+      QuestionStreamValue::KeepAlive => {},
     }
   }
   answer