oaapi 0.2.0

use serde::{Deserialize, Serialize};

use crate::chat::Role;
use crate::chat::ToolType;
use crate::macros::impl_display_for_serialize;

/// The assistant message.
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct AssistantMessage {
    /// The contents of the assistant message.
    /// Required unless tool_calls or function_call is specified.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub content: Option<String>,
    /// The role of the messages author, in this case assistant.
    pub role: Role,
    /// An optional name for the participant.
    /// Provides the model information to differentiate between participants of the same role.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub name: 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>>,
}

impl Default for AssistantMessage {
    fn default() -> Self {
        Self {
            content: None,
            role: Role::Assistant,
            name: None,
            tool_calls: None,
        }
    }
}

impl_display_for_serialize!(AssistantMessage);

impl AssistantMessage {
    pub fn new(
        content: Option<String>,
        name: Option<String>,
        tool_calls: Option<Vec<ToolCall>>,
    ) -> Self {
        Self {
            content,
            role: Role::Assistant,
            name,
            tool_calls,
        }
    }
}

/// The tool call.
#[derive(Debug, Clone, Eq, PartialEq, Default, Serialize, Deserialize)]
pub struct ToolCall {
    /// The ID of the tool call.
    pub id: String,
    /// The type of the tool. Currently, only function is supported.
    #[serde(rename = "type")]
    pub _type: ToolType,
    /// The function that the model called.
    pub function: CalledFunction,
}

impl_display_for_serialize!(ToolCall);

/// The function that the model called by tool.
#[derive(Debug, Clone, Eq, PartialEq, Default, Serialize, Deserialize)]
pub struct CalledFunction {
    /// The name of the function to call.
    pub name: String,
    /// The arguments to call the function with, as generated by the model in JSON format.
    /// Note that the model does not always generate valid JSON, and may hallucinate parameters not defined by your function schema.
    /// Validate the arguments in your code before calling your function.
    pub arguments: String,
}

impl_display_for_serialize!(CalledFunction);

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

    #[test]
    fn deserialize_assistant_message() {
        let json = r#"{
            "content": "Hello, how are you?",
            "role": "assistant",
            "name": "John",
            "tool_calls": [
                {
                    "id": "1",
                    "type": "function",
                    "function": {
                        "name": "my_function",
                        "arguments": "{\"arg1\":\"value1\",\"arg2\":\"value2\"}"
                    }
                }
            ]
        }"#;

        let message: AssistantMessage = serde_json::from_str(json).unwrap();
        assert_eq!(
            message,
            AssistantMessage {
                content: Some("Hello, how are you?".to_string()),
                role: Role::Assistant,
                name: Some("John".to_string()),
                tool_calls: Some(vec![ToolCall {
                    id: "1".to_string(),
                    _type: ToolType::Function,
                    function: CalledFunction {
                        name: "my_function".to_string(),
                        arguments: "{\"arg1\":\"value1\",\"arg2\":\"value2\"}"
                            .to_string(),
                    },
                }]),
            }
        );
    }

    #[test]
    fn deserialize_assistant_message_without_optional() {
        let json = r#"{
            "role": "assistant"
        }"#;

        let message: AssistantMessage = serde_json::from_str(json).unwrap();
        assert_eq!(
            message,
            AssistantMessage {
                content: None,
                role: Role::Assistant,
                name: None,
                tool_calls: None,
            }
        );
    }

    #[test]
    fn serialize_assistant_message() {
        let message = AssistantMessage {
            content: Some("Hello, how are you?".to_string()),
            role: Role::Assistant,
            name: Some("John".to_string()),
            tool_calls: Some(vec![ToolCall {
                id: "1".to_string(),
                _type: ToolType::Function,
                function: CalledFunction {
                    name: "my_function".to_string(),
                    arguments: "{\"arg1\":\"value1\",\"arg2\":\"value2\"}"
                        .to_string(),
                },
            }]),
        };

        let json = serde_json::to_string(&message).unwrap();
        assert_eq!(
            json,
            r#"{"content":"Hello, how are you?","role":"assistant","name":"John","tool_calls":[{"id":"1","type":"function","function":{"name":"my_function","arguments":"{\"arg1\":\"value1\",\"arg2\":\"value2\"}"}}]}"#
        );
    }

    #[test]
    fn serialize_assistant_message_without_optional() {
        let message = AssistantMessage {
            content: None,
            role: Role::Assistant,
            name: None,
            tool_calls: None,
        };

        let json = serde_json::to_string(&message).unwrap();
        assert_eq!(json, r#"{"role":"assistant"}"#);
    }

    #[test]
    fn deserialize_assistant_message_with_content() {
        let json = r#"{
            "role": "assistant",
            "content": "Hello, how are you?"
        }"#;

        let message: AssistantMessage = serde_json::from_str(json).unwrap();
        assert_eq!(
            message,
            AssistantMessage {
                content: Some("Hello, how are you?".to_string()),
                role: Role::Assistant,
                name: None,
                tool_calls: None,
            }
        );
    }

    #[test]
    fn serialize_assistant_message_with_content() {
        let message = AssistantMessage {
            content: Some("Hello, how are you?".to_string()),
            role: Role::Assistant,
            name: None,
            tool_calls: None,
        };

        let json = serde_json::to_string(&message).unwrap();
        assert_eq!(
            json,
            r#"{"content":"Hello, how are you?","role":"assistant"}"#,
        );
    }
}