use crate::utils::current_timestamp;
use hashbrown::HashMap;
use serde::{Deserialize, Serialize};
use crate::tools::result_metadata::ResultMetadata;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ToolEffectiveness {
pub tool_name: String,
pub success_rate: f32,
pub avg_result_quality: f32,
pub usage_count: usize,
pub success_count: usize,
pub last_used_timestamp: u64,
#[serde(default)]
pub failure_modes: Vec<ToolFailureMode>,
#[serde(default)]
pub avg_execution_time_ms: f32,
}
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Hash)]
pub enum ToolFailureMode {
Timeout,
NoResults,
InvalidArgs,
ParseError,
PermissionDenied,
Unknown,
}
impl ToolEffectiveness {
pub fn new(tool_name: String) -> Self {
Self {
tool_name,
success_rate: 0.0,
avg_result_quality: 0.0,
usage_count: 0,
success_count: 0,
last_used_timestamp: 0,
failure_modes: vec![],
avg_execution_time_ms: 0.0,
}
}
pub fn record_success(&mut self, quality: f32, execution_time_ms: f32) {
self.usage_count += 1;
self.success_count += 1;
self.last_used_timestamp = current_timestamp();
self.avg_result_quality = (self.avg_result_quality * (self.success_count - 1) as f32
+ quality)
/ self.success_count as f32;
self.avg_execution_time_ms = (self.avg_execution_time_ms * (self.success_count - 1) as f32
+ execution_time_ms)
/ self.success_count as f32;
self.update_success_rate();
}
pub fn record_failure(&mut self, failure_mode: ToolFailureMode, execution_time_ms: f32) {
self.usage_count += 1;
self.last_used_timestamp = current_timestamp();
if !self.failure_modes.iter().any(|m| m == &failure_mode) {
self.failure_modes.push(failure_mode);
}
let success_count_f = (self.success_count + 1) as f32;
self.avg_execution_time_ms = (self.avg_execution_time_ms
* (self.success_count as f32 / success_count_f))
+ (execution_time_ms / success_count_f);
self.update_success_rate();
}
fn update_success_rate(&mut self) {
if self.usage_count > 0 {
self.success_rate = self.success_count as f32 / self.usage_count as f32;
}
}
pub fn effectiveness_score(&self) -> f32 {
if self.usage_count == 0 {
return 0.5; }
(self.success_rate * 0.6) + (self.avg_result_quality * 0.4)
}
pub fn is_reliable(&self) -> bool {
self.usage_count >= 3 && self.success_rate > 0.7
}
pub fn time_since_last_use_seconds(&self) -> u64 {
if self.last_used_timestamp == 0 {
u64::MAX
} else {
current_timestamp().saturating_sub(self.last_used_timestamp)
}
}
}
#[derive(Debug, Clone)]
pub struct ToolSelectionContext {
pub task_description: String,
pub prior_tools_used: Vec<String>,
pub prior_result_qualities: Vec<f32>,
pub tool_effectiveness: HashMap<String, ToolEffectiveness>,
}
pub trait ToolSelector: Send + Sync {
fn select_tool(&self, context: &ToolSelectionContext, candidates: &[&str]) -> Option<String>;
}
pub struct AdaptiveToolSelector;
impl ToolSelector for AdaptiveToolSelector {
fn select_tool(&self, context: &ToolSelectionContext, candidates: &[&str]) -> Option<String> {
if candidates.is_empty() {
return None;
}
let mut scored: Vec<(String, f32)> = candidates
.iter()
.map(|tool| {
let name = (*tool).to_owned();
let score = score_tool(&name, context);
(name, score)
})
.collect();
scored.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
scored.first().map(|(name, _)| name.clone())
}
}
fn score_tool(tool_name: &str, context: &ToolSelectionContext) -> f32 {
let mut score = 0.5;
if let Some(eff) = context.tool_effectiveness.get(tool_name) {
score += eff.effectiveness_score() * 0.4;
let time_score = 1.0 - (eff.avg_execution_time_ms / 10000.0).min(1.0);
score += time_score * 0.1;
if !eff.failure_modes.is_empty() {
let failure_penalty = (eff.failure_modes.len() as f32) * 0.1;
score -= failure_penalty;
}
}
if context.prior_tools_used.iter().any(|s| s == tool_name) {
score -= 0.15; }
score.clamp(0.0, 1.0)
}
pub struct ToolEffectivenessTracker {
effectiveness: HashMap<String, ToolEffectiveness>,
}
impl ToolEffectivenessTracker {
pub fn new() -> Self {
Self {
effectiveness: HashMap::new(),
}
}
fn get_or_create(&mut self, tool_name: &str) -> &mut ToolEffectiveness {
self.effectiveness
.entry(tool_name.to_owned())
.or_insert_with(|| ToolEffectiveness::new(tool_name.to_owned()))
}
pub fn record_success(
&mut self,
tool_name: &str,
metadata: &ResultMetadata,
execution_time_ms: f32,
) {
let quality = metadata.quality_score();
self.get_or_create(tool_name)
.record_success(quality, execution_time_ms);
}
pub fn record_failure(
&mut self,
tool_name: &str,
mode: ToolFailureMode,
execution_time_ms: f32,
) {
self.get_or_create(tool_name)
.record_failure(mode, execution_time_ms);
}
pub fn snapshot(&self) -> HashMap<String, ToolEffectiveness> {
self.effectiveness.clone()
}
pub fn get(&self, tool_name: &str) -> Option<&ToolEffectiveness> {
self.effectiveness.get(tool_name)
}
pub fn sorted_by_effectiveness(&self) -> Vec<(String, f32)> {
let mut tools: Vec<_> = self
.effectiveness
.iter()
.map(|(name, eff)| (name.clone(), eff.effectiveness_score()))
.collect();
tools.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
tools
}
}
impl Default for ToolEffectivenessTracker {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tool_effectiveness_success() {
let mut eff = ToolEffectiveness::new("grep".to_owned());
eff.record_success(0.9, 100.0);
eff.record_success(0.85, 110.0);
assert_eq!(eff.success_count, 2);
assert_eq!(eff.usage_count, 2);
assert_eq!(eff.success_rate, 1.0);
assert!(eff.avg_result_quality > 0.8);
}
#[test]
fn test_tool_effectiveness_failure() {
let mut eff = ToolEffectiveness::new("find".to_owned());
eff.record_failure(ToolFailureMode::Timeout, 5000.0);
assert_eq!(eff.success_count, 0);
assert_eq!(eff.usage_count, 1);
assert_eq!(eff.success_rate, 0.0);
assert!(eff.failure_modes.contains(&ToolFailureMode::Timeout));
}
#[test]
fn test_adaptive_selector() {
let selector = AdaptiveToolSelector;
let mut effectiveness = HashMap::new();
let mut grep_eff = ToolEffectiveness::new("grep".to_owned());
grep_eff.record_success(0.9, 100.0);
effectiveness.insert("grep".to_owned(), grep_eff);
let mut find_eff = ToolEffectiveness::new("find".to_owned());
find_eff.record_failure(ToolFailureMode::Timeout, 5000.0);
effectiveness.insert("find".to_owned(), find_eff);
let context = ToolSelectionContext {
task_description: "find error patterns".to_owned(),
prior_tools_used: vec![],
prior_result_qualities: vec![],
tool_effectiveness: effectiveness,
};
let selected = selector.select_tool(&context, &["grep", "find"]);
assert_eq!(selected, Some("grep".to_owned()));
}
#[test]
fn test_tool_diversity_penalty() {
let selector = AdaptiveToolSelector;
let effectiveness = HashMap::new();
let context = ToolSelectionContext {
task_description: "find files".to_owned(),
prior_tools_used: vec!["grep".to_owned()],
prior_result_qualities: vec![],
tool_effectiveness: effectiveness,
};
let selected = selector.select_tool(&context, &["grep", "find"]);
assert_eq!(selected, Some("find".to_owned()));
}
#[test]
fn test_effectiveness_tracker() {
let mut tracker = ToolEffectivenessTracker::new();
let meta = ResultMetadata::success(0.8, 0.8);
tracker.record_success("grep", &meta, 100.0);
tracker.record_success("grep", &meta, 110.0);
let sorted = tracker.sorted_by_effectiveness();
assert_eq!(sorted[0].0, "grep");
assert!(sorted[0].1 > 0.7);
}
}