40 lines
1.5 KiB
Python
40 lines
1.5 KiB
Python
# Copyright 2026 zhaoxi826
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
|
|
from pretor.utils.ray_hook import ray_actor_hook
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
|
|
|
|
workflow_router = APIRouter(prefix="/api/v1/workflow", tags=["workflow"])
|
|
|
|
@workflow_router.websocket("/ws/{trace_id}")
|
|
async def get_workflow(websocket: WebSocket, trace_id: str):
|
|
await websocket.accept()
|
|
global_state_machine = ray_actor_hook("global_state_machine")
|
|
try:
|
|
while True:
|
|
await websocket.send(await global_state_machine.get_workflow.remote(trace_id))
|
|
await websocket.send_text(await global_state_machine.get_pending.remote(trace_id))
|
|
response = await websocket.receive_text()
|
|
await global_state_machine.put_received(trace_id, response)
|
|
except WebSocketDisconnect:
|
|
pass
|
|
except RuntimeError as e:
|
|
if "closed" not in str(e) and "GeneratorExit" not in str(e):
|
|
raise
|
|
except Exception:
|
|
pass
|
|
|