Update response_formatter.py

response_formatter.py

@@ -20,52 +20,48 @@ class ToolType:
 class ResponseFormatter:
     @staticmethod
     def format_thought(
-        thought: str, 
-        observation: Optional[str] = None,
+        thought: str,
+        observation: str,
         citations: List[Dict] = None,
-        metadata: Dict = None
+        metadata: Dict = None,
+        tool_outputs: List[Dict] = None
     ) -> Tuple[str, str]:
-        """Format agent thought and observation for both terminal and XML output"""
+        """Format agent thought for both terminal and XML output"""
         # Terminal format
-        terminal_output = thought.strip()
-        if observation:
-            cleaned_obs = ResponseFormatter._clean_markdown(observation)
-            if cleaned_obs:
-                terminal_output += f"\n\nObservation:\n{cleaned_obs}"
-
+        terminal_output = {
+            "type": "agent_thought",
+            "content": thought,
+            "metadata": metadata or {}
+        }
+        
+        if tool_outputs:
+            terminal_output["tool_outputs"] = tool_outputs
+        
         # XML format
         root = ET.Element("agent_response")
         thought_elem = ET.SubElement(root, "thought")
-        thought_elem.text = thought.strip()
+        thought_elem.text = thought
         
         if observation:
             obs_elem = ET.SubElement(root, "observation")
-            
-            # Extract and format tool outputs
-            tool_outputs = ResponseFormatter._extract_tool_outputs(observation)
-            if tool_outputs:
-                tools_elem = ET.SubElement(obs_elem, "tools")
-                for tool_name, tool_data in tool_outputs.items():
-                    tool_elem = ResponseFormatter._create_tool_element(tools_elem, tool_name, tool_data)
-
-        # Add citations if available
+            obs_elem.text = observation
+        
+        if tool_outputs:
+            tools_elem = ET.SubElement(root, "tool_outputs")
+            for tool_output in tool_outputs:
+                tool_elem = ET.SubElement(tools_elem, "tool_output")
+                tool_elem.attrib["type"] = tool_output.get("type", "")
+                tool_elem.text = tool_output.get("content", "")
+        
         if citations:
-            citations_elem = ET.SubElement(root, "citations")
+            cites_elem = ET.SubElement(root, "citations")
             for citation in citations:
-                cite_elem = ET.SubElement(citations_elem, "citation")
+                cite_elem = ET.SubElement(cites_elem, "citation")
                 for key, value in citation.items():
-                    cite_detail = ET.SubElement(cite_elem, key)
-                    cite_detail.text = str(value)
-
-        # Add metadata if available
-        if metadata:
-            metadata_elem = ET.SubElement(root, "metadata")
-            for key, value in metadata.items():
-                meta_detail = ET.SubElement(metadata_elem, key)
-                meta_detail.text = str(value)
-
+                    cite_elem.attrib[key] = str(value)
+        
         xml_output = ET.tostring(root, encoding='unicode')
-        return terminal_output, xml_output
+        return json.dumps(terminal_output), xml_output
 
     @staticmethod
     def _create_tool_element(parent: ET.Element, tool_name: str, tool_data: Dict) -> ET.Element:
@@ -113,16 +109,29 @@ class ResponseFormatter:
 
     @staticmethod
     def _format_mermaid_data(tool_elem: ET.Element, data: Dict) -> None:
-        """Format mermaid diagram data"""
+        """Format mermaid diagram data with improved error handling"""
         try:
             diagram_elem = ET.SubElement(tool_elem, "diagram")
-            if "mermaid_diagram" in data:
-                # Clean the mermaid code
-                mermaid_code = re.sub(r'```mermaid\s*|\s*```', '', data["mermaid_diagram"])
-                diagram_elem.text = mermaid_code.strip()
-        except:
+            
+            # Extract content from data
+            content = ""
+            if isinstance(data, dict):
+                content = data.get("content", data.get("mermaid_diagram", ""))
+            elif isinstance(data, str):
+                content = data
+            
+            # Clean any remaining markdown/JSON formatting
+            content = re.sub(r'```mermaid\s*|\s*```', '', content)
+            content = re.sub(r'tool response:.*?{', '{', content)
+            content = re.sub(r'}\s*\.$', '}', content)
+            
+            # Set cleaned content
+            diagram_elem.text = content.strip()
+            
+        except Exception as e:
+            logger.error(f"Error formatting mermaid data: {e}")
             content_elem = ET.SubElement(tool_elem, "content")
-            content_elem.text = ResponseFormatter._clean_markdown(str(data))
+            content_elem.text = "Error formatting diagram"
 
     @staticmethod
     def _format_health_data(tool_elem: ET.Element, data: Dict) -> None: