kproc-llm 0.6.1

use serde::{Deserialize, Serialize};

/// Speaker
#[derive(Debug)]
pub enum Role
{
  /// User role
  User,
  /// System role
  System,
  /// Assistant role
  Assistant,
  /// Custom role, defined by the string. Often fallback to User.
  Custom(String),
}

/// Message
#[derive(Debug, Serialize, Deserialize)]
pub struct Message
{
  /// Role of the message
  pub role: Role,
  /// Content of the message
  pub content: String,
}

impl Serialize for Role
{
  fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
  where
    S: serde::Serializer,
  {
    let s = match self
    {
      Role::User => "user",
      Role::System => "system",
      Role::Assistant => "assistant",
      Role::Custom(s) => s.as_str(),
    };
    serializer.serialize_str(s)
  }
}

impl<'de> Deserialize<'de> for Role
{
  fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
  where
    D: serde::Deserializer<'de>,
  {
    let s = String::deserialize(deserializer)?;
    Ok(match s.as_str()
    {
      "user" => Role::User,
      "system" => Role::System,
      "assistant" => Role::Assistant,
      other => Role::Custom(other.to_string()),
    })
  }
}