perplexity 0.1.1

use reqwest::{
    header::{HeaderMap, HeaderValue, AUTHORIZATION, CONTENT_TYPE},
    Client as ReqwestClient, Method, RequestBuilder, Url,
};
use serde::{Deserialize, Serialize};

use crate::{config::Config, error::Error};

pub struct Client {
    base_url: Url,
    http_client: ReqwestClient,
}

impl Client {
    pub fn new(config: Config) -> Result<Self, Error> {
        let mut headers = HeaderMap::new();
        headers.insert(
            AUTHORIZATION,
            HeaderValue::from_str(format!("Bearer {}", config.api_key.as_str()).as_str())?,
        );
        headers.insert(CONTENT_TYPE, HeaderValue::from_static("application/json"));

        let http_client = reqwest::Client::builder()
            .default_headers(headers)
            .build()?;

        let base_url =
            Url::parse(&config.base_url).map_err(|err| Error::UrlParse(err.to_string()))?;

        Ok(Self {
            base_url,
            http_client,
        })
    }

    pub async fn create_completion(
        &self,
        payload: CreateChatCompletion,
    ) -> Result<ChatCompletion, Error> {
        let completion = self
            .request(Method::POST, "chat/completions")?
            .json(&payload)
            .send()
            .await?
            .json::<ChatCompletion>()
            .await?;
        Ok(completion)
    }

    fn request(&self, method: Method, path: &str) -> Result<RequestBuilder, Error> {
        let url = self
            .base_url
            .join(path)
            .map_err(|err| Error::UrlParse(err.to_string()))?;
        Ok(self.http_client.request(method, url))
    }
}

#[derive(Debug, Serialize, Deserialize)]
pub struct ChatCompletion {
    /// An ID generated uniquely for each response.
    pub id: String,

    /// The model used to generate the response.
    pub model: Model,

    /// The object type, which always equals **chat.completion**.
    pub object: String,

    /// The Unix timestamp (in seconds) of when the completion was created.
    pub created: u64,

    /// The list of completion choices the model generated for the input prompt.
    pub choices: Vec<CompletionChoice>,

    /// Usage statistics for the completion request.
    pub usage: CompletionUsage,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct CompletionChoice {
    pub index: u64,

    /// The reason the model stopped generating tokens. Possible values include stop if the model hit a natural stopping point, or length if the maximum number of tokens specified in the request was reached.
    pub finish_reason: FinishReason,

    /// The message generated by the model.
    pub message: Message,

    /// The incrementally streamed next tokens. Only meaningful when **stream = true**.
    pub delta: CompletionDelta,
}

#[derive(Debug, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum FinishReason {
    Stop,
    Length,
}

/// The message generated by the model.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
    /// The contents of the message in this turn of conversation.
    pub content: String,

    /// The role of the speaker in this turn of conversation. After the (optional) system message, user and assistant roles should alternate with user then assistant, ending in user.
    pub role: Role,
}

/// The incrementally streamed next tokens. Only meaningful when **stream = true**.
#[derive(Debug, Serialize, Deserialize)]
pub struct CompletionDelta {
    /// The contents of the message in this turn of conversation.
    pub content: String,

