chore: format chat response (#1125)

* chore: chat response with format

* chore: update prompt

* chore: update test

* chore: update test

* chore: fix stress test

* chore: fix test

* chore: test

* chore: test

* chore: fix stress test

* chore: fix test

.github/workflows/stress_test.yml

@@ -4,14 +4,15 @@ on: [ pull_request ]
 
 concurrency:
   group: stress-test-${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
-  cancel-in-progress: false
+  cancel-in-progress: true
 
 env:
+  SQLX_OFFLINE: true
+  RUST_TOOLCHAIN: "1.80"
   POSTGRES_HOST: localhost
   REDIS_HOST: localhost
-  MINIO_HOST: localhost
-  SQLX_OFFLINE: true
-  RUST_TOOLCHAIN: "1.78"
+  LOCALHOST_GOTRUE: http://localhost/gotrue
+  DATABASE_URL: postgres://postgres:password@localhost:5432/postgres
 
 jobs:
   test:
@@ -25,24 +26,34 @@ jobs:
       - name: Install Rust Toolchain
         uses: dtolnay/rust-toolchain@stable
 
-      - name: Copy and Rename deploy.env to .env
-        run: cp deploy.env .env
+      - name: Copy and Rename dev.env to .env
+        run: cp dev.env .env
+
+      - name: Install Prerequisites
+        run: |
+          brew update
+          brew install libpq
+          brew install sqlx-cli
+          brew install protobuf
 
       - name: Replace Values in .env
         run: |
           sed -i '' 's|RUST_LOG=.*|RUST_LOG=debug|' .env
           sed -i '' 's|API_EXTERNAL_URL=.*|API_EXTERNAL_URL=http://localhost:9999|' .env
+          sed -i '' 's|APPFLOWY_INDEXER_ENABLED=.*|APPFLOWY_INDEXER_ENABLED=false|' .env
+          
           sed -i '' 's|APPFLOWY_GOTRUE_BASE_URL=.*|APPFLOWY_GOTRUE_BASE_URL=http://localhost:9999|' .env
+          sed -i '' 's|GOTRUE_MAILER_AUTOCONFIRM=.*|GOTRUE_MAILER_AUTOCONFIRM=false|' .env
+          sed -i '' 's|APPFLOWY_DATABASE_URL=.*|APPFLOWY_DATABASE_URL=postgres://postgres:password@localhost:5432/postgres|' .env
+          
+          cat .env
         shell: bash
 
       - name: Start Docker Compose Services
         run: |
-          docker compose -f docker-compose-stress-test.yml up -d 
-          docker ps -a
-
-      - name: Install Prerequisites
-        run: |
-          brew install protobuf
+          docker compose -f docker-compose-dev.yml up -d 
+          ./script/code_gen.sh
+          cargo sqlx database create && cargo sqlx migrate run
 
       - name: Run Server and Test
         run: |

dev.env

@@ -116,13 +116,13 @@ CLOUDFLARE_TUNNEL_TOKEN=
 APPFLOWY_AI_OPENAI_API_KEY=
 APPFLOWY_AI_SERVER_PORT=5001
 APPFLOWY_AI_SERVER_HOST=localhost
-APPFLOWY_AI_DATABASE_URL=postgresql+psycopg://postgres:password@postgres:5432/postgres
+APPFLOWY_AI_DATABASE_URL=postgresql+psycopg://postgres:password@localhost:5432/postgres
 APPFLOWY_AI_REDIS_URL=redis://redis:6379
 APPFLOWY_LOCAL_AI_TEST_ENABLED=false
 
 # AppFlowy Indexer
 APPFLOWY_INDEXER_ENABLED=true
-APPFLOWY_INDEXER_DATABASE_URL=postgres://postgres:password@postgres:5432/postgres
+APPFLOWY_INDEXER_DATABASE_URL=postgres://postgres:password@localhost:5432/postgres
 APPFLOWY_INDEXER_REDIS_URL=redis://redis:6379
 APPFLOWY_INDEXER_EMBEDDING_BUFFER_SIZE=5000
 

docker-compose-dev.yml

@@ -98,14 +98,6 @@ services:
     ports:
       - 9999:9999
 
-  portainer:
-    restart: on-failure
-    image: portainer/portainer-ce:latest
-    ports:
-      - 9442:9000
-    volumes:
-      - /var/run/docker.sock:/var/run/docker.sock
-
   pgadmin:
     restart: on-failure
     image: dpage/pgadmin4
@@ -119,16 +111,5 @@ services:
     volumes:
       - ./docker/pgadmin/servers.json:/pgadmin4/servers.json
 
-  ai:
-    restart: on-failure
-    image: appflowyinc/appflowy_ai:${APPFLOWY_AI_VERSION:-latest}
-    ports:
-      - 5001:5001
-    environment:
-      - OPENAI_API_KEY=${APPFLOWY_AI_OPENAI_API_KEY}
-      - APPFLOWY_AI_SERVER_PORT=${APPFLOWY_AI_SERVER_PORT}
-      - APPFLOWY_AI_DATABASE_URL=${APPFLOWY_AI_DATABASE_URL}
-      - APPFLOWY_AI_REDIS_URL=${APPFLOWY_AI_REDIS_URL}
-
 volumes:
   postgres_data:

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

@@ -1,8 +1,8 @@
 use crate::dto::{
   AIModel, CalculateSimilarityParams, ChatAnswer, ChatQuestion, CompleteTextResponse,
   CompletionType, CreateChatContext, CustomPrompt, Document, LocalAIConfig, MessageData,
-  RepeatedLocalAIPackage, RepeatedRelatedQuestion, SearchDocumentsRequest, SimilarityResponse,
-  SummarizeRowResponse, TranslateRowData, TranslateRowResponse,
+  RepeatedLocalAIPackage, RepeatedRelatedQuestion, ResponseFormat, SearchDocumentsRequest,
+  SimilarityResponse, SummarizeRowResponse, TranslateRowData, TranslateRowResponse,
 };
 use crate::error::AIError;
 
@@ -187,6 +187,7 @@ impl AppFlowyAIClient {
         rag_ids: vec![],
         message_id: Some(question_id.to_string()),
       },
+      format: Default::default(),
     };
     let url = format!("{}/chat/message", self.url);
     let resp = self
@@ -216,6 +217,7 @@ impl AppFlowyAIClient {
         rag_ids,
         message_id: None,
       },
+      format: Default::default(),
     };
     let url = format!("{}/chat/message/stream", self.url);
     let resp = self
