weavegraph 0.7.0

//! Chat message types for conversation turns and roles.

use serde::{Deserialize, Serialize};
use std::fmt;

/// Participant role in a conversation turn.
///
/// Roles serialize to/from lowercase strings:
/// `User` ↔ `"user"`, `Assistant` ↔ `"assistant"`, `System` ↔ `"system"`,
/// `Tool` ↔ `"tool"`, `Custom("foo")` ↔ `"foo"`.
#[derive(Clone, Debug, PartialEq, Eq, Hash, Default, Serialize, Deserialize)]
#[serde(into = "String", try_from = "String")]
pub enum Role {
    /// Human / end-user turn.
    #[default]
    User,
    /// Model response turn.
    Assistant,
    /// System-level instruction turn.
    System,
    /// Tool or function result turn.
    Tool,
    /// Any role not covered by the standard variants.
    Custom(String),
}

impl Role {
    /// String form of this role.
    #[must_use]
    pub fn as_str(&self) -> &str {
        match self {
            Self::User => "user",
            Self::Assistant => "assistant",
            Self::System => "system",
            Self::Tool => "tool",
            Self::Custom(s) => s.as_str(),
        }
    }

    /// Returns `true` if this role's string form equals `role_str`.
    #[must_use]
    pub fn matches(&self, role_str: &str) -> bool {
        self.as_str() == role_str
    }
}

impl fmt::Display for Role {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(self.as_str())
    }
}

impl From<&str> for Role {
    fn from(s: &str) -> Self {
        match s {
            "user" => Self::User,
            "assistant" => Self::Assistant,
            "system" => Self::System,
            "tool" => Self::Tool,
            other => Self::Custom(other.to_owned()),
        }
    }
}

impl From<String> for Role {
    fn from(s: String) -> Self {
        Self::from(s.as_str())
    }
}

impl From<Role> for String {
    fn from(role: Role) -> Self {
        role.as_str().to_owned()
    }
}

/// A single message in a conversation: a role paired with text content.
#[derive(Clone, Debug, PartialEq, Eq, Default, Serialize, Deserialize)]
pub struct Message {
    /// Who sent the message.
    pub role: Role,
    /// Text body of the message.
    pub content: String,
}

impl Message {
    /// Construct a message with an explicit role and content.
    #[must_use]
    pub fn with_role(role: Role, content: &str) -> Self {
        Self {
            role,
            content: content.to_owned(),
        }
    }

    /// Construct a `User` message.
    #[must_use]
    pub fn user(content: &str) -> Self {
        Self::with_role(Role::User, content)
    }

    /// Construct an `Assistant` message.
    #[must_use]
    pub fn assistant(content: &str) -> Self {
        Self::with_role(Role::Assistant, content)
    }

    /// Construct a `System` message.
    #[must_use]
    pub fn system(content: &str) -> Self {
        Self::with_role(Role::System, content)
    }

    /// Construct a `Tool` message.
    #[must_use]
    pub fn tool(content: &str) -> Self {
        Self::with_role(Role::Tool, content)
    }
}

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

    #[test]
    fn test_role_from_str() {
        assert_eq!(Role::from("user"), Role::User);
        assert_eq!(Role::from("assistant"), Role::Assistant);
        assert_eq!(Role::from("system"), Role::System);
        assert_eq!(Role::from("tool"), Role::Tool);
        assert_eq!(Role::from("custom"), Role::Custom("custom".to_string()));
    }

    #[test]
    fn test_role_as_str() {
        assert_eq!(Role::User.as_str(), "user");
        assert_eq!(Role::Assistant.as_str(), "assistant");
        assert_eq!(Role::System.as_str(), "system");
        assert_eq!(Role::Tool.as_str(), "tool");
        assert_eq!(Role::Custom("foo".into()).as_str(), "foo");
    }

    #[test]
    fn test_message_role_typed_field() {
        let msg = Message::user("hello");
        assert_eq!(msg.role, Role::User);

        let msg = Message::assistant("hi");
        assert_eq!(msg.role, Role::Assistant);

        let msg = Message::with_role(Role::Custom("custom".into()), "data");
        assert_eq!(msg.role, Role::Custom("custom".into()));
    }

    #[test]
    fn test_message_with_role() {
        let msg = Message::with_role(Role::Tool, "result");
        assert_eq!(msg.role, Role::Tool);
        assert_eq!(msg.content, "result");
    }

    #[test]
    fn test_role_serialization() {
        let role = Role::User;
        let json = serde_json::to_string(&role).unwrap();
        assert_eq!(json, "\"user\"");

        let parsed: Role = serde_json::from_str("\"assistant\"").unwrap();
        assert_eq!(parsed, Role::Assistant);

        let custom: Role = serde_json::from_str("\"function\"").unwrap();
        assert_eq!(custom, Role::Custom("function".into()));
    }

    #[test]
    fn test_message_backward_compatibility() {
        let json = r#"{"role": "user", "content": "hello"}"#;
        let msg: Message = serde_json::from_str(json).unwrap();
        assert_eq!(msg.role, Role::User);
    }
}