From e280b4118efa81e4de600d8ea637350717f1801b Mon Sep 17 00:00:00 2001 From: Chris Beck Date: Mon, 8 Dec 2025 20:03:31 -0700 Subject: [PATCH] add a plot tool to claude --- signal-gateway/src/claude/mod.rs | 2 +- signal-gateway/src/claude/tools.rs | 44 ++++++++++++- signal-gateway/src/claude/worker.rs | 25 ++++--- signal-gateway/src/gateway/command_router.rs | 1 - signal-gateway/src/gateway/log_handler.rs | 6 +- signal-gateway/src/gateway/mod.rs | 4 +- signal-gateway/src/prometheus/mod.rs | 69 +++++++++++++++++--- 7 files changed, 125 insertions(+), 26 deletions(-) diff --git a/signal-gateway/src/claude/mod.rs b/signal-gateway/src/claude/mod.rs index e9fc519..dc5ad14 100644 --- a/signal-gateway/src/claude/mod.rs +++ b/signal-gateway/src/claude/mod.rs @@ -3,7 +3,7 @@ mod tools; mod worker; -pub use tools::{Tool, ToolExecutor}; +pub use tools::{Tool, ToolExecutor, ToolResult}; pub use worker::SentBy; use crate::message_handler::AdminMessageResponse; diff --git a/signal-gateway/src/claude/tools.rs b/signal-gateway/src/claude/tools.rs index 6d577dc..3dbcb09 100644 --- a/signal-gateway/src/claude/tools.rs +++ b/signal-gateway/src/claude/tools.rs @@ -3,6 +3,46 @@ use async_trait::async_trait; use serde::Serialize; use serde_json::Value; +use std::path::PathBuf; + +/// Result of executing a tool. +#[derive(Clone, Debug, Default)] +pub struct ToolResult { + /// Text result to return to Claude. + pub text: String, + /// Optional file attachments generated by the tool. + pub attachments: Vec, +} + +impl ToolResult { + /// Create a new tool result with just text. + pub fn new(text: impl Into) -> Self { + Self { + text: text.into(), + attachments: Vec::new(), + } + } + + /// Create a tool result with text and an attachment. + pub fn with_attachment(text: impl Into, path: impl Into) -> Self { + Self { + text: text.into(), + attachments: vec![path.into()], + } + } +} + +impl From for ToolResult { + fn from(text: String) -> Self { + Self::new(text) + } +} + +impl From<&str> for ToolResult { + fn from(text: &str) -> Self { + Self::new(text) + } +} /// Trait for executing tools. Implement this to provide tool capabilities. #[async_trait] @@ -16,8 +56,8 @@ pub trait ToolExecutor: Send + Sync { } /// Execute a tool by name with the given input arguments. - /// Returns the result as a string to be sent back to Claude. - async fn execute(&self, name: &str, input: &Value) -> Result; + /// Returns the result to be sent back to Claude, potentially with attachments. + async fn execute(&self, name: &str, input: &Value) -> Result; } /// A tool definition for the Claude API. diff --git a/signal-gateway/src/claude/worker.rs b/signal-gateway/src/claude/worker.rs index 20d122f..8002b66 100644 --- a/signal-gateway/src/claude/worker.rs +++ b/signal-gateway/src/claude/worker.rs @@ -5,6 +5,7 @@ use crate::message_handler::AdminMessageResponse; use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; use serde_json::Value; +use std::path::PathBuf; use std::sync::Weak; use tokio::sync::{mpsc, oneshot}; use tracing::info; @@ -192,10 +193,14 @@ impl ClaudeWorker { let executor = self.tool_executor.upgrade(); let tools = executor.as_ref().map(|te| te.tools()).unwrap_or_default(); + // Collect attachments from tool results across all iterations + let mut attachments: Vec = Vec::new(); + if let Some(last) = self.messages.last() - && let Some(ContentBlock::Text { text, .. }) = last.content.first() { - info!("Claude request: {}", text); - } + && let Some(ContentBlock::Text { text, .. }) = last.content.first() + { + info!("Claude request: {}", text); + } for iteration in 0..max_iterations { // Check for stop before making API call @@ -248,10 +253,12 @@ impl ClaudeWorker { self.check_stop()?; info!("Claude tool use: {}({})", name, input); - let (result, is_error) = match executor.execute(name, input).await { - Ok(result) => { - info!("Tool result: {}", result); - (result, false) + let (result_text, is_error) = match executor.execute(name, input).await { + Ok(tool_result) => { + info!("Tool result: {}", tool_result.text); + // Collect any attachments from the tool result + attachments.extend(tool_result.attachments); + (tool_result.text, false) } Err(err) => { info!("Tool error: {}", err); @@ -260,7 +267,7 @@ impl ClaudeWorker { }; self.messages.push(MessageContent::tool_result( id.to_string(), - result, + result_text, is_error, )); } @@ -286,7 +293,7 @@ impl ClaudeWorker { .join("\n"); info!("Claude final result: {}", text); - return Ok(AdminMessageResponse::new(text)); + return Ok(AdminMessageResponse::new(text).with_attachments(attachments)); } Err(ClaudeError::TooManyIterations(max_iterations)) diff --git a/signal-gateway/src/gateway/command_router.rs b/signal-gateway/src/gateway/command_router.rs index 32b7421..a9d362d 100644 --- a/signal-gateway/src/gateway/command_router.rs +++ b/signal-gateway/src/gateway/command_router.rs @@ -24,7 +24,6 @@ pub struct CommandRouter { routes: Vec<(String, Handling)>, } - impl CommandRouter { /// Create a builder for constructing a CommandRouter. pub fn builder() -> CommandRouterBuilder { diff --git a/signal-gateway/src/gateway/log_handler.rs b/signal-gateway/src/gateway/log_handler.rs index bacdd73..59d93f0 100644 --- a/signal-gateway/src/gateway/log_handler.rs +++ b/signal-gateway/src/gateway/log_handler.rs @@ -4,7 +4,7 @@ use super::{ route::{Destination, Limit, Route}, }; use crate::{ - claude::{Tool, ToolExecutor}, + claude::{Tool, ToolExecutor, ToolResult}, concurrent_map::LazyMap, log_format::LogFormatConfig, log_message::{LogFilter, LogMessage, Origin}, @@ -308,11 +308,11 @@ impl ToolExecutor for LogHandler { vec![logs_tool()] } - async fn execute(&self, name: &str, input: &serde_json::Value) -> Result { + async fn execute(&self, name: &str, input: &serde_json::Value) -> Result { match name { "logs" => { let filter = input.get("filter").and_then(|v| v.as_str()); - Ok(self.format_logs(filter).await) + Ok(self.format_logs(filter).await.into()) } _ => Err(format!("unknown tool: {name}")), } diff --git a/signal-gateway/src/gateway/mod.rs b/signal-gateway/src/gateway/mod.rs index 10cb20d..9bb170d 100644 --- a/signal-gateway/src/gateway/mod.rs +++ b/signal-gateway/src/gateway/mod.rs @@ -4,7 +4,7 @@ use crate::signal_jsonrpc::connect_ipc; use crate::{ alertmanager::AlertPost, - claude::{ClaudeApi, ClaudeConfig, SentBy, Tool, ToolExecutor}, + claude::{ClaudeApi, ClaudeConfig, SentBy, Tool, ToolExecutor, ToolResult}, log_message::{LogMessage, Origin}, message_handler::{AdminMessage, AdminMessageResponse, Context, MessageHandlerResult}, prometheus::{Prometheus, PrometheusConfig}, @@ -921,7 +921,7 @@ impl ToolExecutor for Gateway { tools } - async fn execute(&self, name: &str, input: &serde_json::Value) -> Result { + async fn execute(&self, name: &str, input: &serde_json::Value) -> Result { if self.log_handler.has_tool(name) { return self.log_handler.execute(name, input).await; } diff --git a/signal-gateway/src/prometheus/mod.rs b/signal-gateway/src/prometheus/mod.rs index 1245604..145cdf1 100644 --- a/signal-gateway/src/prometheus/mod.rs +++ b/signal-gateway/src/prometheus/mod.rs @@ -1,4 +1,4 @@ -use crate::claude::{Tool, ToolExecutor}; +use crate::claude::{Tool, ToolExecutor, ToolResult}; use async_trait::async_trait; use conf::Conf; use prometheus_http_client::{ @@ -299,13 +299,40 @@ fn alerts_tool() -> Tool { } } +fn plot_tool() -> Tool { + Tool { + name: "prometheus_plot", + description: "Create a plot/graph of a Prometheus metric over time. Returns the plot as an image attachment.", + input_schema: serde_json::json!({ + "type": "object", + "properties": { + "query": { + "type": "string", + "description": "A PromQL query expression to plot (e.g., 'rate(http_requests_total[5m])')" + }, + "range": { + "type": "string", + "description": "Time range to plot, as a duration string (e.g., '1h', '30m', '2d'). Defaults to '1h'." + } + }, + "required": ["query"] + }), + } +} + #[async_trait] impl ToolExecutor for Prometheus { fn tools(&self) -> Vec { - vec![query_tool(), series_tool(), labels_tool(), alerts_tool()] + vec![ + query_tool(), + series_tool(), + labels_tool(), + alerts_tool(), + plot_tool(), + ] } - async fn execute(&self, name: &str, input: &serde_json::Value) -> Result { + async fn execute(&self, name: &str, input: &serde_json::Value) -> Result { match name { "prometheus_query" => { let query = input @@ -326,7 +353,7 @@ impl ToolExecutor for Prometheus { .unwrap_or_else(|| "-".to_owned()); writeln!(&mut result, " {:?} = {}", sl, value_str).unwrap(); } - Ok(result) + Ok(result.into()) } Err(err) => Err(format!("prometheus query failed: {err}")), } @@ -351,7 +378,7 @@ impl ToolExecutor for Prometheus { if series.len() > 100 { result.push_str(&format!(" ... and {} more\n", series.len() - 100)); } - Ok(result) + Ok(result.into()) } Err(err) => Err(format!("prometheus series failed: {err}")), } @@ -373,7 +400,7 @@ impl ToolExecutor for Prometheus { for label in &labels { writeln!(&mut result, " {}", label).unwrap(); } - Ok(result) + Ok(result.into()) } Err(err) => Err(format!("prometheus labels failed: {err}")), } @@ -381,7 +408,7 @@ impl ToolExecutor for Prometheus { "prometheus_alerts" => match self.alerts().await { Ok(alerts) => { if alerts.is_empty() { - return Ok("No active alerts".to_owned()); + return Ok("No active alerts".into()); } let mut result = format!("Found {} alerts:\n", alerts.len()); for alert in &alerts { @@ -392,10 +419,36 @@ impl ToolExecutor for Prometheus { ) .unwrap(); } - Ok(result) + Ok(result.into()) } Err(err) => Err(format!("prometheus alerts failed: {err}")), }, + "prometheus_plot" => { + #[cfg(feature = "plot")] + { + let query = input + .get("query") + .and_then(|v| v.as_str()) + .ok_or_else(|| "missing 'query' parameter".to_owned())?; + + // Parse range, default to 1h + let range_str = input.get("range").and_then(|v| v.as_str()).unwrap_or("1h"); + let range = conf_extra::parse_duration(range_str) + .map_err(|e| format!("invalid range '{}': {}", range_str, e))?; + + match self.create_oneoff_plot(query.to_owned(), range).await { + Ok(path) => Ok(ToolResult::with_attachment( + format!("Plot created: {}", path.display()), + path, + )), + Err(err) => Err(format!("prometheus plot failed: {err}")), + } + } + #[cfg(not(feature = "plot"))] + { + Err("prometheus_plot requires the 'plot' feature to be enabled".to_owned()) + } + } _ => Err(format!("unknown tool: {name}")), } }