Merge pull request #88 from dongri/assistant-beta-v2

Assistant beta v2
dongri · Jun 10, 2024 · bfc0ec2 · bfc0ec2
2 parents 3e3844a + a0bde8f
commit bfc0ec2
Show file tree

Hide file tree

Showing 6 changed files with 130 additions and 23 deletions.
diff --git a/examples/assistant.rs b/examples/assistant.rs
@@ -1,6 +1,6 @@
 use openai_api_rs::v1::api::Client;
 use openai_api_rs::v1::assistant::AssistantRequest;
-use openai_api_rs::v1::common::GPT4_1106_PREVIEW;
+use openai_api_rs::v1::common::GPT4_O;
 use openai_api_rs::v1::message::{CreateMessageRequest, MessageRole};
 use openai_api_rs::v1::run::CreateRunRequest;
 use openai_api_rs::v1::thread::CreateThreadRequest;
@@ -13,31 +13,32 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     let mut tools = HashMap::new();
     tools.insert("type".to_string(), "code_interpreter".to_string());
 
-    let req = AssistantRequest::new(GPT4_1106_PREVIEW.to_string());
+    let req = AssistantRequest::new(GPT4_O.to_string());
     let req = req
         .clone()
         .description("this is a test assistant".to_string());
     let req = req.clone().instructions("You are a personal math tutor. When asked a question, write and run Python code to answer the question.".to_string());
     let req = req.clone().tools(vec![tools]);
-    println!("{:?}", req);
+    println!("AssistantRequest: {:?}", req);
 
     let result = client.create_assistant(req)?;
-    println!("{:?}", result.id);
+    println!("Create Assistant Result ID: {:?}", result.id);
 
     let thread_req = CreateThreadRequest::new();
     let thread_result = client.create_thread(thread_req)?;
-    println!("{:?}", thread_result.id.clone());
+    println!("Create Thread Result ID: {:?}", thread_result.id.clone());
 
     let message_req = CreateMessageRequest::new(
         MessageRole::user,
         "`I need to solve the equation 3x + 11 = 14. Can you help me?".to_string(),
     );
 
     let message_result = client.create_message(thread_result.id.clone(), message_req)?;
-    println!("{:?}", message_result.id.clone());
+    println!("Create Message Result ID: {:?}", message_result.id.clone());
 
     let run_req = CreateRunRequest::new(result.id);
     let run_result = client.create_run(thread_result.id.clone(), run_req)?;