@@ -245,12 +247,21 @@ impl AppFlowyAIClient {
         rag_ids,
         message_id: Some(question_id.to_string()),
       },
+      format: ResponseFormat::default(),
     };
+    self.stream_question_v3(model, json).await
+  }
+
+  pub async fn stream_question_v3(
+    &self,
+    model: &AIModel,
+    question: ChatQuestion,
+  ) -> 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(&json)
+      .json(&question)
       .timeout(Duration::from_secs(30))
       .send()
       .await?;

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

@@ -7,15 +7,89 @@ 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";
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct SummarizeRowResponse {
   pub text: String,
 }
 
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct ChatQuestionQuery {
+  pub chat_id: String,
+  pub question_id: i64,
+  pub format: ResponseFormat,
+}
+
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct ChatQuestion {
   pub chat_id: String,
   pub data: MessageData,
+  #[serde(default)]
+  pub format: ResponseFormat,
+}
+
+#[derive(Clone, Default, Debug, Serialize, Deserialize)]
+pub struct ResponseFormat {
+  pub output_layout: OutputLayout,
+  pub output_content: OutputContent,
+  pub output_content_metadata: Option<OutputContentMetadata>,
+}
+
+#[derive(Clone, Debug, Default, Serialize_repr, Deserialize_repr)]
+#[repr(u8)]
+pub enum OutputLayout {
+  #[default]
+  Paragraph = 0,
+  BulletList = 1,
+  NumberedList = 2,
+  SimpleTable = 3,
+}
+
+#[derive(Clone, Debug, Default, Serialize_repr, Deserialize_repr)]
+#[repr(u8)]
+pub enum OutputContent {
+  #[default]
+  TEXT = 0,
+  IMAGE = 1,
+  RichTextImage = 2,
+}
+
+#[derive(Clone, Default, Debug, Serialize, Deserialize)]
+pub struct OutputContentMetadata {
+  /// Custom prompt for image generation.
+  #[serde(default, skip_serializing_if = "Option::is_none")]
+  pub custom_image_prompt: Option<String>,
+
+  /// The image model to use for generation (default: "dall-e-2").
+  #[serde(default = "default_image_model")]
+  pub image_model: String,
+
+  /// Size of the image (default: "256x256").
+  #[serde(
+    default = "default_image_size",
+    skip_serializing_if = "Option::is_none"
+  )]
+  pub size: Option<String>,
+
+  /// Quality of the image (default: "standard").
+  #[serde(
+    default = "default_image_quality",
+    skip_serializing_if = "Option::is_none"
+  )]
+  pub quality: Option<String>,
+}
+
+// Default values for the fields
+fn default_image_model() -> String {
+  "dall-e-2".to_string()
+}
+
+fn default_image_size() -> Option<String> {
+  Some("256x256".to_string())
+}
+
+fn default_image_quality() -> Option<String> {
+  Some("standard".to_string())
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize)]