magi-openai 0.0.10

use std::fmt;

use serde::{Deserialize, Serialize};

use super::request::{FinishReason, ToolCall};

#[derive(Debug, Deserialize, Clone, Default, PartialEq, Serialize)]
pub struct ErrResponse {
    pub error: Err,
}

#[derive(Debug, Deserialize, Clone, Default, PartialEq, Serialize)]
pub struct Err {
    pub message: String,
    pub r#type: String,
    pub param: String,
    pub code: String,
}

#[derive(Debug, Deserialize, Clone, Default, PartialEq, Serialize)]
pub struct Response {
    /// A unique identifier for the chat completion.
    pub id: String,
    /// A list of chat completion choices. Can be more than one if n is greater than 1.
    pub choices: Vec<Choice>,
    /// The Unix timestamp (in seconds) of when the chat completion was created.
    pub created: u64,
    /// The model used for the chat completion.
    pub model: String,
    /// The service tier used for processing the request.
    /// This field is only included if the service_tier parameter is specified in the request.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub service_tier: Option<ServiceTier>,
    /// This fingerprint represents the backend configuration that the model runs with.
    /// Can be used in conjunction with the seed request parameter to understand when backend changes have been
    /// made that might impact determinism.
    pub system_fingerprint: Option<String>,
    /// The object type, which is always chat.completion
    pub object: String,
    /// Usage statistics for the completion request.
    pub usage: Usage,
}

#[derive(Debug, Clone)]
pub struct ResponseBuilder {
    inner: Response,
}

impl Default for ResponseBuilder {
    fn default() -> Self {
        Self::new()
    }
}

impl ResponseBuilder {
    pub fn new() -> Self {
        Self {
            inner: Response {
                object: "chat.completion".to_string(),
                ..Default::default()
            },
        }
    }

    pub fn build(self) -> Response {
        self.inner
    }

    pub fn push_assistant_message(mut self, message: impl Into<String>) -> Self {
        self.inner.choices.push(Choice {
            index: self.inner.choices.len(),
            message: Message {
                role: Role::Assistant,
                content: Some(message.into()),
                ..Default::default()
            },
            finish_reason: None,
            logprobs: None,
        });
        self
    }

    pub fn id(mut self, id: impl Into<String>) -> Self {
        self.inner.id = id.into();
        self
    }

    pub fn model(mut self, model: impl Into<String>) -> Self {
        self.inner.model = model.into();
        self
    }
}

/// Service tier used for processing the request.
///
/// - `auto`: The system will utilize scale tier credits until they are exhausted (if Project is Scale tier enabled), otherwise uses default tier.
/// - `default`: The request will be processed using the default service tier with a lower uptime SLA and no latency guarantee.
/// - `flex`: The request will be processed with the Flex Processing service tier.
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum ServiceTier {
    Auto,
    Default,
    Flex,
}

/// Usage statistics for the completion request.
#[derive(Debug, Deserialize, Serialize, Default, Clone, PartialEq)]
pub struct Usage {
    /// Number of tokens in the generated completion.
    pub completion_tokens: u32,
    /// Number of tokens in the prompt.
    pub prompt_tokens: u32,
    /// Total number of tokens used in the request (prompt + completion).
    pub total_tokens: u32,
    /// Breakdown of tokens used in a completion.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub completion_tokens_details: Option<CompletionTokensDetails>,
    /// Breakdown of tokens used in the prompt.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prompt_tokens_details: Option<PromptTokensDetails>,
}

/// Breakdown of tokens used in a completion.
#[derive(Debug, Deserialize, Serialize, Default, Clone, PartialEq)]
pub struct CompletionTokensDetails {
    /// When using Predicted Outputs, the number of tokens in the prediction that appeared in the completion.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub accepted_prediction_tokens: Option<u32>,
    /// Audio input tokens generated by the model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audio_tokens: Option<u32>,
    /// Tokens generated by the model for reasoning.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_tokens: Option<u32>,
    /// When using Predicted Outputs, the number of tokens in the prediction that did not appear in the completion.
    /// However, like reasoning tokens, these tokens are still counted in the total completion tokens for purposes
    /// of billing, output, and context window limits.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub rejected_prediction_tokens: Option<u32>,
}

