Update json_parser.py

json_parser.py

@@ -39,27 +39,65 @@ class SSEParser:
     def _clean_mermaid_content(self, content: str) -> Optional[str]:
         """Clean and extract mermaid diagram content"""
         try:
-            # Remove tool response prefix/suffix if present
-            if "tool response:" in content:
-                content = content.split("tool response:", 1)[1].strip()
+            # Handle tool output format
+            if isinstance(content, dict) and "tool_output" in content:
+                content = content["tool_output"]
             
+            # Remove tool response prefix/suffix if present
+            if isinstance(content, str):
+                if "tool response:" in content:
+                    content = content.split("tool response:", 1)[1].strip()
+                if content.endswith('.'):
+                    content = content[:-1]
+                
             # Parse JSON if present
             try:
-                data = json.loads(content)
+                if isinstance(content, str):
+                    data = json.loads(content)
+                else:
+                    data = content
+                    
                 # Handle different mermaid output formats
                 if "mermaid_output" in data:
                     content = data["mermaid_output"]
                 elif "mermaid_diagram" in data:
                     content = data["mermaid_diagram"]
+                    
+                # If content is still JSON string, parse again
+                if isinstance(content, str) and content.startswith('{'):
+                    try:
+                        data = json.loads(content)
+                        if "mermaid_output" in data:
+                            content = data["mermaid_output"]
+                        elif "mermaid_diagram" in data:
+                            content = data["mermaid_diagram"]
+                    except:
+                        pass
+                    
             except json.JSONDecodeError:
                 pass
 
             # Clean up markdown formatting
-            content = content.replace("```mermaid\n", "").replace("\n```", "")
+            if isinstance(content, str):
+                content = content.replace("```mermaid\n", "").replace("\n```", "")
+                content = content.strip()
+                
+                # Remove any remaining JSON artifacts
+                if content.startswith('{'):
+                    try:
+                        data = json.loads(content)
+                        if isinstance(data, dict):
+                            content = next(iter(data.values()))
+                    except:
+                        pass
+                    
+                return content
+            
+            return None
             
-            return content.strip()
         except Exception as e:
             self.logger.error(f"Error cleaning mermaid content: {str(e)}")
+            self.logger.debug(f"Original content: {content}")
             return None
 
     def parse_sse_event(self, data: str) -> Optional[Dict]:
@@ -89,9 +127,10 @@ class SSEParser:
                 thought = parsed_data.get("thought", "")
                 observation = parsed_data.get("observation", "")
                 tool = parsed_data.get("tool", "")
+                tool_input = parsed_data.get("tool_input", "")
                 
                 # Handle tool-specific formatting
-                if tool == "mermaid_diagrams":
+                if tool == "mermaid_diagrams" or "mermaid" in str(tool_input).lower():
                     try:
                         cleaned_content = self._clean_mermaid_content(observation)
                         if cleaned_content:
@@ -103,6 +142,7 @@ class SSEParser:
                             }
                     except Exception as e:
                         self.logger.error(f"Failed to parse mermaid data: {str(e)}")
+                        self.logger.debug(f"Raw observation: {observation}")
                 
                 return {
                     "type": "thought",
@@ -125,6 +165,7 @@ class SSEParser:
             
         except Exception as e:
             self.logger.error(f"Parse error: {str(e)}")
+            self.logger.debug(f"Raw data: {data}")
             return None
 
     def _process_observation(self, data: Dict) -> Dict: