File size: 3,137 Bytes
e5034c3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | use std::sync::{Arc, Mutex};
use derive_setters::Setters;
use crate::{ArcSender, ChatResponse, Metrics, TitleFormat, Todo, TodoItem};
/// Provides additional context for tool calls.
#[derive(Debug, Clone, Setters)]
pub struct ToolCallContext {
sender: Option<ArcSender>,
metrics: Arc<Mutex<Metrics>>,
}
impl ToolCallContext {
/// Creates a new ToolCallContext with default values
pub fn new(metrics: Metrics) -> Self {
Self { sender: None, metrics: Arc::new(Mutex::new(metrics)) }
}
/// Send a message through the sender if available
pub async fn send(&self, agent_message: impl Into<ChatResponse>) -> anyhow::Result<()> {
if let Some(sender) = &self.sender {
sender.send(Ok(agent_message.into())).await?
}
Ok(())
}
/// Send tool input title - MUST ONLY be used for presenting tool input
/// information
pub async fn send_tool_input(&self, title: impl Into<TitleFormat>) -> anyhow::Result<()> {
let title = title.into();
self.send(ChatResponse::TaskMessage {
content: crate::ChatResponseContent::ToolInput(title),
})
.await
}
/// Execute a closure with access to the metrics
pub fn with_metrics<F, R>(&self, f: F) -> anyhow::Result<R>
where
F: FnOnce(&mut Metrics) -> R,
{
let mut metrics = self
.metrics
.lock()
.map_err(|_| anyhow::anyhow!("Failed to acquire metrics lock"))?;
Ok(f(&mut metrics))
}
/// Execute a fallible closure with access to the metrics
pub fn try_with_metrics<F, R>(&self, f: F) -> anyhow::Result<R>
where
F: FnOnce(&mut Metrics) -> anyhow::Result<R>,
{
let mut metrics = self
.metrics
.lock()
.map_err(|_| anyhow::anyhow!("Failed to acquire metrics lock"))?;
f(&mut metrics)
}
/// Returns all known todos (active and historical completed todos).
///
/// # Errors
///
/// Returns an error if the metrics lock cannot be acquired.
pub fn get_todos(&self) -> anyhow::Result<Vec<Todo>> {
self.with_metrics(|metrics| metrics.get_todos().to_vec())
}
/// Applies incremental todo changes using content as the matching key.
///
/// # Arguments
///
/// * `changes` - Todo items to add, update, or remove (via `cancelled`
/// status).
///
/// # Errors
///
/// Returns an error if the metrics lock cannot be acquired or todo
/// validation fails.
pub fn update_todos(&self, changes: Vec<TodoItem>) -> anyhow::Result<Vec<Todo>> {
self.try_with_metrics(|metrics| metrics.apply_todo_changes(changes))
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_create_context() {
let metrics = Metrics::default();
let context = ToolCallContext::new(metrics);
assert!(context.sender.is_none());
}
#[test]
fn test_with_sender() {
let metrics = Metrics::default();
let context = ToolCallContext::new(metrics);
assert!(context.sender.is_none());
}
}
|