use anyhow::{anyhow, Context, Result};
use async_trait::async_trait;
use log::{error, info};
use reqwest::{header, multipart, Client};
use serde::{Deserialize, Serialize};
use crate::assistants::{OpenAIAssistantResource, OpenAIAssistantVersion};
use crate::domain::AllmsError;
use crate::files::LLMFiles;
use crate::utils::get_mime_type;
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIFile {
pub id: Option<String>,
debug: bool,
api_key: String,
version: OpenAIAssistantVersion,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIFileResp {
id: String,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIDFileDeleteResp {
id: String,
object: String,
deleted: bool,
}
#[async_trait(?Send)]
impl LLMFiles for OpenAIFile {
fn new(id: Option<String>, open_ai_key: &str) -> Self {
OpenAIFile {
id,
debug: false,
api_key: open_ai_key.to_string(),
version: OpenAIAssistantVersion::V1, }
}
fn debug(mut self) -> Self {
self.debug = true;
self
}
async fn upload(mut self, file_name: &str, file_bytes: Vec<u8>) -> Result<Self> {
let files_url = self.version.get_endpoint(&OpenAIAssistantResource::Files);
if self.debug {
info!("[debug] OpenAI Files Upload API URL: {:#?}", files_url);
}
let mut version_headers = self.version.get_headers(&self.api_key);
version_headers.remove(header::CONTENT_TYPE);
let mime_type = get_mime_type(file_name).ok_or_else(|| anyhow!("Unsupported file type"))?;
let form = multipart::Form::new().text("purpose", "assistants").part(
"file",
multipart::Part::bytes(file_bytes)
.file_name(file_name.to_string())
.mime_str(mime_type)
.context("Failed to set MIME type")?,
);
let client = Client::new();
let response = client
.post(files_url)
.headers(version_headers)
.multipart(form)
.send()
.await?;
let response_status = response.status();
let response_text = response.text().await?;
if self.debug {
info!(
"[debug] OpenAI Files status API response: [{}] {:#?}",
&response_status, &response_text
);
}
let response_deser: OpenAIFileResp =
serde_json::from_str(&response_text).map_err(|error| {
let error = AllmsError {
crate_name: "allms".to_string(),
module: "assistants::openai_file".to_string(),
error_message: format!("Files API response serialization error: {}", error),
error_detail: response_text,
};
error!("{:?}", error);
anyhow!("{:?}", error)
})?;
self.id = Some(response_deser.id);
Ok(self)
}
async fn delete(&self) -> Result<()> {
let file_id = if let Some(id) = &self.id {
id
} else {
return Err(anyhow!(
"[OpenAI][File API] Unable to delete file without an ID."
));
};
let files_resource = OpenAIAssistantResource::File {
file_id: file_id.to_string(),
};
let files_url = self.version.get_endpoint(&files_resource);
let version_headers = self.version.get_headers(&self.api_key);
let client = Client::new();
let response = client
.delete(files_url)
.headers(version_headers)
.send()
.await?;
let response_status = response.status();
let response_text = response.text().await?;
if self.debug {
info!(
"[debug] OpenAI Files status API response: [{}] {:#?}",
&response_status, &response_text
);
}
serde_json::from_str::<OpenAIDFileDeleteResp>(&response_text)
.map_err(|error| {
let error = AllmsError {
crate_name: "allms".to_string(),
module: "assistants::openai_file".to_string(),
error_message: format!(
"Files Delete API response serialization error: {}",
error
),
error_detail: response_text,
};
error!("{:?}", error);
anyhow!("{:?}", error)
})
.and_then(|response| match response.deleted {
true => Ok(()),
false => Err(anyhow!("[OpenAIAssistant] Failed to delete the file.")),
})
}
fn get_id(&self) -> Option<&String> {
self.id.as_ref()
}
fn is_debug(&self) -> bool {
self.debug
}
}
impl OpenAIFile {
pub fn version(mut self, version: OpenAIAssistantVersion) -> Self {
let version = match version {
OpenAIAssistantVersion::V2 => OpenAIAssistantVersion::V1,
_ => version,
};
self.version = version;
self
}
}