autoagents-protocol 0.3.7

use serde::{Deserialize, Serialize};

/// Usage metadata for a chat response.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct Usage {
    pub prompt_tokens: u32,
    pub completion_tokens: u32,
    pub total_tokens: u32,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub completion_tokens_details: Option<CompletionTokensDetails>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub prompt_tokens_details: Option<PromptTokensDetails>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct CompletionTokensDetails {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub reasoning_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audio_tokens: Option<u32>,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PromptTokensDetails {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub cached_tokens: Option<u32>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub audio_tokens: Option<u32>,
}

/// The supported MIME type of an image.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum ImageMime {
    JPEG,
    PNG,
    GIF,
    WEBP,
}

impl ImageMime {
    pub fn mime_type(&self) -> &'static str {
        match self {
            ImageMime::JPEG => "image/jpeg",
            ImageMime::PNG => "image/png",
            ImageMime::GIF => "image/gif",
            ImageMime::WEBP => "image/webp",
        }
    }
}

/// Tool call represents a function call that an LLM wants to make.
#[derive(Debug, Deserialize, Serialize, Clone, Eq, PartialEq)]
pub struct ToolCall {
    pub id: String,
    #[serde(rename = "type")]
    pub call_type: String,
    pub function: FunctionCall,
}

/// FunctionCall contains details about which function to call and with what arguments.
#[derive(Debug, Deserialize, Serialize, Clone, Eq, PartialEq)]
pub struct FunctionCall {
    pub name: String,
    pub arguments: String,
}

/// A streaming chunk that can be either text or a tool call event.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum StreamChunk {
    Text(String),
    ReasoningContent(String),
    ToolUseStart {
        index: usize,
        id: String,
        name: String,
    },
    ToolUseInputDelta {
        index: usize,
        partial_json: String,
    },
    ToolUseComplete {
        index: usize,
        tool_call: ToolCall,
    },
    Done {
        stop_reason: String,
    },
    Usage(Usage),
}

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

    #[test]
    fn stream_chunk_serializes_roundtrip() {
        let chunk = StreamChunk::ToolUseStart {
            index: 1,
            id: "tool_1".to_string(),
            name: "search".to_string(),
        };
        let serialized = serde_json::to_string(&chunk).unwrap();
        let deserialized: StreamChunk = serde_json::from_str(&serialized).unwrap();
        match deserialized {
            StreamChunk::ToolUseStart { id, name, .. } => {
                assert_eq!(id, "tool_1");
                assert_eq!(name, "search");
            }
            _ => panic!("expected ToolUseStart"),
        }
    }

    #[test]
    fn tool_call_serializes_roundtrip() {
        let call = ToolCall {
            id: "call_1".to_string(),
            call_type: "function".to_string(),
            function: FunctionCall {
                name: "lookup".to_string(),
                arguments: "{\"q\":\"value\"}".to_string(),
            },
        };
        let serialized = serde_json::to_string(&call).unwrap();
        let deserialized: ToolCall = serde_json::from_str(&serialized).unwrap();
        assert_eq!(deserialized, call);
    }

    #[test]
    fn image_mime_serializes_roundtrip() {
        let mime = ImageMime::PNG;
        let serialized = serde_json::to_string(&mime).unwrap();
        let deserialized: ImageMime = serde_json::from_str(&serialized).unwrap();
        assert_eq!(deserialized, mime);
    }

    #[test]
    fn image_mime_type_mapping() {
        assert_eq!(ImageMime::JPEG.mime_type(), "image/jpeg");
        assert_eq!(ImageMime::PNG.mime_type(), "image/png");
        assert_eq!(ImageMime::GIF.mime_type(), "image/gif");
        assert_eq!(ImageMime::WEBP.mime_type(), "image/webp");
    }

    #[test]
    fn usage_serializes_with_details() {
        let usage = Usage {
            prompt_tokens: 1,
            completion_tokens: 2,
            total_tokens: 3,
            completion_tokens_details: Some(CompletionTokensDetails {
                reasoning_tokens: Some(1),
                audio_tokens: None,
            }),
            prompt_tokens_details: Some(PromptTokensDetails {
                cached_tokens: Some(2),
                audio_tokens: None,
            }),
        };
        let serialized = serde_json::to_value(&usage).unwrap();
        assert!(serialized.get("completion_tokens_details").is_some());
        assert!(serialized.get("prompt_tokens_details").is_some());
    }
}