use crate::llm::reasoning::ReasoningState;
use alloc::string::{String, ToString};
use serde_json::Value;
#[derive(Debug, Clone, Default, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct Usage {
pub prompt_tokens: Option<u32>,
pub completion_tokens: Option<u32>,
pub total_tokens: Option<u32>,
pub reasoning_tokens: Option<u32>,
pub cache_read_tokens: Option<u32>,
pub cache_write_tokens: Option<u32>,
pub cost_usd: Option<f64>,
pub stop_reason: Option<String>,
}
impl Usage {
#[must_use]
pub const fn new(prompt_tokens: u32, completion_tokens: u32) -> Self {
Self {
prompt_tokens: Some(prompt_tokens),
completion_tokens: Some(completion_tokens),
total_tokens: Some(prompt_tokens + completion_tokens),
reasoning_tokens: None,
cache_read_tokens: None,
cache_write_tokens: None,
cost_usd: None,
stop_reason: None,
}
}
#[must_use]
pub const fn with_reasoning_tokens(mut self, tokens: u32) -> Self {
self.reasoning_tokens = Some(tokens);
self
}
#[must_use]
pub const fn with_cache_tokens(mut self, read: u32, write: u32) -> Self {
self.cache_read_tokens = Some(read);
self.cache_write_tokens = Some(write);
self
}
#[must_use]
pub const fn with_cost(mut self, cost_usd: f64) -> Self {
self.cost_usd = Some(cost_usd);
self
}
#[must_use]
pub fn with_stop_reason(mut self, reason: impl Into<String>) -> Self {
self.stop_reason = Some(reason.into());
self
}
pub fn accumulate(&mut self, other: &Self) {
if let Some(v) = other.prompt_tokens {
*self.prompt_tokens.get_or_insert(0) += v;
}
if let Some(v) = other.completion_tokens {
*self.completion_tokens.get_or_insert(0) += v;
}
if let Some(v) = other.total_tokens {
*self.total_tokens.get_or_insert(0) += v;
}
if let Some(v) = other.reasoning_tokens {
*self.reasoning_tokens.get_or_insert(0) += v;
}
if let Some(v) = other.cache_read_tokens {
*self.cache_read_tokens.get_or_insert(0) += v;
}
if let Some(v) = other.cache_write_tokens {
*self.cache_write_tokens.get_or_insert(0) += v;
}
if let Some(v) = other.cost_usd {
*self.cost_usd.get_or_insert(0.0) += v;
}
if self.stop_reason.is_none() {
self.stop_reason.clone_from(&other.stop_reason);
}
}
}
#[derive(Debug, Clone)]
pub enum Event {
Text(String),
Reasoning(String),
ToolCallDelta {
id: String,
name: String,
arguments_fragment: String,
},
ToolCall(ToolCall),
ReasoningState(ReasoningState),
BuiltInToolResult {
tool: String,
result: String,
},
Usage(Usage),
}
impl Event {
#[must_use]
pub fn text(text: impl Into<String>) -> Self {
Self::Text(text.into())
}
#[must_use]
pub fn reasoning(thought: impl Into<String>) -> Self {
Self::Reasoning(thought.into())
}
#[must_use]
pub fn tool_call_delta(
id: impl Into<String>,
name: impl Into<String>,
arguments_fragment: impl Into<String>,
) -> Self {
Self::ToolCallDelta {
id: id.into(),
name: name.into(),
arguments_fragment: arguments_fragment.into(),
}
}
#[must_use]
pub fn tool_call(id: impl Into<String>, name: impl Into<String>, arguments: Value) -> Self {
Self::ToolCall(ToolCall {
id: id.into(),
name: name.into(),
arguments,
reasoning_state: None,
})
}
#[must_use]
pub fn builtin_result(tool: impl Into<String>, result: impl Into<String>) -> Self {
Self::BuiltInToolResult {
tool: tool.into(),
result: result.into(),
}
}
#[must_use]
pub const fn usage(usage: Usage) -> Self {
Self::Usage(usage)
}
#[must_use]
pub fn as_text(&self) -> Option<&str> {
match self {
Self::Text(s) => Some(s),
_ => None,
}
}
#[must_use]
pub fn as_reasoning(&self) -> Option<&str> {
match self {
Self::Reasoning(s) => Some(s),
_ => None,
}
}
#[must_use]
pub const fn as_tool_call(&self) -> Option<&ToolCall> {
match self {
Self::ToolCall(call) => Some(call),
_ => None,
}
}
#[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 as_usage(&self) -> Option<&Usage> {
match self {
Self::Usage(u) => Some(u),
_ => None,
}
}
#[must_use]
pub const fn is_usage(&self) -> bool {
matches!(self, Self::Usage(_))
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))]
pub struct ToolCall {
pub id: String,
pub name: String,
pub arguments: Value,
#[cfg_attr(
feature = "serde",
serde(default, skip_serializing_if = "Option::is_none")
)]
pub reasoning_state: Option<ReasoningState>,
}
impl ToolCall {
#[must_use]
pub fn new(id: impl Into<String>, name: impl Into<String>, arguments: Value) -> Self {
Self {
id: id.into(),
name: name.into(),
arguments,
reasoning_state: None,
}
}
#[must_use]
pub fn with_reasoning_state(mut self, state: ReasoningState) -> Self {
self.reasoning_state = Some(state);
self
}
#[must_use]
pub fn arguments_json(&self) -> String {
self.arguments.to_string()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_event_constructors() {
let text = Event::text("hello");
assert!(text.is_text());
assert_eq!(text.as_text(), Some("hello"));
let reasoning = Event::reasoning("thinking...");
assert_eq!(reasoning.as_reasoning(), Some("thinking..."));
let tool = Event::tool_call("call_1", "search", serde_json::json!({"query": "rust"}));
assert!(tool.is_tool_call());
let call = tool.as_tool_call().unwrap();
assert_eq!(call.name, "search");
assert_eq!(call.id, "call_1");
}
#[test]
fn test_tool_call_arguments() {
let call = ToolCall::new("id", "test", serde_json::json!({"key": "value"}));
let json = call.arguments_json();
assert!(json.contains("key"));
assert!(json.contains("value"));
}
}