| | |
|---|
| | | import asyncio |
|---|
| | | import websockets |
|---|
| | | from sqlalchemy.orm import Session |
|---|
| | | from app.api import get_current_user_websocket, ResponseList, get_current_user |
|---|
| | | from app.api import get_current_user_websocket, ResponseList, get_current_user, format_file_url, process_files |
|---|
| | | from app.config.config import settings |
|---|
| | | from app.models.agent_model import AgentModel, AgentType |
|---|
| | | from app.models.base_model import get_db |
|---|
| | | from app.models.user_model import UserModel |
|---|
| | | from app.service.bisheng import BishengService |
|---|
| | | from app.service.token import get_bisheng_token |
|---|
| | | from app.service.service_token import get_bisheng_token |
|---|
| | | |
|---|
| | | router = APIRouter() |
|---|
| | | |
|---|
| | |
|---|
| | | return websocket.send_json(ret) |
|---|
| | | |
|---|
| | | token = get_bisheng_token(db, current_user.id) |
|---|
| | | service_uri = f"{settings.bisheng_websocket_url}/api/v1/chat/{agent_id}?type=L1&t=&chat_id={chat_id}" |
|---|
| | | service_uri = f"{settings.sgb_websocket_url}/api/v1/chat/{agent_id}?type=L1&t=&chat_id={chat_id}" |
|---|
| | | headers = {'cookie': f"access_token_cookie={token};"} |
|---|
| | | |
|---|
| | | await websocket.accept() |
|---|
| | |
|---|
| | | |
|---|
| | | # 监听毕昇发来的消息并转发给客户端 |
|---|
| | | async def forward_to_client(): |
|---|
| | | last_message = "step" |
|---|
| | | while True: |
|---|
| | | message = await service_websocket.recv() |
|---|
| | | print(f"Received from bisheng: {message}") |
|---|
| | | data = json.loads(message) |
|---|
| | | files = data.get("files", []) |
|---|
| | | steps = data.get("intermediate_steps", "") |
|---|
| | | if len(files) != 0 or steps != "" or data["type"] == "close": |
|---|
| | | msg = data.get("message", "") |
|---|
| | | category = data.get("category", "") |
|---|
| | | |
|---|
| | | if len(files) != 0 or (msg and category != "answer") or data["type"] == "close": |
|---|
| | | if data["type"] == "close": |
|---|
| | | t = "close" |
|---|
| | | else: |
|---|
| | | t = "stream" |
|---|
| | | result = {"step_message": steps, "type": t, "files": files} |
|---|
| | | process_files(files, agent_id) |
|---|
| | | result = {"message": msg, "type": t, "files": files} |
|---|
| | | await websocket.send_json(result) |
|---|
| | | print(f"Forwarded to client, {chat_id}: {result}") |
|---|
| | | elif steps and last_message == "step": |
|---|
| | | result = {"step_message": steps, "type": "stream", "files": files} |
|---|
| | | await websocket.send_json(result) |
|---|
| | | |
|---|
| | | last_message = "message" if msg else "step" |
|---|
| | | |
|---|
| | | # 启动两个任务，分别处理客户端和服务端的消息 |
|---|
| | | tasks = [ |
|---|
| | |
|---|
| | | task.cancel() |
|---|
| | | try: |
|---|
| | | await task |
|---|
| | | except asyncio.CancelledError: |
|---|
| | | except asyncio.CancelledError as e: |
|---|
| | | print(f"asyncio CancelledError: {e}") |
|---|
| | | pass |
|---|
| | | |
|---|
| | | except WebSocketDisconnect: |
|---|
| | | print(f"Client {chat_id} disconnected") |
|---|
| | | except WebSocketDisconnect as e: |
|---|
| | | print(f"WebSocket connection closed with code {e.code}: {e.reason}") |
|---|
| | | await websocket.close() |
|---|
| | | await service_websocket.close() |
|---|
| | | except Exception as e: |
|---|
| | | print(f"Exception occurred: {e}") |
|---|
| | | finally: |
|---|
| | | print("Cleaning up resources of bisheng report") |
|---|
| | | # 取消所有任务 |
|---|
| | | for task in tasks: |
|---|
| | | if not task.done(): |
|---|
| | | task.cancel() |
|---|
| | | try: |
|---|
| | | await task |
|---|
| | | except asyncio.CancelledError: |
|---|
| | | pass |
|---|
| | | |
|---|
| | | |
|---|
| | | @router.get("/variables/list", response_model=ResponseList) |
|---|
| | |
|---|
| | | agent = db.query(AgentModel).filter(AgentModel.id == agent_id).first() |
|---|
| | | if not agent: |
|---|
| | | return ResponseList(code=404, msg="Agent not found") |
|---|
| | | bisheng_service = BishengService(base_url=settings.bisheng_base_url) |
|---|
| | | bisheng_service = BishengService(base_url=settings.sgb_base_url) |
|---|
| | | try: |
|---|
| | | token = get_bisheng_token(db, current_user.id) |
|---|
| | | result = await bisheng_service.variable_list(token, agent_id) |
|---|
