groqu 0.1.3

/// Groq models. Model definitions are generated by AI because... yeah. Elon ma.
/// But don't worry, I still home-made those `impl`'s.

use std::{ marker::PhantomData, ops::Index };

use ijson::{ ijson, IValue };
use serde::{ Deserialize, Serialize };
use base64::{ engine::general_purpose, Engine as _ };

#[derive(Debug, Clone, Default, Serialize)]
pub struct ChatCompletionRequest {
    /// ID of the model to use
    pub model: IValue,
    /// A list of messages comprising the conversation so far
    pub messages: IValue,
    /// What sampling temperature to use, between 0 and 2
    #[serde(skip_serializing_if = "Option::is_none")]
    pub temperature: Option<f32>,
    /// An alternative to sampling with temperature, called nucleus sampling
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_p: Option<f32>,
    /// How many chat completion choices to generate for each input message
    #[serde(skip_serializing_if = "Option::is_none")]
    pub n: Option<u32>,
    /// Up to 4 sequences where the API will stop generating further tokens
    #[serde(skip_serializing_if = "Option::is_none")]
    pub stop: Option<StopSequence>,
    /// The maximum number of tokens that can be generated in the chat completion
    #[serde(skip_serializing_if = "Option::is_none")]
    pub max_completion_tokens: Option<u32>,
    /// Number between -2.0 and 2.0 affecting frequency penalty
    #[serde(skip_serializing_if = "Option::is_none")]
    pub frequency_penalty: Option<f32>,
    /// Number between -2.0 and 2.0 affecting presence penalty
    #[serde(skip_serializing_if = "Option::is_none")]
    pub presence_penalty: Option<f32>,
    /// Specifies how to output reasoning tokens
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_format: Option<String>,
    /// An object specifying the format that the model must output
    #[serde(skip_serializing_if = "Option::is_none")]
    pub response_format: Option<ResponseFormat>,
    /// For deterministic sampling
    #[serde(skip_serializing_if = "Option::is_none")]
    pub seed: Option<u64>,
    /// A unique identifier representing your end-user
    #[serde(skip_serializing_if = "Option::is_none")]
    pub user: Option<String>,
}

impl<'a> ChatCompletionRequest {
    pub fn builder() -> _ChatCompletionRequestBuilder<'a> {
        _ChatCompletionRequestBuilder {
            ph: Default::default(),
            rq: ChatCompletionRequest::default(),
        }
    }
}

#[derive(Debug)]
pub struct _ChatCompletionRequestBuilder<'a> {
    ph: PhantomData<&'a ()>,
    rq: ChatCompletionRequest,
}

