use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "role")]
#[non_exhaustive]
pub enum Message {
Human(HumanMessage),
Ai(AiMessage),
System(SystemMessage),
Tool(ToolMessage),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HumanMessage {
pub content: MessageContent,
#[serde(default)]
pub id: Option<String>,
#[serde(default)]
pub name: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct AiMessage {
pub content: MessageContent,
#[serde(default)]
pub tool_calls: Vec<ToolCall>,
#[serde(default)]
pub invalid_tool_calls: Vec<InvalidToolCall>,
#[serde(default)]
pub usage_metadata: Option<UsageMetadata>,
#[serde(default)]
pub response_metadata: HashMap<String, serde_json::Value>,
#[serde(default)]
pub id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SystemMessage {
pub content: String,
#[serde(default)]
pub id: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolMessage {
pub content: String,
pub tool_call_id: String,
#[serde(default)]
pub name: Option<String>,
#[serde(default)]
pub id: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub metadata: Option<HashMap<String, Value>>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub duration_ms: Option<u64>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(untagged)]
#[non_exhaustive]
pub enum MessageContent {
Text(String),
Blocks(Vec<ContentBlock>),
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
#[non_exhaustive]
pub enum ContentBlock {
Text { text: String },
Image { url: String },
Audio { data: String },
File { path: String, mime_type: String },
Reasoning { content: String },
Citation { source: String, quote: String },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolCall {
pub id: String,
pub name: String,
pub args: serde_json::Value,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct InvalidToolCall {
pub id: String,
pub name: String,
pub args: String,
pub error: String,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct UsageMetadata {
pub input_tokens: u32,
pub output_tokens: u32,
pub total_tokens: u32,
#[serde(default)]
pub input_token_details: Option<InputTokenDetails>,
#[serde(default)]
pub output_token_details: Option<OutputTokenDetails>,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct InputTokenDetails {
#[serde(default)]
pub cache_read: u32,
#[serde(default)]
pub cache_write: u32,
}
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct OutputTokenDetails {
#[serde(default)]
pub reasoning_tokens: u32,
}
pub const REMOVE_ALL_MESSAGES: &str = "__REMOVE_ALL_MESSAGES__";
impl Message {
pub fn human(content: impl Into<String>) -> Self {
Message::Human(HumanMessage {
content: MessageContent::Text(content.into()),
id: None,
name: None,
})
}
pub fn system(content: impl Into<String>) -> Self {
Message::System(SystemMessage {
content: content.into(),
id: None,
})
}
pub fn ai(content: impl Into<String>) -> Self {
Message::Ai(AiMessage {
content: MessageContent::Text(content.into()),
tool_calls: vec![],
invalid_tool_calls: vec![],
usage_metadata: None,
response_metadata: HashMap::new(),
id: None,
})
}
pub fn tool(content: impl Into<String>, tool_call_id: impl Into<String>) -> Self {
Message::Tool(ToolMessage {
content: content.into(),
tool_call_id: tool_call_id.into(),
name: None,
id: None,
metadata: None,
duration_ms: None,
})
}
pub fn tool_with_metadata(
content: impl Into<String>,
tool_call_id: impl Into<String>,
metadata: Option<HashMap<String, Value>>,
duration_ms: Option<u64>,
) -> Self {
Message::Tool(ToolMessage {
content: content.into(),
tool_call_id: tool_call_id.into(),
name: None,
id: None,
metadata,
duration_ms,
})
}
pub fn id(&self) -> Option<&str> {
match self {
Message::Human(m) => m.id.as_deref(),
Message::Ai(m) => m.id.as_deref(),
Message::System(m) => m.id.as_deref(),
Message::Tool(m) => m.id.as_deref(),
}
}
}
impl MessageContent {
pub fn as_text(&self) -> Option<&str> {
match self {
MessageContent::Text(t) => Some(t),
MessageContent::Blocks(_) => None,
}
}
}