+    println!("Create Run Result ID: {:?}", run_result.id.clone());
 
     loop {
         let run_result = client

diff --git a/src/v1/api.rs b/src/v1/api.rs
@@ -93,7 +93,7 @@ impl Client {
             request = request.with_header("openai-organization", organization);
         }
         if is_beta {
-            request = request.with_header("OpenAI-Beta", "assistants=v1");
+            request = request.with_header("OpenAI-Beta", "assistants=v2");
         }
         if let Some(proxy) = &self.proxy {
             request = request.with_proxy(minreq::Proxy::new(proxy).unwrap());

diff --git a/src/v1/assistant.rs b/src/v1/assistant.rs
@@ -15,7 +15,7 @@ pub struct AssistantRequest {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub tools: Option<Vec<HashMap<String, String>>>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub file_ids: Option<Vec<String>>,
+    pub tool_resources: Option<ToolResource>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub metadata: Option<HashMap<String, String>>,
 }
@@ -28,7 +28,7 @@ impl AssistantRequest {
             description: None,
             instructions: None,
             tools: None,
-            file_ids: None,
+            tool_resources: None,
             metadata: None,
         }
     }
@@ -40,7 +40,7 @@ impl_builder_methods!(
     description: String,
     instructions: String,
     tools: Vec<HashMap<String, String>>,
-    file_ids: Vec<String>,
+    tool_resources: ToolResource,
     metadata: HashMap<String, String>
 );
 
@@ -57,11 +57,44 @@ pub struct AssistantObject {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub instructions: Option<String>,
     pub tools: Vec<HashMap<String, String>>,
-    pub file_ids: Vec<String>,
-    pub metadata: HashMap<String, String>,
+    #[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)]
+pub struct ToolResource {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub code_interpreter: Option<CodeInterpreter>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub file_search: Option<FileSearch>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct CodeInterpreter {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub file_ids: Option<Vec<String>>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct FileSearch {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub vector_store_ids: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub vector_stores: Option<VectorStores>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct VectorStores {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub file_ids: Option<Vec<String>>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub chunking_strategy: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub metadata: Option<HashMap<String, String>>,
+}
+
 #[derive(Debug, Deserialize, Serialize)]
 pub struct DeletionStatus {
     pub id: String,

diff --git a/src/v1/message.rs b/src/v1/message.rs
@@ -8,7 +8,7 @@ pub struct CreateMessageRequest {
     pub role: MessageRole,
     pub content: String,
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub file_ids: Option<Vec<String>>,
+    pub attachments: Option<Vec<Attachment>>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub metadata: Option<HashMap<String, String>>,
 }
@@ -18,15 +18,15 @@ impl CreateMessageRequest {
         Self {
             role,
             content,
-            file_ids: None,
+            attachments: None,
             metadata: None,
         }
     }
 }
 
 impl_builder_methods!(
     CreateMessageRequest,
-    file_ids: Vec<String>,
+    attachments: Vec<Attachment>,
     metadata: HashMap<String, String>
 );
 
@@ -65,11 +65,23 @@ pub struct MessageObject {
     pub assistant_id: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub run_id: Option<String>,
-    pub file_ids: Vec<String>,
-    pub metadata: HashMap<String, String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub attachments: Option<Vec<Attachment>>,
+    pub metadata: Option<HashMap<String, String>>,
     pub headers: Option<HashMap<String, String>>,
 }
 
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct Attachment {
+    pub file_id: Option<String>,
+    pub tools: Vec<Tool>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct Tool {
+    pub r#type: String,
+}
+
 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]
 #[allow(non_camel_case_types)]
 pub enum MessageRole {

diff --git a/src/v1/run.rs b/src/v1/run.rs
@@ -85,7 +85,6 @@ pub struct RunObject {
     pub model: String,
     pub instructions: Option<String>,
     pub tools: Vec<HashMap<String, String>>,
-    pub file_ids: Vec<String>,
     pub metadata: HashMap<String, String>,
     pub headers: Option<HashMap<String, String>>,
 }

diff --git a/src/v1/thread.rs b/src/v1/thread.rs
@@ -8,13 +8,40 @@ pub struct CreateThreadRequest {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub messages: Option<Vec<Message>>,
     #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_resources: Option<ToolResource>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub metadata: Option<HashMap<String, String>>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct ToolResource {
+    pub code_interpreter: Option<CodeInterpreter>,
+    pub file_search: Option<FileSearch>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct CodeInterpreter {
+    pub file_ids: Option<Vec<String>>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct FileSearch {
+    pub vector_store_ids: Option<Vec<String>>,
+    pub vector_stores: Option<VectorStores>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct VectorStores {
+    pub file_ids: Option<Vec<String>>,
+    pub chunking_strategy: Option<String>,
     pub metadata: Option<HashMap<String, String>>,
 }
 
 impl CreateThreadRequest {
     pub fn new() -> Self {
         Self {
             messages: None,
+            tool_resources: None,
             metadata: None,
         }
     }
@@ -27,9 +54,9 @@ impl Default for CreateThreadRequest {
 }
 
 impl_builder_methods!(
-  CreateThreadRequest,
-  messages: Vec<Message>,
-  metadata: HashMap<String, String>
+    CreateThreadRequest,
+    messages: Vec<Message>,
+    tool_resources: ToolResource
 );
 
 #[derive(Debug, Deserialize, Serialize)]
@@ -38,17 +65,52 @@ pub struct ThreadObject {
     pub object: String,
     pub created_at: i64,
     pub metadata: HashMap<String, String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub tool_resources: Option<ToolResource>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub headers: Option<HashMap<String, String>>,
 }
 
 #[derive(Debug, Deserialize, Serialize, Clone)]
 pub struct Message {
+    pub id: String,
+    pub object: String,
+    pub created_at: i64,
+    pub thread_id: String,
     pub role: MessageRole,
-    pub content: String,
+    pub content: Vec<Content>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub assistant_id: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
-    pub file_ids: Option<String>,
+    pub run_id: Option<String>,
     #[serde(skip_serializing_if = "Option::is_none")]
+    pub attachments: Option<Vec<Attachment>>,
     pub metadata: Option<HashMap<String, String>>,
+    pub headers: Option<HashMap<String, String>>,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct Content {
+    #[serde(rename = "type")]
+    pub content_type: String,
+    pub text: ContentText,
+}
+
+#[derive(Debug, Deserialize, Serialize, Clone)]
+pub struct ContentText {
+    pub value: String,
+    pub annotations: Vec<String>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct Attachment {
+    pub file_id: String,
+    pub tools: Vec<Tool>,
+}
+
+#[derive(Serialize, Deserialize, Debug, Clone)]
+pub struct Tool {
+    pub r#type: String,
 }
 
 #[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq)]