use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct StepResult {
pub step_id: String,
pub success: bool,
pub output: StepOutput,
pub confidence: f64,
pub duration_ms: u64,
pub tokens: TokenUsage,
pub error: Option<String>,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
#[derive(Default)]
pub enum StepOutput {
Text {
content: String,
},
List {
items: Vec<ListItem>,
},
Structured {
data: HashMap<String, serde_json::Value>,
},
Score {
value: f64,
},
Boolean {
value: bool,
reason: Option<String>,
},
#[default]
Empty,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ListItem {
pub content: String,
#[serde(default)]
pub confidence: Option<f64>,
#[serde(default)]
pub metadata: HashMap<String, serde_json::Value>,
}
impl ListItem {
pub fn new(content: impl Into<String>) -> Self {
Self {
content: content.into(),
confidence: None,
metadata: HashMap::new(),
}
}
pub fn with_confidence(content: impl Into<String>, confidence: f64) -> Self {
Self {
content: content.into(),
confidence: Some(confidence),
metadata: HashMap::new(),
}
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
#[cfg_attr(
feature = "performance",
derive(rkyv::Archive, rkyv::Serialize, rkyv::Deserialize)
)]
pub struct TokenUsage {
pub input_tokens: u32,
pub output_tokens: u32,
pub total_tokens: u32,
pub cost_usd: f64,
}
impl TokenUsage {
pub fn new(input: u32, output: u32, cost: f64) -> Self {
Self {
input_tokens: input,
output_tokens: output,
total_tokens: input + output,
cost_usd: cost,
}
}
pub fn add(&mut self, other: &TokenUsage) {
self.input_tokens += other.input_tokens;
self.output_tokens += other.output_tokens;
self.total_tokens += other.total_tokens;
self.cost_usd += other.cost_usd;
}
}
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub enum OutputFormat {
#[default]
Raw,
Json,
List,
KeyValue,
Numeric,
}
impl StepResult {
pub fn success(step_id: impl Into<String>, output: StepOutput, confidence: f64) -> Self {
Self {
step_id: step_id.into(),
success: true,
output,
confidence,
duration_ms: 0,
tokens: TokenUsage::default(),
error: None,
}
}
pub fn failure(step_id: impl Into<String>, error: impl Into<String>) -> Self {
Self {
step_id: step_id.into(),
success: false,
output: StepOutput::Empty,
confidence: 0.0,
duration_ms: 0,
tokens: TokenUsage::default(),
error: Some(error.into()),
}
}
pub fn with_duration(mut self, ms: u64) -> Self {
self.duration_ms = ms;
self
}
pub fn with_tokens(mut self, tokens: TokenUsage) -> Self {
self.tokens = tokens;
self
}
pub fn meets_threshold(&self, threshold: f64) -> bool {
self.success && self.confidence >= threshold
}
pub fn as_text(&self) -> Option<&str> {
match &self.output {
StepOutput::Text { content } => Some(content),
_ => None,
}
}
pub fn as_list(&self) -> Option<&[ListItem]> {
match &self.output {
StepOutput::List { items } => Some(items),
_ => None,
}
}
pub fn as_score(&self) -> Option<f64> {
match &self.output {
StepOutput::Score { value } => Some(*value),
_ => None,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_step_result_success() {
let result = StepResult::success(
"test_step",
StepOutput::Text {
content: "Hello world".to_string(),
},
0.85,
);
assert!(result.success);
assert_eq!(result.confidence, 0.85);
assert_eq!(result.as_text(), Some("Hello world"));
}
#[test]
fn test_step_result_failure() {
let result = StepResult::failure("test_step", "Something went wrong");
assert!(!result.success);
assert_eq!(result.confidence, 0.0);
assert_eq!(result.error, Some("Something went wrong".to_string()));
}
#[test]
fn test_token_usage_add() {
let mut usage1 = TokenUsage::new(100, 50, 0.001);
let usage2 = TokenUsage::new(200, 100, 0.002);
usage1.add(&usage2);
assert_eq!(usage1.input_tokens, 300);
assert_eq!(usage1.output_tokens, 150);
assert_eq!(usage1.total_tokens, 450);
}
#[test]
fn test_list_item() {
let item = ListItem::with_confidence("Test item", 0.9);
assert_eq!(item.content, "Test item");
assert_eq!(item.confidence, Some(0.9));
}
}