traitclaw_core/traits/
output_transformer.rs

1//! Async tool output transformation.
2//!
3//! [`OutputTransformer`] provides context-aware, async tool output processing.
4//! It adds context-awareness (tool name and agent state) and supports async operations
5//! such as LLM-powered summarization of tool output.
6//!
7//! # Example
8//!
9//! ```rust
10//! use traitclaw_core::traits::output_transformer::OutputTransformer;
11//! use traitclaw_core::types::agent_state::AgentState;
12//! use async_trait::async_trait;
13//!
14//! struct BudgetAwareTransformer {
15//!     max_chars: usize,
16//! }
17//!
18//! #[async_trait]
19//! impl OutputTransformer for BudgetAwareTransformer {
20//!     async fn transform(
21//!         &self,
22//!         output: String,
23//!         tool_name: &str,
24//!         state: &AgentState,
25//!     ) -> String {
26//!         // Truncate more aggressively when context is nearly full
27//!         let budget = if state.context_utilization() > 0.8 {
28//!             self.max_chars / 2
29//!         } else {
30//!             self.max_chars
31//!         };
32//!         if output.len() > budget {
33//!             format!("{}...\n[truncated]", &output[..budget])
34//!         } else {
35//!             output
36//!         }
37//!     }
38//! }
39//! ```
40
41use async_trait::async_trait;
42
43use crate::types::agent_state::AgentState;
44
45/// Async trait for context-aware tool output transformation.
46///
47/// Called after each tool execution to process the output before adding it
48/// to the message context. Supports async operations such as LLM-powered
49/// summarization.
50#[async_trait]
51pub trait OutputTransformer: Send + Sync {
52    /// Transform tool output, optionally using context about which tool
53    /// produced it and the current agent state.
54    ///
55    /// `tool_name` identifies the tool that produced the output.
56    /// `state` provides runtime context (token usage, iteration count, etc.).
57    async fn transform(&self, output: String, tool_name: &str, state: &AgentState) -> String;
58
59    /// Estimate token count for a given output string.
60    ///
61    /// Default: 4-characters ≈ 1-token approximation.
62    fn estimate_output_tokens(&self, output: &str) -> usize {
63        output.len() / 4 + 1
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use super::*;
70    use crate::types::model_info::ModelTier;
71    use std::sync::Arc;
72
73    // ── Object safety: confirm Arc<dyn OutputTransformer> compiles ───────
74    #[test]
75    fn test_output_transformer_is_object_safe() {
76        struct Dummy;
77
78        #[async_trait]
79        impl OutputTransformer for Dummy {
80            async fn transform(
81                &self,
82                output: String,
83                _tool_name: &str,
84                _state: &AgentState,
85            ) -> String {
86                output
87            }
88        }
89
90        let _: Arc<dyn OutputTransformer> = Arc::new(Dummy);
91    }
92
93    // ── Default estimate_output_tokens() ────────────────────────────────
94    #[test]
95    fn test_default_estimate_output_tokens() {
96        struct Dummy;
97
98        #[async_trait]
99        impl OutputTransformer for Dummy {
100            async fn transform(
101                &self,
102                output: String,
103                _tool_name: &str,
104                _state: &AgentState,
105            ) -> String {
106                output
107            }
108        }
109
110        let t = Dummy;
111        // 400 chars → 400/4 + 1 = 101 tokens
112        assert_eq!(t.estimate_output_tokens(&"a".repeat(400)), 101);
113        // empty → 0/4 + 1 = 1 token
114        assert_eq!(t.estimate_output_tokens(""), 1);
115    }
116
117    // ── Context-aware transformer test ──────────────────────────────────
118    #[tokio::test]
119    async fn test_context_aware_transformer() {
120        struct ToolAwareTransformer;
121
122        #[async_trait]
123        impl OutputTransformer for ToolAwareTransformer {
124            async fn transform(
125                &self,
126                output: String,
127                tool_name: &str,
128                state: &AgentState,
129            ) -> String {
130                format!(
131                    "[tool={}, util={:.0}%] {}",
132                    tool_name,
133                    state.context_utilization() * 100.0,
134                    output
135                )
136            }
137        }
138
139        let t = ToolAwareTransformer;
140        let mut state = AgentState::new(ModelTier::Medium, 1000);
141        state.total_context_tokens = 750;
142
143        let result = t.transform("data".to_string(), "search", &state).await;
144        assert!(result.contains("tool=search"), "should include tool name");
145        assert!(result.contains("75%"), "should include utilization");
146        assert!(result.contains("data"), "should include original output");
147    }
148}
traitclaw_core/traits/output_transformer.rs

traitclaw_core/traits/
output_transformer.rs
