Spaces:
Sleeping
Sleeping
Update json_parser.py
Browse files- json_parser.py +48 -7
json_parser.py
CHANGED
@@ -39,27 +39,65 @@ class SSEParser:
|
|
39 |
def _clean_mermaid_content(self, content: str) -> Optional[str]:
|
40 |
"""Clean and extract mermaid diagram content"""
|
41 |
try:
|
42 |
-
#
|
43 |
-
if
|
44 |
-
content = content
|
45 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
46 |
# Parse JSON if present
|
47 |
try:
|
48 |
-
|
|
|
|
|
|
|
|
|
49 |
# Handle different mermaid output formats
|
50 |
if "mermaid_output" in data:
|
51 |
content = data["mermaid_output"]
|
52 |
elif "mermaid_diagram" in data:
|
53 |
content = data["mermaid_diagram"]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
54 |
except json.JSONDecodeError:
|
55 |
pass
|
56 |
|
57 |
# Clean up markdown formatting
|
58 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
59 |
|
60 |
-
return content.strip()
|
61 |
except Exception as e:
|
62 |
self.logger.error(f"Error cleaning mermaid content: {str(e)}")
|
|
|
63 |
return None
|
64 |
|
65 |
def parse_sse_event(self, data: str) -> Optional[Dict]:
|
@@ -89,9 +127,10 @@ class SSEParser:
|
|
89 |
thought = parsed_data.get("thought", "")
|
90 |
observation = parsed_data.get("observation", "")
|
91 |
tool = parsed_data.get("tool", "")
|
|
|
92 |
|
93 |
# Handle tool-specific formatting
|
94 |
-
if tool == "mermaid_diagrams":
|
95 |
try:
|
96 |
cleaned_content = self._clean_mermaid_content(observation)
|
97 |
if cleaned_content:
|
@@ -103,6 +142,7 @@ class SSEParser:
|
|
103 |
}
|
104 |
except Exception as e:
|
105 |
self.logger.error(f"Failed to parse mermaid data: {str(e)}")
|
|
|
106 |
|
107 |
return {
|
108 |
"type": "thought",
|
@@ -125,6 +165,7 @@ class SSEParser:
|
|
125 |
|
126 |
except Exception as e:
|
127 |
self.logger.error(f"Parse error: {str(e)}")
|
|
|
128 |
return None
|
129 |
|
130 |
def _process_observation(self, data: Dict) -> Dict:
|
|
|
39 |
def _clean_mermaid_content(self, content: str) -> Optional[str]:
|
40 |
"""Clean and extract mermaid diagram content"""
|
41 |
try:
|
42 |
+
# Handle tool output format
|
43 |
+
if isinstance(content, dict) and "tool_output" in content:
|
44 |
+
content = content["tool_output"]
|
45 |
|
46 |
+
# Remove tool response prefix/suffix if present
|
47 |
+
if isinstance(content, str):
|
48 |
+
if "tool response:" in content:
|
49 |
+
content = content.split("tool response:", 1)[1].strip()
|
50 |
+
if content.endswith('.'):
|
51 |
+
content = content[:-1]
|
52 |
+
|
53 |
# Parse JSON if present
|
54 |
try:
|
55 |
+
if isinstance(content, str):
|
56 |
+
data = json.loads(content)
|
57 |
+
else:
|
58 |
+
data = content
|
59 |
+
|
60 |
# Handle different mermaid output formats
|
61 |
if "mermaid_output" in data:
|
62 |
content = data["mermaid_output"]
|
63 |
elif "mermaid_diagram" in data:
|
64 |
content = data["mermaid_diagram"]
|
65 |
+
|
66 |
+
# If content is still JSON string, parse again
|
67 |
+
if isinstance(content, str) and content.startswith('{'):
|
68 |
+
try:
|
69 |
+
data = json.loads(content)
|
70 |
+
if "mermaid_output" in data:
|
71 |
+
content = data["mermaid_output"]
|
72 |
+
elif "mermaid_diagram" in data:
|
73 |
+
content = data["mermaid_diagram"]
|
74 |
+
except:
|
75 |
+
pass
|
76 |
+
|
77 |
except json.JSONDecodeError:
|
78 |
pass
|
79 |
|
80 |
# Clean up markdown formatting
|
81 |
+
if isinstance(content, str):
|
82 |
+
content = content.replace("```mermaid\n", "").replace("\n```", "")
|
83 |
+
content = content.strip()
|
84 |
+
|
85 |
+
# Remove any remaining JSON artifacts
|
86 |
+
if content.startswith('{'):
|
87 |
+
try:
|
88 |
+
data = json.loads(content)
|
89 |
+
if isinstance(data, dict):
|
90 |
+
content = next(iter(data.values()))
|
91 |
+
except:
|
92 |
+
pass
|
93 |
+
|
94 |
+
return content
|
95 |
+
|
96 |
+
return None
|
97 |
|
|
|
98 |
except Exception as e:
|
99 |
self.logger.error(f"Error cleaning mermaid content: {str(e)}")
|
100 |
+
self.logger.debug(f"Original content: {content}")
|
101 |
return None
|
102 |
|
103 |
def parse_sse_event(self, data: str) -> Optional[Dict]:
|
|
|
127 |
thought = parsed_data.get("thought", "")
|
128 |
observation = parsed_data.get("observation", "")
|
129 |
tool = parsed_data.get("tool", "")
|
130 |
+
tool_input = parsed_data.get("tool_input", "")
|
131 |
|
132 |
# Handle tool-specific formatting
|
133 |
+
if tool == "mermaid_diagrams" or "mermaid" in str(tool_input).lower():
|
134 |
try:
|
135 |
cleaned_content = self._clean_mermaid_content(observation)
|
136 |
if cleaned_content:
|
|
|
142 |
}
|
143 |
except Exception as e:
|
144 |
self.logger.error(f"Failed to parse mermaid data: {str(e)}")
|
145 |
+
self.logger.debug(f"Raw observation: {observation}")
|
146 |
|
147 |
return {
|
148 |
"type": "thought",
|
|
|
165 |
|
166 |
except Exception as e:
|
167 |
self.logger.error(f"Parse error: {str(e)}")
|
168 |
+
self.logger.debug(f"Raw data: {data}")
|
169 |
return None
|
170 |
|
171 |
def _process_observation(self, data: Dict) -> Dict:
|