ai-video-fullstack/backend/services/background_tasks/native.py

"""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