mod types;
use crate::chat::{ChatMessage, ChatRole};
pub use types::*;
impl ContentBlock {
pub fn as_text(&self) -> Option<&str> {
match self {
ContentBlock::Text(text) => Some(text),
_ => None,
}
}
}
fn concat_text(content: &[ContentBlock]) -> String {
content
.iter()
.filter_map(ContentBlock::as_text)
.collect::<Vec<_>>()
.join("")
}
impl Message {
pub fn system(content: impl Into<String>) -> Self {
Message::System(SystemMessage {
content: vec![ContentBlock::Text(content.into())],
})
}
pub fn user(content: impl Into<String>) -> Self {
Message::User(UserMessage {
content: vec![ContentBlock::Text(content.into())],
})
}
pub fn assistant(content: impl Into<String>) -> Self {
Message::Assistant(AssistantMessage {
id: None,
content: vec![ContentBlock::Text(content.into())],
tool_calls: Vec::new(),
usage: None,
})
}
pub fn tool(tool_call_id: impl Into<String>, content: impl Into<String>) -> Self {
Message::Tool(ToolMessage {
tool_call_id: tool_call_id.into(),
content: vec![ContentBlock::Text(content.into())],
})
}
pub fn text(&self) -> String {
match self {
Message::System(m) => concat_text(&m.content),
Message::User(m) => concat_text(&m.content),
Message::Assistant(m) => concat_text(&m.content),
Message::Tool(m) => concat_text(&m.content),
}
}
pub fn to_chat(&self) -> ChatMessage {
match self {
Message::System(_) => ChatMessage::system(self.text()),
Message::User(_) => ChatMessage::user(self.text()),
Message::Assistant(_) => ChatMessage::assistant(self.text()),
Message::Tool(m) => ChatMessage::tool(m.tool_call_id.clone(), self.text()),
}
}
}
impl From<ChatMessage> for Message {
fn from(msg: ChatMessage) -> Self {
match msg.role {
ChatRole::System => Message::system(msg.content),
ChatRole::User => Message::user(msg.content),
ChatRole::Assistant => Message::assistant(msg.content),
ChatRole::Tool => Message::tool(msg.name.unwrap_or_default(), msg.content),
}
}
}
#[cfg(test)]
mod test;