#![allow(dead_code)]
use colored::Colorize;
use serde::{Deserialize, Serialize};
use std::fmt;
#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "lowercase")]
pub enum Role {
System,
User,
Assistant,
}
#[derive(Serialize, Deserialize, Debug)]
pub struct Message {
pub role: Role,
pub content: String,
}
impl Message {
pub fn system(content: String) -> Self {
Self {
role: Role::System,
content,
}
}
pub fn user(content: String) -> Self {
Self {
role: Role::User,
content,
}
}
pub fn assistant(content: String) -> Self {
Self {
role: Role::Assistant,
content,
}
}
}
#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct ErrorRoot {
pub error: Error,
}
#[derive(Serialize, Deserialize, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Error {
pub message: String,
#[serde(rename = "type")]
pub type_field: String,
pub param: Option<String>,
pub code: Option<String>,
}
impl fmt::Display for Error {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(
f,
"{} ({:?}): {:?}",
self.type_field.red(),
self.code,
self.message
)
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct Request {
pub model: String,
pub messages: Vec<Message>,
pub n: i32,
pub temperature: f64,
pub frequency_penalty: f64,
}
impl Request {
pub fn new(
model: String,
messages: Vec<Message>,
n: i32,
temperature: f64,
frequency_penalty: f64,
) -> Self {
Self {
model,
messages,
n,
temperature,
frequency_penalty,
}
}
}
#[derive(Serialize, Deserialize, Debug)]
pub struct Response {
pub id: String,
pub object: String,
pub created: i64,
pub model: String,
pub choices: Vec<Choice>,
pub usage: Usage,
}
#[derive(Serialize, Deserialize, Debug)]
pub struct Choice {
pub index: i64,
pub finish_reason: Option<String>,
pub message: Message,
}
#[derive(Serialize, Deserialize, Debug)]
pub struct Usage {
pub prompt_tokens: i64,
pub completion_tokens: i64,
pub total_tokens: i64,
}
pub fn count_token(s: &str) -> anyhow::Result<usize> {
let bpe = tiktoken_rs::cl100k_base()?;
let tokens = bpe.encode_with_special_tokens(s);
Ok(tokens.len())
}
pub fn cost(token: i64) -> f64 {
token as f64 * (PRICE / 1000.0)
}
const PRICE: f64 = 0.002;