Update api.py

api.py

@@ -119,6 +119,7 @@ class AgentProcessor:
             parser = SSEParser()
             citations = []
             metadata = {}
+            tool_outputs = []
             
             try:
                 async with self.client.stream(
@@ -145,72 +146,93 @@ class AgentProcessor:
                                 data = line.split("data:", 1)[1].strip()
                                 parsed = json.loads(data)
                                 
+                                # Enhanced mermaid diagram handling
+                                if parsed.get("observation"):
+                                    try:
+                                        observation = parsed["observation"]
+                                        if isinstance(observation, str):
+                                            if "mermaid_diagram" in observation:
+                                                try:
+                                                    # Clean and extract diagram content
+                                                    cleaned_content = parser.clean_mermaid_content(
+                                                        observation
+                                                    )
+                                                    
+                                                    # Create tool output without extra wrapping
+                                                    tool_output = {
+                                                        "type": "mermaid_diagram",
+                                                        "content": cleaned_content
+                                                    }
+                                                    tool_outputs.append(tool_output)
+                                                    
+                                                    # Send clean event
+                                                    yield (
+                                                        "event: tool_output\n"
+                                                        f"data: {json.dumps(tool_output)}\n\n"
+                                                    )
+                                                except Exception as e:
+                                                    self.logger.error(
+                                                        f"Failed to process mermaid diagram: {e}"
+                                                    )
+                                    except Exception as e:
+                                        self.logger.error(
+                                            f"Error processing observation: {e}"
+                                        )
+                                
                                 if parsed.get("event") == "message_end":
                                     citations = parsed.get("retriever_resources", [])
                                     metadata = parsed.get("metadata", {})
+                                    metadata["tool_outputs"] = tool_outputs
                                     self.logger.debug(
                                         f"Message end event:\n"
                                         f"Citations: {citations}\n"
-                                        f"Metadata: {metadata}"
+                                        f"Metadata: {metadata}\n"
+                                        f"Tool outputs: {tool_outputs}"
                                     )
                                 
                                 formatted = self.format_terminal_output(
                                     parsed,
                                     citations=citations,
-                                    metadata=metadata
+                                    metadata=metadata,
+                                    tool_outputs=tool_outputs
                                 )
                                 if formatted:
                                     self.logger.info(formatted)
+                                    yield f"data: {formatted}\n\n"
+                                
                             except Exception as e:
                                 await self.log_error(
                                     e, 
                                     {"line": line, "event": "parse_data"}
                                 )
-                        
-                        buffer += line + "\n"
-                        
-                        if line.startswith("data:") or buffer.strip().endswith("}"):
-                            try:
-                                processed_response = parser.parse_sse_event(buffer)
-                                if processed_response and isinstance(processed_response, dict):
-                                    cleaned_response = self.clean_response(processed_response)
-                                    if cleaned_response:
-                                        xml_content = cleaned_response.get("content", "")
-                                        yield f"data: {xml_content}\n\n"
-                            except Exception as parse_error:
-                                await self.log_error(
-                                    parse_error,
-                                    {"buffer": buffer, "event": "process_buffer"}
-                                )
-                                error_xml = (
-                                    f"<agent_response>"
-                                    f"<error>{str(parse_error)}</error>"
-                                    f"</agent_response>"
-                                )
-                                yield f"data: {error_xml}\n\n"
-                            finally:
-                                buffer = ""
-            
-            except httpx.ConnectError as e:
-                await self.log_error(e, {"event": "connection_error"})
-                error_xml = (
-                    f"<agent_response>"
-                    f"<error>Connection error: {str(e)}</error>"
-                    f"</agent_response>"
-                )
-                yield f"data: {error_xml}\n\n"
+                    
+                    buffer += line + "\n"
+                    
+                    if line.startswith("data:") or buffer.strip().endswith("}"):
+                        try:
+                            processed_response = parser.parse_sse_event(buffer)
+                            if processed_response and isinstance(processed_response, dict):
+                                cleaned_response = self.clean_response(processed_response)
+                                if cleaned_response:
+                                    xml_content = cleaned_response.get("content", "")
+                                    yield f"data: {xml_content}\n\n"
+                        except Exception as parse_error:
+                            await self.log_error(
+                                parse_error,
+                                {"buffer": buffer, "event": "process_buffer"}
+                            )
+                            error_xml = (
+                                f"<agent_response>"
+                                f"<error>{str(parse_error)}</error>"
+                                f"</agent_response>"
+                            )
+                            yield f"data: {error_xml}\n\n"
+                        finally:
+                            buffer = ""
+                            
             except Exception as e:
-                await self.log_error(e, {"event": "stream_error"})
-                error_xml = (
-                    f"<agent_response>"
-                    f"<error>Streaming error: {str(e)}</error>"
-                    f"</agent_response>"
-                )
-                yield f"data: {error_xml}\n\n"
-            finally:
-                end_time = datetime.now()
-                duration = (end_time - start_time).total_seconds()
-                self.logger.info(f"Request completed in {duration:.2f} seconds")
+                self.logger.error(f"Stream processing error: {str(e)}")
+                yield f"data: {{'error': '{str(e)}'}}\n\n"
 
         return StreamingResponse(
             event_generator(),
@@ -227,7 +249,8 @@ class AgentProcessor:
         self, 
         response: Dict, 
         citations: List[Dict] = None,
-        metadata: Dict = None
+        metadata: Dict = None,
+        tool_outputs: List[Dict] = None
     ) -> Optional[str]:
         """Format response for terminal output"""
         event_type = response.get("event")
@@ -239,7 +262,8 @@ class AgentProcessor:
                 thought, 
                 observation,
                 citations=citations,
-                metadata=metadata
+                metadata=metadata,
+                tool_outputs=tool_outputs
             )
             return terminal_output
             