    /// The role of the speaker in this turn of conversation. After the (optional) system message, user and assistant roles should alternate with user then assistant, ending in user.
    pub role: Role,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    System,
    User,
    Assistant,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct CompletionUsage {
    /// The number of tokens provided in the request prompt.
    pub prompt_tokens: u64,

    /// The number of tokens generated in the response output.
    pub completion_tokens: u64,

    /// The total number of tokens used in the chat completion (prompt + completion).
    pub total_tokens: u64,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CreateChatCompletion {
    /// The name of the model that will complete your prompt.
    pub model: Model,

    /// The conversation messages.
    pub messages: Vec<Message>,

    /// The maximum number of completion tokens returned by the API. The total number of tokens requested in **max_tokens** plus the number of prompt tokens sent in messages must not exceed the context window token limit of model requested. If left unspecified, then the model will generate tokens until either it reaches its stop token or the end of its context window.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_tokens: Option<u64>,

    /// The amount of randomness in the response, valued between 0 inclusive and 2 exclusive. Higher values are more random, and lower values are more deterministic.
    ///
    /// Defaults to **0.2**.
    pub temperature: f32,

    /// The nucleus sampling threshold, valued between 0 and 1 inclusive. For each subsequent token, the model considers the results of the tokens with **top_p** probability mass.
    ///
    /// Defaults to **0.9**.
    pub top_p: f32,

    /// Determines whether or not a request to an online model should return citations.
    ///
    /// Defaults to **false**.
    pub return_citations: bool,

    /// Given a list of domains, limit the citations used by the online model to URLs from the specified domains. Currently limited to only 3 domains for whitelisting and blacklisting. For blacklisting add a **-** to the beginning of the domain string.
    pub search_domain_filter: Option<Vec<String>>,

    /// Determines whether or not a request to an online model should return images.
    ///
    /// Defaults to **false**.
    pub return_images: bool,

    /// Determines whether or not a request to an online model should return related questions.
    ///
    /// Defaults to **false**.
    pub return_related_questions: bool,

    /// Returns search results within the specified time interval - does not apply to images.
    pub search_recency_filter: RecencyFilter,

    /// The number of tokens to keep for highest **top-k** filtering, specified as an integer between **0** and **2048** inclusive. If set to **0**, **top-k** filtering is disabled.
    ///
    /// Defaults to **0**.
    pub top_k: u16,

    /// Determines whether or not to incrementally stream the response with server-sent events.
    ///
    /// Defaults to **false**.
    pub stream: bool,

    /// A value between **-2.0** and **2.0**. Positive values penalize new tokens based on whether they appear in the text so far, increasing the model's likelihood to talk about new topics. Incompatible with **frequency_penalty**.
    ///
    /// Defaults to **0**.
    pub presence_penalty: u8,

    /// A multiplicative penalty greater than **0**. Values greater than **1.0** penalize new tokens based on their existing frequency in the text so far, decreasing the model's likelihood to repeat the same line verbatim. A value of **1.0** means no penalty. Incompatible with **presence_penalty**.
    ///
    /// Defaults to **1**.
    pub frequency_penalty: f32,
}

impl CreateChatCompletion {
    pub fn new(model: Model, messages: Vec<Message>) -> Self {
        Self {
            model,
            messages,
            ..Default::default()
        }
    }

    pub fn with_model(mut self, model: Model) -> Self {
        self.model = model;
        self
    }

    pub fn with_stream(mut self, stream: bool) -> Self {
        self.stream = stream;
        self
    }

    pub fn with_max_tokens(mut self, max_tokens: u64) -> Self {
        self.max_tokens = Some(max_tokens);
        self
    }

    pub fn with_temperature(mut self, temperature: f32) -> Self {
        self.temperature = temperature;
        self
    }

    pub fn with_top_k(mut self, top_k: u16) -> Self {
        self.top_k = top_k;
        self
    }

    pub fn with_top_p(mut self, top_p: f32) -> Self {
        self.top_p = top_p;
        self
    }

    pub fn with_frequency_penalty(mut self, frequency_penalty: f32) -> Self {
        self.frequency_penalty = frequency_penalty;
        self
    }

    pub fn with_presence_penalty(mut self, presence_penalty: u8) -> Self {
        self.presence_penalty = presence_penalty;
        self
    }

    pub fn with_images(mut self, return_images: bool) -> Self {
        self.return_images = return_images;
        self
    }

    pub fn with_citations(mut self, return_citations: bool) -> Self {
        self.return_citations = return_citations;
        self
    }

    pub fn with_return_related_questions(mut self, return_related_questions: bool) -> Self {
        self.return_related_questions = return_related_questions;
        self
    }

    pub fn with_search_domain_filter(mut self, search_domain_filter: Vec<String>) -> Self {
        self.search_domain_filter = Some(search_domain_filter);
        self
    }

    pub fn with_recency_filter(mut self, recency_filter: RecencyFilter) -> Self {
        self.search_recency_filter = recency_filter;
        self
    }
}

impl Default for CreateChatCompletion {
    fn default() -> Self {
        Self {
            model: Model::Llama31SonarLargeOnline,
            messages: Vec::new(),
            stream: false,
            max_tokens: None,
            temperature: 0.2,
            top_p: 0.9,
            top_k: 0,
            frequency_penalty: 1.0,
            presence_penalty: 0,
            return_images: false,
            return_citations: false,
            return_related_questions: false,
            search_domain_filter: None,
            search_recency_filter: RecencyFilter::default(),
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Model {
    /// **8 Billion** parameters and **127,072** context length model.
    #[serde(rename = "llama-3.1-sonar-small-128k-online")]
    Llama31SonarSmallOnline,

    /// **70 Billion** parameters and **127,072** context length model.
    #[serde(rename = "llama-3.1-sonar-large-128k-online")]
    Llama31SonarLargeOnline,

    /// **405 Billion** parameters and **127,072** context length model.
    #[serde(rename = "llama-3.1-sonar-huge-128k-online")]
    Llama31SonarHugeOnline,

    /// **8 Billion** parameters and **127,072** context length model.
    #[serde(rename = "llama-3.1-sonar-small-128k-chat")]
    Llama31SonarSmallChat,

    /// **70 Billion** parameters and **127,072** context length model.
    #[serde(rename = "llama-3.1-sonar-large-128k-chat")]
    Llama31SonarLargeChat,

    /// **8 Billion** parameters and **131,072** context length model.
    #[serde(rename = "llama-3.1-8b-instruct")]
    Llama31InstructSmall,

    /// **70 Billion** parameters and **131,072** context length model.
    #[serde(rename = "llama-3.1-70b-instruct")]
    Llama31InstructLarge,
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum RecencyFilter {
    Hour,
    Day,
    Week,
    Month,
}

impl Default for RecencyFilter {
    fn default() -> Self {
        Self::Month
    }
}