use serde::Serialize;
use crate::llm::Usage;
#[derive(Debug, Clone, Serialize)]
pub struct AgentTrace {
pub started_at: String,
pub user_message: String,
pub history: Vec<TraceMessage>,
pub events: Vec<TraceEvent>,
pub usage: Option<TraceUsage>,
}
#[derive(Debug, Clone, Serialize)]
pub struct TraceMessage {
pub role: String,
pub content: String,
}
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type")]
pub enum TraceEvent {
#[serde(rename = "tool_call")]
ToolCall {
name: String,
arguments: String,
},
#[serde(rename = "tool_result")]
ToolResult {
name: String,
content: String,
},
#[serde(rename = "text_delta")]
TextDelta {
text: String,
},
#[serde(rename = "build_mutation")]
BuildMutation {
label: String,
},
}
#[derive(Debug, Clone, Serialize)]
pub struct TraceUsage {
pub input_tokens: u32,
pub output_tokens: u32,
pub cached_tokens: u32,
pub total_tokens: u32,
}
impl From<&Usage> for TraceUsage {
fn from(u: &Usage) -> Self {
Self {
input_tokens: u.input_tokens,
output_tokens: u.output_tokens,
cached_tokens: u.cached_tokens(),
total_tokens: u.total_tokens,
}
}
}
pub(crate) struct TraceBuilder {
started_at: String,
user_message: String,
history: Vec<TraceMessage>,
events: Vec<TraceEvent>,
}
impl TraceBuilder {
pub fn new(user_message: &str, history: impl IntoIterator<Item = TraceMessage>) -> Self {
Self {
started_at: now_iso8601(),
user_message: user_message.to_owned(),
history: history.into_iter().collect(),
events: Vec::new(),
}
}
pub fn tool_call(&mut self, name: &str, arguments: &str) {
self.events.push(TraceEvent::ToolCall {
name: name.to_owned(),
arguments: arguments.to_owned(),
});
}
pub fn tool_result(&mut self, name: &str, content: &str) {
self.events.push(TraceEvent::ToolResult {
name: name.to_owned(),
content: content.to_owned(),
});
}
pub fn text_delta(&mut self, text: &str) {
self.events.push(TraceEvent::TextDelta {
text: text.to_owned(),
});
}
pub fn build_mutation(&mut self, label: &str) {
self.events.push(TraceEvent::BuildMutation {
label: label.to_owned(),
});
}
pub fn finish(self, usage: &Usage) -> AgentTrace {
AgentTrace {
started_at: self.started_at,
user_message: self.user_message,
history: self.history,
events: self.events,
usage: Some(TraceUsage::from(usage)),
}
}
}
fn now_iso8601() -> String {
let now = std::time::SystemTime::now();
let duration = now
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default();
let secs = duration.as_secs();
let days = secs / 86400;
let time_secs = secs % 86400;
let hours = time_secs / 3600;
let minutes = (time_secs % 3600) / 60;
let seconds = time_secs % 60;
let (y, m, d) = days_to_ymd(days);
format!("{y:04}-{m:02}-{d:02}T{hours:02}:{minutes:02}:{seconds:02}Z")
}
fn days_to_ymd(days: u64) -> (u64, u64, u64) {
let z = days + 719468;
let era = z / 146097;
let doe = z - era * 146097;
let yoe = (doe - doe / 1460 + doe / 36524 - doe / 146096) / 365;
let y = yoe + era * 400;
let doy = doe - (365 * yoe + yoe / 4 - yoe / 100);
let mp = (5 * doy + 2) / 153;
let d = doy - (153 * mp + 2) / 5 + 1;
let m = if mp < 10 { mp + 3 } else { mp - 9 };
let y = if m <= 2 { y + 1 } else { y };
(y, m, d)
}