From 85c3611afbbe8df30dfc7512cc381ed304ce4024 Mon Sep 17 00:00:00 2001
From: Nick Anderson <nicholasanderson5@gmail.com>
Date: Tue, 14 May 2024 22:20:41 -0500
Subject: [PATCH] feat(chat): add response_format for JSON return values

---
 src/v1/chat.rs | 36 ++++++++++++++++++++++++++++++++++--
 1 file changed, 34 insertions(+), 2 deletions(-)

diff --git a/src/v1/chat.rs b/src/v1/chat.rs
index 80c32aa..ebb309f 100644
--- a/src/v1/chat.rs
+++ b/src/v1/chat.rs
@@ -38,6 +38,20 @@ pub enum ChatMessageRole {
     User,
 }
 
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ResponseFormat {
+    #[serde(rename = "type")]
+    pub type_: String,
+}
+
+impl ResponseFormat {
+    pub fn json_object() -> Self {
+        Self {
+            type_: "json_object".to_string(),
+        }
+    }
+}
+
 // -----------------------------------------------------------------------------
 // Request
 
@@ -50,6 +64,7 @@ pub struct ChatParams {
     pub tool_choice: Option<tool::ToolChoice>,
     pub tools: Option<Vec<tool::Tool>>,
     pub top_p: Option<f32>,
+    pub response_format: Option<ResponseFormat>,
 }
 impl Default for ChatParams {
     fn default() -> Self {
@@ -61,6 +76,21 @@ impl Default for ChatParams {
             tool_choice: None,
             tools: None,
             top_p: None,
+            response_format: None,
+        }
+    }
+}
+impl ChatParams {
+    pub fn json_default() -> Self {
+        Self {
+            max_tokens: None,
+            random_seed: None,
+            safe_prompt: None,
+            temperature: None,
+            tool_choice: None,
+            tools: None,
+            top_p: None,
+            response_format: Some(ResponseFormat::json_object()),
         }
     }
 }
@@ -87,8 +117,8 @@ pub struct ChatRequest {
     pub top_p: Option<f32>,
     // TODO Check this prop (seen in official Python client but not in API doc).
     // pub tool_choice: Option<String>,
-    // TODO Check this prop (seen in official Python client but not in API doc).
-    // pub response_format: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub response_format: Option<ResponseFormat>,
 }
 impl ChatRequest {
     pub fn new(
@@ -105,6 +135,7 @@ impl ChatRequest {
             tool_choice,
             tools,
             top_p,
+            response_format,
         } = options.unwrap_or_default();
 
         Self {
@@ -119,6 +150,7 @@ impl ChatRequest {
             tool_choice,
             tools,
             top_p,
+            response_format,
         }
     }
 }