websocket 简单service框架可以接收和传输消息

import asyncio
import time

from fastapi import FastAPI, WebSocket
import json
from pydantic import BaseModel
from typing import Set
from test_post_data import *

session_id = "your_session_id"
finishtraceId = False

class ContenVideo(BaseModel):
    content:str

class SessionID(BaseModel):
    session_id:str

app = FastAPI()

from fastapi.middleware.cors import CORSMiddleware
# 配置 CORS
app.add_middleware(
CORSMiddleware,
allow_origins=["*"], # 允许所有来源
allow_credentials=True,
allow_methods=["*"], # 允许所有方法
allow_headers=["*"], # 允许所有头部
)

class MyService:
    def __init__(self):
        self.websocket = None
        # 存储 WebSocket 连接的集合
        self.connections: Set[WebSocket] = set()

    async def close_websocket(self):
        if self.websocket:
            await self.websocket.close()
            print("Disconnected from WebSocket server")
            # 将连接从集合中移除
            self.connections.remove(self.websocket)

    async def websocket_endpoint(self, websocket: WebSocket):
        self.websocket = websocket
        await self.websocket.accept()
        # 将连接添加到集合中
        self.connections.add(self.websocket)
        while True:
            data = await self.websocket.receive_text()
            if data:
                print(data)
                if 'finish' in str(data):
                    # 处理接收到的消息
                    global finishtraceId
                    finishtraceId = True

            # time.sleep(1)
            # if data:
            #     # 处理接收到的消息
            #     await self.websocket_send({"message": "Received: " + data})

    async def websocket_send(self, item):
        if self.websocket is None or self.websocket.client_state !=  "WebSocketState.CONNECTED":
            # 客户端未连接或已断开，直接返回
            return

        json_data = json.dumps(item, ensure_ascii=False)
        await self.websocket.send_json(json_data)

    async def broadcast_message(self, item: dict):
        for connection in self.connections.copy():
            if str(connection.client_state) != "WebSocketState.CONNECTED":
                # 连接已断开，从集合中移除
                self.connections.remove(connection)
            else:
                json_data = json.dumps(item, ensure_ascii=False)
                await connection.send_json(json_data)

service = MyService()

@app.websocket("/ws_chat")
async def websocket_route(websocket: WebSocket):
    await service.websocket_endpoint(websocket)


@app.post("/trigger_status")
async def entrance__status_trigger(query:ContenVideo):
    sned_status={}
    sned_status['status']=query.content
    await service.broadcast_message(sned_status)
    return {"message": "已触发status发送"}



if __name__ == "__main__":
    import uvicorn

    uvicorn.run('new_service:app', host='0.0.0.0', port=7999, log_level='info')