mirror of
https://github.com/docling-project/docling-serve.git
synced 2025-11-29 08:33:50 +00:00
77 lines
2.9 KiB
Python
77 lines
2.9 KiB
Python
from fastapi import WebSocket
|
|
|
|
from docling_jobkit.datamodel.task_meta import TaskStatus
|
|
from docling_jobkit.orchestrators.base_notifier import BaseNotifier
|
|
from docling_jobkit.orchestrators.base_orchestrator import BaseOrchestrator
|
|
|
|
from docling_serve.datamodel.responses import (
|
|
MessageKind,
|
|
TaskStatusResponse,
|
|
WebsocketMessage,
|
|
)
|
|
|
|
|
|
class WebsocketNotifier(BaseNotifier):
|
|
def __init__(self, orchestrator: BaseOrchestrator):
|
|
super().__init__(orchestrator)
|
|
self.task_subscribers: dict[str, set[WebSocket]] = {}
|
|
|
|
async def add_task(self, task_id: str):
|
|
self.task_subscribers[task_id] = set()
|
|
|
|
async def remove_task(self, task_id: str):
|
|
if task_id in self.task_subscribers:
|
|
for websocket in self.task_subscribers[task_id]:
|
|
await websocket.close()
|
|
|
|
del self.task_subscribers[task_id]
|
|
|
|
async def notify_task_subscribers(self, task_id: str):
|
|
if task_id not in self.task_subscribers:
|
|
raise RuntimeError(f"Task {task_id} does not have a subscribers list.")
|
|
|
|
try:
|
|
# Get task status from Redis or RQ directly instead of in-memory registry
|
|
task = await self.orchestrator.task_status(task_id=task_id)
|
|
task_queue_position = await self.orchestrator.get_queue_position(task_id)
|
|
msg = TaskStatusResponse(
|
|
task_id=task.task_id,
|
|
task_type=task.task_type,
|
|
task_status=task.task_status,
|
|
task_position=task_queue_position,
|
|
task_meta=task.processing_meta,
|
|
)
|
|
for websocket in self.task_subscribers[task_id]:
|
|
await websocket.send_text(
|
|
WebsocketMessage(
|
|
message=MessageKind.UPDATE, task=msg
|
|
).model_dump_json()
|
|
)
|
|
if task.is_completed():
|
|
await websocket.close()
|
|
except Exception as e:
|
|
# Log the error but don't crash the notifier
|
|
import logging
|
|
|
|
_log = logging.getLogger(__name__)
|
|
_log.error(f"Error notifying subscribers for task {task_id}: {e}")
|
|
|
|
async def notify_queue_positions(self):
|
|
"""Notify all subscribers of pending tasks about queue position updates."""
|
|
for task_id in self.task_subscribers.keys():
|
|
try:
|
|
# Check task status directly from Redis or RQ
|
|
task = await self.orchestrator.task_status(task_id)
|
|
|
|
# Notify only pending tasks
|
|
if task.task_status == TaskStatus.PENDING:
|
|
await self.notify_task_subscribers(task_id)
|
|
except Exception as e:
|
|
# Log the error but don't crash the notifier
|
|
import logging
|
|
|
|
_log = logging.getLogger(__name__)
|
|
_log.error(
|
|
f"Error checking task {task_id} status for queue position notification: {e}"
|
|
)
|