claude-flows 0.2.1

use http_req::{
    request::{Method, Request},
    uri::Uri,
};
use serde::{Deserialize, Serialize, Serializer};
use urlencoding::encode;

use crate::Retry;

/// Models
#[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
    }
}

/// struct for setting the chat options.
///
/// For more detail about parameters, please refer to
/// [Claude docs](https://docs.anthropic.com/claude/reference/complete_post)
#[derive(Debug, Serialize)]
pub struct ChatOptions<'a> {
    /// The ID or name of the model to use for completion.
    pub model: ClaudeModel,

    /// When true, a new conversation will be created.
    pub restart: bool,

    /// The prompt of the system role.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub system_prompt: Option<&'a str>,

    /// The prompt that will be prepended to user's prompt without saving in history.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub pre_prompt: Option<&'a str>,

    /// The prompt that will be appended to user's prompt without saving in history.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub post_prompt: Option<&'a str>,

    /// The maximum number of tokens to generate before stopping.
    pub max_tokens_to_sample: u32,

    /// Sequences that will cause the model to stop generating completion text.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop_sequences: Option<Vec<String>>,

    /// Amount of randomness injected into the response.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,

    /// Use nucleus sampling.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f32>,

    /// Only sample from the top K options for each subsequent token.
    #[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 {
    /// Create chat completion with the provided sentence.
    /// It uses Claude's [completion](https://docs.anthropic.com/claude/reference/complete_post) API to make a conversation.
    ///
    /// `conversation_id` is the identifier of the conversation.
    /// The history will be fetched and attached to the `sentence` as a whole prompt for Claude.
    ///
    /// `sentence` is a String that reprensents the current utterance of the conversation.
    ///
    /// If you haven't connected your Claude account with [Flows.network platform](https://flows.network),
    /// you will receive an error in the flow's building log or running log.
    ///
    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 => {
                                // 409 TryAgain 429 RateLimitError
                                // 503 ServiceUnavailable
                                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,
}

/// Fetch the question history of conversation_id
/// Result will be an array of string whose length is
/// restricted by limit.
/// When limit is 0, all history will be returned.
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,
        }
    }
}