use anyhow::{anyhow, Result};
use log::{error, info};
use reqwest::{header, Client};
use schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use crate::enums::{OpenAIAssistantRole, OpenAIRunStatus, OpenAIToolTypes};
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAPICompletionsResponse {
pub id: Option<String>,
pub object: Option<String>,
pub created: Option<u32>,
pub model: Option<String>,
pub choices: Option<Vec<OpenAPICompletionsChoices>>,
pub usage: Option<OpenAPIUsage>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAPICompletionsChoices {
pub text: Option<String>,
pub index: Option<u32>,
pub logprobs: Option<u32>,
pub finish_reason: Option<String>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAPIChatResponse {
pub id: Option<String>,
pub object: Option<String>,
pub created: Option<u32>,
pub model: Option<String>,
pub choices: Option<Vec<OpenAPIChatChoices>>,
pub usage: Option<OpenAPIUsage>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAPIChatChoices {
pub message: OpenAPIChatMessage,
pub index: Option<u32>,
pub finish_reason: Option<String>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAPIChatMessage {
pub role: String,
pub content: Option<String>,
pub function_call: Option<OpenAPIChatFunctionCall>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAPIChatFunctionCall {
pub(crate) name: String,
pub(crate) arguments: String,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAPIUsage {
prompt_tokens: Option<u32>,
completion_tokens: Option<u32>,
total_tokens: Option<u32>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIRateLimit {
pub tpm: usize, pub rpm: usize, }
#[derive(Deserialize, Serialize, Debug, Clone, JsonSchema)]
pub struct OpenAIDataResponse<T: JsonSchema> {
pub data: T,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIAssistantResp {
pub id: String,
pub object: String,
pub created_at: u32,
pub name: Option<String>,
pub description: Option<String>,
pub instructions: Option<String>,
pub model: String,
pub tools: Vec<OpenAITools>,
pub file_ids: Vec<String>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAITools {
#[serde(rename(deserialize = "type", serialize = "type"))]
tool_type: OpenAIToolTypes,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIThreadResp {
pub id: String,
pub object: String,
pub created_at: u32,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIMessageListResp {
pub object: String,
pub data: Vec<OpenAIMessageResp>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIMessageResp {
pub id: String,
pub object: String,
pub created_at: u32,
pub thread_id: String,
pub role: OpenAIAssistantRole,
pub content: Vec<OpenAIContent>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIContent {
#[serde(rename(deserialize = "type", serialize = "type"))]
pub content_type: String,
pub text: Option<OpenAIContentText>,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIContentText {
pub value: String,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIRunResp {
pub id: String,
pub object: String,
pub created_at: u32,
pub status: OpenAIRunStatus,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIFile {
pub id: String,
debug: bool,
api_key: String,
}
#[derive(Deserialize, Serialize, Debug, Clone)]
pub struct OpenAIFileResp {
id: String,
}
impl OpenAIFile {
pub async fn new(file_bytes: Vec<u8>, open_ai_key: &str, debug: bool) -> Result<Self> {
let mut new_file = OpenAIFile {
id: "this-will-be-overwritten".to_string(),
debug,
api_key: open_ai_key.to_string(),
};
new_file.upload_file(file_bytes).await?;
Ok(new_file)
}
async fn upload_file(&mut self, file_bytes: Vec<u8>) -> Result<()> {
let files_url = "https://api.openai.com/v1/files";
let form = reqwest::multipart::Form::new()
.text("purpose", "assistants")
.part(
"file",
reqwest::multipart::Part::bytes(file_bytes).file_name("file.pdf"),
);
let client = Client::new();
let response = client
.post(files_url)
.header(header::CONTENT_TYPE, "application/json")
.header("OpenAI-Beta", "assistants=v1")
.bearer_auth(&self.api_key)
.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| {
error!(
"[OpenAIAssistant] Files API response serialization error: {}",
&error
);
anyhow!("Error: {}", error)
})?;
self.id = response_deser.id;
Ok(())
}
}