impl _ChatCompletionRequestBuilder<'_> {
    /// ID of the model to use
    pub fn model<K: AsRef<str>>(mut self, name: K) -> Self {
        self.rq.model = IValue::from(name.as_ref());
        self
    }

    /// A list of messages comprising the conversation so far
    pub fn messages(mut self, messages: &[ChatMessage]) -> Self {
        self.rq.messages = ijson!(messages);
        self
    }

    /// What sampling temperature to use, between 0 and 2
    pub fn temperature(mut self, tmp: f32) -> Self {
        self.rq.temperature = Some(tmp);
        self
    }

    /// An alternative to sampling with temperature, called nucleus sampling
    pub fn top_p(mut self, p: f32) -> Self {
        self.rq.top_p = Some(p);
        self
    }

    /// How many chat completion choices to generate for each input message
    pub fn n(mut self, x: u32) -> Self {
        self.rq.n = Some(x);
        self
    }

    /// Up to 4 sequences where the API will stop generating further tokens
    pub fn stop(mut self, seq: StopSequence) -> Self {
        self.rq.stop = Some(seq);
        self
    }

    /// Up to 4 sequences where the API will stop generating further tokens
    ///
    /// This method targets a single token.
    pub fn stop_single(mut self, token: String) -> Self {
        self.rq.stop = Some(StopSequence::Single(token));
        self
    }

    /// Up to 4 sequences where the API will stop generating further tokens
    ///
    /// This method targets multiple tokens.
    pub fn stop_multiple(mut self, tokens: Vec<String>) -> Self {
        self.rq.stop = Some(StopSequence::Multiple(tokens));
        self
    }

    /// The maximum number of tokens that can be generated in the chat completion
    pub fn max_completion_tokens(mut self, n: u32) -> Self {
        self.rq.max_completion_tokens = Some(n);
        self
    }

    /// Number between -2.0 and 2.0 affecting frequency penalty
    pub fn frequency_penalty(mut self, x: f32) -> Self {
        self.rq.frequency_penalty = Some(x);
        self
    }

    /// Number between -2.0 and 2.0 affecting presence penalty
    pub fn presence_penalty(mut self, x: f32) -> Self {
        self.rq.presence_penalty = Some(x);
        self
    }

    /// Specifies how to output reasoning tokens
    ///
    /// See https://console.groq.com/docs/api-reference#chat-create
    pub fn reasoning_format(mut self, fmt: String) -> Self {
        self.rq.reasoning_format = Some(fmt);
        self
    }

    /// Type of the format that the model must output
    ///
    /// You do not need to pass an object here.
    /// See https://console.groq.com/docs/api-reference#chat-create
    pub fn response_format(mut self, fmt: String) -> Self {
        self.rq.response_format = Some(ResponseFormat { r#type: fmt });
        self
    }

    /// For deterministic sampling
    pub fn seed(mut self, x: u64) -> Self {
        self.rq.seed = Some(x);
        self
    }

    /// A unique identifier representing your end-user
    pub fn user(mut self, id: String) -> Self {
        self.rq.user = Some(id);
        self
    }

    /// Builds the request
    ///
    /// **You must use this** in order to build a request object.
    pub fn build(self) -> ChatCompletionRequest {
        self.rq
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
    /// The role of the message author
    pub role: ChatRole,
    /// The content of the message
    pub content: ContentType,
    /// Optional name for the message author
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: Option<String>,
}

macro_rules! messagerole {
    ($name:ident, $e:expr) => {
        pub fn $name(content: impl Into<ContentType>, name: Option<String>) -> Self {
            Self { role: $e, content: content.into(), name }
        }
    };
}

impl ChatMessage {
    messagerole!(system, ChatRole::System);
    messagerole!(assistant, ChatRole::Assistant);
    messagerole!(user, ChatRole::User);
}

#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ContentType {
    /// Text content type.
    Text(String),
    /// Array of content parts.
    Parts(Vec<ContentPart>),
}

impl ContentType {
    /// Creates a new instance, pure text only.
    pub fn text(x: String) -> Self {
        Self::Text(x)
    }

    /// Creates a part with text and image.
    pub fn text_and_image<T: Into<ContentPart>>(a: T, b: T) -> Self {
        Self::Parts(vec![a.into(), b.into()])
    }

    /// Used for constructed instances. Gets the text if this is a text part.
    pub fn get_text(self) -> String {
        match self {
            Self::Text(t) => t,
            _ => panic!("Expected text"),
        }
    }

    /// Used for constructed instances. Gets the text as a referenced str.
    pub fn get_text_as_str(&self) -> &str {
        match self {
            Self::Text(t) => t,
            _ => panic!("Expected text"),
        }
    }
}

impl From<String> for ContentType {
    fn from(value: String) -> Self {
        Self::Text(value)
    }
}

impl From<&str> for ContentType {
    fn from(value: &str) -> Self {
        Self::Text(value.to_string())
    }
}

impl<T: Into<ContentPart>> From<(T, T)> for ContentType {
    fn from((a, b): (T, T)) -> Self {
        Self::Parts(vec![a.into(), b.into()])
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum ContentPart {
    /// A text content part.
    /// This can be paired with a `ContentPart::Image`.
    Text {
        r#type: String,
        /// The text.
        text: String,
    },
    /// An image content part. Requires an [`ImageUrl`] object.
    Image {
        r#type: String,
        /// The image.
        image_url: ImageUrl,
    },
}

impl ContentPart {
    /// Creates a text content part.
    pub fn text(text: String) -> Self {
        Self::Text {
            r#type: String::from("text"),
            text,
        }
    }

    /// Creates an image content part.
    ///
    /// This method is not recommended, if you're lazy like me.
    pub fn image(image_url: ImageUrl) -> Self {
        Self::Image { r#type: String::from("image_url"), image_url }
    }

    /// Creates an image content part from an image URL.
    pub fn image_url(url: String) -> Self {
        Self::Image { r#type: String::from("image_url"), image_url: ImageUrl::new_string(url) }
    }

    /// Creates an image content part from bytes of a PNG.
    pub fn image_png(image: &[u8]) -> Self {
        Self::Image { r#type: String::from("image_url"), image_url: ImageUrl::new_png(image) }
    }

    /// Creates an image content part from bytes of any content type.
    ///
    /// For instance, for a PNG, you could pass `"image/png"` to the `dt` parameter.
    pub fn image_any<K: AsRef<str>>(dt: K, image: &[u8]) -> Self {
        Self::Image { r#type: String::from("image_url"), image_url: ImageUrl::new_any(dt, image) }
    }

    /// Creates an image content part directly from a `String`.
    pub fn image_string(url_s: String) -> Self {
        Self::Image { r#type: String::from("image_url"), image_url: ImageUrl::new_string(url_s) }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageUrl {
    pub url: String,
}

impl ImageUrl {
    /// Image URL from a PNG.
    pub fn new_png(image: &[u8]) -> Self {
        Self {
            url: format!("data:image/png;base64,{}", general_purpose::STANDARD.encode(image)),
        }
    }

    /// Image URL from any data type.
    ///
    /// For instance, for a PNG, you could pass `"image/png"` to the `dt` parameter.
    pub fn new_any<K: AsRef<str>>(dt: K, image: &[u8]) -> Self {
        Self {
            url: format!("data:{};base64,{}", dt.as_ref(), general_purpose::STANDARD.encode(image)),
        }
    }

    /// ...or you could do it manually, of course.
    ///
    /// I'm not sure if external URLs are supported.
    pub fn new_string(url: String) -> Self {
        Self { url }
    }
}

/// Represents a chat role.
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum ChatRole {
    System,
    Assistant,
    User,
}

/// Represents a stop sequence.
#[derive(Debug, Clone, Serialize)]
#[serde(untagged)]
pub enum StopSequence {
    Single(String),
    Multiple(Vec<String>),
}

impl From<String> for StopSequence {
    fn from(value: String) -> Self {
        Self::Single(value)
    }
}

impl From<&[String]> for StopSequence {
    fn from(value: &[String]) -> Self {
        Self::Multiple(value.to_vec())
    }
}

impl From<Vec<String>> for StopSequence {
    fn from(value: Vec<String>) -> Self {
        Self::Multiple(value)
    }
}

#[derive(Debug, Clone, Serialize)]
pub struct ResponseFormat {
    /// The type of response format
    pub r#type: String,
}

#[derive(Debug, Clone, Deserialize, Default)]
pub struct ChatCompletionResponse {
    /// The list of completion choices
    pub choices: Option<Vec<ChatCompletionChoice>>,
}

impl ChatCompletionResponse {
    /// Get the last choice. This is the most commonly used one.
    pub fn get_choice(self) -> ChatCompletionChoice {
        let mut choices = self.get_choices();
        choices.pop().unwrap()
    }

    /// Get all the choices.
    pub fn get_choices(self) -> Vec<ChatCompletionChoice> {
        self.choices.unwrap()
    }
}

impl<K: AsRef<str>> Index<K> for ChatCompletionResponse {
    type Output = Vec<ChatCompletionChoice>;
    fn index(&self, index: K) -> &Self::Output {
        assert_eq!(index.as_ref(), "choices");
        if let Some(item) = &self.choices {
            item
        } else {
            panic!("Key error: choices")
        }
    }
}

#[derive(Debug, Clone, Deserialize)]
pub struct ChatCompletionChoice {
    /// The index of this completion choice
    pub index: usize,
    /// The message produced by the model
    pub message: ChatMessage,
    /// The reason the model stopped generating tokens
    pub finish_reason: String,
}