Update api.py

api.py

@@ -10,7 +10,7 @@ from dotenv import load_dotenv
 from dify_client_python.dify_client import models
 from sse_starlette.sse import EventSourceResponse
 import httpx
-from json_parser import SSEParser
+from json_parser import SSEParser, MessageState
 from logger_config import setup_logger
 from fastapi.responses import StreamingResponse
 from fastapi.responses import JSONResponse
@@ -57,7 +57,24 @@ class AgentProcessor:
         self.formatter = ResponseFormatter()
         self.client = httpx.AsyncClient(timeout=60.0)
         self.logger = setup_logger("agent_processor")
+        
+        # Initialize headers
+        self.headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "Accept": "text/event-stream"
+        }
 
+    def prepare_request(self, request: AgentRequest) -> Dict:
+        """Prepare request payload for API"""
+        return {
+            "query": request.query,
+            "inputs": request.inputs,
+            "response_mode": "streaming" if request.stream else "blocking",
+            "user": request.user,
+            "conversation_id": request.conversation_id,
+            "files": request.files
+        }
 
     async def log_request_details(
         self, 
@@ -97,153 +114,67 @@ class AgentProcessor:
         await self.client.aclose()
 
     async def process_stream(self, request: AgentRequest):
-        start_time = datetime.now()
-        await self.log_request_details(request, start_time)
-        
-        headers = {
-            "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json",
-            "Accept": "text/event-stream"
-        }
+        """Process streaming request and format for frontend"""
+        parser = SSEParser()
+        formatter = ResponseFormatter()
         
-        chat_request = {
-            "query": request.query,
-            "inputs": request.inputs,
-            "response_mode": "streaming" if request.stream else "blocking",
-            "user": request.user,
-            "conversation_id": request.conversation_id,
-            "files": request.files
-        }
-
         async def event_generator():
-            parser = SSEParser()
-            citations = []
-            metadata = {}
-            tool_outputs = []
-            buffer = ""
-            
             try:
                 async with self.client.stream(
                     "POST",
                     f"{self.api_base}/chat-messages",
-                    headers=headers,
-                    json=chat_request
+                    headers=self.headers,
+                    json=self.prepare_request(request)
                 ) as response:
-                    self.logger.debug(
-                        f"Stream connection established\n"
-                        f"Status: {response.status_code}\n"
-                        f"Headers: {dict(response.headers)}"
-                    )
-                    
                     async for line in response.aiter_lines():
                         if not line.strip():
                             continue
                         
-                        self.logger.debug(f"Raw SSE line: {line}")
-                        buffer += line + "\n"
+                        # Parse the event
+                        parsed = parser.parse_sse_event(line)
+                        if not parsed:
+                            continue
                         
-                        if "data:" in line:
-                            try:
-                                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}\n"
-                                        f"Tool outputs: {tool_outputs}"
-                                    )
-                                
-                                formatted = self.format_terminal_output(
-                                    parsed,
-                                    citations=citations,
-                                    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"}
-                                )
+                        event_type = parsed.get("type")
+                        
+                        # Format based on type
+                        if event_type == "message":
+                            formatted = formatter.format_message(
+                                message=parsed["content"],
+                                message_id=parsed["message_id"]
+                            )
+                        elif event_type == "thought":
+                            formatted = formatter.format_thought(
+                                thought=parsed["content"]["thought"],
+                                observation=parsed["content"]["observation"],
+                                message_id=parsed["message_id"]
+                            )
+                        elif event_type == "tool_output":
+                            # Special handling for tool outputs
+                            formatted = formatter.format_thought(
+                                thought="",
+                                observation="",
+                                tool_outputs=[{
+                                    "type": parsed["tool"],
+                                    "content": parsed["content"]
+                                }],
+                                message_id=parsed["message_id"]
+                            )
+                        else:
+                            continue
+                        
+                        if formatted:
+                            _, xml_output = formatted
+                            yield f"data: {xml_output}\n\n"
                         
-                        # Process buffer when we have a complete message
-                        if 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:
-                self.logger.error(f"Stream processing error: {str(e)}")
-                yield f"data: {{'error': '{str(e)}'}}\n\n"
-
-        return StreamingResponse(
-            event_generator(),
-            media_type="text/event-stream",
-            headers={
-                "Cache-Control": "no-cache",
-                "Connection": "keep-alive",
-                "X-Accel-Buffering": "no",
-                "Access-Control-Allow-Origin": "*"
-            }
-        )
+                self.logger.error(f"Stream error: {str(e)}")
+                error_msg = formatter.format_error(str(e))
+                if error_msg:
+                    _, xml_output = error_msg
+                    yield f"data: {xml_output}\n\n"
+        
+        return EventSourceResponse(event_generator())
 
     def format_terminal_output(
         self,