agtrace_engine/
export.rs

1use agtrace_types::{AgentEvent, EventPayload};
2use std::str::FromStr;
3
4#[derive(Debug, Clone, Copy)]
5pub enum ExportStrategy {
6    Raw,
7    Clean,
8    Reasoning,
9}
10
11impl FromStr for ExportStrategy {
12    type Err = String;
13
14    fn from_str(s: &str) -> Result<Self, Self::Err> {
15        match s {
16            "raw" => Ok(ExportStrategy::Raw),
17            "clean" => Ok(ExportStrategy::Clean),
18            "reasoning" => Ok(ExportStrategy::Reasoning),
19            _ => Err(format!("Unknown export strategy: {}", s)),
20        }
21    }
22}
23
24pub fn transform(events: &[AgentEvent], strategy: ExportStrategy) -> Vec<AgentEvent> {
25    match strategy {
26        ExportStrategy::Raw => events.to_vec(),
27        ExportStrategy::Clean => apply_clean_strategy(events),
28        ExportStrategy::Reasoning => apply_reasoning_strategy(events),
29    }
30}
31
32fn apply_clean_strategy(events: &[AgentEvent]) -> Vec<AgentEvent> {
33    let mut cleaned = Vec::new();
34    let mut skip_until_next_success = false;
35
36    for event in events.iter() {
37        match &event.payload {
38            EventPayload::ToolResult(result) => {
39                if result.is_error {
40                    skip_until_next_success = true;
41                    continue;
42                } else {
43                    skip_until_next_success = false;
44                }
45            }
46            EventPayload::Message(msg) => {
47                let text_lower = msg.text.to_lowercase();
48                if text_lower.contains("i apologize")
49                    || text_lower.contains("my mistake")
50                    || text_lower.contains("sorry")
51                {
52                    continue;
53                }
54            }
55            _ => {}
56        }
57
58        if !skip_until_next_success {
59            let mut cleaned_event = event.clone();
60
61            // Truncate long text in payloads
62            match &mut cleaned_event.payload {
63                EventPayload::Message(msg) => {
64                    if msg.text.len() > 5000 {
65                        msg.text = format!(
66                            "{}...<truncated_output_for_training>",
67                            msg.text.chars().take(1000).collect::<String>()
68                        );
69                    }
70                }
71                EventPayload::ToolResult(result) => {
72                    if result.output.len() > 5000 {
73                        result.output = format!(
74                            "{}...<truncated_output_for_training>",
75                            result.output.chars().take(1000).collect::<String>()
76                        );
77                    }
78                }
79                EventPayload::Reasoning(reasoning) => {
80                    if reasoning.text.len() > 5000 {
81                        reasoning.text = format!(
82                            "{}...<truncated_output_for_training>",
83                            reasoning.text.chars().take(1000).collect::<String>()
84                        );
85                    }
86                }
87                _ => {}
88            }
89
90            cleaned.push(cleaned_event);
91        }
92    }
93
94    cleaned
95}
96
97fn apply_reasoning_strategy(events: &[AgentEvent]) -> Vec<AgentEvent> {
98    let mut reasoning_pairs = Vec::new();
99
100    for i in 0..events.len() {
101        if matches!(events[i].payload, EventPayload::Reasoning(_)) {
102            reasoning_pairs.push(events[i].clone());
103
104            if let Some(next) = events.get(i + 1)
105                && matches!(next.payload, EventPayload::ToolCall(_))
106            {
107                reasoning_pairs.push(next.clone());
108            }
109        }
110    }
111
112    reasoning_pairs
113}
agtrace_engine/export.rs

agtrace_engine/
export.rs
