use http_req::{
request::{Method, Request},
uri::Uri,
};
use serde::{Deserialize, Serialize, Serializer};
use urlencoding::encode;
use crate::Retry;
#[derive(Debug)]
pub enum ClaudeModel {
Claude1,
Claude2,
ClaudeInstant1,
}
impl Serialize for ClaudeModel {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
match self {
ClaudeModel::Claude1 => serializer.serialize_str("claude-1"),
ClaudeModel::Claude2 => serializer.serialize_str("claude-2"),
ClaudeModel::ClaudeInstant1 => serializer.serialize_str("claude-instant-1"),
}
}
}
impl Default for ClaudeModel {
fn default() -> ClaudeModel {
ClaudeModel::Claude2
}
}
#[derive(Debug, Serialize)]
pub struct ChatOptions<'a> {
pub model: ClaudeModel,
pub restart: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub system_prompt: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
pub pre_prompt: Option<&'a str>,
#[serde(skip_serializing_if = "Option::is_none")]
pub post_prompt: Option<&'a str>,
pub max_tokens_to_sample: u32,
#[serde(skip_serializing_if = "Option::is_none")]
pub stop_sequences: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_k: Option<u32>,
}
impl<'a> Default for ChatOptions<'a> {
fn default() -> ChatOptions<'a> {
ChatOptions {
model: Default::default(),
restart: false,
system_prompt: None,
pre_prompt: None,
post_prompt: None,
max_tokens_to_sample: 256,
stop_sequences: None,
temperature: None,
top_p: None,
top_k: None,
}
}
}
impl crate::ClaudeFlows {
pub async fn chat_completion(
&self,
conversation_id: &str,
sentence: &str,
options: &ChatOptions<'_>,
) -> Result<String, String> {
self.keep_trying(|account| {
chat_completion_inner(account, conversation_id, sentence, options)
})
}
}
fn chat_completion_inner(
account: &str,
conversation_id: &str,
sentence: &str,
options: &ChatOptions,
) -> Retry<String> {
unsafe {
let flows_user = crate::_get_flows_user();
let flow_id = crate::_get_flow_id();
let mut writer = Vec::new();
let uri = format!(
"{}/{}/{}/chat_completion?account={}&conversation={}",
crate::CLAUDE_API_PREFIX.as_str(),
flows_user,
flow_id,
encode(account),
encode(conversation_id),
);
let uri = Uri::try_from(uri.as_str()).unwrap();
let body = serde_json::to_vec(&serde_json::json!({
"sentence": sentence,
"params": options
}))
.unwrap_or_default();
match Request::new(&uri)
.method(Method::POST)
.header("Content-Type", "application/json")
.header("Content-Length", &body.len())
.body(&body)
.send(&mut writer)
{
Ok(res) => {
match res.status_code().is_success() {
true => Retry::No(Ok(String::from_utf8_lossy(&writer).into_owned())),
false => {
match res.status_code().into() {
409 | 429 | 503 => {
Retry::Yes(String::from_utf8_lossy(&writer).into_owned())
}
_ => Retry::No(Err(String::from_utf8_lossy(&writer).into_owned())),
}
}
}
}
Err(e) => Retry::No(Err(e.to_string())),
}
}
}
#[derive(Debug, Deserialize)]
pub enum ChatRole {
Human,
Assistant,
}
#[derive(Debug, Deserialize)]
pub struct ChatMessage {
pub role: ChatRole,
pub content: String,
}
pub fn chat_history(conversation_id: &str, limit: u8) -> Option<Vec<ChatMessage>> {
unsafe {
let flows_user = crate::_get_flows_user();
let flow_id = crate::_get_flow_id();
let mut writer = Vec::new();
let uri = format!(
"{}/{}/{}/chat_history?conversation={}&limit={}",
crate::CLAUDE_API_PREFIX.as_str(),
flows_user,
flow_id,
encode(conversation_id),
limit
);
let uri = Uri::try_from(uri.as_str()).unwrap();
match Request::new(&uri).method(Method::GET).send(&mut writer) {
Ok(res) => match res.status_code().is_success() {
true => serde_json::from_slice::<Vec<ChatMessage>>(&writer).ok(),
false => None,
},
Err(_) => None,
}
}
}