use serde::{Deserialize, Serialize};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum Role {
User,
Assistant,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum ContentBlock {
Text(String),
Thinking {
text: String,
signature: Option<String>,
},
ToolUse {
id: String,
name: String,
input: serde_json::Value,
},
Opaque {
kind: String,
raw: serde_json::Value,
},
Image {
media_type: String,
source: ImageSource,
},
ToolResult {
tool_use_id: String,
content: String,
is_error: bool,
},
}
impl ContentBlock {
pub fn answer_text(&self) -> Option<&str> {
match self {
ContentBlock::Text(text) => Some(text),
_ => None,
}
}
pub fn reasoning(&self) -> Option<&str> {
match self {
ContentBlock::Thinking { text, .. } => Some(text),
_ => None,
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum ImageSource {
Bytes(Vec<u8>),
Url(String),
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Message {
pub role: Role,
pub content: Vec<ContentBlock>,
}
impl Message {
pub fn user(text: impl Into<String>) -> Message {
Message {
role: Role::User,
content: vec![ContentBlock::Text(text.into())],
}
}
pub fn assistant(text: impl Into<String>) -> Message {
Message {
role: Role::Assistant,
content: vec![ContentBlock::Text(text.into())],
}
}
pub fn text(&self) -> String {
self.content
.iter()
.filter_map(ContentBlock::answer_text)
.collect::<Vec<_>>()
.join("\n")
}
pub fn tool_calls(&self) -> Vec<&ContentBlock> {
self.content
.iter()
.filter(|b| matches!(b, ContentBlock::ToolUse { .. }))
.collect()
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub enum StopReason {
EndTurn,
ToolUse,
StopSequence,
MaxTokens,
Refusal,
PauseTurn,
ContextWindowExceeded,
Interrupted,
Other,
}
impl StopReason {
pub fn is_complete(self) -> bool {
matches!(
self,
StopReason::EndTurn | StopReason::ToolUse | StopReason::StopSequence
)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn a_truncated_answer_is_not_complete() {
assert!(!StopReason::MaxTokens.is_complete());
assert!(!StopReason::Refusal.is_complete());
}
#[test]
fn a_paused_turn_is_not_a_finished_one() {
assert!(!StopReason::PauseTurn.is_complete());
assert!(!StopReason::ContextWindowExceeded.is_complete());
}
#[test]
fn a_block_this_crate_does_not_know_survives_a_round_trip() {
let block = ContentBlock::Opaque {
kind: "redacted_thinking".into(),
raw: serde_json::json!({ "type": "redacted_thinking", "data": "EroBCkY..." }),
};
let json = serde_json::to_string(&block).unwrap_or_default();
let back: ContentBlock =
serde_json::from_str(&json).unwrap_or(ContentBlock::Text("lost".into()));
assert_eq!(block, back);
assert_eq!(
block.answer_text(),
None,
"an opaque block is not an answer"
);
}
#[test]
fn waiting_for_a_tool_is_complete() {
assert!(StopReason::ToolUse.is_complete());
}
#[test]
fn reasoning_is_not_an_answer() {
let thinking = ContentBlock::Thinking {
text: "the user probably means".into(),
signature: None,
};
assert_eq!(thinking.answer_text(), None);
assert_eq!(thinking.reasoning(), Some("the user probably means"));
let answer = ContentBlock::Text("Here it is.".into());
assert_eq!(answer.answer_text(), Some("Here it is."));
assert_eq!(answer.reasoning(), None);
}
#[test]
fn joining_text_leaves_out_the_reasoning() {
let turn = Message {
role: Role::Assistant,
content: vec![
ContentBlock::Thinking {
text: "let me work this out".into(),
signature: None,
},
ContentBlock::Text("Four.".into()),
],
};
assert_eq!(turn.text(), "Four.");
}
#[test]
fn joining_text_leaves_out_the_tool_calls() {
let turn = Message {
role: Role::Assistant,
content: vec![
ContentBlock::Text("Looking that up.".into()),
ContentBlock::ToolUse {
id: "call-1".into(),
name: "search".into(),
input: serde_json::json!({ "q": "rust" }),
},
],
};
assert_eq!(turn.text(), "Looking that up.");
assert_eq!(turn.tool_calls().len(), 1);
}
#[test]
fn a_thinking_block_keeps_its_signature_through_a_round_trip() {
let block = ContentBlock::Thinking {
text: "considering".into(),
signature: Some("sig-abc".into()),
};
let json = serde_json::to_string(&block).unwrap_or_default();
let back: ContentBlock = serde_json::from_str(&json)
.unwrap_or(ContentBlock::Text("deserialization failed".into()));
assert_eq!(block, back);
}
}