/// Breakdown of tokens used in the prompt.
#[derive(Debug, Deserialize, Serialize, Default, Clone, PartialEq)]
pub struct PromptTokensDetails {
    /// Audio input tokens present in the prompt.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audio_tokens: Option<u32>,
    /// Cached tokens present in the prompt.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cached_tokens: Option<u32>,
}

#[derive(Debug, Deserialize, Serialize, Default, Clone, PartialEq)]
pub struct Message {
    /// The contents of the message.
    pub content: Option<String>,

    /// OpenRouter-style gateways may send this alongside `content`; official OpenAI Chat Completions
    /// message objects do not document this field.
    #[serde(
        default,
        skip_serializing_if = "Option::is_none",
        alias = "reasoning_content",
        alias = "thinking",
    )]
    pub reasoning: Option<String>,

    /// The tool calls generated by the model, such as function calls.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tool_calls: Option<Vec<ToolCall>>,

    /// The role of the author of this message.
    pub role: Role,

    /// The refusal message generated by the model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub refusal: Option<String>,

    /// Annotations for the message, when applicable, as when using the web search tool.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub annotations: Option<Vec<Annotation>>,

    /// If the audio output modality is requested, this object contains data about the audio response from the model.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audio: Option<ChatCompletionAudio>,
}

/// The role of the author of a message.
#[derive(Debug, Serialize, Deserialize, Clone, Copy, Default, PartialEq)]
#[serde(rename_all = "lowercase")]
pub enum Role {
    /// Developer-provided instructions (replaces system in some contexts).
    Developer,
    /// System instructions.
    System,
    #[default]
    /// User input.
    User,
    /// Assistant response.
    Assistant,
    /// Tool result.
    Tool,
    /// Deprecated: Function result.
    #[deprecated(note = "Use Tool instead")]
    Function,
}

impl fmt::Display for Role {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            Role::Developer => write!(f, "developer"),
            Role::System => write!(f, "system"),
            Role::User => write!(f, "user"),
            Role::Assistant => write!(f, "assistant"),
            Role::Tool => write!(f, "tool"),
            #[allow(deprecated)]
            Role::Function => write!(f, "function"),
        }
    }
}

impl Role {
    pub fn as_str(&self) -> &'static str {
        match *self {
            Role::Developer => "developer",
            Role::System => "system",
            Role::User => "user",
            Role::Assistant => "assistant",
            Role::Tool => "tool",
            #[allow(deprecated)]
            Role::Function => "function",
        }
    }
}

#[derive(Debug, Deserialize, Default, Serialize, Clone, PartialEq)]
pub struct Choice {
    pub index: usize,
    pub message: Message,
    /// The reason the model stopped generating tokens. This will be `stop` if the model hit a natural stop point or a provided stop sequence,
    /// `length` if the maximum number of tokens specified in the request was reached,
    /// `content_filter` if content was omitted due to a flag from our content filters,
    /// `tool_calls` if the model called a tool, or `function_call` (deprecated) if the model called a function.
    pub finish_reason: Option<FinishReason>,
    /// Log probability information for the choice.
    pub logprobs: Option<Logprobs>,
}

#[derive(Debug, Default, Deserialize, Serialize, Clone, PartialEq)]
pub struct Logprobs {
    /// A list of message content tokens with log probability information.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<Vec<LogprobContent>>,
    /// A list of message refusal tokens with log probability information.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub refusal: Option<Vec<LogprobContent>>,
}

#[derive(Debug, Default, Deserialize, Serialize, Clone, PartialEq)]
pub struct LogprobContent {
    /// The token.
    pub token: String,
    /// The log probability of this token, if it is within the top 20 most likely tokens.
    /// Otherwise, the value -9999.0 is used to signify that the token is very unlikely.
    pub logprob: f64,
    /// A list of integers representing the UTF-8 bytes representation of the token.
    /// Useful in instances where characters are represented by multiple tokens and their
    /// byte representations must be combined to generate the correct text representation.
    /// Can be null if there is no bytes representation for the token.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bytes: Option<Vec<u8>>,
    /// List of the most likely tokens and their log probability, at this token position.
    /// In rare cases, there may be fewer than the number of requested top_logprobs returned.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub top_logprobs: Option<Vec<TopLogprobs>>,
}

