huangfu
a6c2027caa
- 添加 start_all.sh 一键启动脚本，支持启动llama-server和FastAPI服务 - 修改启动脚本使用venv虚拟环境替代conda环境 - 更新README.md，添加一键启动脚本使用说明 - 更新py_tree_generator.py，添加final_prompt返回字段 - 禁用Qwen3模型的思考功能 - 添加RAG检索结果的终端打印 - 移除ROS2相关代码（ros2_client.py已删除）
50 lines
1.8 KiB
Python
50 lines
1.8 KiB
Python
import asyncio
|
|
from typing import List
|
|
from fastapi import WebSocket
|
|
import logging
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self.active_connections: List[WebSocket] = []
|
|
self.loop: asyncio.AbstractEventLoop | None = None
|
|
|
|
def set_loop(self, loop: asyncio.AbstractEventLoop):
|
|
"""Sets the asyncio event loop."""
|
|
self.loop = loop
|
|
|
|
async def connect(self, websocket: WebSocket):
|
|
await websocket.accept()
|
|
self.active_connections.append(websocket)
|
|
|
|
def disconnect(self, websocket: WebSocket):
|
|
self.active_connections.remove(websocket)
|
|
|
|
def broadcast(self, message: str):
|
|
"""
|
|
Thread-safely broadcasts a message to all active WebSocket connections.
|
|
This method is designed to be called from a different thread.
|
|
(Note: ROS2 callback support has been removed as the project is decoupled from ROS2)
|
|
"""
|
|
if not self.loop:
|
|
logging.error("Event loop not set in ConnectionManager. Cannot broadcast.")
|
|
return
|
|
|
|
# Schedule the coroutine to be executed in the event loop
|
|
self.loop.call_soon_threadsafe(self._broadcast_in_loop, message)
|
|
|
|
def _broadcast_in_loop(self, message: str):
|
|
"""
|
|
Helper to run the broadcast coroutine in the correct event loop.
|
|
"""
|
|
asyncio.ensure_future(self._broadcast_async(message), loop=self.loop)
|
|
|
|
async def _broadcast_async(self, message: str):
|
|
"""
|
|
The actual async method that sends messages.
|
|
"""
|
|
tasks = [connection.send_text(message) for connection in self.active_connections]
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
# Create a single instance of the manager to be used across the application
|
|
websocket_manager = ConnectionManager()
|