diff --git a/libs/appflowy-ai-client/src/dto.rs b/libs/appflowy-ai-client/src/dto.rs
index 6f4e2adc..8e354fb2 100644
--- a/libs/appflowy-ai-client/src/dto.rs
+++ b/libs/appflowy-ai-client/src/dto.rs
@@ -447,7 +447,16 @@ pub struct SimilarityResponse {
 
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct CompletionMetadata {
+  /// A unique identifier for the object.
   pub object_id: String,
+  /// The workspace identifier.
+  ///
+  /// This field must be provided when generating images.
+  pub workspace_id: Option<String>,
+  /// A list of relevant document IDs.
+  ///
+  /// When using completions for document-related tasks, this should include the document ID.
+  /// In some cases, `object_id` may be the same as the document ID.
   pub rag_ids: Option<Vec<String>>,
 }
 
@@ -459,6 +468,8 @@ pub struct CompleteTextParams {
   #[serde(default)]
   #[serde(skip_serializing_if = "Option::is_none")]
   pub metadata: Option<CompletionMetadata>,
+  #[serde(default)]
+  pub format: ResponseFormat,
 }
 
 impl CompleteTextParams {
@@ -472,6 +483,7 @@ impl CompleteTextParams {
       completion_type: Some(completion_type),
       custom_prompt: None,
       metadata,
+      format: Default::default(),
     }
   }
 }
diff --git a/libs/appflowy-ai-client/tests/chat_test/completion_test.rs b/libs/appflowy-ai-client/tests/chat_test/completion_test.rs
index 2efa1965..d07bf182 100644
--- a/libs/appflowy-ai-client/tests/chat_test/completion_test.rs
+++ b/libs/appflowy-ai-client/tests/chat_test/completion_test.rs
@@ -1,6 +1,35 @@
 use crate::appflowy_ai_client;
 use appflowy_ai_client::client::collect_stream_text;
-use appflowy_ai_client::dto::{CompleteTextParams, CompletionType};
+use appflowy_ai_client::dto::{
+  CompleteTextParams, CompletionMetadata, CompletionType, OutputContent, OutputLayout,
+  ResponseFormat,
+};
+
+#[tokio::test]
+async fn completion_image_test() {
+  let client = appflowy_ai_client();
+  let params = CompleteTextParams {
+    text: "A yellow cat".to_string(),
+    completion_type: Some(CompletionType::ImproveWriting),
+    custom_prompt: None,
+    metadata: Some(CompletionMetadata {
+      object_id: uuid::Uuid::new_v4().to_string(),
+      workspace_id: Some(uuid::Uuid::new_v4().to_string()),
+      rag_ids: None,
+    }),
+    format: ResponseFormat {
+      output_content: OutputContent::IMAGE,
+      ..Default::default()
+    },
+  };
+  let stream = client
+    .stream_completion_text(params, "gpt-4o-mini")
+    .await
+    .unwrap();
+  let text = collect_stream_text(stream).await;
+  assert!(text.contains("http://localhost:8000"));
+}
+
 #[tokio::test]
 async fn continue_writing_test() {
   let client = appflowy_ai_client();
@@ -9,6 +38,10 @@ async fn continue_writing_test() {
     completion_type: Some(CompletionType::ImproveWriting),
     custom_prompt: None,
     metadata: None,
+    format: ResponseFormat {
+      output_layout: OutputLayout::SimpleTable,
+      ..Default::default()
+    },
   };
   let stream = client
     .stream_completion_text(params, "gpt-4o-mini")
@@ -27,6 +60,7 @@ async fn improve_writing_test() {
     completion_type: Some(CompletionType::ImproveWriting),
     custom_prompt: None,
     metadata: None,
+    format: ResponseFormat::default(),
   };
   let stream = client
     .stream_completion_text(params, "gpt-4o-mini")
@@ -47,6 +81,7 @@ async fn make_text_shorter_text() {
         completion_type: Some(CompletionType::MakeShorter),
         custom_prompt: None,
         metadata: None,
+        format: ResponseFormat::default(),
     };
   let stream = client
     .stream_completion_text(params, "gpt-4o-mini")