From 2bc68e859950ffde9d2450fc1534c593758b407c Mon Sep 17 00:00:00 2001 From: dal Date: Sun, 26 Jan 2025 19:31:51 -0700 Subject: [PATCH] refactor: Move ToolExecutor trait to tools module - Relocated `ToolExecutor` trait from `agent/types.rs` to `tools/mod.rs` - Updated import paths in `agent/agent.rs` to reflect new trait location - Simplified `types.rs` by removing unused imports and trait definition - Prepared for more comprehensive tool management in the tools module --- api/src/utils/agent/agent.rs | 6 +- api/src/utils/agent/types.rs | 16 +---- api/src/utils/tools/bulk_modify_files.rs | 87 ++++++++++++++++++++++++ api/src/utils/tools/create_files.rs | 50 ++++++++++++++ api/src/utils/tools/mod.rs | 28 ++++++++ api/src/utils/tools/open_files.rs | 45 ++++++++++++ api/src/utils/tools/search_datasets.rs | 45 ++++++++++++ api/src/utils/tools/search_files.rs | 45 ++++++++++++ 8 files changed, 304 insertions(+), 18 deletions(-) create mode 100644 api/src/utils/tools/bulk_modify_files.rs create mode 100644 api/src/utils/tools/create_files.rs create mode 100644 api/src/utils/tools/open_files.rs create mode 100644 api/src/utils/tools/search_datasets.rs create mode 100644 api/src/utils/tools/search_files.rs diff --git a/api/src/utils/agent/agent.rs b/api/src/utils/agent/agent.rs index ced9a8d84..afa2ee682 100644 --- a/api/src/utils/agent/agent.rs +++ b/api/src/utils/agent/agent.rs @@ -1,9 +1,9 @@ -use crate::utils::clients::ai::litellm::{ChatCompletionRequest, LiteLLMClient, Message, Tool}; +use crate::utils::{clients::ai::litellm::{ChatCompletionRequest, LiteLLMClient, Message, Tool}, tools::ToolExecutor}; use anyhow::Result; use std::collections::HashMap; use tokio::sync::mpsc; -use super::types::{AgentThread, ToolExecutor}; +use super::types::{AgentThread}; /// The Agent struct is responsible for managing conversations with the LLM /// and coordinating tool executions. It maintains a registry of available tools @@ -178,4 +178,4 @@ impl Agent { Ok(rx) } -} +} \ No newline at end of file diff --git a/api/src/utils/agent/types.rs b/api/src/utils/agent/types.rs index 167c7b8c5..b52cfcc2a 100644 --- a/api/src/utils/agent/types.rs +++ b/api/src/utils/agent/types.rs @@ -1,9 +1,6 @@ -use anyhow::Result; -use async_trait::async_trait; use serde::{Deserialize, Serialize}; -use serde_json::Value; -use crate::utils::clients::ai::litellm::{Message, ToolCall}; +use crate::utils::clients::ai::litellm::Message; /// A Thread represents a conversation between a user and the AI agent. /// It contains a sequence of messages in chronological order. @@ -14,14 +11,3 @@ pub struct AgentThread { /// Ordered sequence of messages in the conversation pub messages: Vec, } - -/// A trait that defines how tools should be implemented. -/// Any struct that wants to be used as a tool must implement this trait. -#[async_trait] -pub trait ToolExecutor: Send + Sync { - /// Execute the tool with given arguments and return a result - async fn execute(&self, tool_call: &ToolCall) -> Result; - - /// Return the JSON schema that describes this tool's interface - fn get_schema(&self) -> serde_json::Value; -} diff --git a/api/src/utils/tools/bulk_modify_files.rs b/api/src/utils/tools/bulk_modify_files.rs new file mode 100644 index 000000000..539dc8d33 --- /dev/null +++ b/api/src/utils/tools/bulk_modify_files.rs @@ -0,0 +1,87 @@ +use anyhow::Result; +use async_trait::async_trait; +use serde_json::Value; +use serde::{Deserialize, Serialize}; + +use crate::utils::{clients::ai::litellm::ToolCall, tools::ToolExecutor}; + +#[derive(Debug, Serialize, Deserialize)] +struct Modification { + new_content: String, + line_numbers: Vec, +} + +#[derive(Debug, Serialize, Deserialize)] +struct FileModification { + file: String, + modifications: Vec, +} + +#[derive(Debug, Serialize, Deserialize)] +struct BulkModifyFilesParams { + files_with_modifications: Vec, +} + +pub struct BulkModifyFilesTool; + +#[async_trait] +impl ToolExecutor for BulkModifyFilesTool { + async fn execute(&self, tool_call: &ToolCall) -> Result { + let params: BulkModifyFilesParams = serde_json::from_str(&tool_call.function.arguments.clone())?; + // TODO: Implement actual file modification logic + Ok(Value::Array(vec![])) + } + + fn get_schema(&self) -> Value { + serde_json::json!({ + "name": "bulk_modify_files", + "strict": true, + "parameters": { + "type": "object", + "required": ["files_with_modifications"], + "properties": { + "files_with_modifications": { + "type": "array", + "items": { + "type": "object", + "required": ["file", "modifications"], + "properties": { + "file": { + "type": "string", + "description": "The path to a yml file that needs to be modified." + }, + "modifications": { + "type": "array", + "items": { + "type": "object", + "required": ["new_content", "line_numbers"], + "properties": { + "new_content": { + "type": "string", + "description": "The new content that will replace the existing lines. If continuous line changes are made, then you should keep them together." + }, + "line_numbers": { + "type": "array", + "items": { + "type": "number", + "description": "Line numbers in the yml file." + }, + "description": "Array of line numbers that need to be replaced, includes start and end line. If continuous lines are being edited please keep them together. i.e. [20, 34]" + } + }, + "additionalProperties": false + }, + "description": "List of modifications to be made to the file." + } + }, + "additionalProperties": false + }, + "description": "Array of objects containing file paths and their respective modifications." + } + }, + "additionalProperties": false + }, + "description": "Makes multiple line-level modifications to one or more existing YAML files in a single call. If you need to update SQL, chart config, or other sections within a file, use this." + }) + } +} \ No newline at end of file diff --git a/api/src/utils/tools/create_files.rs b/api/src/utils/tools/create_files.rs new file mode 100644 index 000000000..321441a4f --- /dev/null +++ b/api/src/utils/tools/create_files.rs @@ -0,0 +1,50 @@ +use anyhow::Result; +use async_trait::async_trait; +use serde_json::Value; +use serde::{Deserialize, Serialize}; + +use crate::utils::{clients::ai::litellm::ToolCall, tools::ToolExecutor}; + +#[derive(Debug, Serialize, Deserialize)] +struct CreateFilesParams { + names: Vec, + yml_content: String, +} + +pub struct CreateFilesTool; + +#[async_trait] +impl ToolExecutor for CreateFilesTool { + async fn execute(&self, tool_call: &ToolCall) -> Result { + let params: CreateFilesParams = serde_json::from_str(&tool_call.function.arguments.clone())?; + // TODO: Implement actual file creation logic + Ok(Value::Array(vec![])) + } + + fn get_schema(&self) -> Value { + serde_json::json!({ + "name": "create_files", + "strict": true, + "parameters": { + "type": "object", + "required": ["names", "yml_content"], + "properties": { + "names": { + "type": "array", + "items": { + "type": "string", + "description": "The name of a metric or dashboard file to be created" + }, + "description": "An array of names for the metric or dashboard files to be created" + }, + "yml_content": { + "type": "string", + "description": "The YAML content to be included in the created files" + } + }, + "additionalProperties": false + }, + "description": "Creates **new** metric or dashboard files by name. Use this if no existing file can fulfill the user's needs. This will automatically open the metric/dashboard for the user." + }) + } +} \ No newline at end of file diff --git a/api/src/utils/tools/mod.rs b/api/src/utils/tools/mod.rs index e69de29bb..7ed1ef2b2 100644 --- a/api/src/utils/tools/mod.rs +++ b/api/src/utils/tools/mod.rs @@ -0,0 +1,28 @@ +use anyhow::Result; +use async_trait::async_trait; +use serde_json::Value; + +use crate::utils::clients::ai::litellm::ToolCall; + +mod search_files; +mod create_files; +mod bulk_modify_files; +mod search_datasets; +mod open_files; + +pub use search_files::SearchFilesTool; +pub use create_files::CreateFilesTool; +pub use bulk_modify_files::BulkModifyFilesTool; +pub use search_datasets::SearchDatasetsTool; +pub use open_files::OpenFilesTool; + +/// A trait that defines how tools should be implemented. +/// Any struct that wants to be used as a tool must implement this trait. +#[async_trait] +pub trait ToolExecutor: Send + Sync { + /// Execute the tool with given arguments and return a result + async fn execute(&self, tool_call: &ToolCall) -> Result; + + /// Return the JSON schema that describes this tool's interface + fn get_schema(&self) -> serde_json::Value; +} diff --git a/api/src/utils/tools/open_files.rs b/api/src/utils/tools/open_files.rs new file mode 100644 index 000000000..97cf49bdd --- /dev/null +++ b/api/src/utils/tools/open_files.rs @@ -0,0 +1,45 @@ +use anyhow::Result; +use async_trait::async_trait; +use serde_json::Value; +use serde::{Deserialize, Serialize}; + +use crate::utils::{clients::ai::litellm::ToolCall, tools::ToolExecutor}; + +#[derive(Debug, Serialize, Deserialize)] +struct OpenFilesParams { + file_names: Vec, +} + +pub struct OpenFilesTool; + +#[async_trait] +impl ToolExecutor for OpenFilesTool { + async fn execute(&self, tool_call: &ToolCall) -> Result { + let params: OpenFilesParams = serde_json::from_str(&tool_call.function.arguments.clone())?; + // TODO: Implement actual file opening logic + Ok(Value::Array(vec![])) + } + + fn get_schema(&self) -> Value { + serde_json::json!({ + "name": "open_files", + "strict": true, + "parameters": { + "type": "object", + "required": ["file_names"], + "properties": { + "file_names": { + "type": "array", + "items": { + "type": "string", + "description": "The name of a file to be opened" + }, + "description": "List of file names to be opened" + } + }, + "additionalProperties": false + }, + "description": "Opens one or more files in read mode and displays **their entire contents** to the user. If you use this, the user will actually see the metric/dashboard you open." + }) + } +} \ No newline at end of file diff --git a/api/src/utils/tools/search_datasets.rs b/api/src/utils/tools/search_datasets.rs new file mode 100644 index 000000000..a82e67168 --- /dev/null +++ b/api/src/utils/tools/search_datasets.rs @@ -0,0 +1,45 @@ +use anyhow::Result; +use async_trait::async_trait; +use serde_json::Value; +use serde::{Deserialize, Serialize}; + +use crate::utils::{clients::ai::litellm::ToolCall, tools::ToolExecutor}; + +#[derive(Debug, Serialize, Deserialize)] +struct SearchDatasetsParams { + search_terms: Vec, +} + +pub struct SearchDatasetsTool; + +#[async_trait] +impl ToolExecutor for SearchDatasetsTool { + async fn execute(&self, tool_call: &ToolCall) -> Result { + let params: SearchDatasetsParams = serde_json::from_str(&tool_call.function.arguments.clone())?; + // TODO: Implement actual dataset search logic + Ok(Value::Array(vec![])) + } + + fn get_schema(&self) -> Value { + serde_json::json!({ + "name": "search_datasets", + "strict": true, + "parameters": { + "type": "object", + "required": ["search_terms"], + "properties": { + "search_terms": { + "type": "array", + "items": { + "type": "string", + "description": "A search term for finding relevant datasets" + }, + "description": "Array of strings representing the terms to search for" + } + }, + "additionalProperties": false + }, + "description": "Searches for relevant datasets or tables you can query. If you need to write SQL but don't know which dataset to reference, call this with relevant search terms (e.g., \"orders,\" \"customers,\" \"sales transactions\")." + }) + } +} \ No newline at end of file diff --git a/api/src/utils/tools/search_files.rs b/api/src/utils/tools/search_files.rs new file mode 100644 index 000000000..b56bb8c9f --- /dev/null +++ b/api/src/utils/tools/search_files.rs @@ -0,0 +1,45 @@ +use anyhow::Result; +use async_trait::async_trait; +use serde_json::Value; +use serde::{Deserialize, Serialize}; + +use crate::utils::{clients::ai::litellm::ToolCall, tools::ToolExecutor}; + +#[derive(Debug, Serialize, Deserialize)] +struct SearchFilesParams { + query_params: Vec, +} + +pub struct SearchFilesTool; + +#[async_trait] +impl ToolExecutor for SearchFilesTool { + async fn execute(&self, tool_call: &ToolCall) -> Result { + let params: SearchFilesParams = serde_json::from_str(&tool_call.function.arguments.clone())?; + // TODO: Implement actual file search logic + Ok(Value::Array(vec![])) + } + + fn get_schema(&self) -> Value { + serde_json::json!({ + "name": "search_files", + "strict": true, + "parameters": { + "type": "object", + "required": ["query_params"], + "properties": { + "query_params": { + "type": "array", + "items": { + "type": "string", + "description": "A single search query represented as a string." + }, + "description": "An array of natural language queries used to search for files." + } + }, + "additionalProperties": false + }, + "description": "Searches for metric and dashboard files using natural-language queries. Typically used if you suspect there might already be a relevant metric or dashboard in the repository. If results are found, you can then decide whether to open them with `open_files`." + }) + } +} \ No newline at end of file