openai-fork 1.0.0-alpha.15

//! Upload, download, list, and delete files in openapi platform. Usually used for fine-tuning files.
//!
//! See the [Files API for OpenAI](https://platform.openai.com/docs/api-reference/files) for
//! more information.
//!
//! # Examples
//!
//! All examples and tests require the `OPENAI_KEY` environment variable
//! be set with your personal openai platform API key.
//!
//! Upload a new file. [Reference API](https://platform.openai.com/docs/api-reference/files/upload)
//! ```
//!use openai::files::File;
//!use openai::ApiResponseOrError;
//!use dotenvy::dotenv;
//!use std::env;
//!use openai::set_key;
//!
//!#[tokio::main]
//!async fn main() -> ApiResponseOrError<()> {
//!     dotenv().ok();
//!     set_key(env::var("OPENAI_KEY").unwrap());
//!     let uploaded_file = File::builder()
//!         .file_name("test_data/file_upload_test1.jsonl") // local file path to upload.
//!         .purpose("fine-tune")
//!         .create()
//!         .await?;
//!     assert_eq!(uploaded_file.filename, "file_upload_test1.jsonl");
//!     Ok(())
//!}
//! ```
//!
//! List files. [Reference API](https://platform.openai.com/docs/api-reference/files/list)
//! ```
//!use openai::files::Files;
//!use openai::ApiResponseOrError;
//!use dotenvy::dotenv;
//!use std::env;
//!use openai::set_key;
//!
//!#[tokio::main]
//!async fn main() -> ApiResponseOrError<()> {
//!     dotenv().ok();
//!     set_key(env::var("OPENAI_KEY").unwrap());
//!     let openai_files = Files::list().await?;
//!     let file_count = openai_files.len();
//!     println!("Listing {} files", file_count);
//!     for openai_file in openai_files.into_iter() {
//!         println!("  id: {}, file: {}, size: {}", openai_file.id, openai_file.filename, openai_file.bytes)
//!     }
//!     Ok(())
//!}
//! ```
//!
//! Retrieve a file (json metadata only). [Reference API](https://platform.openai.com/docs/api-reference/files/retrieve)
//!
//! ```no_run
//!use openai::files::File;
//!use openai::ApiResponseOrError;
//!use dotenvy::dotenv;
//!use std::env;
//!use openai::set_key;
//!
//!#[tokio::main]
//!async fn main() -> ApiResponseOrError<()> {
//!     dotenv().ok();
//!     set_key(env::var("OPENAI_KEY").unwrap());
//!     let file_id = "file-XjGxS3KTG0uNmNOK362iJua3"; // Use a real file id.
//!     let file = File::get(file_id).await?;
//!     println!("id: {}, file: {}, size: {}", file.id, file.filename, file.bytes);
//!     Ok(())
//!}
//! ```
//!
//! Download to a local file. [Reference API](https://platform.openai.com/docs/api-reference/files/retrieve-content)
//!
//! ```no_run
//!use openai::files::File;
//!use openai::ApiResponseOrError;
//!use dotenvy::dotenv;
//!use std::env;
//!use openai::set_key;
//!
//!#[tokio::main]
//!async fn main() -> ApiResponseOrError<()> {
//!     dotenv().ok();
//!     set_key(env::var("OPENAI_KEY").unwrap());
//!     let test_file = "test_file.jsonl";
//!     let file_id = "file-XjGxS3KTG0uNmNOK362iJua3"; // Use a real file id.
//!     File::download_content_to_file(file_id, test_file).await?;
//!     Ok(())
//!}
//! ```
//!
//! Delete a file. [Reference API](https://platform.openai.com/docs/api-reference/files/delete)
//!
//! ```no_run
//!use openai::files::File;
//!use openai::ApiResponseOrError;
//!use dotenvy::dotenv;
//!use std::env;
//!use openai::set_key;
//!
//!#[tokio::main]
//!async fn main() -> ApiResponseOrError<()> {
//!     dotenv().ok();
//!     set_key(env::var("OPENAI_KEY").unwrap());
//!     let file_id = "file-XjGxS3KTG0uNmNOK362iJua3"; // Use a real file id.
//!     File::delete(file_id).await?;
//!     Ok(())
//!}
//! ```
//!
//! For more examples see the files tests.
//!

use std::io::Write;
use std::path::Path;

use bytes::{BufMut, BytesMut};
use derive_builder::Builder;
use futures_util::StreamExt;
use reqwest::multipart::{Form, Part};
use reqwest::Method;
use serde::{Deserialize, Serialize};

