use std::path::{Path, PathBuf};
use bytes::Bytes;
use reqwest::multipart::{Form, Part};
use crate::{client::Client, error::Result, Envelope, UploadedFile};
#[derive(Clone, Debug)]
pub enum FileUpload {
Bytes {
bytes: Bytes,
filename: String,
content_type: Option<String>,
},
Path {
path: PathBuf,
filename: Option<String>,
content_type: Option<String>,
},
}
impl FileUpload {
pub fn from_bytes(bytes: impl Into<Bytes>, filename: impl Into<String>) -> Self {
Self::Bytes {
bytes: bytes.into(),
filename: filename.into(),
content_type: None,
}
}
pub fn from_bytes_with_content_type(
bytes: impl Into<Bytes>,
filename: impl Into<String>,
content_type: impl Into<String>,
) -> Self {
Self::Bytes {
bytes: bytes.into(),
filename: filename.into(),
content_type: Some(content_type.into()),
}
}
pub fn from_path(path: impl Into<PathBuf>) -> Self {
Self::Path {
path: path.into(),
filename: None,
content_type: None,
}
}
pub fn filename(mut self, filename: impl Into<String>) -> Self {
match &mut self {
Self::Bytes {
filename: current, ..
} => *current = filename.into(),
Self::Path {
filename: current, ..
} => *current = Some(filename.into()),
}
self
}
pub fn content_type(mut self, content_type: impl Into<String>) -> Self {
match &mut self {
Self::Bytes {
content_type: current,
..
} => *current = Some(content_type.into()),
Self::Path {
content_type: current,
..
} => *current = Some(content_type.into()),
}
self
}
async fn into_part(self) -> Result<Part> {
match self {
Self::Bytes {
bytes,
filename,
content_type,
} => {
let content_type = content_type.unwrap_or_else(|| infer_content_type(&filename));
Ok(Part::bytes(bytes.to_vec())
.file_name(filename)
.mime_str(&content_type)?)
}
Self::Path {
path,
filename,
content_type,
} => {
let bytes = tokio::fs::read(&path).await?;
let filename = filename.unwrap_or_else(|| file_name(&path));
let content_type = content_type.unwrap_or_else(|| infer_content_type(&filename));
Ok(Part::bytes(bytes)
.file_name(filename)
.mime_str(&content_type)?)
}
}
}
}
#[derive(Clone, Debug, Default)]
pub struct UploadOptions {
pub file_type: Option<String>,
pub model: Option<String>,
pub mode: Option<String>,
}
impl UploadOptions {
pub fn new() -> Self {
Self::default()
}
pub fn file_type(mut self, file_type: impl Into<String>) -> Self {
self.file_type = Some(file_type.into());
self
}
pub fn model(mut self, model: impl Into<String>) -> Self {
self.model = Some(model.into());
self
}
pub fn mode(mut self, mode: impl Into<String>) -> Self {
self.mode = Some(mode.into());
self
}
}
#[derive(Clone, Debug)]
pub struct FilesService {
client: Client,
}
impl FilesService {
pub(crate) fn new(client: Client) -> Self {
Self { client }
}
pub async fn upload(
&self,
file: FileUpload,
options: UploadOptions,
) -> Result<Envelope<UploadedFile>> {
let mut form = Form::new().part("file", file.into_part().await?);
if let Some(file_type) = options.file_type {
form = form.text("file_type", file_type);
}
if let Some(model) = options.model {
form = form.text("model", model);
}
if let Some(mode) = options.mode {
form = form.text("mode", mode);
}
self.client.send_multipart("/api/v1/files", form).await
}
}
fn file_name(path: &Path) -> String {
path.file_name()
.and_then(|value| value.to_str())
.unwrap_or("file")
.to_string()
}
fn infer_content_type(filename: &str) -> String {
mime_guess::from_path(filename)
.first_raw()
.unwrap_or("application/octet-stream")
.to_string()
}