Update json_parser.py

json_parser.py

@@ -60,7 +60,7 @@ class SSEParser:
         self.logger = setup_logger("sse_parser")
 
     def parse_sse_event(self, data: str) -> Optional[Dict]:
-        """Parse SSE event data and return cleaned dictionary"""
+        """Parse SSE event data with improved mermaid handling"""
         self.logger.debug("Parsing SSE event")
         
         try:
@@ -71,19 +71,32 @@ class SSEParser:
             # Parse JSON data
             parsed_data = json.loads(data)
             
-            # Clean tool outputs if present
+            # Enhanced mermaid diagram handling
             if "observation" in parsed_data:
                 try:
                     observation = parsed_data["observation"]
                     if observation and isinstance(observation, str):
-                        tool_data = json.loads(observation)
-                        # Extract relevant tool output
-                        for key, value in tool_data.items():
-                            if isinstance(value, str) and "llm_result" in value:
-                                tool_result = json.loads(value)["llm_result"]
-                                parsed_data["observation"] = self.clean_tool_output(tool_result)
-                except:
-                    pass  # Keep original observation if parsing fails
+                        if "mermaid_diagram" in observation:
+                            try:
+                                tool_data = json.loads(observation)
+                                if isinstance(tool_data, dict):
+                                    mermaid_content = tool_data.get(
+                                        "mermaid_diagram", ""
+                                    )
+                                    if mermaid_content:
+                                        # Clean and format mermaid content
+                                        cleaned_content = self.clean_mermaid_content(
+                                            mermaid_content
+                                        )
+                                        parsed_data["observation"] = json.dumps({
+                                            "mermaid_diagram": cleaned_content
+                                        })
+                            except json.JSONDecodeError:
+                                self.logger.warning(
+                                    "Failed to parse mermaid diagram content"
+                                )
+                except Exception as e:
+                    self.logger.error(f"Error processing observation: {str(e)}")
             
             return parsed_data
             
@@ -94,12 +107,35 @@ class SSEParser:
             self.logger.error(f"Parse error: {str(e)}")
             return None
 
-    def clean_tool_output(self, output: str) -> str:
-        """Clean tool output by removing markdown and other formatting"""
-        # Remove markdown code blocks
-        output = re.sub(r'```.*?```', '', output, flags=re.DOTALL)
-        # Remove other markdown formatting
-        output = re.sub(r'[*_`#]', '', output)
-        # Clean up whitespace
-        output = re.sub(r'\n{3,}', '\n\n', output.strip())
-        return output
+    def clean_mermaid_content(self, content: str) -> str:
+        """Clean and format mermaid diagram content"""
+        try:
+            # If content is JSON string, parse it
+            if isinstance(content, str) and content.strip().startswith('{'):
+                content_dict = json.loads(content)
+                if "mermaid_diagram" in content_dict:
+                    content = content_dict["mermaid_diagram"]
+            
+            # Remove markdown code blocks
+            content = re.sub(r'```mermaid\s*|\s*```', '', content)
+            
+            # Remove "tool response:" and any JSON wrapper
+            content = re.sub(r'tool response:.*?{', '{', content)
+            content = re.sub(r'}\s*\.$', '}', content)
+            
+            # If still JSON, extract mermaid content
+            if content.strip().startswith('{'):
+                try:
+                    content_dict = json.loads(content)
+                    if "mermaid_diagram" in content_dict:
+                        content = content_dict["mermaid_diagram"]
+                except:
+                    pass
+            
+            # Final cleanup
+            content = re.sub(r'\s+', ' ', content.strip())
+            return content
+            
+        except Exception as e:
+            self.logger.error(f"Error cleaning mermaid content: {e}")
+            return content