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
This commit is contained in:
dal 2025-01-26 19:31:51 -07:00
parent 692f8f7a1d
commit 2bc68e8599
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
8 changed files with 304 additions and 18 deletions

View File

@ -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)
}
}
}

View File

@ -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<Message>,
}
/// 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<Value>;
/// Return the JSON schema that describes this tool's interface
fn get_schema(&self) -> serde_json::Value;
}

View File

@ -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<i64>,
}
#[derive(Debug, Serialize, Deserialize)]
struct FileModification {
file: String,
modifications: Vec<Modification>,
}
#[derive(Debug, Serialize, Deserialize)]
struct BulkModifyFilesParams {
files_with_modifications: Vec<FileModification>,
}
pub struct BulkModifyFilesTool;
#[async_trait]
impl ToolExecutor for BulkModifyFilesTool {
async fn execute(&self, tool_call: &ToolCall) -> Result<Value> {
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."
})
}
}

View File

@ -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<String>,
yml_content: String,
}
pub struct CreateFilesTool;
#[async_trait]
impl ToolExecutor for CreateFilesTool {
async fn execute(&self, tool_call: &ToolCall) -> Result<Value> {
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."
})
}
}

View File

@ -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<Value>;
/// Return the JSON schema that describes this tool's interface
fn get_schema(&self) -> serde_json::Value;
}

View File

@ -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<String>,
}
pub struct OpenFilesTool;
#[async_trait]
impl ToolExecutor for OpenFilesTool {
async fn execute(&self, tool_call: &ToolCall) -> Result<Value> {
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."
})
}
}

View File

@ -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<String>,
}
pub struct SearchDatasetsTool;
#[async_trait]
impl ToolExecutor for SearchDatasetsTool {
async fn execute(&self, tool_call: &ToolCall) -> Result<Value> {
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\")."
})
}
}

View File

@ -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<String>,
}
pub struct SearchFilesTool;
#[async_trait]
impl ToolExecutor for SearchFilesTool {
async fn execute(&self, tool_call: &ToolCall) -> Result<Value> {
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`."
})
}
}