From 854d26fbd2b3ba02426c5a64793e7ada090a607a Mon Sep 17 00:00:00 2001 From: dal Date: Fri, 14 Feb 2025 16:01:37 -0700 Subject: [PATCH] added in the send user a message tool --- api/.cursor/rules/tools.mdc | 184 ++++++++++++++++++ .../post_thread/agent_message_transformer.rs | 68 +++++++ .../post_thread/agent_thread.rs | 14 +- api/src/utils/tools/file_tools/mod.rs | 4 +- ...{send_to_user.rs => send_files_to_user.rs} | 6 +- api/src/utils/tools/interaction_tools/mod.rs | 3 + .../interaction_tools/send_message_to_user.rs | 57 ++++++ api/src/utils/tools/mod.rs | 1 + 8 files changed, 329 insertions(+), 8 deletions(-) create mode 100644 api/.cursor/rules/tools.mdc rename api/src/utils/tools/file_tools/{send_to_user.rs => send_files_to_user.rs} (93%) create mode 100644 api/src/utils/tools/interaction_tools/mod.rs create mode 100644 api/src/utils/tools/interaction_tools/send_message_to_user.rs diff --git a/api/.cursor/rules/tools.mdc b/api/.cursor/rules/tools.mdc new file mode 100644 index 000000000..1546c8251 --- /dev/null +++ b/api/.cursor/rules/tools.mdc @@ -0,0 +1,184 @@ +--- +description: Helpful for working with and building tools +globs: */tools/* +--- +# Tools Documentation and Guidelines + +## Overview +This document outlines the architecture, patterns, and best practices for building tools in our system. Tools are modular, reusable components that provide specific functionality to our AI agents and application. + +## Core Architecture + +### ToolExecutor Trait +The foundation of our tools system is the `ToolExecutor` trait. Any struct that wants to be used as a tool must implement this trait: + +```rust +#[async_trait] +pub trait ToolExecutor: Send + Sync { + type Output: Serialize + Send; + async fn execute(&self, tool_call: &ToolCall) -> Result; + fn get_schema(&self) -> serde_json::Value; + fn get_name(&self) -> String; +} +``` + +Key components: +- `Output`: The return type of your tool (must be serializable) +- `execute()`: The main function that implements your tool's logic +- `get_schema()`: Returns the JSON schema describing the tool's interface +- `get_name()`: Returns the tool's unique identifier + +## Tool Categories + +### 1. File Tools +Our file tools provide a robust example of well-structured tool implementation. They handle: +- File creation and modification +- File searching and cataloging +- File type-specific operations +- User interaction with files + +Key patterns from file tools: +- Modular organization by functionality +- Clear separation of concerns +- Type-safe file operations +- Consistent error handling + +### 2. Interaction Tools +Tools that manage user and system interactions. + +## Best Practices + +### 1. Tool Structure +- Create a new module for each tool category +- Implement the `ToolExecutor` trait +- Use meaningful types for `Output` +- Provide comprehensive error handling + +### 2. Schema Design +- Document all parameters clearly +- Use descriptive names for properties +- Include example values where helpful +- Validate input parameters + +### 3. Error Handling +- Use `anyhow::Result` for flexible error handling +- Provide meaningful error messages +- Handle edge cases gracefully +- Implement proper error propagation + +### 4. Testing +- Write unit tests for each tool +- Test edge cases and error conditions +- Mock external dependencies +- Ensure thread safety for async operations + +## Creating New Tools + +### Step 1: Define Your Tool +```rust +pub struct MyNewTool { + // Tool-specific fields +} + +#[async_trait] +impl ToolExecutor for MyNewTool { + type Output = YourOutputType; + + async fn execute(&self, tool_call: &ToolCall) -> Result { + // Implementation + } + + fn get_schema(&self) -> Value { + // Schema definition + } + + fn get_name(&self) -> String { + "my_new_tool".to_string() + } +} +``` + +### Step 2: Schema Definition +```json +{ + "name": "my_new_tool", + "description": "Clear description of what the tool does", + "parameters": { + "type": "object", + "properties": { + // Tool parameters + }, + "required": ["param1", "param2"] + } +} +``` + +### Step 3: Integration +1. Add your tool to the appropriate module +2. Register it in the tool registry +3. Add necessary tests +4. Document usage examples + +## Common Patterns + +### Value Conversion +Use `IntoValueTool` trait when you need to convert tool output to generic JSON: +```rust +my_tool.into_value_tool() +``` + +### File Operations +For tools that modify files: +- Implement `FileModificationTool` trait +- Use `add_line_numbers` for better output formatting +- Handle file permissions appropriately + +## Security Considerations +1. Validate all input parameters +2. Check file permissions before operations +3. Sanitize file paths +4. Handle sensitive data appropriately + +## Examples + +### File Tool Example +```rust +pub struct ReadFileTool { + base_path: PathBuf, +} + +#[async_trait] +impl ToolExecutor for ReadFileTool { + type Output = String; + + async fn execute(&self, tool_call: &ToolCall) -> Result { + // Implementation + } +} +``` + +### Interaction Tool Example +```rust +pub struct UserPromptTool; + +#[async_trait] +impl ToolExecutor for UserPromptTool { + type Output = UserResponse; + + async fn execute(&self, tool_call: &ToolCall) -> Result { + // Implementation + } +} +``` + +## Troubleshooting +1. Check tool registration +2. Verify schema correctness +3. Ensure proper error handling +4. Validate async operations + +## Future Considerations +1. Tool versioning +2. Performance optimization +3. Enhanced type safety +4. Extended testing frameworks \ No newline at end of file diff --git a/api/src/routes/ws/threads_and_messages/post_thread/agent_message_transformer.rs b/api/src/routes/ws/threads_and_messages/post_thread/agent_message_transformer.rs index 04ef525b3..37cf52f3c 100644 --- a/api/src/routes/ws/threads_and_messages/post_thread/agent_message_transformer.rs +++ b/api/src/routes/ws/threads_and_messages/post_thread/agent_message_transformer.rs @@ -15,6 +15,7 @@ use crate::utils::tools::file_tools::modify_files::ModifyFilesParams; use crate::utils::tools::file_tools::open_files::OpenFilesOutput; use crate::utils::tools::file_tools::search_data_catalog::SearchDataCatalogOutput; use crate::utils::tools::file_tools::search_files::SearchFilesOutput; +use crate::utils::tools::interaction_tools::send_message_to_user::{SendMessageToUserInput, SendMessageToUserOutput}; struct StreamingParser { buffer: String, @@ -431,6 +432,7 @@ fn transform_tool_message( "create_files" => tool_create_file(id, content, progress), "modify_files" => tool_modify_file(id, content, progress), "open_files" => tool_open_files(id, content, progress), + "send_message_to_user" => tool_send_message_to_user(id, content, progress), _ => Err(anyhow::anyhow!("Unsupported tool name")), }?; @@ -464,6 +466,7 @@ fn transform_assistant_tool_message( "create_files" => assistant_create_file(id, tool_calls, progress), "modify_files" => assistant_modify_file(id, tool_calls, progress), "open_files" => assistant_open_files(id, progress, initial), + "send_message_to_user" => assistant_send_message_to_user(id, tool_calls, progress), _ => Err(anyhow::anyhow!("Unsupported tool name")), }?; @@ -1061,3 +1064,68 @@ fn tool_modify_file( Err(anyhow::anyhow!("Tool modify file requires progress.")) } } + +fn assistant_send_message_to_user( + id: Option, + tool_calls: Vec, + progress: Option, +) -> Result> { + if let Some(progress) = progress { + if let Some(tool_call) = tool_calls.first() { + // Try to parse the message from the tool call arguments + if let Ok(input) = serde_json::from_str::(&tool_call.function.arguments) { + match progress { + MessageProgress::InProgress => { + Ok(vec![BusterThreadMessage::ChatMessage(BusterChatMessage { + id: id.unwrap_or_else(|| Uuid::new_v4().to_string()), + message_type: "text".to_string(), + message: None, + message_chunk: Some(input.message), + })]) + } + _ => Err(anyhow::anyhow!( + "Assistant send message to user only supports in progress." + )), + } + } else { + Err(anyhow::anyhow!("Failed to parse send message to user input")) + } + } else { + Err(anyhow::anyhow!("No tool call found")) + } + } else { + Err(anyhow::anyhow!( + "Assistant send message to user requires progress." + )) + } +} + +fn tool_send_message_to_user( + id: Option, + content: String, + progress: Option, +) -> Result> { + if let Some(progress) = progress { + // Parse the output to get the message + let output = match serde_json::from_str::(&content) { + Ok(result) => result, + Err(_) => return Ok(vec![]), // Silently ignore parsing errors + }; + + match progress { + MessageProgress::Complete => { + Ok(vec![BusterThreadMessage::ChatMessage(BusterChatMessage { + id: id.unwrap_or_else(|| Uuid::new_v4().to_string()), + message_type: "text".to_string(), + message: Some(output.message), + message_chunk: None, + })]) + } + _ => Err(anyhow::anyhow!( + "Tool send message to user only supports complete." + )), + } + } else { + Err(anyhow::anyhow!("Tool send message to user requires progress.")) + } +} diff --git a/api/src/routes/ws/threads_and_messages/post_thread/agent_thread.rs b/api/src/routes/ws/threads_and_messages/post_thread/agent_thread.rs index 8669727a2..d49f4c744 100644 --- a/api/src/routes/ws/threads_and_messages/post_thread/agent_thread.rs +++ b/api/src/routes/ws/threads_and_messages/post_thread/agent_thread.rs @@ -1,7 +1,7 @@ use anyhow::{Error, Result}; use chrono::Utc; use diesel::{insert_into, ExpressionMethods, QueryDsl}; -use diesel_async::{RunQueryDsl}; +use diesel_async::RunQueryDsl; use handlers::messages::types::{ThreadMessage, ThreadUserMessage}; use handlers::threads::types::ThreadWithMessages; use serde::{Deserialize, Serialize}; @@ -11,6 +11,7 @@ use tokio::sync::mpsc::Receiver; use tracing; use uuid::Uuid; +use crate::utils::tools::interaction_tools::SendMessageToUser; use crate::{ database::{ enums::Verification, @@ -35,7 +36,7 @@ use crate::{ tools::{ file_tools::{ CreateFilesTool, ModifyFilesTool, OpenFilesTool, SearchDataCatalogTool, - SearchFilesTool, SendToUserTool, + SearchFilesTool, SendFilesToUserTool, }, IntoValueTool, ToolExecutor, }, @@ -71,7 +72,8 @@ impl AgentThreadHandler { let modify_files_tool = ModifyFilesTool; let create_files_tool = CreateFilesTool; let open_files_tool = OpenFilesTool; - let send_to_user_tool = SendToUserTool; + let send_to_user_tool = SendFilesToUserTool; + let send_message_to_user_tool = SendMessageToUser; agent.add_tool( search_data_catalog_tool.get_name(), @@ -97,6 +99,10 @@ impl AgentThreadHandler { send_to_user_tool.get_name(), send_to_user_tool.into_value_tool(), ); + agent.add_tool( + send_message_to_user_tool.get_name(), + send_message_to_user_tool.into_value_tool(), + ); Ok(Self { agent }) } @@ -504,6 +510,8 @@ const AGENT_PROMPT: &str = r##" You are an expert analytics/data engineer helping non-technical users get answers to their analytics questions quickly and accurately. You primarily do this by creating or returning metrics and dashboards that already exist or can be built from available datasets. +You should always start by sending a message to the user basically confirming their request. + ## Core Responsibilities - Only open (and show) files that clearly fulfill the user's request - Search data catalog if you can't find solutions to verify you can build what's needed diff --git a/api/src/utils/tools/file_tools/mod.rs b/api/src/utils/tools/file_tools/mod.rs index 415b32b99..8f431f55f 100644 --- a/api/src/utils/tools/file_tools/mod.rs +++ b/api/src/utils/tools/file_tools/mod.rs @@ -4,14 +4,14 @@ pub mod file_types; pub mod open_files; pub mod search_data_catalog; pub mod search_files; -pub mod send_to_user; +pub mod send_files_to_user; pub use modify_files::ModifyFilesTool; pub use create_files::CreateFilesTool; pub use open_files::OpenFilesTool; pub use search_data_catalog::SearchDataCatalogTool; pub use search_files::SearchFilesTool; -pub use send_to_user::SendToUserTool; +pub use send_files_to_user::SendFilesToUserTool; use crate::utils::tools::ToolExecutor; diff --git a/api/src/utils/tools/file_tools/send_to_user.rs b/api/src/utils/tools/file_tools/send_files_to_user.rs similarity index 93% rename from api/src/utils/tools/file_tools/send_to_user.rs rename to api/src/utils/tools/file_tools/send_files_to_user.rs index 11edf234c..66ceb811a 100644 --- a/api/src/utils/tools/file_tools/send_to_user.rs +++ b/api/src/utils/tools/file_tools/send_files_to_user.rs @@ -16,16 +16,16 @@ pub struct SendToUserOutput { success: bool, } -pub struct SendToUserTool; +pub struct SendFilesToUserTool; -impl SendToUserTool { +impl SendFilesToUserTool { pub fn new() -> Self { Self } } #[async_trait] -impl ToolExecutor for SendToUserTool { +impl ToolExecutor for SendFilesToUserTool { type Output = SendToUserOutput; fn get_name(&self) -> String { diff --git a/api/src/utils/tools/interaction_tools/mod.rs b/api/src/utils/tools/interaction_tools/mod.rs new file mode 100644 index 000000000..3579f5d6a --- /dev/null +++ b/api/src/utils/tools/interaction_tools/mod.rs @@ -0,0 +1,3 @@ +pub mod send_message_to_user; + +pub use send_message_to_user::{SendMessageToUser, SendMessageToUserOutput}; diff --git a/api/src/utils/tools/interaction_tools/send_message_to_user.rs b/api/src/utils/tools/interaction_tools/send_message_to_user.rs new file mode 100644 index 000000000..a4ca41aa5 --- /dev/null +++ b/api/src/utils/tools/interaction_tools/send_message_to_user.rs @@ -0,0 +1,57 @@ +use anyhow::Result; +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::utils::tools::ToolCall; +use crate::utils::tools::ToolExecutor; + +#[derive(Debug, Serialize, Deserialize)] +pub struct SendMessageToUserOutput { + pub success: bool, + pub message: String, +} + +#[derive(Debug, Deserialize)] +pub struct SendMessageToUserInput { + pub message: String, +} + +pub struct SendMessageToUser; + +#[async_trait] +impl ToolExecutor for SendMessageToUser { + type Output = SendMessageToUserOutput; + + async fn execute(&self, tool_call: &ToolCall) -> Result { + let input: SendMessageToUserInput = serde_json::from_str(&tool_call.function.arguments)?; + + // For now, we'll consider it a success if we can parse the message + // In a real implementation, we might have additional delivery logic + Ok(SendMessageToUserOutput { + success: true, + message: input.message, + }) + } + + fn get_schema(&self) -> Value { + serde_json::json!({ + "name": "send_message_to_user", + "description": "Sends a text message to the user", + "parameters": { + "type": "object", + "properties": { + "message": { + "type": "string", + "description": "The message to send to the user" + } + }, + "required": ["message"] + } + }) + } + + fn get_name(&self) -> String { + "send_message_to_user".to_string() + } +} diff --git a/api/src/utils/tools/mod.rs b/api/src/utils/tools/mod.rs index 7ec00a459..cc5647c08 100644 --- a/api/src/utils/tools/mod.rs +++ b/api/src/utils/tools/mod.rs @@ -6,6 +6,7 @@ use litellm::ToolCall; pub mod file_tools; +pub mod interaction_tools; /// A trait that defines how tools should be implemented. /// Any struct that wants to be used as a tool must implement this trait.