Max Novich
75 lines
1.8 KiB
Rust
75 lines
1.8 KiB
Rust
use serde::{Deserialize, Serialize};
|
|
use std::collections::HashMap;
|
|
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct ToolCall {
|
|
name: String,
|
|
parameters: serde_json::Value,
|
|
}
|
|
|
|
impl ToolCall {
|
|
pub fn new(name: String, parameters: serde_json::Value) -> Self {
|
|
Self { name, parameters }
|
|
}
|
|
|
|
fn matches(&self, other: &ToolCall) -> bool {
|
|
self.name == other.name && self.parameters == other.parameters
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct ToolMonitor {
|
|
max_repetitions: Option<u32>,
|
|
last_call: Option<ToolCall>,
|
|
repeat_count: u32,
|
|
call_counts: HashMap<String, u32>,
|
|
}
|
|
|
|
impl ToolMonitor {
|
|
pub fn new(max_repetitions: Option<u32>) -> Self {
|
|
Self {
|
|
max_repetitions,
|
|
last_call: None,
|
|
repeat_count: 0,
|
|
call_counts: HashMap::new(),
|
|
}
|
|
}
|
|
|
|
pub fn check_tool_call(&mut self, tool_call: ToolCall) -> bool {
|
|
let total_calls = self.call_counts.entry(tool_call.name.clone()).or_insert(0);
|
|
*total_calls += 1;
|
|
|
|
if self.max_repetitions.is_none() {
|
|
self.last_call = Some(tool_call);
|
|
self.repeat_count = 1;
|
|
return true;
|
|
}
|
|
|
|
if let Some(last) = &self.last_call {
|
|
if last.matches(&tool_call) {
|
|
self.repeat_count += 1;
|
|
if self.repeat_count > self.max_repetitions.unwrap() {
|
|
return false;
|
|
}
|
|
} else {
|
|
self.repeat_count = 1;
|
|
}
|
|
} else {
|
|
self.repeat_count = 1;
|
|
}
|
|
|
|
self.last_call = Some(tool_call);
|
|
true
|
|
}
|
|
|
|
pub fn get_stats(&self) -> HashMap<String, u32> {
|
|
self.call_counts.clone()
|
|
}
|
|
|
|
pub fn reset(&mut self) {
|
|
self.last_call = None;
|
|
self.repeat_count = 0;
|
|
self.call_counts.clear();
|
|
}
|
|
}
|