#[derive(Debug, Default, Deserialize, Serialize, Clone, PartialEq)]
pub struct TopLogprobs {
    /// The token.
    pub token: String,
    /// The log probability of this token, if it is within the top 20 most likely tokens.
    /// Otherwise, the value -9999.0 is used to signify that the token is very unlikely.
    pub logprob: f64,
    /// A list of integers representing the UTF-8 bytes representation of the token.
    /// Useful in instances where characters are represented by multiple tokens and their
    /// byte representations must be combined to generate the correct text representation.
    /// Can be null if there is no bytes representation for the token.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub bytes: Option<Vec<u8>>,
}

/// Annotation for a message, such as a URL citation when using web search.
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
pub struct Annotation {
    /// The type of the annotation. Currently, only 'url_citation' is supported.
    #[serde(rename = "type")]
    pub annotation_type: String,
    /// A URL citation when using web search.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub url_citation: Option<AnnotationURLCitation>,
}

/// A URL citation when using web search.
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
pub struct AnnotationURLCitation {
    /// The index of the last character of the URL citation in the message.
    pub end_index: usize,
    /// The index of the first character of the URL citation in the message.
    pub start_index: usize,
    /// The title of the web resource.
    pub title: String,
    /// The URL of the web resource.
    pub url: String,
}

/// Audio response from the model when audio output modality is requested.
#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
pub struct ChatCompletionAudio {
    /// Unique identifier for this audio response.
    pub id: String,
    /// Base64 encoded audio bytes generated by the model, in the format specified in the request.
    pub data: String,
    /// The Unix timestamp (in seconds) for when this audio response will no longer be accessible on the server for use in multi-turn conversations.
    pub expires_at: u64,
    /// Transcript of the audio generated by the model.
    pub transcript: String,
}

impl Response {
    pub fn first_assistant_message(&self) -> Option<&Message> {
        self.choices
            .iter()
            .find(|choice| choice.message.role == Role::Assistant)
            .map(|choice| &choice.message)
    }

    pub fn first_assistant_message_text(&self) -> Option<String> {
        self.first_assistant_message()
            .and_then(|message| message.content.to_owned())
    }

    pub fn first_assistant_message_reasoning_text(&self) -> Option<String> {
        self.first_assistant_message()
            .and_then(|message| message.reasoning.to_owned())
    }
}

#[cfg(test)]
mod tests {
    use crate::completions::request::{ToolCallFunction, ToolCallFunctionObj};

    use super::*;

