add a plot tool to claude
This commit is contained in:
@@ -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;
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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}")),
|
||||
}
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
|
||||
@@ -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}")),
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user