fix: more robust json parse for follow ups and locate

Author: humbertoyusta

refact-agent/engine/src/agentic/generate_follow_up_message.rs

@@ -2,10 +2,12 @@ use std::sync::Arc;
 use serde::Deserialize;
 use tokio::sync::{RwLock as ARwLock, Mutex as AMutex};
 
+use crate::custom_error::MapErrToString;
 use crate::global_context::GlobalContext;
 use crate::at_commands::at_commands::AtCommandsContext;
 use crate::subchat::subchat_single;
 use crate::call_validation::{ChatContent, ChatMessage};
+use crate::json_utils;
 
 const PROMPT: &str = r#"
 Your task is to do two things for a conversation between a user and an assistant:
@@ -117,6 +119,7 @@ pub async fn generate_follow_up_message(
 
     tracing::info!("follow-up model says {:?}", response);
 
-    let response: FollowUpResponse = serde_json::from_str(&response).map_err(|e| e.to_string())?;
+    let response: FollowUpResponse = json_utils::extract_json_object(&response)
+        .map_err_with_prefix("Failed to parse json:")?;
     Ok(response)
 }

refact-agent/engine/src/json_utils.rs

@@ -0,0 +1,75 @@
+use crate::custom_error::MapErrToString;
+
+pub fn extract_json_object<T: for<'de> serde::Deserialize<'de>>(text: &str) -> Result<T, String> {
+    let start = text.find('{').ok_or_else(|| "No opening brace '{' found".to_string())?;
+    let end = text.rfind('}').ok_or_else(|| "No closing brace '}' found".to_string())?;
+    
+    if end <= start {
+        return Err("Closing brace appears before opening brace".to_string());
+    }
+
+    serde_json::from_str(&text[start..=end]).map_err_to_string()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use serde_json::json;
+    use serde::Deserialize;
+    use indexmap::IndexMap;
+
+    #[test]
+    fn test_extract_json_clean_input() {
+        let input = r#"{"key": "value", "number": 42}"#;
+        let result: Result<serde_json::Value, _> = extract_json_object(input);
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), json!({"key": "value", "number": 42}));
+    }
+
+    #[test]
+    fn test_extract_json_with_text_before_after() {
+        let input = "Some text before\n {\"key\": \"value\",\n \"number\": 42}\n\n and text after";
+        let result: Result<serde_json::Value, _> = extract_json_object(input);
+        assert!(result.is_ok());
+        assert_eq!(result.unwrap(), json!({"key": "value", "number": 42}));
+    }
+
+    #[test]
+    fn test_extract_json_nested() {
+        let input = r#"LLM response: {"FOUND": {"file1.rs": "symbol1,symbol2"}, "SIMILAR": {"file2.rs": "symbol3"}}"#;
+        let result: Result<IndexMap<String, IndexMap<String, String>>, _> = extract_json_object(input);
+        assert!(result.is_ok());
+        
+        let map = result.unwrap();
+        assert_eq!(map.len(), 2);
+        assert_eq!(map.get("FOUND").unwrap().get("file1.rs").unwrap(), "symbol1,symbol2");
+        assert_eq!(map.get("SIMILAR").unwrap().get("file2.rs").unwrap(), "symbol3");
+    }
+    
+    #[derive(Deserialize, Debug, PartialEq)]
+    struct FollowUpResponse {
+        pub follow_ups: Vec<String>,
+        pub topic_changed: bool,
+    }
+    
+    #[test]
+    fn test_follow_up_response_format() {
+        let input = r#"
+        Here are the follow up questions:
+        ```json
+        {
+          "follow_ups": ["How?", "More examples", "Thank you"],
+          "topic_changed": false
+        }
+        ```
+        "#;
+        
+        let result: Result<FollowUpResponse, _> = extract_json_object(input);
+        
+        assert!(result.is_ok());
+        let response = result.unwrap();
+        
+        assert_eq!(response.follow_ups, vec!["How?", "More examples", "Thank you"]);
+        assert_eq!(response.topic_changed, false);
+    }
+}

refact-agent/engine/src/main.rs

@@ -28,6 +28,7 @@ mod telemetry;
 mod global_context;
 mod background_tasks;
 mod yaml_configs;
+mod json_utils;
 
 mod file_filter;
 mod files_in_workspace;

refact-agent/engine/src/tools/tool_locate_search.rs

@@ -219,7 +219,7 @@ async fn find_relevant_files_with_search(
     crate::tools::tool_relevant_files::update_usage_from_message(&mut usage, &last_message);
     assert!(last_message.role == "assistant");
 
-    let assistant_output1 = serde_json::from_str::<IndexMap<String, serde_json::Value>>(last_message.content.content_text_only().as_str()).map_err(|e| {
+    let assistant_output1 = crate::json_utils::extract_json_object::<IndexMap<String, serde_json::Value>>(last_message.content.content_text_only().as_str()).map_err(|e| {
         tracing::warn!("\n{}\nUnable to parse JSON: {:?}", last_message.content.content_text_only(), e);
         format!("Unable to parse JSON: {:?}", e)
     })?;
@@ -229,7 +229,7 @@ async fn find_relevant_files_with_search(
         return Ok((results, usage, serde_json::to_string_pretty(&assistant_output1).unwrap(), cd_instruction));
     }
 
-    let assistant_output2 = serde_json::from_str::<IndexMap<String, IndexMap<String, String>>>(last_message.content.content_text_only().as_str()).map_err(|e| {
+    let assistant_output2 = crate::json_utils::extract_json_object::<IndexMap<String, IndexMap<String, String>>>(last_message.content.content_text_only().as_str()).map_err(|e| {
         tracing::warn!("\n{}\nUnable to parse JSON: {:?}", last_message.content.content_text_only(), e);
         format!("Unable to parse JSON: {:?}", e)
     })?;