magi-openai 0.1.0

use serde::{Deserialize, Serialize};

/// Response returned by the Images API.
#[derive(Debug, Deserialize, Clone, Default, PartialEq, Serialize)]
pub struct Response {
    /// Unix timestamp in seconds when the image response was created.
    pub created: u64,

    /// Background mode used for generated images.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub background: Option<String>,

    /// Generated image payloads.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data: Option<Vec<ImageData>>,

    /// Output image format.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub output_format: Option<String>,

    /// Image quality used for generation.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub quality: Option<String>,

    /// Generated image size.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub size: Option<String>,

    /// Token usage for the image generation request.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub usage: Option<Usage>,
}

/// Image data returned by the Images API.
#[derive(Debug, Deserialize, Clone, Default, PartialEq, Serialize)]
pub struct ImageData {
    /// Base64 encoded image bytes.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub b64_json: Option<String>,

    /// URL for image bytes when the API returns a hosted image.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub url: Option<String>,

    /// Revised prompt used by DALL-E 3.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub revised_prompt: Option<String>,
}

/// Token usage for an image generation request.
#[derive(Debug, Deserialize, Serialize, Default, Clone, PartialEq)]
pub struct Usage {
    /// Number of tokens present in the input.
    pub input_tokens: u32,

    /// Detailed input token usage statistics.
    pub input_tokens_details: InputTokensDetails,

    /// Number of tokens generated in the output.
    pub output_tokens: u32,

    /// Detailed output token usage statistics.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub output_tokens_details: Option<OutputTokensDetails>,

    /// Total number of tokens used for the request.
    pub total_tokens: u32,
}

/// Detailed input token usage statistics.
#[derive(Debug, Deserialize, Serialize, Default, Clone, PartialEq)]
pub struct InputTokensDetails {
    /// Image input tokens present in the prompt.
    pub image_tokens: u32,

    /// Text input tokens present in the prompt.
    pub text_tokens: u32,
}

/// Detailed output token usage statistics.
#[derive(Debug, Deserialize, Serialize, Default, Clone, PartialEq)]
pub struct OutputTokensDetails {
    /// Image tokens generated by the model.
    pub image_tokens: u32,

    /// Text tokens generated by the model.
    pub text_tokens: u32,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn serde_gpt_image_2_response() {
        let json = r#"{
            "created": 1782974226,
            "background": "opaque",
            "data": [
                {
                    "b64_json": "iVBORw0"
                }
            ],
            "output_format": "png",
            "quality": "low",
            "size": "1536x656",
            "usage": {
                "input_tokens": 162,
                "input_tokens_details": {
                    "image_tokens": 0,
                    "text_tokens": 162
                },
                "output_tokens": 85,
                "output_tokens_details": {
                    "image_tokens": 85,
                    "text_tokens": 0
                },
                "total_tokens": 247
            }
        }"#;

        let expected = Response {
            created: 1782974226,
            background: Some("opaque".to_string()),
            data: Some(vec![ImageData {
                b64_json: Some("iVBORw0".to_string()),
                url: None,
                revised_prompt: None,
            }]),
            output_format: Some("png".to_string()),
            quality: Some("low".to_string()),
            size: Some("1536x656".to_string()),
            usage: Some(Usage {
                input_tokens: 162,
                input_tokens_details: InputTokensDetails {
                    image_tokens: 0,
                    text_tokens: 162,
                },
                output_tokens: 85,
                output_tokens_details: Some(OutputTokensDetails {
                    image_tokens: 85,
                    text_tokens: 0,
                }),
                total_tokens: 247,
            }),
        };

        let actual: Response = serde_json::from_str(json).unwrap();
        assert_eq!(actual, expected);

        let serialized = serde_json::to_value(&expected).unwrap();
        let source: serde_json::Value = serde_json::from_str(json).unwrap();
        assert_eq!(serialized, source);

        let roundtrip: Response = serde_json::from_value(serialized).unwrap();
        assert_eq!(roundtrip, expected);
    }

    #[test]
    fn serde_official_reference_response_example() {
        let json = r#"{
            "created": 1713833628,
            "data": [
                {
                    "b64_json": "..."
                }
            ],
            "usage": {
                "total_tokens": 100,
                "input_tokens": 50,
                "output_tokens": 50,
                "input_tokens_details": {
                    "text_tokens": 10,
                    "image_tokens": 40
                }
            }
        }"#;

        let expected = Response {
            created: 1713833628,
            background: None,
            data: Some(vec![ImageData {
                b64_json: Some("...".to_string()),
                url: None,
                revised_prompt: None,
            }]),
            output_format: None,
            quality: None,
            size: None,
            usage: Some(Usage {
                input_tokens: 50,
                input_tokens_details: InputTokensDetails {
                    image_tokens: 40,
                    text_tokens: 10,
                },
                output_tokens: 50,
                output_tokens_details: None,
                total_tokens: 100,
            }),
        };

        let actual: Response = serde_json::from_str(json).unwrap();
        assert_eq!(actual, expected);

        let serialized = serde_json::to_value(&expected).unwrap();
        let source: serde_json::Value = serde_json::from_str(json).unwrap();
        assert_eq!(serialized, source);

        let roundtrip: Response = serde_json::from_value(serialized).unwrap();
        assert_eq!(roundtrip, expected);
    }
}