use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::fmt;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Message {
pub role: Role,
pub parts: Vec<MessagePart>,
}
impl Message {
pub fn user(text: impl Into<String>) -> Self {
Self {
role: Role::User,
parts: vec![MessagePart::text(text)],
}
}
pub fn assistant(text: impl Into<String>) -> Self {
Self {
role: Role::Assistant,
parts: vec![MessagePart::text(text)],
}
}
#[must_use]
pub const fn new(role: Role, parts: Vec<MessagePart>) -> Self {
Self { role, parts }
}
#[must_use]
pub fn text_content(&self) -> String {
self.parts
.iter()
.filter_map(MessagePart::as_text)
.collect::<Vec<_>>()
.join("")
}
#[must_use]
pub fn tool_call_parts(&self) -> Vec<(&str, &str, &serde_json::Value)> {
self.parts
.iter()
.filter_map(|part| match part {
MessagePart::ToolCall { id, name, input } => {
Some((id.as_str(), name.as_str(), input))
}
_ => None,
})
.collect()
}
}
impl fmt::Display for Message {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut chunks = Vec::new();
for part in &self.parts {
match part {
MessagePart::Text { text } => {
chunks.push(text.clone());
}
MessagePart::ToolCall { name, input, .. } => {
if let Ok(input_str) = serde_json::to_string(input) {
chunks.push(format!("[Tool: {name} with input: {input_str}]"));
}
}
MessagePart::ToolResult { output, .. } => {
chunks.push(format!("[Tool Result: {output}]"));
}
MessagePart::Image { source } => {
chunks.push(format!("[Image: {}]", source.media_type));
}
}
}
write!(f, "{}", chunks.join("\n"))
}
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum Role {
User,
Assistant,
System,
}
impl fmt::Display for Role {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::User => write!(f, "user"),
Self::Assistant => write!(f, "assistant"),
Self::System => write!(f, "system"),
}
}
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type")]
pub enum MessagePart {
#[serde(rename = "text")]
Text {
text: String,
},
#[serde(rename = "image")]
Image {
source: ImageSource,
},
#[serde(rename = "tool_call")]
ToolCall {
id: String,
name: String,
input: Value,
},
#[serde(rename = "tool_result")]
ToolResult {
call_id: String,
#[serde(default)]
name: String,
output: ToolContent,
is_error: Option<bool>,
},
}
impl MessagePart {
pub fn text(text: impl Into<String>) -> Self {
Self::Text { text: text.into() }
}
pub fn tool_call(id: impl Into<String>, name: impl Into<String>, input: Value) -> Self {
Self::ToolCall {
id: id.into(),
name: name.into(),
input,
}
}
pub fn tool_result(
call_id: impl Into<String>,
name: impl Into<String>,
output: impl Into<ToolContent>,
is_error: bool,
) -> Self {
Self::ToolResult {
call_id: call_id.into(),
name: name.into(),
output: output.into(),
is_error: Some(is_error),
}
}
#[must_use]
pub const fn is_text(&self) -> bool {
matches!(self, Self::Text { .. })
}
#[must_use]
pub const fn is_tool_call(&self) -> bool {
matches!(self, Self::ToolCall { .. })
}
#[must_use]
pub const fn is_tool_result(&self) -> bool {
matches!(self, Self::ToolResult { .. })
}
#[must_use]
pub fn as_text(&self) -> Option<&str> {
match self {
Self::Text { text } => Some(text),
_ => None,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ImageSource {
pub encoding: String,
pub media_type: String,
pub data: String,
}
impl ImageSource {
#[must_use]
pub fn new_base64(media_type: impl Into<String>, data: impl Into<String>) -> Self {
Self {
encoding: "base64".into(),
media_type: media_type.into(),
data: data.into(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum ToolContent {
Text(String),
Multipart(Vec<ToolContentPart>),
}
impl ToolContent {
#[must_use]
pub fn from_string(s: impl Into<String>) -> Self {
Self::Text(s.into())
}
#[must_use]
pub const fn from_multipart(parts: Vec<ToolContentPart>) -> Self {
Self::Multipart(parts)
}
#[must_use]
pub const fn is_string(&self) -> bool {
matches!(self, Self::Text(_))
}
}
impl Default for ToolContent {
fn default() -> Self {
Self::Text(String::new())
}
}
impl From<String> for ToolContent {
fn from(s: String) -> Self {
Self::Text(s)
}
}
impl From<&str> for ToolContent {
fn from(s: &str) -> Self {
Self::Text(s.to_owned())
}
}
impl fmt::Display for ToolContent {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Text(s) => write!(f, "{s}"),
Self::Multipart(parts) => {
let texts: Vec<&str> = parts
.iter()
.filter_map(|part| {
if let ToolContentPart::Text { text } = part {
Some(text.as_str())
} else {
None
}
})
.collect();
write!(f, "{}", texts.join("\n"))
}
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ToolContentPart {
Image {
source: ImageSource,
},
Text {
text: String,
},
}
impl ToolContentPart {
pub fn text(text: impl Into<String>) -> Self {
Self::Text { text: text.into() }
}
#[must_use]
pub const fn image(source: ImageSource) -> Self {
Self::Image { source }
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_message_user_shortcut() {
let msg = Message::user("Hello");
assert_eq!(msg.role, Role::User);
assert_eq!(msg.parts.len(), 1);
assert_eq!(msg.parts[0].as_text(), Some("Hello"));
}
#[test]
fn test_message_assistant_shortcut() {
let msg = Message::assistant("Hi there!");
assert_eq!(msg.role, Role::Assistant);
assert_eq!(msg.parts.len(), 1);
}
#[test]
fn test_message_display() {
let msg = Message::user("Hello world");
assert_eq!(msg.to_string(), "Hello world");
}
#[test]
fn test_message_display_with_tool_call() {
let msg = Message {
role: Role::Assistant,
parts: vec![MessagePart::tool_call(
"id1",
"read_file",
serde_json::json!({"path": "/tmp/test.txt"}),
)],
};
let display = msg.to_string();
assert!(display.contains("Tool: read_file"));
}
#[test]
fn test_role_display() {
assert_eq!(Role::User.to_string(), "user");
assert_eq!(Role::Assistant.to_string(), "assistant");
}
#[test]
fn test_part_helpers() {
let text = MessagePart::text("hello");
assert!(text.is_text());
assert!(!text.is_tool_call());
assert_eq!(text.as_text(), Some("hello"));
let tool_call = MessagePart::tool_call("id", "tool", serde_json::json!({}));
assert!(tool_call.is_tool_call());
assert!(!tool_call.is_text());
assert!(tool_call.as_text().is_none());
let tool_result = MessagePart::tool_result("id", "tool", "ok", false);
assert!(tool_result.is_tool_result());
}
#[test]
fn test_image_source() {
let src = ImageSource::new_base64("image/png", "iVBOR...");
assert_eq!(src.encoding, "base64");
assert_eq!(src.media_type, "image/png");
}
#[test]
fn test_tool_result_from_string() {
let result: ToolContent = "hello".into();
assert!(result.is_string());
assert_eq!(result.to_string(), "hello");
}
#[test]
fn tool_result_deserializes_without_name_field() {
let json = r#"{"type":"tool_result","call_id":"tc_1","output":"ok","is_error":false}"#;
let part: MessagePart =
serde_json::from_str(json).expect("old data without name must parse");
match part {
MessagePart::ToolResult { call_id, name, .. } => {
assert_eq!(call_id, "tc_1");
assert_eq!(name, "");
}
other => panic!("expected ToolResult, got {other:?}"),
}
}
#[test]
fn test_tool_result_default() {
let result = ToolContent::default();
assert!(result.is_string());
}
#[test]
fn test_tool_result_part_text() {
let part = ToolContentPart::text("output");
match &part {
ToolContentPart::Text { text } => assert_eq!(text, "output"),
ToolContentPart::Image { .. } => panic!("expected text part"),
}
}
#[test]
fn test_message_serialization() {
let msg = Message::user("test");
let json = serde_json::to_string(&msg).unwrap();
let deserialized: Message = serde_json::from_str(&json).unwrap();
assert_eq!(msg.role, deserialized.role);
}
#[test]
fn test_part_serialization_roundtrip() {
let parts = vec![
MessagePart::text("hello"),
MessagePart::tool_call("id1", "my_tool", serde_json::json!({"key": "value"})),
];
let json = serde_json::to_string(&parts).unwrap();
let back: Vec<MessagePart> = serde_json::from_str(&json).unwrap();
assert_eq!(parts.len(), back.len());
}
#[test]
fn test_tool_result_multipart_display() {
let result = ToolContent::from_multipart(vec![
ToolContentPart::text("line 1"),
ToolContentPart::text("line 2"),
]);
assert_eq!(result.to_string(), "line 1\nline 2");
}
#[test]
fn test_role_system_serializes_snake_case() {
let s = serde_json::to_string(&Role::System).expect("System serializes");
assert_eq!(s, "\"system\"");
}
#[test]
fn test_role_system_display() {
assert_eq!(Role::System.to_string(), "system");
}
#[test]
fn test_role_system_round_trips() {
let parsed: Role =
serde_json::from_str("\"system\"").expect("\"system\" deserializes to a Role");
assert_eq!(parsed, Role::System);
}
}