@@ -266,6 +290,60 @@ class AgentProcessor:
             if event_type == "agent_thought":
                 thought = response.get("thought", "")
                 observation = response.get("observation", "")
+                tool = response.get("tool", "")
+                
+                # Handle mermaid diagram observations
+                if tool == "mermaid_diagram" and observation:
+                    try:
+                        # First check if observation is error message
+                        if isinstance(observation, str):
+                            obs_data = json.loads(observation)
+                            if "mermaid_diagram" in obs_data:
+                                if obs_data["mermaid_diagram"].startswith("tool invoke error"):
+                                    self.logger.warning(
+                                        f"Mermaid diagram tool error: {obs_data['mermaid_diagram']}"
+                                    )
+                                    return None
+                        
+                        # Handle successful mermaid diagram
+                        if isinstance(observation, dict):
+                            mermaid_data = observation.get("mermaid_diagram", "")
+                        else:
+                            obs_data = json.loads(observation)
+                            mermaid_data = obs_data.get("mermaid_diagram", "")
+
+                        if mermaid_data:
+                            # Handle nested JSON structure
+                            if isinstance(mermaid_data, str):
+                                mermaid_data = json.loads(mermaid_data)
+                            
+                            # Extract diagram from either format
+                            if isinstance(mermaid_data, dict):
+                                diagram = mermaid_data.get("mermaid_diagram", "")
+                            else:
+                                diagram = mermaid_data
+
+                            # Clean up the diagram code
+                            if isinstance(diagram, str):
+                                if "tool response:" in diagram:
+                                    diagram = diagram.split("tool response:")[0]
+                                if diagram.startswith('{"mermaid_diagram": "'):
+                                    diagram = json.loads(diagram)["mermaid_diagram"]
+                                if diagram.startswith("```mermaid\n"):
+                                    diagram = diagram[10:]
+                                if diagram.endswith("\n```"):
+                                    diagram = diagram[:-4]
+                                
+                                return {
+                                    "type": "mermaid_diagram",
+                                    "content": diagram.strip()
+                                }
+                    except (json.JSONDecodeError, KeyError) as e:
+                        self.logger.error(f"Failed to parse mermaid diagram data: {str(e)}")
+                        self.logger.debug(f"Raw observation: {observation}")
+                        return None
+                
+                # Handle regular thought
                 _, xml_output = self.formatter.format_thought(thought, observation)
                 return {
                     "type": "thought",