    #[test]
    fn serde() {
        let tests = vec![
            (
                "default",
                r#"{
                "id": "chatcmpl-123",
                "object": "chat.completion",
                "created": 1677652288,
                "model": "gpt-3.5-turbo-0613",
                "system_fingerprint": "fp_44709d6fcb",
                "choices": [{
                  "index": 0,
                  "message": {
                    "role": "assistant",
                    "content": "\n\nHello there, how may I assist you today?"
                  },
                  "logprobs": null,
                  "finish_reason": "stop"
                }],
                "usage": {
                  "prompt_tokens": 9,
                  "completion_tokens": 12,
                  "total_tokens": 21
                }
              }"#,
                Response {
                    id: "chatcmpl-123".to_string(),
                    object: "chat.completion".to_string(),
                    created: 1677652288,
                    model: "gpt-3.5-turbo-0613".to_string(),
                    system_fingerprint: Some("fp_44709d6fcb".to_string()),
                    choices: vec![Choice {
                        index: 0,
                        message: Message {
                            role: Role::Assistant,
                            content: Some(
                                "\n\nHello there, how may I assist you today?".to_string(),
                            ),
                            reasoning: None,
                            tool_calls: None,
                            refusal: None,
                            annotations: None,
                            audio: None,
                        },
                        logprobs: None,
                        finish_reason: Some(FinishReason::Stop),
                    }],
                    usage: Usage {
                        prompt_tokens: 9,
                        completion_tokens: 12,
                        total_tokens: 21,
                        ..Default::default()
                    },
                    service_tier: None,
                },
            ),
            (
                "function",
                r#"{
                    "id": "chatcmpl-abc123",
                    "object": "chat.completion",
                    "created": 1699896916,
                    "model": "gpt-3.5-turbo-0613",
                    "system_fingerprint": "fp_6b68a8204b",
                    "choices": [
                      {
                        "index": 0,
                        "message": {
                          "role": "assistant",
                          "content": null,
                          "tool_calls": [
                            {
                              "id": "call_abc123",
                              "type": "function",
                              "function": {
                                "name": "get_current_weather",
                                "arguments": "{\n\"location\": \"Boston, MA\"\n}"
                              }
                            }
                          ]
                        },
                        "logprobs": null,
                        "finish_reason": "tool_calls"
                      }
                    ],
                    "usage": {
                      "prompt_tokens": 82,
                      "completion_tokens": 17,
                      "total_tokens": 99
                    }
                  }"#,
                Response {
                    id: "chatcmpl-abc123".to_string(),
                    object: "chat.completion".to_string(),
                    created: 1699896916,
                    model: "gpt-3.5-turbo-0613".to_string(),
                    system_fingerprint: Some("fp_6b68a8204b".to_string()),
                    choices: vec![Choice {
                        index: 0,
                        message: Message {
                            role: Role::Assistant,
                            content: None,
                            tool_calls: Some(vec![ToolCall::Function(ToolCallFunction {
                                id: "call_abc123".to_string(),
                                function: ToolCallFunctionObj {
                                    name: "get_current_weather".to_string(),
                                    arguments: "{\n\"location\": \"Boston, MA\"\n}".to_string(),
                                },
                            })]),
                            refusal: None,
                            reasoning: None,
                            annotations: None,
                            audio: None,
                        },
                        logprobs: None,
                        finish_reason: Some(FinishReason::ToolCalls),
                    }],
                    usage: Usage {
                        prompt_tokens: 82,
                        completion_tokens: 17,
                        total_tokens: 99,
                        ..Default::default()
                    },
                    service_tier: None,
                },
            ),
            (
                "logprobs",
                r#"{
                    "id": "chatcmpl-123",
                    "object": "chat.completion",
                    "created": 1702685778,
                    "model": "gpt-3.5-turbo-0613",
                    "choices": [
                      {
                        "index": 0,
                        "message": {
                          "role": "assistant",
                          "content": "Hello! How can I assist you today?"
                        },
                        "logprobs": {
                          "content": [
                            {
                              "token": "Hello",
                              "logprob": -0.31725305,
                              "bytes": [72, 101, 108, 108, 111],
                              "top_logprobs": [
                                {
                                  "token": "Hello",
                                  "logprob": -0.31725305,
                                  "bytes": [72, 101, 108, 108, 111]
                                },
                                {
                                  "token": "Hi",
                                  "logprob": -1.3190403,
                                  "bytes": [72, 105]
                                }
                              ]
                            },
                            {
                              "token": "!",
                              "logprob": -0.02380986,
                              "bytes": [33],
                              "top_logprobs": [
                                {
                                  "token": "!",
                                  "logprob": -0.02380986,
                                  "bytes": [33]
                                },
                                {
                                  "token": " there",
                                  "logprob": -3.787621,
                                  "bytes": [32, 116, 104, 101, 114, 101]
                                }
                              ]
                            },
                            {
                              "token": " How",
                              "logprob": -0.000054669687,
                              "bytes": [32, 72, 111, 119],
                              "top_logprobs": [
                                {
                                  "token": " How",
                                  "logprob": -0.000054669687,
                                  "bytes": [32, 72, 111, 119]
                                },
                                {
                                  "token": "<|end|>",
                                  "logprob": -10.953937,
                                  "bytes": null
                                }
                              ]
                            }
                          ]
                        },
                        "finish_reason": "stop"
                      }
                    ],
                    "usage": {
                      "prompt_tokens": 9,
                      "completion_tokens": 9,
                      "total_tokens": 18
                    },
                    "system_fingerprint": "fp_44709d6fcb"
                  }"#,
                Response {
                    id: "chatcmpl-123".to_string(),
                    object: "chat.completion".to_string(),
                    created: 1702685778,
                    model: "gpt-3.5-turbo-0613".to_string(),
                    system_fingerprint: Some("fp_44709d6fcb".to_string()),
                    choices: vec![Choice {
                        index: 0,
                        message: Message {
                            role: Role::Assistant,
                            content: Some("Hello! How can I assist you today?".to_string()),
                            reasoning: None,
                            tool_calls: None,
                            refusal: None,
                            annotations: None,
                            audio: None,
                        },
                        logprobs: Some(Logprobs {
                            content: Some(vec![
                                LogprobContent {
                                    token: "Hello".to_string(),
                                    logprob: -0.31725305,
                                    bytes: Some(vec![72, 101, 108, 108, 111]),
                                    top_logprobs: Some(vec![
                                        TopLogprobs {
                                            token: "Hello".to_string(),
                                            logprob: -0.31725305,
                                            bytes: Some(vec![72, 101, 108, 108, 111]),
                                        },
                                        TopLogprobs {
                                            token: "Hi".to_string(),
                                            logprob: -1.3190403,
                                            bytes: Some(vec![72, 105]),
                                        },
                                    ]),
                                },
                                LogprobContent {
                                    token: "!".to_string(),
                                    logprob: -0.02380986,
                                    bytes: Some(vec![33]),
                                    top_logprobs: Some(vec![
                                        TopLogprobs {
                                            token: "!".to_string(),
                                            logprob: -0.02380986,
                                            bytes: Some(vec![33]),
                                        },
                                        TopLogprobs {
                                            token: " there".to_string(),
                                            logprob: -3.787621,
                                            bytes: Some(vec![32, 116, 104, 101, 114, 101]),
                                        },
                                    ]),
                                },
                                LogprobContent {
                                    token: " How".to_string(),
                                    logprob: -0.000054669687,
                                    bytes: Some(vec![32, 72, 111, 119]),
                                    top_logprobs: Some(vec![
                                        TopLogprobs {
                                            token: " How".to_string(),
                                            logprob: -0.000054669687,
                                            bytes: Some(vec![32, 72, 111, 119]),
                                        },
                                        TopLogprobs {
                                            token: "<|end|>".to_string(),
                                            logprob: -10.953937,
                                            bytes: None,
                                        },
                                    ]),
                                },
                            ]),
                            refusal: None,
                        }),
                        finish_reason: Some(FinishReason::Stop),
                    }],
                    usage: Usage {
                        prompt_tokens: 9,
                        completion_tokens: 9,
                        total_tokens: 18,
                        completion_tokens_details: None,
                        prompt_tokens_details: None,
                    },
                    service_tier: None,
                },
            ),
            (
                "refusal",
                r#"{
                    "id": "chatcmpl-123456",
                    "object": "chat.completion",
                    "created": 1728933352,
                    "model": "gpt-4o-2024-08-06",
                    "choices": [
                        {
                            "index": 0,
                            "message": {
                                "role": "assistant",
                                "content": "Hi there! How can I assist you today?"
                            },
                            "logprobs": null,
                            "finish_reason": "stop"
                        }
                    ],
                    "usage": {
                        "prompt_tokens": 19,
                        "completion_tokens": 10,
                        "total_tokens": 29,
                        "prompt_tokens_details": {
                            "cached_tokens": 0
                        },
                        "completion_tokens_details": {
                            "reasoning_tokens": 0,
                            "accepted_prediction_tokens": 0,
                            "rejected_prediction_tokens": 0
                        }
                    },
                    "system_fingerprint": "fp_6b68a8204b"
                }"#,
                Response {
                    id: "chatcmpl-123456".to_string(),
                    object: "chat.completion".to_string(),
                    created: 1728933352,
                    model: "gpt-4o-2024-08-06".to_string(),
                    system_fingerprint: Some("fp_6b68a8204b".to_string()),
                    choices: vec![Choice {
                        index: 0,
                        message: Message {
                            role: Role::Assistant,
                            content: Some("Hi there! How can I assist you today?".to_string()),
                            reasoning: None,
                            tool_calls: None,
                            refusal: None,
                            annotations: None,
                            audio: None,
                        },
                        logprobs: None,
                        finish_reason: Some(FinishReason::Stop),
                    }],
                    usage: Usage {
                        prompt_tokens: 19,
                        completion_tokens: 10,
                        total_tokens: 29,
                        prompt_tokens_details: Some(PromptTokensDetails {
                            audio_tokens: None,
                            cached_tokens: Some(0),
                        }),
                        completion_tokens_details: Some(CompletionTokensDetails {
                            reasoning_tokens: Some(0),
                            accepted_prediction_tokens: Some(0),
                            rejected_prediction_tokens: Some(0),
                            audio_tokens: None,
                        }),
                    },
                    service_tier: None,
                },
            ),
        ];
        for (name, json, expected) in tests {
            //test deserialize
            let actual: Response = serde_json::from_str(json).unwrap();
            assert_eq!(actual, expected, "deserialize test failed: {}", name);
            //test serialize
            let serialized = serde_json::to_string(&expected).unwrap();
            let actual: Response = serde_json::from_str(&serialized).unwrap();
            assert_eq!(actual, expected, "serialize test failed: {}", name);
        }
    }
}