add a plot tool to claude

This commit is contained in:
Chris Beck
2025-12-08 20:03:31 -07:00
parent 429c4beef9
commit e280b4118e
7 changed files with 125 additions and 26 deletions
+1 -1
View File
@@ -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;
+42 -2
View File
@@ -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<PathBuf>,
}
impl ToolResult {
/// Create a new tool result with just text.
pub fn new(text: impl Into<String>) -> Self {
Self {
text: text.into(),
attachments: Vec::new(),
}
}
/// Create a tool result with text and an attachment.
pub fn with_attachment(text: impl Into<String>, path: impl Into<PathBuf>) -> Self {
Self {
text: text.into(),
attachments: vec![path.into()],
}
}
}
impl From<String> 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<String, String>;
/// Returns the result to be sent back to Claude, potentially with attachments.
async fn execute(&self, name: &str, input: &Value) -> Result<ToolResult, String>;
}
/// A tool definition for the Claude API.
+16 -9
View File
@@ -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<PathBuf> = 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))
@@ -24,7 +24,6 @@ pub struct CommandRouter {
routes: Vec<(String, Handling)>,
}
impl CommandRouter {
/// Create a builder for constructing a CommandRouter.
pub fn builder() -> CommandRouterBuilder {
+3 -3
View File
@@ -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<String, String> {
async fn execute(&self, name: &str, input: &serde_json::Value) -> Result<ToolResult, String> {
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}")),
}
+2 -2
View File
@@ -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<String, String> {
async fn execute(&self, name: &str, input: &serde_json::Value) -> Result<ToolResult, String> {
if self.log_handler.has_tool(name) {
return self.log_handler.execute(name, input).await;
}
+61 -8
View File
@@ -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<Tool> {
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<String, String> {
async fn execute(&self, name: &str, input: &serde_json::Value) -> Result<ToolResult, String> {
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}")),
}
}