vtcode_core/tools/
improvements_registry_ext.rs

1//! Real integration with ToolRegistry
2//!
3//! Extends ToolRegistry with improvement capabilities:
4//! - Tool effectiveness tracking
5//! - Intelligent tool selection based on patterns
6//! - Result caching and optimization (migrated to UnifiedCache)
7//! - Observability integration
8
9use crate::cache::{CacheKey, DEFAULT_CACHE_TTL, EvictionPolicy, UnifiedCache};
10use crate::tools::{
11    improvements_errors::ObservabilityContext,
12    pattern_engine::{ExecutionEvent, PatternEngine},
13};
14use std::sync::Arc;
15
16/// Cache key for tool results
17#[derive(Debug, Clone, Hash, PartialEq, Eq)]
18struct ToolResultKey(String);
19
20impl CacheKey for ToolResultKey {
21    fn to_cache_key(&self) -> String {
22        self.0.clone()
23    }
24}
25
26/// Tool effectiveness metrics
27#[derive(Clone, Debug)]
28pub struct ToolMetrics {
29    pub name: String,
30    pub total_calls: usize,
31    pub successful_calls: usize,
32    pub total_duration_ms: u64,
33    pub avg_quality: f32,
34}
35
36impl ToolMetrics {
37    pub fn success_rate(&self) -> f32 {
38        if self.total_calls == 0 {
39            0.0
40        } else {
41            self.successful_calls as f32 / self.total_calls as f32
42        }
43    }
44
45    pub fn avg_duration_ms(&self) -> u64 {
46        if self.total_calls == 0 {
47            0
48        } else {
49            self.total_duration_ms / self.total_calls as u64
50        }
51    }
52}
53
54/// ToolRegistry improvement extension (migrated to UnifiedCache)
55pub struct ToolRegistryImprovement {
56    pattern_engine: Arc<PatternEngine>,
57    tool_metrics: Arc<parking_lot::RwLock<hashbrown::HashMap<String, ToolMetrics>>>,
58    result_cache: Arc<parking_lot::RwLock<UnifiedCache<ToolResultKey, String>>>,
59    obs_context: Arc<ObservabilityContext>,
60}
61
62impl ToolRegistryImprovement {
63    /// Create new registry extension (now using UnifiedCache)
64    pub fn new(obs_context: Arc<ObservabilityContext>) -> Self {
65        Self {
66            pattern_engine: Arc::new(PatternEngine::new(1000, 20)),
67            tool_metrics: Arc::new(parking_lot::RwLock::new(hashbrown::HashMap::new())),
68            result_cache: Arc::new(parking_lot::RwLock::new(UnifiedCache::new(
69                10000,
70                DEFAULT_CACHE_TTL,
71                EvictionPolicy::Lru,
72            ))),
73            obs_context,
74        }
75    }
76
77    /// Record tool execution
78    pub fn record_execution(
79        &self,
80        tool_name: &str,
81        arguments: &str,
82        success: bool,
83        quality_score: f32,
84        duration_ms: u64,
85    ) {
86        // Update metrics
87        {
88            let mut metrics = self.tool_metrics.write();
89            let entry = if let Some(m) = metrics.get_mut(tool_name) {
90                m
91            } else {
92                metrics.entry(tool_name.to_string()).or_insert_with(|| ToolMetrics {
93                    name: tool_name.to_string(),
94                    total_calls: 0,
95                    successful_calls: 0,
96                    total_duration_ms: 0,
97                    avg_quality: 0.0,
98                })
99            };
100
101            entry.total_calls += 1;
102            if success {
103                entry.successful_calls += 1;
104            }
105            entry.total_duration_ms += duration_ms;
106            entry.avg_quality =
107                (entry.avg_quality * (entry.total_calls as f32 - 1.0) + quality_score) / entry.total_calls as f32;
108        }
109
110        // Record pattern event
111        let now = std::time::SystemTime::now()
112            .duration_since(std::time::UNIX_EPOCH)
113            .unwrap_or_default()
114            .as_secs();
115
116        self.pattern_engine.record(ExecutionEvent {
117            tool_name: tool_name.to_string(),
118            arguments: arguments.to_string(),
119            success,
120            quality_score,
121            duration_ms,
122            timestamp: now,
123        });
124
125        // Log metric - avoid format! if possible, but for key names it's needed
126        let metric_key = format!("{tool_name}_success_rate");
127        self.obs_context.metric("tool_effectiveness", &metric_key, {
128            let metrics = self.tool_metrics.read();
129            metrics.get(tool_name).map(|m| m.success_rate()).unwrap_or(0.0)
130        });
131    }
132
133    /// Get tool metrics
134    pub fn get_tool_metrics(&self, tool_name: &str) -> Option<ToolMetrics> {
135        self.tool_metrics.read().get(tool_name).cloned()
136    }
137
138    /// Get all tool metrics
139    pub fn get_all_metrics(&self) -> Vec<ToolMetrics> {
140        self.tool_metrics.read().values().cloned().collect()
141    }
142
143    /// Get execution summary
144    pub fn get_summary(&self) -> crate::tools::pattern_engine::ExecutionSummary {
145        self.pattern_engine.summary()
146    }
147
148    /// Predict next tool based on patterns
149    pub fn predict_next_tool(&self) -> Option<String> {
150        self.pattern_engine.predict_next_tool()
151    }
152
153    /// Cache result for tool execution (migrated to UnifiedCache)
154    pub fn cache_result(&self, tool: &str, args: &str, result: &str) {
155        let key = ToolResultKey(format!("{tool}::{args}"));
156        let size = result.len() as u64;
157        self.result_cache.write().insert(key, result.to_string(), size);
158    }
159
160    /// Try to get cached result (migrated to UnifiedCache)
161    pub fn get_cached_result(&self, tool: &str, args: &str) -> Option<String> {
162        let key = ToolResultKey(format!("{tool}::{args}"));
163        self.result_cache.write().get_owned(&key)
164    }
165
166    /// Clear cache (migrated to UnifiedCache)
167    pub fn clear_cache(&self) {
168        self.result_cache.write().clear();
169    }
170
171    /// Get cache stats (migrated to UnifiedCache)
172    pub fn cache_stats(&self) -> crate::cache::CacheStats {
173        self.result_cache.read().stats().clone()
174    }
175
176    /// Rank tools by effectiveness
177    pub fn rank_tools(&self) -> Vec<(String, f32)> {
178        let metrics = self.tool_metrics.read();
179        let mut tools: Vec<_> = metrics
180            .values()
181            .map(|m| {
182                let score = (m.success_rate() * 0.6) + ((1.0 - (m.avg_duration_ms() as f32 / 5000.0).min(1.0)) * 0.4);
183                (&m.name, score)
184            })
185            .collect();
186
187        tools.sort_unstable_by(|a, b| b.1.total_cmp(&a.1));
188        tools.into_iter().map(|(n, s)| (n.clone(), s)).collect()
189    }
190}
191
192#[cfg(test)]
193mod tests {
194    use super::*;
195    use crate::config::constants::tools;
196
197    #[test]
198    fn test_record_execution() {
199        let obs = Arc::new(ObservabilityContext::noop());
200        let ext = ToolRegistryImprovement::new(obs);
201
202        ext.record_execution(tools::CODE_SEARCH, "query", true, 0.8, 100);
203
204        let metrics = ext.get_tool_metrics(tools::CODE_SEARCH);
205        assert!(metrics.is_some());
206        assert!((metrics.unwrap().success_rate() - 1.0).abs() < f32::EPSILON);
207    }
208
209    #[test]
210    fn test_cache_result() {
211        let obs = Arc::new(ObservabilityContext::noop());
212        let ext = ToolRegistryImprovement::new(obs);
213
214        ext.cache_result(tools::CODE_SEARCH, "query", "result");
215        assert_eq!(ext.get_cached_result(tools::CODE_SEARCH, "query"), Some("result".to_owned()));
216    }
217
218    #[test]
219    fn test_rank_tools() {
220        let obs = Arc::new(ObservabilityContext::noop());
221        let ext = ToolRegistryImprovement::new(obs);
222
223        ext.record_execution("tool1", "arg", true, 0.9, 100);
224        ext.record_execution("tool2", "arg", false, 0.3, 50);
225
226        let ranked = ext.rank_tools();
227        assert_eq!(ranked[0].0, "tool1"); // Higher success rate
228    }
229}
vtcode_core/tools/improvements_registry_ext.rs

vtcode_core/tools/
improvements_registry_ext.rs
