use std::str::FromStr; use chrono::{DateTime, Utc}; use derive_more::derive::Display; use derive_setters::Setters; use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::{Context, Error, Metrics, Result, TokenCount}; #[derive(Debug, Default, Display, Serialize, Deserialize, Clone, PartialEq, Eq, Hash)] #[serde(transparent)] pub struct ConversationId(Uuid); impl Copy for ConversationId {} impl ConversationId { pub fn generate() -> Self { Self(Uuid::new_v4()) } pub fn into_string(&self) -> String { self.0.to_string() } pub fn parse(value: impl ToString) -> Result { Ok(Self( Uuid::parse_str(&value.to_string()).map_err(Error::ConversationId)?, )) } } impl FromStr for ConversationId { type Err = Error; fn from_str(s: &str) -> Result { Self::parse(s) } } #[derive(Debug, Setters, Serialize, Deserialize, Clone)] #[setters(into)] pub struct Conversation { pub id: ConversationId, pub title: Option, pub context: Option, pub metrics: Metrics, pub metadata: MetaData, } #[derive(Debug, Setters, Serialize, Deserialize, Clone)] #[setters(into)] pub struct MetaData { pub created_at: DateTime, pub updated_at: Option>, } impl MetaData { pub fn new(created_at: DateTime) -> Self { Self { created_at, updated_at: None } } } impl Conversation { pub fn new(id: ConversationId) -> Self { let created_at = Utc::now(); let metrics = Metrics::default().started_at(created_at); Self { id, metrics, metadata: MetaData::new(created_at), title: None, context: None, } } /// Creates a new conversation with a new conversation ID. /// /// This is a convenience constructor that automatically generates a unique /// conversation ID, making it easy to create new conversations without /// having to manually create the ID. pub fn generate() -> Self { Self::new(ConversationId::generate()) } /// Generates an HTML representation of the conversation /// /// This method uses Handlebars to render the conversation as HTML /// from the template file, including all agents, events, and variables. /// /// # Errors /// - If the template file cannot be found or read /// - If the Handlebars template registration fails /// - If the template rendering fails pub fn to_html(&self) -> String { // Instead of using Handlebars, we now use our Element DSL crate::conversation_html::render_conversation_html(self) } /// Generates an HTML representation with related agent conversations /// /// Creates a single HTML document containing the main conversation /// and all related agent conversations with anchor links for navigation. pub fn to_html_with_related(&self, related: &[Conversation]) -> String { crate::conversation_html::render_conversation_html_with_related(self, related) } /// Returns a vector of user messages, selecting the first message from /// each consecutive sequence of user messages. pub fn first_user_messages(&self) -> Vec<&crate::ContextMessage> { self.context .as_ref() .map(|ctx| ctx.first_user_messages()) .unwrap_or_default() } /// Returns the total token usage across all messages in the conversation. /// /// This is a convenience method that aggregates usage from the context, /// if available. pub fn accumulated_usage(&self) -> Option { self.context.as_ref().and_then(|ctx| ctx.accumulate_usage()) } pub fn usage(&self) -> Option { self.context .iter() .flat_map(|ctx| ctx.messages.iter()) .flat_map(|msg| msg.usage.into_iter()) .last() } pub fn token_count(&self) -> Option { self.context.as_ref().map(|ctx| ctx.token_count()) } pub fn accumulated_cost(&self) -> Option { self.accumulated_usage().and_then(|usage| usage.cost) } /// Calculates the total cost including related conversations. /// /// This method sums the costs of the current conversation and all related /// conversations provided. It's useful for computing the total cost of a /// conversation tree when agent tools spawn child conversations. /// /// # Arguments /// /// * `conversations` - A slice of conversations to include in cost /// calculation pub fn total_cost(conversations: &[Conversation]) -> Option { let costs: Vec = conversations .iter() .filter_map(|conv| conv.accumulated_cost()) .collect(); if costs.is_empty() { return None; } Some(costs.iter().sum()) } /// Returns the number of messages in the conversation context. /// /// Returns `0` if the context has not been initialized yet. pub fn len(&self) -> usize { self.context .as_ref() .map(|ctx| ctx.messages.len()) .unwrap_or(0) } /// Returns `true` if the conversation context has no messages. pub fn is_empty(&self) -> bool { self.len() == 0 } /// Extracts all related conversation IDs from agent tool calls. /// /// This method scans through all tool results in the conversation's context /// and collects conversation IDs from AI tool values, which are created /// when agent tools are called and trigger new conversations. pub fn related_conversation_ids(&self) -> Vec { self.context .as_ref() .map(|ctx| { ctx.messages .iter() .filter_map(|msg| msg.as_tool_result()) .flat_map(|result| &result.output.values) .filter_map(|value| { if let crate::ToolValue::AI { conversation_id, .. } = value { Some(conversation_id) } else { None } }) .copied() .collect() }) .unwrap_or_default() } } #[cfg(test)] mod tests { use pretty_assertions::assert_eq; use super::*; use crate::{Context, ContextMessage, ToolOutput, ToolResult, ToolValue}; #[test] fn test_related_conversation_ids_empty() { let conversation = Conversation::generate(); let actual = conversation.related_conversation_ids(); assert_eq!(actual, vec![]); } #[test] fn test_related_conversation_ids_with_agent_calls() { let agent_conv_id_1 = ConversationId::generate(); let agent_conv_id_2 = ConversationId::generate(); let context = Context::default() .add_message(ContextMessage::user("Test task", None)) .add_message(ContextMessage::assistant("Working on it", None, None, None)) .add_message(ContextMessage::Tool(ToolResult::new("agent_tool").output( Ok(ToolOutput { is_error: false, values: vec![ToolValue::AI { value: "Agent result".to_string(), conversation_id: agent_conv_id_1, }], }), ))) .add_message(ContextMessage::assistant("Continuing", None, None, None)) .add_message(ContextMessage::Tool( ToolResult::new("another_agent").output(Ok(ToolOutput { is_error: false, values: vec![ToolValue::AI { value: "Another agent result".to_string(), conversation_id: agent_conv_id_2, }], })), )); let conversation = Conversation::generate().context(context); let actual = conversation.related_conversation_ids(); assert_eq!(actual.len(), 2); assert!(actual.contains(&agent_conv_id_1)); assert!(actual.contains(&agent_conv_id_2)); } #[test] fn test_related_conversation_ids_mixed_tool_results() { let agent_conv_id = ConversationId::generate(); let context = Context::default() .add_message(ContextMessage::user("Test task", None)) .add_message(ContextMessage::Tool( ToolResult::new("regular_tool").output(Ok(ToolOutput::text("Regular result"))), )) .add_message(ContextMessage::Tool(ToolResult::new("agent_tool").output( Ok(ToolOutput { is_error: false, values: vec![ToolValue::AI { value: "Agent result".to_string(), conversation_id: agent_conv_id, }], }), ))) .add_message(ContextMessage::Tool( ToolResult::new("another_regular_tool") .output(Ok(ToolOutput::text("Another regular result"))), )); let conversation = Conversation::generate().context(context); let actual = conversation.related_conversation_ids(); assert_eq!(actual, vec![agent_conv_id]); } #[test] fn test_total_cost() { use crate::{MessageEntry, Usage}; // Create main conversation with cost // Create main conversation with cost let main_usage = Usage { cost: Some(0.01), ..Usage::default() }; let main_entry: MessageEntry = ContextMessage::assistant("Response", None, None, None).into(); let main_context = Context::default() .add_message(ContextMessage::user("Test", None)) .add_entry(main_entry.usage(main_usage)); let main_conv = Conversation::generate().context(main_context); // Create related conversations with costs let related_usage_1 = Usage { cost: Some(0.02), ..Usage::default() }; let related_entry_1: MessageEntry = ContextMessage::assistant("Result 1", None, None, None).into(); let related_context_1 = Context::default() .add_message(ContextMessage::user("Task 1", None)) .add_entry(related_entry_1.usage(related_usage_1)); let related_conv_1 = Conversation::generate().context(related_context_1); let related_usage_2 = Usage { cost: Some(0.03), ..Usage::default() }; let related_entry_2: MessageEntry = ContextMessage::assistant("Result 2", None, None, None).into(); let related_context_2 = Context::default() .add_message(ContextMessage::user("Task 2", None)) .add_entry(related_entry_2.usage(related_usage_2)); let related_conv_2 = Conversation::generate().context(related_context_2); let actual = Conversation::total_cost(&[main_conv, related_conv_1, related_conv_2]); // Check that cost is approximately 0.06 (accounting for floating point // precision) assert!(actual.is_some()); assert!((actual.unwrap() - 0.06).abs() < 0.0001); } #[test] fn test_total_cost_no_costs() { let main_conv = Conversation::generate(); let related_conv = Conversation::generate(); let actual = Conversation::total_cost(&[main_conv, related_conv]); let expected = None; assert_eq!(actual, expected); } #[test] fn test_total_cost_partial_costs() { use crate::{MessageEntry, Usage}; // Main conversation has no cost let main_conv = Conversation::generate(); // Related conversation has cost let related_usage = Usage { cost: Some(0.05), ..Usage::default() }; let related_entry: MessageEntry = ContextMessage::assistant("Result", None, None, None).into(); let related_context = Context::default() .add_message(ContextMessage::user("Task", None)) .add_entry(related_entry.usage(related_usage)); let related_conv = Conversation::generate().context(related_context); let actual = Conversation::total_cost(&[main_conv, related_conv]); let expected = Some(0.05); assert_eq!(actual, expected); } }