use crate::{openai_delete, openai_get, openai_post_multipart, openai_request};

use super::ApiResponseOrError;

/// Upload, download and delete a file from the openai platform.
#[derive(Deserialize, Serialize, Clone)]
pub struct File {
    /// The unique id for this uploaded the in the openai platform.
    /// This id is generated by openai for each uploaded file.
    pub id: String,
    /// The object type uploaded. ie: "file"
    pub object: String,
    /// The size in bytes of the uploaded file.
    pub bytes: usize,
    /// Unix timestamp, seconds since epoch, of when the file was uploaded.
    pub created_at: usize,
    /// The name of the file uploaded.
    pub filename: String,
    /// The purpose of the file. ie: "fine-tine"
    pub purpose: String,
}

#[derive(Deserialize, Serialize, Clone)]
pub struct DeletedFile {
    pub id: String,
    pub object: String,
    pub deleted: bool,
}

/// List files in the openai platform.
#[derive(Deserialize, Serialize, Clone)]
pub struct Files {
    data: Vec<File>,
    pub object: String,
}

#[derive(Serialize, Builder, Debug, Clone)]
#[builder(pattern = "owned")]
#[builder(name = "FileUploadBuilder")]
#[builder(setter(strip_option, into))]
pub struct FileUploadRequest {
    file_name: String,
    purpose: String,
}

impl File {
    async fn create(request: &FileUploadRequest) -> ApiResponseOrError<Self> {
        let purpose = request.purpose.clone();
        let upload_file_path = Path::new(request.file_name.as_str());
        let upload_file_path = upload_file_path.canonicalize()?;
        let simple_name = upload_file_path
            .file_name()
            .unwrap()
            .to_str()
            .unwrap()
            .to_string()
            .clone();
        let async_file = tokio::fs::File::open(upload_file_path).await?;
        let file_part = Part::stream(async_file)
            .file_name(simple_name)
            .mime_str("application/jsonl")?;
        let form = Form::new().part("file", file_part).text("purpose", purpose);
        openai_post_multipart("files", form).await
    }

    /// New FileUploadBuilder
    pub fn builder() -> FileUploadBuilder {
        FileUploadBuilder::create_empty()
    }

    /// Delete a file from openai platform by id.
    pub async fn delete(id: &str) -> ApiResponseOrError<DeletedFile> {
        openai_delete(format!("files/{}", id).as_str()).await
    }

    /// Get a file from openai platform by id.
    pub async fn get(id: &str) -> ApiResponseOrError<File> {
        openai_get(format!("files/{}", id).as_str()).await
    }

    /// Download a file as bytes into memory by id.
    pub async fn get_content_bytes(id: &str) -> ApiResponseOrError<Vec<u8>> {
        let route = format!("files/{}/content", id);
        let response = openai_request(Method::GET, route.as_str(), |request| request).await?;
        let content_len = response.content_length().unwrap_or(1024) as usize;
        let mut file_bytes = BytesMut::with_capacity(content_len);
        let mut bytes_stream = response.bytes_stream();
        while let Some(Ok(bytes)) = bytes_stream.next().await {
            file_bytes.put(bytes);
        }
        Ok(file_bytes.to_vec())
    }

    /// Download a file to a new local file by id.
    pub async fn download_content_to_file(id: &str, file_path: &str) -> ApiResponseOrError<()> {
        let mut output_file = std::fs::File::create(file_path)?;
        let route = format!("files/{}/content", id);
        let response = openai_request(Method::GET, route.as_str(), |request| request).await?;
        let mut bytes_stream = response.bytes_stream();
        while let Some(Ok(bytes)) = bytes_stream.next().await {
            output_file.write_all(bytes.as_ref())?;
        }
        Ok(())
    }
}

impl FileUploadBuilder {
    /// Upload the file to the openai platform.
    pub async fn create(self) -> ApiResponseOrError<File> {
        File::create(&self.build().unwrap()).await
    }
}

impl Files {
    /// Get a list of all uploaded files in the openai platform.
    pub async fn list() -> ApiResponseOrError<Files> {
        openai_get("files").await
    }
    pub fn len(&self) -> usize {
        self.data.len()
    }
}

impl<'a> IntoIterator for &'a Files {
    type Item = &'a File;
    type IntoIter = core::slice::Iter<'a, File>;

    fn into_iter(self) -> Self::IntoIter {
        self.data.as_slice().iter()
    }
}

