Merge pull request #118 from dongri/refactor-function-types

refactoring function
dongri · Oct 15, 2024 · e46412c · e46412c
2 parents 34aae91 + 5867baa
commit e46412c
Show file tree

Hide file tree

Showing 6 changed files with 106 additions and 58 deletions.
diff --git a/examples/function_call.rs b/examples/function_call.rs
@@ -1,6 +1,7 @@
 use openai_api_rs::v1::api::OpenAIClient;
 use openai_api_rs::v1::chat_completion::{self, ChatCompletionRequest};
 use openai_api_rs::v1::common::GPT4_O;
+use openai_api_rs::v1::types;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::{env, vec};
@@ -21,8 +22,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let mut properties = HashMap::new();
     properties.insert(
         "coin".to_string(),
-        Box::new(chat_completion::JSONSchemaDefine {
-            schema_type: Some(chat_completion::JSONSchemaType::String),
+        Box::new(types::JSONSchemaDefine {
+            schema_type: Some(types::JSONSchemaType::String),
             description: Some("The cryptocurrency to get the price of".to_string()),
             ..Default::default()
         }),
@@ -40,11 +41,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     )
     .tools(vec![chat_completion::Tool {
         r#type: chat_completion::ToolType::Function,
-        function: chat_completion::Function {
+        function: types::Function {
             name: String::from("get_coin_price"),
             description: Some(String::from("Get the price of a cryptocurrency")),
-            parameters: chat_completion::FunctionParameters {
-                schema_type: chat_completion::JSONSchemaType::Object,
+            parameters: types::FunctionParameters {
+                schema_type: types::JSONSchemaType::Object,
                 properties: Some(properties),
                 required: Some(vec![String::from("coin")]),
             },

diff --git a/examples/function_call_role.rs b/examples/function_call_role.rs
@@ -1,6 +1,7 @@
 use openai_api_rs::v1::api::OpenAIClient;
 use openai_api_rs::v1::chat_completion::{self, ChatCompletionRequest};
 use openai_api_rs::v1::common::GPT4_O;
+use openai_api_rs::v1::types;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::{env, vec};
@@ -21,8 +22,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let mut properties = HashMap::new();
     properties.insert(
         "coin".to_string(),
-        Box::new(chat_completion::JSONSchemaDefine {
-            schema_type: Some(chat_completion::JSONSchemaType::String),
+        Box::new(types::JSONSchemaDefine {
+            schema_type: Some(types::JSONSchemaType::String),
             description: Some("The cryptocurrency to get the price of".to_string()),
             ..Default::default()
         }),
@@ -40,11 +41,11 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     )
     .tools(vec![chat_completion::Tool {
         r#type: chat_completion::ToolType::Function,
-        function: chat_completion::Function {
+        function: types::Function {
             name: String::from("get_coin_price"),
             description: Some(String::from("Get the price of a cryptocurrency")),
-            parameters: chat_completion::FunctionParameters {
-                schema_type: chat_completion::JSONSchemaType::Object,
+            parameters: types::FunctionParameters {
+                schema_type: types::JSONSchemaType::Object,
                 properties: Some(properties),
                 required: Some(vec![String::from("coin")]),
             },

diff --git a/src/v1/assistant.rs b/src/v1/assistant.rs
@@ -1,6 +1,7 @@
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
+use super::types;
 use crate::impl_builder_methods;
 
 #[derive(Debug, Serialize, Clone)]
@@ -56,13 +57,53 @@ pub struct AssistantObject {
     pub model: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub instructions: Option<String>,
-    pub tools: Vec<HashMap<String, String>>,
+    pub tools: Vec<Tools>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub tool_resources: Option<ToolResource>,
     pub metadata: Option<HashMap<String, String>>,
     pub headers: Option<HashMap<String, String>>,
 }
 
+#[derive(Debug, Deserialize, Serialize, Clone)]
+#[serde(tag = "type")]
+#[serde(rename_all = "snake_case")]
+pub enum Tools {
+    CodeInterpreter,
+    FileSearch(ToolsFileSearch),
+    Function(ToolsFunction),
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct ToolsFileSearch {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub file_search: Option<ToolsFileSearchObject>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct ToolsFunction {
+    pub function: types::Function,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct ToolsFileSearchObject {
+    pub max_num_results: Option<u8>,
+    pub ranking_options: Option<FileSearchRankingOptions>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct FileSearchRankingOptions {
+    pub ranker: Option<FileSearchRanker>,
+    pub score_threshold: Option<f32>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub enum FileSearchRanker {
+    #[serde(rename = "auto")]
+    Auto,
+    #[serde(rename = "default_2024_08_21")]
+    Default2024_08_21,
+}
+
 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct ToolResource {
     #[serde(skip_serializing_if = "Option::is_none")]

diff --git a/src/v1/chat_completion.rs b/src/v1/chat_completion.rs
@@ -1,5 +1,6 @@
+use super::{common, types};
 use crate::impl_builder_methods;
-use crate::v1::common;
+
 use serde::de::{self, MapAccess, SeqAccess, Visitor};
 use serde::ser::SerializeMap;
 use serde::{Deserialize, Deserializer, Serialize, Serializer};
@@ -185,6 +186,7 @@ impl<'de> Deserialize<'de> for Content {
         deserializer.deserialize_any(ContentVisitor)
     }
 }
+
 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
 #[allow(non_camel_case_types)]
 pub enum ContentType {
@@ -251,51 +253,6 @@ pub struct ChatCompletionResponse {
     pub headers: Option<HashMap<String, String>>,
 }
 
-#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
-pub struct Function {
-    pub name: String,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub description: Option<String>,
-    pub parameters: FunctionParameters,
-}
-
-#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
-#[serde(rename_all = "lowercase")]
-pub enum JSONSchemaType {
-    Object,
-    Number,
-    String,
-    Array,
-    Null,
-    Boolean,
-}
-
-#[derive(Debug, Deserialize, Serialize, Clone, Default, PartialEq, Eq)]
-pub struct JSONSchemaDefine {
-    #[serde(rename = "type")]
-    pub schema_type: Option<JSONSchemaType>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub description: Option<String>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub enum_values: Option<Vec<String>>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub properties: Option<HashMap<String, Box<JSONSchemaDefine>>>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub required: Option<Vec<String>>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub items: Option<Box<JSONSchemaDefine>>,
-}
-
-#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
-pub struct FunctionParameters {
-    #[serde(rename = "type")]
-    pub schema_type: JSONSchemaType,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub properties: Option<HashMap<String, Box<JSONSchemaDefine>>>,
-    #[serde(skip_serializing_if = "Option::is_none")]
-    pub required: Option<Vec<String>>,
-}
-
 #[derive(Debug, Deserialize, Serialize, PartialEq, Eq)]
 #[allow(non_camel_case_types)]
 pub enum FinishReason {
@@ -352,7 +309,7 @@ where
 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
 pub struct Tool {
     pub r#type: ToolType,
-    pub function: Function,
+    pub function: types::Function,
 }
 
 #[derive(Debug, Deserialize, Serialize, Copy, Clone, PartialEq, Eq)]

diff --git a/src/v1/mod.rs b/src/v1/mod.rs
@@ -1,5 +1,6 @@
 pub mod common;
 pub mod error;
+pub mod types;
 
 pub mod audio;
 pub mod batch;

diff --git a/src/v1/types.rs b/src/v1/types.rs
@@ -0,0 +1,47 @@
+use serde::{Deserialize, Serialize};
+use std::collections::HashMap;
+
+#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
+pub struct Function {
+    pub name: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub description: Option<String>,
+    pub parameters: FunctionParameters,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
+pub struct FunctionParameters {
+    #[serde(rename = "type")]
+    pub schema_type: JSONSchemaType,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub properties: Option<HashMap<String, Box<JSONSchemaDefine>>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub required: Option<Vec<String>>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
+#[serde(rename_all = "lowercase")]
+pub enum JSONSchemaType {
+    Object,
+    Number,
+    String,
+    Array,
+    Null,
+    Boolean,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone, Default, PartialEq, Eq)]
+pub struct JSONSchemaDefine {
+    #[serde(rename = "type")]
+    pub schema_type: Option<JSONSchemaType>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub description: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub enum_values: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub properties: Option<HashMap<String, Box<JSONSchemaDefine>>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub required: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub items: Option<Box<JSONSchemaDefine>>,
+}