84 lines
3.1 KiB
Python
84 lines
3.1 KiB
Python
"""A long-lived owner for native Pipecat jobs, independent of media calls."""
|
|
|
|
import asyncio
|
|
from collections.abc import Awaitable, Callable
|
|
|
|
from pipecat.pipeline.job_context import JobParams, JobStatus
|
|
from pipecat.pipeline.job_decorator import job
|
|
from pipecat.workers.base_worker import BaseWorker
|
|
from pipecat.workers.runner import WorkerRunner
|
|
|
|
JobCallback = Callable[[object], Awaitable[None]]
|
|
Execute = Callable[[object, BaseWorker], Awaitable[dict]]
|
|
|
|
|
|
class CoordinatorWorker(BaseWorker):
|
|
def __init__(self, on_update: JobCallback, on_result: JobCallback):
|
|
super().__init__("background-task-coordinator")
|
|
self._update = on_update
|
|
self._result = on_result
|
|
|
|
async def on_job_update(self, message):
|
|
await self._update(message)
|
|
|
|
async def on_job_response(self, message):
|
|
await self._result(message)
|
|
|
|
|
|
class RemoteAgentWorker(BaseWorker):
|
|
def __init__(self, execute: Execute, cancel: JobCallback):
|
|
super().__init__("background-task-executor")
|
|
self._execute = execute
|
|
self._cancel = cancel
|
|
|
|
@job(name="execute", sequential=True)
|
|
async def execute(self, message):
|
|
try:
|
|
result = await self._execute(message, self)
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception:
|
|
# Error details belong in the repository; don't publish credentials.
|
|
await self.send_job_response(message.job_id, status=JobStatus.ERROR)
|
|
else:
|
|
await self.send_job_response(message.job_id, result)
|
|
|
|
async def on_job_cancelled(self, message):
|
|
await self._cancel(message)
|
|
|
|
|
|
class NativeTaskRuntime:
|
|
"""Only wiring: scheduling, deadlines and cancellation belong to Pipecat."""
|
|
|
|
def __init__(self, execute: Execute, cancel: JobCallback,
|
|
on_update: JobCallback, on_result: JobCallback):
|
|
self.coordinator = CoordinatorWorker(on_update, on_result)
|
|
self.executor = RemoteAgentWorker(execute, cancel)
|
|
self.runner: WorkerRunner | None = None
|
|
self._run: asyncio.Task | None = None
|
|
|
|
async def start(self):
|
|
self.runner = WorkerRunner(handle_sigint=False, check_dangling_tasks=False)
|
|
ready = asyncio.Event()
|
|
async def on_ready(_runner):
|
|
ready.set()
|
|
self.runner.add_event_handler("on_ready", on_ready)
|
|
await self.runner.add_workers(self.coordinator, self.executor)
|
|
self._run = asyncio.create_task(self.runner.run(auto_end=False), name="background-jobs")
|
|
await asyncio.wait_for(ready.wait(), 10)
|
|
|
|
async def submit(self, task_id: str, timeout: float) -> str:
|
|
return await self.coordinator.request_job(
|
|
self.executor.name,
|
|
params=JobParams(name="execute", payload={"task_id": task_id}, timeout=timeout),
|
|
)
|
|
|
|
async def cancel(self, job_id: str):
|
|
return await self.coordinator.request_cancel_job_group(job_id, reason="user requested stop")
|
|
|
|
async def close(self):
|
|
if self.runner:
|
|
await self.runner.cancel(reason="application shutdown")
|
|
if self._run:
|
|
await self._run
|