79 lines
2.1 KiB
Python
79 lines
2.1 KiB
Python
|
import json
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
from fastapi import WebSocket
|
||
|
|
|
||
|
|
|
||
|
|
class ConnectionManager:
|
||
|
|
def __init__(self):
|
||
|
|
self._connections: dict[str, list[WebSocket]] = {}
|
||
|
|
|
||
|
|
async def connect(self, websocket: WebSocket, session_id: str) -> None:
|
||
|
|
await websocket.accept()
|
||
|
|
if session_id not in self._connections:
|
||
|
|
self._connections[session_id] = []
|
||
|
|
self._connections[session_id].append(websocket)
|
||
|
|
|
||
|
|
def disconnect(self, websocket: WebSocket, session_id: str) -> None:
|
||
|
|
if session_id in self._connections:
|
||
|
|
self._connections[session_id].remove(websocket)
|
||
|
|
if not self._connections[session_id]:
|
||
|
|
del self._connections[session_id]
|
||
|
|
|
||
|
|
async def send_message(self, session_id: str, message: dict[str, Any]) -> None:
|
||
|
|
if session_id in self._connections:
|
||
|
|
for connection in self._connections[session_id]:
|
||
|
|
await connection.send_json(message)
|
||
|
|
|
||
|
|
async def broadcast(self, message: dict[str, Any]) -> None:
|
||
|
|
for connections in self._connections.values():
|
||
|
|
for connection in connections:
|
||
|
|
await connection.send_json(message)
|
||
|
|
|
||
|
|
|
||
|
|
manager = ConnectionManager()
|
||
|
|
|
||
|
|
|
||
|
|
async def send_progress(
|
||
|
|
session_id: str,
|
||
|
|
current_agent: str | None,
|
||
|
|
status: str,
|
||
|
|
progress: float,
|
||
|
|
message: str,
|
||
|
|
) -> None:
|
||
|
|
await manager.send_message(
|
||
|
|
session_id,
|
||
|
|
{
|
||
|
|
"type": "progress",
|
||
|
|
"session_id": session_id,
|
||
|
|
"data": {
|
||
|
|
"current_agent": current_agent,
|
||
|
|
"status": status,
|
||
|
|
"progress": progress,
|
||
|
|
"message": message,
|
||
|
|
},
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
async def send_result(session_id: str, result: dict[str, Any]) -> None:
|
||
|
|
await manager.send_message(
|
||
|
|
session_id,
|
||
|
|
{
|
||
|
|
"type": "result",
|
||
|
|
"session_id": session_id,
|
||
|
|
"data": result,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
async def send_error(session_id: str, error: str) -> None:
|
||
|
|
await manager.send_message(
|
||
|
|
session_id,
|
||
|
|
{
|
||
|
|
"type": "error",
|
||
|
|
"session_id": session_id,
|
||
|
|
"data": {"error": error},
|
||
|
|
},
|
||
|
|
)
|