Severian commited on
Commit
1ea16aa
·
verified ·
1 Parent(s): 2531f2f

Update json_parser.py

Browse files
Files changed (1) hide show
  1. 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
- # Remove tool response prefix/suffix if present
43
- if "tool response:" in content:
44
- content = content.split("tool response:", 1)[1].strip()
45
 
 
 
 
 
 
 
 
46
  # Parse JSON if present
47
  try:
48
- data = json.loads(content)
 
 
 
 
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
- content = content.replace("```mermaid\n", "").replace("\n```", "")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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: