Added Line Number Formatting on File outputs

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(),
                     ));
                 }

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<FileEnum> {
 
     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"));
+    }
+}

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<T: Serialize>(&self, output: &T) -> Result<String, serde_json::Error> {
+        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::<Vec<_>>()
+        .join("\n")
+}
+

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<FileModification>,
 }
 
 #[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<Self::Output> {
-        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",

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;