From 0e9075ca2cb6f9f7a08c6c69a172b36e90d47648 Mon Sep 17 00:00:00 2001 From: dal Date: Fri, 7 Feb 2025 00:29:11 -0700 Subject: [PATCH] Added Line Number Formatting on File outputs --- api/src/utils/agent/agent.rs | 5 +- .../utils/tools/file_tools/create_files.rs | 42 ++++++++++++++++- api/src/utils/tools/file_tools/mod.rs | 47 ++++++++++++++++++- .../{bulk_modify_files.rs => modify_files.rs} | 21 +++++---- api/src/utils/tools/file_tools/open_files.rs | 3 ++ 5 files changed, 104 insertions(+), 14 deletions(-) rename api/src/utils/tools/file_tools/{bulk_modify_files.rs => modify_files.rs} (89%) diff --git a/api/src/utils/agent/agent.rs b/api/src/utils/agent/agent.rs index 68e86df45..81d575938 100644 --- a/api/src/utils/agent/agent.rs +++ b/api/src/utils/agent/agent.rs @@ -1,6 +1,7 @@ use crate::utils::{ clients::ai::litellm::{ChatCompletionRequest, LiteLLMClient, Message, Tool}, tools::ToolExecutor, + tools::file_tools::FileModificationTool, }; use anyhow::Result; use serde_json::Value; @@ -112,9 +113,9 @@ impl Agent { for tool_call in tool_calls { if let Some(tool) = self.tools.get(&tool_call.function.name) { let result = tool.execute(tool_call).await?; - // Create a message for the tool's response + let result_str = serde_json::to_string(&result)?; results.push(Message::tool( - serde_json::to_string(&result).unwrap(), + result_str, tool_call.id.clone(), )); } diff --git a/api/src/utils/tools/file_tools/create_files.rs b/api/src/utils/tools/file_tools/create_files.rs index a0b1da059..fe3c56979 100644 --- a/api/src/utils/tools/file_tools/create_files.rs +++ b/api/src/utils/tools/file_tools/create_files.rs @@ -17,7 +17,7 @@ use crate::{ utils::{clients::ai::litellm::ToolCall, tools::ToolExecutor}, }; -use super::file_types::{dashboard_yml::DashboardYml, file::FileEnum, metric_yml::MetricYml}; +use super::{file_types::{dashboard_yml::DashboardYml, file::FileEnum, metric_yml::MetricYml}, FileModificationTool}; #[derive(Debug, Serialize, Deserialize, Clone)] struct FileParams { @@ -39,6 +39,8 @@ pub struct CreateFilesOutput { pub struct CreateFilesTool; +impl FileModificationTool for CreateFilesTool {} + #[async_trait] impl ToolExecutor for CreateFilesTool { type Output = CreateFilesOutput; @@ -258,3 +260,41 @@ async fn create_dashboard_file(file: FileParams) -> Result { Ok(FileEnum::Dashboard(dashboard_yml)) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_create_files_serialization() { + let tool = CreateFilesTool; + let output = CreateFilesOutput { + message: "Test message".to_string(), + files: vec![], + }; + + // Use the custom serialization + let result = tool.serialize_output(&output); + assert!(result.is_ok()); + } + + #[test] + fn test_create_files_with_content() { + let tool = CreateFilesTool; + let yml_content = "name: test\ntype: metric\ndescription: A test metric"; + let metric = MetricYml::new(yml_content.to_string()).unwrap(); + + let output = CreateFilesOutput { + message: "Test message".to_string(), + files: vec![FileEnum::Metric(metric)], + }; + + // Use the custom serialization + let result = tool.serialize_output(&output).unwrap(); + + // Verify line numbers were added + assert!(result.contains("1 | name: test")); + assert!(result.contains("2 | type: metric")); + assert!(result.contains("3 | description: A test metric")); + } +} diff --git a/api/src/utils/tools/file_tools/mod.rs b/api/src/utils/tools/file_tools/mod.rs index 5d60e2a4d..db1621653 100644 --- a/api/src/utils/tools/file_tools/mod.rs +++ b/api/src/utils/tools/file_tools/mod.rs @@ -1,4 +1,4 @@ -mod bulk_modify_files; +mod modify_files; mod create_files; pub mod file_types; mod open_files; @@ -6,10 +6,53 @@ mod search_data_catalog; mod search_files; mod send_to_user; -pub use bulk_modify_files::BulkModifyFilesTool; +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; +use crate::utils::tools::ToolExecutor; +use serde::Serialize; +use serde_json::Value; + +/// Trait to mark tools that should have line numbers added to their file content output +pub trait FileModificationTool: ToolExecutor { + /// Process the output Value before it's serialized to add line numbers to file content + fn process_output_value(&self, value: &mut Value) { + if let Some(obj) = value.as_object_mut() { + // Process any string fields that might contain file content + for (_, v) in obj.iter_mut() { + if let Some(content) = v.as_str() { + *v = Value::String(add_line_numbers(content)); + } else if let Some(arr) = v.as_array_mut() { + // Handle arrays of objects that might contain file content + for item in arr.iter_mut() { + if let Some(content) = item.as_str() { + *item = Value::String(add_line_numbers(content)); + } + } + } + } + } + } + + /// Custom serialization for tool output that needs line numbers + fn serialize_output(&self, output: &T) -> Result { + let mut value = serde_json::to_value(output)?; + self.process_output_value(&mut value); + serde_json::to_string(&value) + } +} + +/// Adds line numbers to content in the format "line_number | content" +pub fn add_line_numbers(content: &str) -> String { + content + .lines() + .enumerate() + .map(|(i, line)| format!("{:>4} | {}", i + 1, line)) + .collect::>() + .join("\n") +} + diff --git a/api/src/utils/tools/file_tools/bulk_modify_files.rs b/api/src/utils/tools/file_tools/modify_files.rs similarity index 89% rename from api/src/utils/tools/file_tools/bulk_modify_files.rs rename to api/src/utils/tools/file_tools/modify_files.rs index a939f10fb..a8a5b3129 100644 --- a/api/src/utils/tools/file_tools/bulk_modify_files.rs +++ b/api/src/utils/tools/file_tools/modify_files.rs @@ -4,6 +4,7 @@ use serde_json::Value; use serde::{Deserialize, Serialize}; use crate::utils::{clients::ai::litellm::ToolCall, tools::ToolExecutor}; +use super::FileModificationTool; #[derive(Debug, Serialize, Deserialize)] struct Modification { @@ -18,29 +19,31 @@ struct FileModification { } #[derive(Debug, Serialize, Deserialize)] -struct BulkModifyFilesParams { +struct ModifyFilesParams { files_with_modifications: Vec, } #[derive(Debug, Serialize)] -pub struct BulkModifyFilesOutput { +pub struct ModifyFilesOutput { success: bool, } -pub struct BulkModifyFilesTool; +pub struct ModifyFilesTool; + +impl FileModificationTool for ModifyFilesTool {} #[async_trait] -impl ToolExecutor for BulkModifyFilesTool { - type Output = BulkModifyFilesOutput; +impl ToolExecutor for ModifyFilesTool { + type Output = ModifyFilesOutput; fn get_name(&self) -> String { - "bulk_modify_files".to_string() + "modify_files".to_string() } async fn execute(&self, tool_call: &ToolCall) -> Result { - let params: BulkModifyFilesParams = serde_json::from_str(&tool_call.function.arguments.clone())?; + let params: ModifyFilesParams = serde_json::from_str(&tool_call.function.arguments.clone())?; // TODO: Implement actual file modification logic - let output = BulkModifyFilesOutput { + let output = ModifyFilesOutput { success: true, }; @@ -49,7 +52,7 @@ impl ToolExecutor for BulkModifyFilesTool { fn get_schema(&self) -> Value { serde_json::json!({ - "name": "bulk_modify_files", + "name": "modify_files", "strict": true, "parameters": { "type": "object", diff --git a/api/src/utils/tools/file_tools/open_files.rs b/api/src/utils/tools/file_tools/open_files.rs index 5b1f098f5..f53e7e8b6 100644 --- a/api/src/utils/tools/file_tools/open_files.rs +++ b/api/src/utils/tools/file_tools/open_files.rs @@ -22,6 +22,7 @@ use crate::{ tools::ToolExecutor, }, }; +use super::FileModificationTool; #[derive(Debug, Serialize, Deserialize)] struct FileRequest { @@ -42,6 +43,8 @@ pub struct OpenFilesOutput { pub struct OpenFilesTool; +impl FileModificationTool for OpenFilesTool {} + #[async_trait] impl ToolExecutor for OpenFilesTool { type Output = OpenFilesOutput;