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, Credentials};
use super::ApiResponseOrError;
#[derive(Deserialize, Serialize, Clone)]
pub struct File {
pub id: String,
pub object: String,
pub bytes: usize,
pub created_at: usize,
pub filename: String,
pub purpose: String,
}
#[derive(Deserialize, Serialize, Clone)]
pub struct DeletedFile {
pub id: String,
pub object: String,
pub deleted: bool,
}
#[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,
#[serde(skip_serializing)]
#[builder(default)]
pub credentials: Option<Credentials>,
}
impl File {
async fn create(request: FileUploadRequest) -> ApiResponseOrError<Self> {
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", request.purpose);
openai_post_multipart("files", form, request.credentials).await
}
pub fn builder() -> FileUploadBuilder {
FileUploadBuilder::create_empty()
}
pub async fn delete(id: &str, credentials: Credentials) -> ApiResponseOrError<DeletedFile> {
openai_delete(format!("files/{}", id).as_str(), Some(credentials)).await
}
#[deprecated(since = "1.0.0-alpha.16", note = "use `fetch` instead")]
pub async fn get(id: &str) -> ApiResponseOrError<File> {
openai_get(format!("files/{}", id).as_str(), None).await
}
pub async fn fetch(id: &str, credentials: Credentials) -> ApiResponseOrError<File> {
openai_get(format!("files/{}", id).as_str(), Some(credentials)).await
}
#[deprecated(since = "1.0.0-alpha.16", note = "use `fetch_content_bytes` instead")]
pub async fn get_content_bytes(id: &str) -> ApiResponseOrError<Vec<u8>> {
Self::fetch_content_bytes_with_credentials_opt(id, None).await
}
pub async fn fetch_content_bytes(
id: &str,
credentials: Credentials,
) -> ApiResponseOrError<Vec<u8>> {
Self::fetch_content_bytes_with_credentials_opt(id, Some(credentials)).await
}
async fn fetch_content_bytes_with_credentials_opt(
id: &str,
credentials_opt: Option<Credentials>,
) -> ApiResponseOrError<Vec<u8>> {
let route = format!("files/{}/content", id);
let response = openai_request(
Method::GET,
route.as_str(),
|request| request,
credentials_opt,
)
.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())
}
pub async fn download_content_to_file(
id: &str,
file_path: &str,
credentials: Credentials,
) -> 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,
Some(credentials),
)
.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 {
pub async fn create(self) -> ApiResponseOrError<File> {
File::create(self.build().unwrap()).await
}
}
impl Files {
pub async fn list(credentials: Credentials) -> ApiResponseOrError<Files> {
openai_get("files", Some(credentials)).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::DEFAULT_CREDENTIALS;
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();
let credentials = Credentials::from_env();
let file_upload = test_upload_builder()
.credentials(credentials)
.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();
let credentials = Credentials::from_env();
let test_builder = File::builder()
.file_name("test_data/missing_file.jsonl")
.credentials(credentials)
.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();
let credentials = Credentials::from_env();
test_upload_builder().create().await.unwrap();
let openai_files = Files::list(credentials).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();
let credentials = Credentials::from_env();
test_upload_builder().create().await.unwrap();
tokio::time::sleep(Duration::from_secs(7)).await;
let openai_files = Files::list(credentials).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(),
DEFAULT_CREDENTIALS.read().unwrap().clone(),
)
.await
.unwrap();
assert!(deleted_file.deleted);
println!("deleted: {} {}", deleted_file.id, deleted_file.deleted)
}
}
#[tokio::test]
async fn get_file_and_contents() {
dotenv().ok();
let credentials = Credentials::from_env();
let file = test_upload_builder()
.credentials(credentials.clone())
.create()
.await
.unwrap();
let file_get = File::fetch(file.id.as_str(), credentials.clone())
.await
.unwrap();
assert_eq!(file.id, file_get.id);
let body_bytes = File::fetch_content_bytes(file.id.as_str(), credentials.clone())
.await
.unwrap();
assert_eq!(body_bytes.len(), file.bytes);
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(), credentials)
.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.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)
}
}