From bbe3d5cb43765253744e4481c3f41eb5d43bc215 Mon Sep 17 00:00:00 2001 From: rauner <29773634+rauner@users.noreply.github.com> Date: Tue, 30 Jan 2024 12:14:40 +0100 Subject: [PATCH] refactor files specific methods into Files struct --- src/assistant.rs | 223 +++++++++++++++++++++++------------------------ src/main.rs | 17 +++- 2 files changed, 123 insertions(+), 117 deletions(-) diff --git a/src/assistant.rs b/src/assistant.rs index 6f1445d..4ee6de3 100644 --- a/src/assistant.rs +++ b/src/assistant.rs @@ -45,11 +45,6 @@ impl From for AssistantError { AssistantError::OpenAIError(err.to_string()) } } -// Define the response type for the file upload response. -#[derive(Deserialize)] -struct FileUploadResponse { - id: String, -} // Define the response type for attaching files to an assistant. #[derive(Serialize)] struct AttachFilesRequest { @@ -115,85 +110,18 @@ pub struct AssistantChatResponse { pub messages: Vec, } -/// A struct representing an OpenAI assistant. -/// The tools are currently hardcoded as a code_interpreter. -pub struct Assistant { - pub id: String, - name: String, - model: String, - instructions: String, +// Define the response type for the file upload response. +#[derive(Deserialize)] +struct FileUploadResponse { + id: String, +} + +pub struct Files { folder_path: String, file_ids: Vec, scrape_urls: Vec, } -impl Assistant { - /// create an OpenAI assistant and set the assistant's ID - pub async fn initialize(&mut self) -> Result<(), AssistantError> { - let client = Client::new(); - let api_key = env::var("OPENAI_API_KEY") - .map_err(|_| AssistantError::OpenAIError("OPENAI_API_KEY not set".to_string()))?; - let payload = json!({ - "instructions": self.instructions, - "name": self.name, - "tools": [{"type": "code_interpreter"}], - "model": self.model, - }); - let response = client - .post("https://api.openai.com/v1/assistants") - .header("OpenAI-Beta", "assistants=v1") - .bearer_auth(&api_key) - .json(&payload) - .send() - .await; - match response { - Ok(res) if res.status().is_success() => match res.json::().await { - Ok(assistant_response) => { - if let Some(id) = assistant_response["id"].as_str() { - self.id = id.to_string(); - Ok(()) - } else { - Err(AssistantError::OpenAIError( - "Failed to extract assistant ID from response".to_string(), - )) - } - } - Err(_) => Err(AssistantError::OpenAIError( - "Failed to parse response from OpenAI".to_string(), - )), - }, - Ok(res) => { - let error_message = res.text().await.unwrap_or_default(); - Err(AssistantError::OpenAIError(error_message)) - } - Err(e) => Err(AssistantError::OpenAIError(format!( - "Failed to send request to OpenAI: {}", - e - ))), - } - } - /// Delete the OpenAI assistant with the given ID - pub async fn delete(&self) -> Result<(), AssistantError> { - let api_key = env::var("OPENAI_API_KEY") - .map_err(|_| AssistantError::OpenAIError("OPENAI_API_KEY not set".to_string()))?; - let client = Client::new(); - let response = client - .delete(format!("https://api.openai.com/v1/assistants/{}", self.id)) - .header("OpenAI-Beta", "assistants=v1") - .bearer_auth(&api_key) - .send() - .await; - match response { - Ok(res) if res.status().is_success() => Ok(()), - Ok(res) => { - let error_message = res.text().await.unwrap_or_default(); - Err(AssistantError::OpenAIError(error_message)) - } - Err(e) => Err(AssistantError::OpenAIError(format!( - "Failed to send DELETE request to OpenAI: {}", - e - ))), - } - } +impl Files { pub async fn scrape_context(&self) -> Result<(), AssistantError> { let client = Client::new(); let folder_path = Path::new(&self.folder_path); @@ -229,7 +157,6 @@ impl Assistant { } Ok(()) } - /// upload a all files in a folder and return a vector of file IDs pub async fn upload_files(&mut self) -> Result<(), AssistantError> { let api_key = env::var("OPENAI_API_KEY") .map_err(|_| AssistantError::OpenAIError("OPENAI_API_KEY not set".to_string()))?; @@ -241,10 +168,8 @@ impl Assistant { .map_err(|e| AssistantError::DatabaseError(e.to_string()))? .path(); if path.is_file() { - // Read the file content into memory let file_content = fs::read(&path).map_err(|e| AssistantError::OpenAIError(e.to_string()))?; - // Get the file name as an owned String let file_name = path .file_name() .ok_or_else(|| { @@ -257,15 +182,10 @@ impl Assistant { ) })? .to_owned(); - // Create a Part from the file content let part = Part::bytes(file_content) - .file_name(file_name) // Set the file name for the part - .mime_str("application/octet-stream")?; // Set the MIME type for the part - // Create a Form and add the Part to it - let form = Form::new() - .part("file", part) - .text("purpose", "assistants"); - // Send the multipart form as the body of a POST request + .file_name(file_name) + .mime_str("application/octet-stream")?; + let form = Form::new().part("file", part).text("purpose", "assistants"); let response = client .post("https://api.openai.com/v1/files") .header("OpenAI-Beta", "assistants=v1") @@ -293,12 +213,10 @@ impl Assistant { } Ok(()) } - /// delelte all files that where uploaded when creating the assistant pub async fn delete_files(&mut self) -> Result<(), AssistantError> { let api_key = env::var("OPENAI_API_KEY") .map_err(|_| AssistantError::OpenAIError("OPENAI_API_KEY not set".to_string()))?; let client = Client::new(); - // Iterate over the file IDs and delete each file for file_id in &self.file_ids { let response = client .delete(format!("https://api.openai.com/v1/files/{}", file_id)) @@ -306,15 +224,12 @@ impl Assistant { .send() .await; match response { - Ok(res) if res.status().is_success() => { - } + Ok(res) if res.status().is_success() => {} Ok(res) => { - // API returned an error status, handle it let error_message = res.text().await.unwrap_or_default(); return Err(AssistantError::OpenAIError(error_message)); } Err(e) => { - // Network or other error occurred, handle it return Err(AssistantError::OpenAIError(format!( "Failed to send DELETE request to OpenAI: {}", e @@ -322,17 +237,92 @@ impl Assistant { } } } - // Clear the file_ids vector since all files have been deleted self.file_ids.clear(); Ok(()) } +} +/// A struct representing an OpenAI assistant. +/// The tools are currently hardcoded as a code_interpreter. +pub struct Assistant { + pub id: String, + name: String, + model: String, + instructions: String, +} +impl Assistant { + /// create an OpenAI assistant and set the assistant's ID + pub async fn initialize(&mut self) -> Result<(), AssistantError> { + let client = Client::new(); + let api_key = env::var("OPENAI_API_KEY") + .map_err(|_| AssistantError::OpenAIError("OPENAI_API_KEY not set".to_string()))?; + let payload = json!({ + "instructions": self.instructions, + "name": self.name, + "tools": [{"type": "code_interpreter"}], + "model": self.model, + }); + let response = client + .post("https://api.openai.com/v1/assistants") + .header("OpenAI-Beta", "assistants=v1") + .bearer_auth(&api_key) + .json(&payload) + .send() + .await; + match response { + Ok(res) if res.status().is_success() => match res.json::().await { + Ok(assistant_response) => { + if let Some(id) = assistant_response["id"].as_str() { + self.id = id.to_string(); + Ok(()) + } else { + Err(AssistantError::OpenAIError( + "Failed to extract assistant ID from response".to_string(), + )) + } + } + Err(_) => Err(AssistantError::OpenAIError( + "Failed to parse response from OpenAI".to_string(), + )), + }, + Ok(res) => { + let error_message = res.text().await.unwrap_or_default(); + Err(AssistantError::OpenAIError(error_message)) + } + Err(e) => Err(AssistantError::OpenAIError(format!( + "Failed to send request to OpenAI: {}", + e + ))), + } + } + /// Delete the OpenAI assistant with the given ID + pub async fn delete(&self) -> Result<(), AssistantError> { + let api_key = env::var("OPENAI_API_KEY") + .map_err(|_| AssistantError::OpenAIError("OPENAI_API_KEY not set".to_string()))?; + let client = Client::new(); + let response = client + .delete(format!("https://api.openai.com/v1/assistants/{}", self.id)) + .header("OpenAI-Beta", "assistants=v1") + .bearer_auth(&api_key) + .send() + .await; + match response { + Ok(res) if res.status().is_success() => Ok(()), + Ok(res) => { + let error_message = res.text().await.unwrap_or_default(); + Err(AssistantError::OpenAIError(error_message)) + } + Err(e) => Err(AssistantError::OpenAIError(format!( + "Failed to send DELETE request to OpenAI: {}", + e + ))), + } + } - /// Attach the files with IDs stored in the file_ids field to the assistant. - pub async fn attach_files(&self) -> Result<(), AssistantError> { + pub async fn attach_files(&self, file_ids: &[String]) -> Result<(), AssistantError> { let api_key = env::var("OPENAI_API_KEY") .map_err(|_| AssistantError::OpenAIError("OPENAI_API_KEY not set".to_string()))?; let client = Client::new(); - for file_id in &self.file_ids { + for file_id in file_ids { let payload = json!({ "file_id": file_id }); let response = client .post(format!( @@ -356,32 +346,37 @@ impl Assistant { Ok(()) } } + +pub async fn create_files(folder_path: &str, scrape_urls: Vec) -> Result { + // Initialize the Files struct directly + let mut files = Files { + folder_path: folder_path.to_string(), + file_ids: Vec::new(), // Initially empty, will be filled during file upload + scrape_urls, // Provided scrape URLs + }; + // Scrape the context from the provided URLs + files.scrape_context().await?; + // Upload the scraped files to OpenAI + files.upload_files().await?; + Ok(files) +} pub async fn create_assistant( assistant_name: &str, model: &str, instructions: &str, - folder_path: &str, + files: Files, // Add this parameter to accept a Files struct ) -> Result { let mut assistant = Assistant { id: String::new(), name: assistant_name.to_string(), model: model.to_string(), instructions: instructions.to_string(), - folder_path: folder_path.to_string(), - scrape_urls: Vec::new(), - file_ids: Vec::new(), // Make sure this field exists in the Assistant struct }; // Initialize the assistant by creating it on the OpenAI platform assistant.initialize().await?; info!("Assistant created with ID: {}", assistant.id); - // Scrape context if needed (optional, based on whether scrape_urls are provided) - if !assistant.scrape_urls.is_empty() { - assistant.scrape_context().await?; - } - // Upload files from the folder_path to OpenAI and store the file IDs in the assistant - assistant.upload_files().await?; - // Attach the uploaded files to the assistant using the stored file IDs - assistant.attach_files().await?; + // Attach the uploaded files to the assistant using the file IDs from the Files struct + assistant.attach_files(&files.file_ids).await?; Ok(assistant) } @@ -391,8 +386,6 @@ pub async fn teardown_assistant( // Delete the assistant on the OpenAI platform assistant.delete().await?; info!("Assistant with ID: {} has been deleted", assistant.id); - assistant.delete_files().await?; - info!("All files uploaded for assistant with ID: {} deleted", assistant.id); Ok(()) } diff --git a/src/main.rs b/src/main.rs index ffb15f5..9f7997b 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,5 +1,5 @@ mod assistant; -use assistant::create_assistant; +use assistant::{create_assistant, create_files}; use assistant::{assistant_chat_handler, DB}; use axum::{extract::Extension, routing::post, Router}; use dotenv::dotenv; @@ -16,12 +16,25 @@ async fn app(db_pool: SqlitePool, assistant_id: String) -> Router { async fn main() { env_logger::init(); dotenv().ok(); +// Create the files for the assistant. + let files = match create_files( + "context", + Vec::new(), + ) + .await + { + Ok(files) => files, + Err(e) => { + eprintln!("Failed to create files: {:?}", e); + return; + } + }; // Create an assistant outside of the main function. let assistant = match create_assistant( "My Assistant", "gpt-4-0125-preview", "Your instructions here", - "context", + files, ) .await {