Helper for streaming responses.
| 41 | |
| 42 | |
| 43 | class StreamingResponse: |
| 44 | """Helper for streaming responses.""" |
| 45 | |
| 46 | def __init__(self, websocket: WebSocket, message_id: str): |
| 47 | self.websocket = websocket |
| 48 | self.message_id = message_id |
| 49 | self.chunks = [] |
| 50 | self.start_time = datetime.utcnow() |
| 51 | |
| 52 | async def send_chunk(self, chunk: str): |
| 53 | """Send a chunk of the response.""" |
| 54 | self.chunks.append(chunk) |
| 55 | await self.websocket.send_json({ |
| 56 | "type": MessageType.RESPONSE, |
| 57 | "message_id": self.message_id, |
| 58 | "chunk": chunk, |
| 59 | "is_streaming": True, |
| 60 | "timestamp": datetime.utcnow().isoformat() |
| 61 | }) |
| 62 | |
| 63 | async def finish(self): |
| 64 | """Finish streaming and send complete response.""" |
| 65 | complete_response = "".join(self.chunks) |
| 66 | duration = (datetime.utcnow() - self.start_time).total_seconds() |
| 67 | |
| 68 | await self.websocket.send_json({ |
| 69 | "type": MessageType.RESPONSE, |
| 70 | "message_id": self.message_id, |
| 71 | "content": complete_response, |
| 72 | "is_streaming": True, # Keep as streaming to avoid duplicate handling |
| 73 | "is_complete": True, |
| 74 | "duration": duration, |
| 75 | "timestamp": datetime.utcnow().isoformat() |
| 76 | }) |
| 77 | |
| 78 | |
| 79 | class EnhancedWebSocketManager: |
no outgoing calls
no test coverage detected