#[cfg(test)]
mod tests {
    use std::env;
    use std::io::Read;
    use std::time::Duration;

    use dotenvy::dotenv;

    use crate::set_key;

    use super::*;

    fn test_upload_builder() -> FileUploadBuilder {
        File::builder()
            .file_name("test_data/file_upload_test1.jsonl")
            .purpose("fine-tune")
    }

    fn test_upload_request() -> FileUploadRequest {
        test_upload_builder().build().unwrap()
    }

    #[tokio::test]
    async fn upload_file() {
        dotenv().ok();
        set_key(env::var("OPENAI_KEY").unwrap());
        let file_upload = test_upload_builder().create().await.unwrap();
        println!(
            "upload: {}",
            serde_json::to_string_pretty(&file_upload).unwrap()
        );
        assert_eq!(file_upload.id.as_bytes()[..5], *"file-".as_bytes())
    }

    #[tokio::test]
    async fn missing_file() {
        dotenv().ok();
        set_key(env::var("OPENAI_KEY").unwrap());
        let test_builder = File::builder()
            .file_name("test_data/missing_file.jsonl")
            .purpose("fine-tune");
        let response = test_builder.create().await;
        assert!(response.is_err());
        let openapi_err = response.err().unwrap();
        assert_eq!(openapi_err.error_type, "io");
        assert_eq!(
            openapi_err.message,
            "No such file or directory (os error 2)"
        )
    }

    #[tokio::test]
    async fn list_files() {
        dotenv().ok();
        set_key(env::var("OPENAI_KEY").unwrap());
        // ensure at least one file exists
        test_upload_builder().create().await.unwrap();
        let openai_files = Files::list().await.unwrap();
        let file_count = openai_files.len();
        assert!(file_count > 0);
        for openai_file in openai_files.into_iter() {
            assert_eq!(openai_file.id.as_bytes()[..5], *"file-".as_bytes())
        }
        println!(
            "files [{}]: {}",
            file_count,
            serde_json::to_string_pretty(&openai_files).unwrap()
        );
    }

    #[tokio::test]
    async fn delete_files() {
        dotenv().ok();
        set_key(env::var("OPENAI_KEY").unwrap());
        // ensure at least one file exists
        test_upload_builder().create().await.unwrap();
        // wait to avoid recent upload still processing error
        tokio::time::sleep(Duration::from_secs(7)).await;
        let openai_files = Files::list().await.unwrap();
        assert!(openai_files.data.len() > 0);
        let mut files = openai_files.data;
        files.sort_by(|a, b| a.created_at.cmp(&b.created_at));
        for file in files {
            let deleted_file = File::delete(file.id.as_str()).await.unwrap();
            assert!(deleted_file.deleted);
            println!("deleted: {} {}", deleted_file.id, deleted_file.deleted)
        }
    }

    #[tokio::test]
    async fn get_file_and_contents() {
        dotenv().ok();
        set_key(env::var("OPENAI_KEY").unwrap());

        let file = test_upload_builder().create().await.unwrap();
        let file_get = File::get(file.id.as_str()).await.unwrap();
        assert_eq!(file.id, file_get.id);

        // get file as bytes
        let body_bytes = File::get_content_bytes(file.id.as_str()).await.unwrap();
        assert_eq!(body_bytes.len(), file.bytes);

        // download file to a file
        let manifest_dir = env::var("CARGO_MANIFEST_DIR").unwrap();
        let test_dir = format!("{}/{}", manifest_dir, "target/files-test");
        std::fs::create_dir_all(test_dir.as_str()).unwrap();
        let test_file_save_path = format!("{}/{}", test_dir.as_str(), file.filename);
        File::download_content_to_file(file.id.as_str(), test_file_save_path.as_str())
            .await
            .unwrap();
        let mut local_file = std::fs::File::open(test_file_save_path.as_str()).unwrap();
        let mut local_bytes: Vec<u8> = Vec::new();
        local_file.read_to_end(&mut local_bytes).unwrap();
        assert_eq!(body_bytes, local_bytes)
    }

    #[test]
    fn file_name_path_test() {
        let request = test_upload_request();
        let file_upload_path = Path::new(request.file_name.as_str());
        let file_name = file_upload_path
            .clone()
            .file_name()
            .unwrap()
            .to_str()
            .unwrap();
        assert_eq!(file_name, "file_upload_test1.jsonl");
        let file_upload_path = file_upload_path.canonicalize().unwrap();
        let file_exists = file_upload_path.exists();
        assert!(file_exists)
    }
}