From 2ff445e2cfd95f6b6c425d174d3dd4ddad763323 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Domen=20Ko=C5=BEar?= <domen@dev.si>
Date: Mon, 12 Aug 2024 22:17:30 +0100
Subject: [PATCH] Support structured outputs

---
 async-openai/src/types/assistant.rs | 28 ++++++++++++++++++++++++-
 async-openai/src/types/chat.rs      | 32 ++++++++++++++++++++++++++++-
 2 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/async-openai/src/types/assistant.rs b/async-openai/src/types/assistant.rs
index c610734a..3792f589 100644
--- a/async-openai/src/types/assistant.rs
+++ b/async-openai/src/types/assistant.rs
@@ -129,8 +129,33 @@ pub enum AssistantsApiResponseFormatOption {
 /// An object describing the expected output of the model. If `json_object` only `function` type `tools` are allowed to be passed to the Run. If `text` the model can return text or any value needed.
 #[derive(Clone, Serialize, Debug, Deserialize, PartialEq, Default)]
 pub struct AssistantsApiResponseFormat {
-    /// Must be one of `text` or `json_object`.
+    /// Must be one of `text`, `json_object` or `json_schema`.
     pub r#type: AssistantsApiResponseFormatType,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub json_schema: Option<AssistantJsonSchemaObject>,
+}
+
+#[derive(Clone, Serialize, Default, Debug, Deserialize, Builder, PartialEq)]
+#[builder(name = "AssistantJsonSchemaObjectArgs")]
+#[builder(pattern = "mutable")]
+#[builder(setter(into, strip_option), default)]
+#[builder(derive(Debug))]
+#[builder(build_fn(error = "OpenAIError"))]
+pub struct AssistantJsonSchemaObject {
+    /// The name of the JSONSchema. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 64.
+    pub name: String,
+    /// A description of what the function does, used by the model to choose when and how to call the function.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub description: Option<String>,
+    /// The parameters the functions accepts, described as a JSON Schema object. See the [guide](https://platform.openai.com/docs/guides/text-generation/function-calling) for examples, and the [JSON Schema reference](https://json-schema.org/understanding-json-schema/) for documentation about the format.
+    ///
+    /// Omitting `parameters` defines a function with an empty parameter list.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub schema: Option<serde_json::Value>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub strict: Option<bool>,
 }
 
 #[derive(Clone, Serialize, Debug, Deserialize, PartialEq, Default)]
@@ -139,6 +164,7 @@ pub enum AssistantsApiResponseFormatType {
     #[default]
     Text,
     JsonObject,
+    JsonSchema,
 }
 
 /// Retrieval tool
diff --git a/async-openai/src/types/chat.rs b/async-openai/src/types/chat.rs
index 2dc65d16..d2fa2e7f 100644
--- a/async-openai/src/types/chat.rs
+++ b/async-openai/src/types/chat.rs
@@ -263,6 +263,8 @@ pub struct ChatCompletionResponseMessage {
     /// The role of the author of this message.
     pub role: Role,
 
+    pub refusal: Option<String>,
+
     /// Deprecated and replaced by `tool_calls`.
     /// The name and arguments of a function that should be called, as generated by the model.
     #[deprecated]
@@ -312,8 +314,33 @@ pub struct FunctionObject {
 pub enum ChatCompletionResponseFormatType {
     Text,
     JsonObject,
+    JsonSchema,
+}
+
+
+#[derive(Clone, Serialize, Default, Debug, Deserialize, Builder, PartialEq)]
+#[builder(name = "ChatJsonSchemaObjectArgs")]
+#[builder(pattern = "mutable")]
+#[builder(setter(into, strip_option), default)]
+#[builder(derive(Debug))]
+#[builder(build_fn(error = "OpenAIError"))]
+pub struct ChatJsonSchemaObject {
+    /// The name of the JSONSchema. Must be a-z, A-Z, 0-9, or contain underscores and dashes, with a maximum length of 64.
+    pub name: String,
+    /// A description of what the function does, used by the model to choose when and how to call the function.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub description: Option<String>,
+    /// The parameters the functions accepts, described as a JSON Schema object. See the [guide](https://platform.openai.com/docs/guides/text-generation/function-calling) for examples, and the [JSON Schema reference](https://json-schema.org/understanding-json-schema/) for documentation about the format.
+    ///
+    /// Omitting `parameters` defines a function with an empty parameter list.
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub schema: Option<serde_json::Value>,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub strict: Option<bool>,
 }
 
+
 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
 pub struct ChatCompletionResponseFormat {
     /// Setting to `json_object` enables JSON mode. This guarantees that the message the model generates is valid JSON.
@@ -323,8 +350,11 @@ pub struct ChatCompletionResponseFormat {
     /// content may be partial (i.e. cut off) if `finish_reason="length"`, which indicates the generation
     /// exceeded `max_tokens` or the conversation exceeded the max context length.
     ///
-    /// Must be one of `text` or `json_object`.
+    /// Must be one of `text`, `json_schema` or `json_object`.
     pub r#type: ChatCompletionResponseFormatType,
+
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub json_schema: Option<ChatJsonSchemaObject>,
 }
 
 #[derive(Clone, Serialize, Default, Debug, Deserialize, PartialEq)]