grok-client 0.3.0

pub mod api {
    // ####################
    // AI API TYPES
    // ####################
    use crate::ai;
    use serde::{Deserialize, Serialize};

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct ApiKey {
        pub redacted_api_key: String,
        pub user_id: String,
        pub name: String,
        pub create_time: String,
        pub modify_time: String,
        pub modified_by: String,
        pub team_id: String,
        pub acls: Vec<String>,
        pub api_key_id: String,
        pub team_blocked: bool,
        pub api_key_blocked: bool,
        pub api_key_disabled: bool,
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct Model {
        pub id: String,
        pub created: i64,
        pub object: String,
        pub owned_by: String,
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct LanguageModel {
        pub id: String,
        pub fingerprint: String,
        pub created: i64,
        pub object: String,
        pub owned_by: String,
        pub version: String,
        pub cached_prompt_text_token_price: i64,
        pub completion_text_token_price: i64,
        pub prompt_image_token_price: i64,
        pub prompt_text_token_price: i64,
        pub input_modalities: Vec<String>,
        pub output_modalities: Vec<String>,
        pub aliases: Vec<String>,
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct LanguageModels {
        pub models: Vec<LanguageModel>,
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct ImageModel {
        pub id: String,
        pub fingerprint: String,
        pub max_prompt_length: i64,
        pub created: i64,
        pub object: String,
        pub owned_by: String,
        pub version: String,
        pub image_price: i64,
        pub input_modalities: Vec<String>,
        pub output_modalities: Vec<String>,
        pub aliases: Vec<String>,
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct ImageModels {
        pub models: Vec<ImageModel>,
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct TokenizeRequest {
        pub model: String,
        pub text: String,
    }

    impl TokenizeRequest {
        pub fn init(model: ai::LanguageModel, text: String) -> Self {
            Self {
                model: model.to_string(),
                text,
            }
        }
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct Token {
        pub token_id: i64,
        pub string_token: String,
        pub token_bytes: Vec<i64>,
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct TokenizeResponse {
        pub token_ids: Vec<Token>,
    }
}

pub mod chat {
    // ####################
    // AI CHAT API TYPES
    // ####################
    use crate::ai;
    use serde::{Deserialize, Serialize};

    pub mod stream {
        use crate::types::chat::Usage;
        use serde::{Deserialize, Serialize};

        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
        pub struct Options {
            pub include_usage: Option<bool>,
        }

        impl Options {
            pub fn default() -> Self {
                Self {
                    include_usage: Some(true),
                }
            }
        }

        #[derive(Debug, Serialize, Deserialize)]
        pub struct ChatCompletionChunk {
            pub id: String,
            pub object: String,
            pub created: i64,
            pub model: String,
            pub choices: Vec<Choice>,
            pub usage: Option<Usage>,
            pub citations: Option<Vec<String>>,
            pub system_fingerprint: String,
        }

        #[derive(Debug, Serialize, Deserialize)]
        pub struct Choice {
            pub index: i32,
            pub delta: Delta,
        }

        #[derive(Debug, Serialize, Deserialize)]
        pub struct Delta {
            pub content: Option<String>,
            pub reasoning_content: Option<String>,
            pub role: Option<String>,
        }
    }

    pub mod search {
        use serde::{Deserialize, Serialize};

        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
        pub struct Parameters {
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub from_date: Option<String>,
            #[serde(
                default = "default_max_search_results",
                skip_serializing_if = "Option::is_none"
            )]
            pub max_search_results: Option<i32>,
            #[serde(default = "default_mode", skip_serializing_if = "Option::is_none")]
            pub mode: Option<String>,
            #[serde(
                default = "default_return_citations",
                skip_serializing_if = "Option::is_none"
            )]
            pub return_citations: Option<bool>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub sources: Option<Vec<Source>>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub to_date: Option<String>,
        }

        impl Parameters {
            pub fn default() -> Self {
                Self {
                    from_date: None,
                    max_search_results: default_max_search_results(),
                    mode: default_mode(),
                    return_citations: default_return_citations(),
                    sources: None,
                    to_date: None,
                }
            }
        }

        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
        pub struct Source {
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub excluded_x_handles: Option<Vec<String>>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub included_x_handles: Option<Vec<String>>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub x_handles: Option<Vec<String>>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub post_favorite_count: Option<i32>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub post_view_count: Option<i32>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub allowed_websites: Option<Vec<String>>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub country: Option<String>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub excluded_websites: Option<Vec<String>>,
            #[serde(
                default = "default_safe_search",
                skip_serializing_if = "Option::is_none"
            )]
            pub safe_search: Option<bool>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub links: Option<Vec<String>>,
            pub r#type: String,
        }

        fn default_max_search_results() -> Option<i32> {
            Some(15)
        }

        fn default_mode() -> Option<String> {
            Some(String::from("auto"))
        }

        fn default_return_citations() -> Option<bool> {
            Some(true)
        }

        fn default_safe_search() -> Option<bool> {
            Some(false)
        }
    }

    pub mod tool {
        use serde::{Deserialize, Serialize};

        // "none" "auto" "required" or a single tool::Tool
        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
        pub enum Choice {
            None,
            Auto,
            Required,
            Instance(Tool),
        }

        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
        pub struct Tool {
            pub r#type: String, // Using `r#type` to match the schema's `type` field
            pub function: Function,

            // CALL
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub id: Option<String>,
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub index: Option<i32>,
        }

        #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
        pub struct Function {
            pub name: String,

            // DEFINITION
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub parameters: Option<serde_json::Value>, // Using Value to handle arbitrary JSON schema
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub description: Option<String>,

            // CALL
            // #[serde(
            //     default,
            //     skip_serializing_if = "Option::is_none",
            //     deserialize_with = "deserialize_to_object"
            // )]
            #[serde(default, skip_serializing_if = "Option::is_none")]
            pub arguments: Option<serde_json::Value>, // Using Value to handle arbitrary JSON schema
        }
    }

    // Struct for chat completion request
    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct ChatCompletionRequest {
        pub model: String,
        pub messages: Vec<Message>,
        // STREAM
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub stream: Option<bool>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub stream_options: Option<stream::Options>,
        // DEFER
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub deferred: Option<bool>,
        // SEARCH
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub search_parameters: Option<search::Parameters>,
        // TOOLS
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub parallel_tool_calls: Option<bool>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub tools: Option<Vec<tool::Tool>>,
        // #[serde(default, skip_serializing_if = "Option::is_none")]
        // pub tool_choice: Option<tool::Choice>, // "none" "auto" "required" or a single tool::Tool
        // #[serde(default, skip_serializing_if = "Option::is_none")]
        // pub n: Option<i32>,
    }

    impl ChatCompletionRequest {
        pub fn init(model: ai::LanguageModel, messages: Vec<Message>) -> Self {
            Self {
                model: model.to_string(),
                messages,
                deferred: Some(false),
                stream: Some(false),
                stream_options: None,
                search_parameters: None,
                parallel_tool_calls: None,
                tools: None,
                // tool_choice: None,
            }
        }
    }

    // Message object (used in both request and response)
    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct Message {
        pub role: String,
        pub content: String, // Could also be and Vec<Struct>
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub reasoning_content: Option<String>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub refusal: Option<String>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub tool_calls: Option<Vec<tool::Tool>>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub tool_call_id: Option<String>,
    }

    impl Message {
        pub fn init(role: ai::Role, content: String) -> Self {
            Self {
                role: role.to_string(),
                content,
                reasoning_content: None,
                refusal: None,
                tool_calls: None,
                tool_call_id: None,
            }
        }
        pub fn tool(tool_call_id: String, content: String) -> Self {
            Self {
                role: "tool".to_string(),
                tool_call_id: Some(tool_call_id),
                content,
                reasoning_content: None,
                refusal: None,
                tool_calls: None,
            }
        }
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct DeferredChatCompletionResponse {
        pub request_id: String,
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct ChatCompletionResponse {
        pub id: String,
        pub object: String, // Always "chat.response"
        pub created: i64,   // Unix timestamp
        pub model: String,
        pub choices: Vec<Choice>, // Length corresponds to 'n' in request
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub usage: Option<Usage>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub citations: Option<Vec<String>>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub system_fingerprint: Option<String>,
    }

    impl ChatCompletionResponse {
        pub fn new(i: i64) -> Self {
            ChatCompletionResponse {
                id: format!("chatcmpl-{}", i),
                object: "chat.response".to_string(),
                created: 1697056000 + i as i64,
                model: format!("grok-beta-{}", i),
                choices: vec![],
                usage: None,
                citations: None,
                system_fingerprint: None,
            }
        }

        pub fn mock(count: i64) -> Vec<Self> {
            (0..count).map(|i| ChatCompletionResponse::new(i)).collect()
        }
    }

    // Choice object within the response
    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct Choice {
        pub index: i32,
        pub message: Message,
        pub finish_reason: String, // "stop" "tool_calls" "length" "end_turn"
    }

    // Usage object
    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct Usage {
        pub prompt_tokens: i32,
        pub prompt_tokens_details: PromptTokensDetails,
        pub completion_tokens: i32,
        pub completion_tokens_details: CompletionTokensDetails,
        pub total_tokens: i32,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub num_sources_used: Option<i32>, // Optional for backward compatibility
    }

    // Details for prompt tokens
    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct PromptTokensDetails {
        pub text_tokens: i32,
        pub audio_tokens: i32,
        pub image_tokens: i32,
        pub cached_tokens: i32,
    }

    // Details for completion tokens
    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct CompletionTokensDetails {
        pub reasoning_tokens: i32,
        pub audio_tokens: i32,
        pub accepted_prediction_tokens: i32,
        pub rejected_prediction_tokens: i32,
    }
}

pub mod image {
    // ####################
    // AI IMAGE API TYPES
    // ####################
    use serde::{Deserialize, Serialize};
    use strum::{Display, EnumIter, EnumString};

    #[derive(
        Clone, Copy, Debug, Display, PartialEq, EnumIter, EnumString, Serialize, Deserialize,
    )]
    pub enum ImageResponseFormat {
        #[strum(serialize = "url")]
        URL,
        #[strum(serialize = "b64_json")]
        B64JSON,
    }

    impl ImageResponseFormat {
        pub fn err_ivalid_response_format(format: &String) -> String {
            format!("Invalid response format '{format}'")
        }
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct ImageRequest {
        prompt: String,
        model: String,
        response_format: String,
        n: u32,
    }

    impl ImageRequest {
        pub fn init(
            prompt: String,
            model: String,
            response_format: ImageResponseFormat,
            n: u32,
        ) -> Self {
            Self {
                prompt,
                model,
                response_format: response_format.to_string(),
                n,
            }
        }
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct ImageData {
        pub revised_prompt: String,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub url: Option<String>,
        #[serde(default, skip_serializing_if = "Option::is_none")]
        pub b64_json: Option<String>,
    }

    #[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
    pub struct ImageResponse {
        pub data: Vec<ImageData>,
    }
}