From 7d28c46a5d3b8c52ac44dee873a090e2e11ace9d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Aleix=20Conchillo=20Flaqu=C3=A9?= Date: Wed, 13 May 2026 19:15:20 -0700 Subject: [PATCH] Add tasks package with LLMTask, LLMContextTask, and proxy stubs Adds `pipecat.tasks.llm` with `LLMTask` (LLM pipeline + `@tool` collection + tool-call deferral via `PipelineFlushFrame`), `LLMContextTask` (LLM + `LLMContextAggregatorPair`), and the `@tool` decorator. Also includes `pipecat.tasks.proxy.websocket` client/server stubs that need a follow-up port to the new `BaseTask` lifecycle. --- src/pipecat/tasks/llm/__init__.py | 18 + src/pipecat/tasks/llm/llm_context_task.py | 119 ++++++ src/pipecat/tasks/llm/llm_task.py | 338 ++++++++++++++++++ src/pipecat/tasks/llm/tool_decorator.py | 62 ++++ src/pipecat/tasks/proxy/__init__.py | 17 + src/pipecat/tasks/proxy/websocket/__init__.py | 15 + src/pipecat/tasks/proxy/websocket/client.py | 214 +++++++++++ src/pipecat/tasks/proxy/websocket/server.py | 231 ++++++++++++ 8 files changed, 1014 insertions(+) create mode 100644 src/pipecat/tasks/llm/__init__.py create mode 100644 src/pipecat/tasks/llm/llm_context_task.py create mode 100644 src/pipecat/tasks/llm/llm_task.py create mode 100644 src/pipecat/tasks/llm/tool_decorator.py create mode 100644 src/pipecat/tasks/proxy/__init__.py create mode 100644 src/pipecat/tasks/proxy/websocket/__init__.py create mode 100644 src/pipecat/tasks/proxy/websocket/client.py create mode 100644 src/pipecat/tasks/proxy/websocket/server.py diff --git a/src/pipecat/tasks/llm/__init__.py b/src/pipecat/tasks/llm/__init__.py new file mode 100644 index 000000000..dc6304c09 --- /dev/null +++ b/src/pipecat/tasks/llm/__init__.py @@ -0,0 +1,18 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""LLM agent and tool decorator.""" + +from pipecat.tasks.llm.llm_context_task import LLMContextTask +from pipecat.tasks.llm.llm_task import LLMTask, LLMTaskActivationArgs +from pipecat.tasks.llm.tool_decorator import tool + +__all__ = [ + "LLMTask", + "LLMTaskActivationArgs", + "LLMContextTask", + "tool", +] diff --git a/src/pipecat/tasks/llm/llm_context_task.py b/src/pipecat/tasks/llm/llm_context_task.py new file mode 100644 index 000000000..265c7b537 --- /dev/null +++ b/src/pipecat/tasks/llm/llm_context_task.py @@ -0,0 +1,119 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""LLM task with a built-in `LLMContext` and aggregator pair. + +Provides the `LLMContextTask` class that extends `LLMTask` with a +self-contained conversation context, removing the need for subclasses +to manually wire `LLMContextAggregatorPair`. +""" + +from pipecat.bus import TaskBus +from pipecat.pipeline.pipeline import Pipeline +from pipecat.processors.aggregators.llm_context import LLMContext +from pipecat.processors.aggregators.llm_response_universal import ( + LLMAssistantAggregator, + LLMAssistantAggregatorParams, + LLMContextAggregatorPair, + LLMUserAggregator, + LLMUserAggregatorParams, +) +from pipecat.services.llm_service import LLMService +from pipecat.tasks.llm.llm_task import LLMTask + + +class LLMContextTask(LLMTask): + """LLM task that owns an `LLMContext` and a context aggregator pair. + + Useful for tasks that need to track their own conversation history, + typically workers that run their own LLM pipeline outside of a shared + transport pipeline. Subclasses do not need to instantiate the context + or aggregators themselves; the pipeline is built as + ``[user_aggregator, llm, assistant_aggregator]`` automatically. + + Example:: + + task = LLMContextTask( + "worker", + bus=bus, + llm=OpenAILLMService(...), + ) + + @task.assistant_aggregator.event_handler("on_assistant_turn_stopped") + async def _on_stopped(aggregator, message): + ... + """ + + def __init__( + self, + name: str, + *, + bus: TaskBus, + llm: LLMService, + active: bool = False, + bridged: tuple[str, ...] | None = None, + defer_tool_frames: bool = True, + context: LLMContext | None = None, + user_params: LLMUserAggregatorParams | None = None, + assistant_params: LLMAssistantAggregatorParams | None = None, + ): + """Initialize the LLMContextTask. + + Args: + name: Unique name for this task. + bus: The `TaskBus` for inter-task communication. + llm: The LLM service. + active: Whether the task starts active. Defaults to False. + bridged: Bridge configuration forwarded to ``PipelineTask``. + Pass ``()`` to wrap the pipeline with bus edges so it + can exchange frames with another bridged task. + defer_tool_frames: Whether to defer frames queued during + tool execution until all tools complete. Defaults to True. + context: Optional pre-built `LLMContext`. When omitted, a + fresh empty context is created. + user_params: Optional parameters for the user aggregator. + assistant_params: Optional parameters for the assistant + aggregator. + """ + self._context = context or LLMContext() + self._aggregators = LLMContextAggregatorPair( + self._context, + user_params=user_params, + assistant_params=assistant_params, + ) + + pipeline = Pipeline( + [ + self._aggregators.user(), + llm, + self._aggregators.assistant(), + ] + ) + + super().__init__( + name, + bus=bus, + llm=llm, + pipeline=pipeline, + active=active, + bridged=bridged, + defer_tool_frames=defer_tool_frames, + ) + + @property + def context(self) -> LLMContext: + """The `LLMContext` owned by this task.""" + return self._context + + @property + def user_aggregator(self) -> LLMUserAggregator: + """The user-side context aggregator.""" + return self._aggregators.user() + + @property + def assistant_aggregator(self) -> LLMAssistantAggregator: + """The assistant-side context aggregator.""" + return self._aggregators.assistant() diff --git a/src/pipecat/tasks/llm/llm_task.py b/src/pipecat/tasks/llm/llm_task.py new file mode 100644 index 000000000..8c770be2d --- /dev/null +++ b/src/pipecat/tasks/llm/llm_task.py @@ -0,0 +1,338 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""LLM task with tool registration. + +Provides the `LLMTask` class that extends `PipelineTask` with an LLM +pipeline and automatic tool registration. +""" + +import asyncio +import functools +from collections import deque +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +from pipecat.adapters.schemas.tools_schema import ToolsSchema +from pipecat.bus import TaskBus +from pipecat.frames.frames import ( + ControlFrame, + Frame, + FunctionCallResultProperties, + LLMMessagesAppendFrame, + LLMSetToolsFrame, + UninterruptibleFrame, +) +from pipecat.pipeline.base_task import TaskActivationArgs +from pipecat.pipeline.pipeline import Pipeline +from pipecat.pipeline.task import PipelineParams, PipelineTask +from pipecat.processors.frame_processor import FrameDirection +from pipecat.services.llm_service import LLMService +from pipecat.tasks.llm.tool_decorator import _collect_tools + +FunctionCallResultCallback = Callable[..., Any] + + +@dataclass +class PipelineFlushFrame(ControlFrame, UninterruptibleFrame): + """Probe frame used to flush all in-flight frames from the pipeline.""" + + pass + + +@dataclass +class LLMTaskActivationArgs(TaskActivationArgs): + """Activation arguments for LLM tasks. + + Attributes: + messages: LLM context messages to inject on activation. + run_llm: Whether to run the LLM after appending messages. + Defaults to True when ``messages`` is set. + """ + + messages: list | None = None + run_llm: bool | None = None + + +class LLMTask(PipelineTask): + """Task with an LLM pipeline and automatic tool registration. + + Methods decorated with ``@tool`` are registered as direct functions + on the LLM and tracked so that frames queued during tool execution + can be deferred until all tools complete. + + Example:: + + class MyTask(LLMTask): + @tool + async def my_function(self, params, arg: str): + ... + + task = MyTask("worker", bus=bus, llm=OpenAILLMService(api_key="...")) + """ + + def __init__( + self, + name: str, + *, + bus: TaskBus, + llm: LLMService, + active: bool = False, + bridged: tuple[str, ...] | None = None, + defer_tool_frames: bool = True, + ): + """Initialize the LLMTask. + + Args: + name: Unique name for this task. + bus: The `TaskBus` for inter-task communication. + llm: The LLM service. ``@tool`` decorated methods are + automatically registered on it. + active: Whether the task starts active. Defaults to False. + bridged: Bridge configuration forwarded to ``PipelineTask``. + Pass ``()`` to wrap the LLM pipeline with bus edge + processors so it can exchange frames with another + bridged task. + defer_tool_frames: Whether to defer frames queued during + tool execution until all tools complete. Defaults to True. + """ + # State referenced by tool wrapper closures; must be set before + # _register_tools wraps any handlers. + self._defer_tool_frames = defer_tool_frames + self._tool_call_inflight: int = 0 + self._deferred_frames: deque[tuple[Frame, FrameDirection]] = deque() + self._closing: bool = False + self._flush_done: asyncio.Event = asyncio.Event() + + self._llm = llm + self._register_tools(llm) + + pipeline = Pipeline([self._llm]) + + super().__init__( + pipeline, + name=name, + bus=bus, + bridged=bridged, + exclude_frames=(PipelineFlushFrame,), + enable_rtvi=False, + idle_timeout_secs=None, + params=PipelineParams( + enable_metrics=True, + enable_usage_metrics=True, + ), + ) + # PipelineTask's __init__ doesn't accept active; configure after. + self._active = active + self._pending_activation = active + + # Pipeline flush wiring: the probe frame travels up to the source + # and back down to the sink, signalling that all in-flight frames + # have been processed. + self.add_reached_upstream_filter((PipelineFlushFrame,)) + self.add_reached_downstream_filter((PipelineFlushFrame,)) + + @self.event_handler("on_frame_reached_upstream") + async def _on_flush_upstream(task, frame): + if isinstance(frame, PipelineFlushFrame): + await super().queue_frame(PipelineFlushFrame()) + + @self.event_handler("on_frame_reached_downstream") + async def _on_flush_downstream(task, frame): + if isinstance(frame, PipelineFlushFrame): + self._flush_done.set() + + @property + def llm(self) -> LLMService: + """The LLM service this task wraps.""" + return self._llm + + @property + def tool_call_active(self) -> bool: + """True when one or more ``@tool`` methods are executing.""" + return self._tool_call_inflight > 0 + + async def on_activated(self, args: dict | None) -> None: + """Configure the LLM with tools and activation messages. + + Args: + args: Optional activation arguments with messages to append. + """ + await super().on_activated(args) + + activation = LLMTaskActivationArgs.from_dict(args) if args else LLMTaskActivationArgs() + + tools = self.build_tools() + if tools: + await self.queue_frame(LLMSetToolsFrame(tools=ToolsSchema(standard_tools=tools))) + + if activation.messages: + run_llm = activation.run_llm if activation.run_llm is not None else True + await self.queue_frame( + LLMMessagesAppendFrame(messages=activation.messages, run_llm=run_llm) + ) + + async def queue_frame( + self, frame: Frame, direction: FrameDirection = FrameDirection.DOWNSTREAM + ) -> None: + """Queue a frame, deferring delivery until all tools complete (if any). + + When tool calls are in progress, the frame is held in an internal + queue and delivered automatically once the last tool finishes. + When no tools are active, the frame is queued immediately. + + Args: + frame: Any ``Frame`` to deliver. + direction: Direction the frame should travel. Defaults to + ``FrameDirection.DOWNSTREAM``. + """ + if self._defer_tool_frames and self._tool_call_inflight > 0 and not self._closing: + self._deferred_frames.append((frame, direction)) + else: + await super().queue_frame(frame, direction) + + def build_tools(self) -> list: + """Return the tools for this task's LLM. + + By default, returns all methods decorated with ``@tool``. + Override to provide additional or different tools. + + Returns: + List of tool functions. + """ + return _collect_tools(self) + + async def end( + self, + *, + reason: str | None = None, + messages: list | None = None, + result_callback: FunctionCallResultCallback | None = None, + ) -> None: + """Request a graceful end of the session. + + When called from a ``@tool`` handler, pass ``params.result_callback`` to + ensure any pending LLM output is fully delivered before ending. + + Args: + reason: Optional human-readable reason for ending. + messages: Optional LLM messages to inject and speak before + ending. The LLM runs immediately so the output is + delivered before the session terminates. + result_callback: The ``result_callback`` from + `FunctionCallParams`. + """ + self._closing = True + await self._finish_function_call(result_callback, messages=messages) + await super().end(reason=reason) + + async def handoff_to( + self, + task_name: str, + *, + activation_args: TaskActivationArgs | None = None, + messages: list | None = None, + result_callback: FunctionCallResultCallback | None = None, + ) -> None: + """Hand off to another task. + + When called from a ``@tool`` handler, pass ``params.result_callback`` to + ensure any pending LLM output is fully delivered before handing off. + + Args: + task_name: The name of the task to hand off to. + activation_args: Optional arguments forwarded to the target + task's ``on_activated`` handler. + messages: Optional LLM messages to inject and speak before + handing off. The LLM runs immediately so the output is + delivered before the transfer completes. + result_callback: The ``result_callback`` from `FunctionCallParams`. + """ + await self._finish_function_call(result_callback, messages=messages) + await super().handoff_to(task_name, activation_args=activation_args) + + async def process_deferred_tool_frames( + self, frames: list[tuple[Frame, FrameDirection]] + ) -> list[tuple[Frame, FrameDirection]]: + """Process deferred frames before they are flushed. + + Called after all in-flight tools complete, before the deferred + frames are queued into the pipeline. Override to inspect, modify, + reorder, or filter the frames. + + Args: + frames: The deferred frames collected during tool execution. + + Returns: + The frames to queue. Return the list as-is for default behavior. + """ + return frames + + def _register_tools(self, llm: LLMService) -> None: + """Register ``@tool`` methods on the LLM in place.""" + for method in _collect_tools(self): + tracked = self._track_tool_call(method) + llm.register_direct_function( + tracked, + cancel_on_interruption=method.cancel_on_interruption, + timeout_secs=method.timeout, + ) + + def _track_tool_call(self, method: Callable) -> Callable: + @functools.wraps(method) + async def wrapper(params, *args, **kwargs): + self._tool_call_inflight += 1 + try: + return await method(params, *args, **kwargs) + finally: + self._tool_call_inflight = max(0, self._tool_call_inflight - 1) + if not self._closing and self._tool_call_inflight == 0: + await self._flush_deferred_frames() + + return wrapper + + async def _flush_deferred_frames(self) -> None: + # Wait until the function result frame is really processed. + await self._flush_pipeline() + + frames = list(self._deferred_frames) + self._deferred_frames.clear() + for frame, direction in await self.process_deferred_tool_frames(frames): + await self.queue_frame(frame, direction) + + async def _flush_pipeline(self) -> None: + self._flush_done.clear() + # Bypass our deferral override by going to PipelineTask directly. + await super().queue_frame(PipelineFlushFrame(), FrameDirection.UPSTREAM) + await self._flush_done.wait() + + async def _finish_function_call( + self, + result_callback: FunctionCallResultCallback | None, + *, + messages: list | None = None, + ) -> None: + """Finish an in-progress function call before taking action. + + Optionally injects LLM messages and flushes the pipeline so the + output is fully delivered before handing off or ending. + + Args: + result_callback: The callback from `FunctionCallParams`, or None. + messages: Optional LLM messages to inject before completing. + """ + if messages: + await self.queue_frame(LLMMessagesAppendFrame(messages=messages, run_llm=True)) + await self._flush_pipeline() + + if not result_callback: + return + + await result_callback(None, properties=FunctionCallResultProperties(run_llm=False)) + + # Wait until the function result frame is really processed. + await self._flush_pipeline() diff --git a/src/pipecat/tasks/llm/tool_decorator.py b/src/pipecat/tasks/llm/tool_decorator.py new file mode 100644 index 000000000..6cd9f7076 --- /dev/null +++ b/src/pipecat/tasks/llm/tool_decorator.py @@ -0,0 +1,62 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""Decorator for marking agent methods as tools.""" + + +def tool(fn=None, *, cancel_on_interruption=True, timeout=None): + """Mark a method as a tool. + + On ``LLMTask`` subclasses, decorated methods are automatically + registered with the LLM via ``register_direct_function`` and + included in ``build_tools()``. + + Can be used with or without arguments:: + + @tool + async def my_tool(self, params, arg: str): + ... + + @tool(cancel_on_interruption=False, timeout=60) + async def my_tool(self, params, arg: str): + ... + + Args: + fn: The function to decorate (when used without arguments). + cancel_on_interruption: Whether to cancel this tool call when + an interruption occurs. Defaults to True. Only applies to + ``LLMTask`` tools. + timeout: Optional timeout in seconds for this tool call. + Defaults to None (uses the LLM service default). + """ + + def decorator(fn): + fn.is_agent_tool = True + fn.cancel_on_interruption = cancel_on_interruption + fn.timeout = timeout + return fn + + if fn is not None: + return decorator(fn) + return decorator + + +def _collect_tools(obj) -> list: + """Collect all ``@tool`` decorated bound methods from an object. + + Walks the MRO so that overridden methods in subclasses take + precedence over base-class definitions. + """ + seen: set[str] = set() + tools = [] + for cls in type(obj).__mro__: + for name, val in cls.__dict__.items(): + if name in seen: + continue + seen.add(name) + if callable(val) and getattr(val, "is_agent_tool", False): + tools.append(getattr(obj, name)) + return tools diff --git a/src/pipecat/tasks/proxy/__init__.py b/src/pipecat/tasks/proxy/__init__.py new file mode 100644 index 000000000..7d4fa4557 --- /dev/null +++ b/src/pipecat/tasks/proxy/__init__.py @@ -0,0 +1,17 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""Proxy agents for forwarding bus messages over network transports.""" + +from pipecat.tasks.proxy.websocket import ( + WebSocketProxyClientAgent, + WebSocketProxyServerAgent, +) + +__all__ = [ + "WebSocketProxyClientAgent", + "WebSocketProxyServerAgent", +] diff --git a/src/pipecat/tasks/proxy/websocket/__init__.py b/src/pipecat/tasks/proxy/websocket/__init__.py new file mode 100644 index 000000000..43ca6e122 --- /dev/null +++ b/src/pipecat/tasks/proxy/websocket/__init__.py @@ -0,0 +1,15 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""WebSocket proxy agents for forwarding bus messages.""" + +from pipecat.tasks.proxy.websocket.client import WebSocketProxyClientAgent +from pipecat.tasks.proxy.websocket.server import WebSocketProxyServerAgent + +__all__ = [ + "WebSocketProxyClientAgent", + "WebSocketProxyServerAgent", +] diff --git a/src/pipecat/tasks/proxy/websocket/client.py b/src/pipecat/tasks/proxy/websocket/client.py new file mode 100644 index 000000000..957b4a28f --- /dev/null +++ b/src/pipecat/tasks/proxy/websocket/client.py @@ -0,0 +1,214 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""WebSocket client proxy that forwards bus messages to a remote server.""" + +import asyncio + +from loguru import logger + +from pipecat.bus import BusMessage, BusTaskRegistryMessage, TaskBus +from pipecat.bus.messages import BusLocalMessage +from pipecat.bus.serializers import JSONMessageSerializer +from pipecat.bus.serializers.base import MessageSerializer +from pipecat.pipeline.base_task import BaseTask + +try: + import websockets + from websockets.asyncio.client import connect +except ModuleNotFoundError as e: + logger.error(f"Exception: {e}") + logger.error( + "In order to use WebSocketProxyClientAgent, you need to `pip install pipecat-ai-subagents[websocket]`." + ) + raise Exception(f"Missing module: {e}") + + +class WebSocketProxyClientAgent(BaseTask): + """Forwards bus messages to a remote agent over WebSocket. + + Connects to a WebSocket URL and forwards messages between a local + agent and a remote agent. Only messages targeted at the remote agent + are sent. Only messages targeted at the local agent are accepted. + + Event handlers available: + + - on_connected: Fired when the WebSocket connection is established. + - on_disconnected: Fired when the WebSocket connection is closed. + + Example:: + + proxy = WebSocketProxyClientAgent( + "proxy", + bus=runner.bus, + url="ws://remote-server:8765/ws", + remote_agent_name="worker", + local_agent_name="voice", + ) + + @proxy.event_handler("on_connected") + async def on_connected(agent, websocket): + logger.info("Connected to remote server") + + @proxy.event_handler("on_disconnected") + async def on_disconnected(agent, websocket): + logger.info("Disconnected from remote server") + + await runner.add_task(proxy) + """ + + def __init__( + self, + name: str, + *, + bus: TaskBus, + url: str, + remote_agent_name: str, + local_agent_name: str, + forward_messages: tuple[type[BusMessage], ...] = (), + headers: dict[str, str] | None = None, + serializer: MessageSerializer | None = None, + ): + """Initialize the WebSocketProxyClientAgent. + + Args: + name: Unique name for this agent. + bus: The `TaskBus` for inter-agent communication. + url: The WebSocket URL to connect to. + remote_agent_name: Name of the agent on the remote server. + Only messages targeted at this agent are forwarded. + local_agent_name: Name of the local agent that should + receive responses. Only inbound messages targeted at + this agent are accepted. + forward_messages: Additional message types to forward from + the local agent (e.g. ``(BusFrameMessage,)`` for frame + routing). These are forwarded based on source agent name + only, regardless of target. + headers: Optional HTTP headers sent with the WebSocket + handshake (e.g. for authentication). + serializer: Serializer for bus messages. Defaults to + `JSONMessageSerializer`. + """ + super().__init__(name, bus=bus) + self._url = url + self._remote_agent_name = remote_agent_name + self._local_agent_name = local_agent_name + self._forward_messages = forward_messages + self._headers = headers or {} + self._serializer = serializer or JSONMessageSerializer() + self._ws = None + self._receive_task: asyncio.Task | None = None + + self._register_event_handler("on_connected") + self._register_event_handler("on_disconnected") + + async def cleanup(self): + """Cancel the receive loop task and release resources.""" + await super().cleanup() + if self._receive_task: + await self.cancel_task(self._receive_task) + self._receive_task = None + + async def on_activated(self, args: dict | None) -> None: + """Connect to the remote WebSocket server.""" + await super().on_activated(args) + + logger.debug(f"Agent '{self}': connecting to {self._url}") + + self._ws = await connect(self._url, additional_headers=self._headers) + + logger.debug(f"Agent '{self}': connected to {self._url}") + + await self._call_event_handler("on_connected", self._ws) + + self._receive_task = self.create_task(self._receive_loop(), f"{self.name}::ws_receive") + + # Schedule task right away. + await asyncio.sleep(0) + + async def on_bus_message(self, message: BusMessage) -> None: + """Forward messages targeted at the remote agent. + + Args: + message: The bus message to process. + """ + await super().on_bus_message(message) + + if not self._ws: + return + + if isinstance(message, BusLocalMessage): + return + + # Forward targeted messages to the remote agent + if message.target == self._remote_agent_name: + await self._send_ws(message) + # Forward additional message types from the local agent + elif isinstance(message, self._forward_messages): + if message.source == self._local_agent_name: + await self._send_ws(message) + + async def _stop(self) -> None: + """Close the WebSocket connection and stop.""" + if self._ws: + await self._ws.close() + logger.debug(f"Agent '{self}': WebSocket connection closed") + await super()._stop() + + async def _send_ws(self, message: BusMessage) -> None: + """Serialize and send a message over the WebSocket.""" + if not self._ws: + return + try: + data = self._serializer.serialize(message) + await self._ws.send(data) + logger.trace(f"Agent '{self}': sent {message}") + except websockets.exceptions.ConnectionClosed: + logger.warning(f"Agent '{self}': connection closed, stopping forwarding") + ws = self._ws + self._ws = None + await self._call_event_handler("on_disconnected", ws) + + async def _receive_loop(self) -> None: + """Read messages from the WebSocket and put them on the local bus.""" + try: + async for data in self._ws: + try: + message = self._serializer.deserialize(data) + if not message: + continue + + # Accept registry messages (target=None) for agent discovery + if isinstance(message, BusTaskRegistryMessage): + logger.trace( + f"Agent '{self}': received registry from remote: {message.agents}" + ) + await self.send_message(message) + continue + + # Accept additional message types (e.g. BusFrameMessage) + if self._forward_messages and isinstance(message, self._forward_messages): + logger.trace(f"Agent '{self}': received {message} from remote") + await self.send_message(message) + continue + + # Only accept other messages targeted at the local agent + if message.target != self._local_agent_name: + logger.warning( + f"Agent '{self}': dropped inbound message with " + f"unexpected target '{message.target}'" + ) + continue + + logger.trace(f"Agent '{self}': received {message} from remote") + await self.send_message(message) + except Exception: + logger.exception(f"Agent '{self}': failed to deserialize remote message") + except websockets.exceptions.ConnectionClosed: + logger.warning(f"Agent '{self}': WebSocket connection closed") + ws = self._ws + self._ws = None + await self._call_event_handler("on_disconnected", ws) diff --git a/src/pipecat/tasks/proxy/websocket/server.py b/src/pipecat/tasks/proxy/websocket/server.py new file mode 100644 index 000000000..ea6b61a84 --- /dev/null +++ b/src/pipecat/tasks/proxy/websocket/server.py @@ -0,0 +1,231 @@ +# +# Copyright (c) 2026, Daily +# +# SPDX-License-Identifier: BSD 2-Clause License +# + +"""WebSocket server proxy that receives bus messages from a remote client.""" + +import asyncio + +from loguru import logger + +from pipecat.bus import BusMessage, BusTaskRegistryMessage, TaskBus +from pipecat.bus.messages import BusLocalMessage +from pipecat.bus.serializers import JSONMessageSerializer +from pipecat.bus.serializers.base import MessageSerializer +from pipecat.pipeline.base_task import BaseTask +from pipecat.registry.types import TaskReadyData, TaskRegistryEntry + +try: + from starlette.websockets import WebSocket, WebSocketDisconnect, WebSocketState +except ModuleNotFoundError as e: + logger.error(f"Exception: {e}") + logger.error( + "In order to use WebSocketProxyServerAgent, you need to `pip install pipecat-ai-subagents[websocket]`." + ) + raise Exception(f"Missing module: {e}") + + +class WebSocketProxyServerAgent(BaseTask): + """Receives bus messages from a remote client over WebSocket. + + Accepts a FastAPI/Starlette WebSocket connection and forwards + messages between the remote client and a local agent. Only messages + from the local agent targeted at the remote agent are sent. Only + inbound messages targeted at the local agent are accepted. + + Event handlers available: + + - on_client_connected: Fired when the WebSocket client connects and the proxy is ready. + - on_client_disconnected: Fired when the WebSocket client disconnects. + + Example:: + + @app.websocket("/ws") + async def websocket_endpoint(websocket: WebSocket): + await websocket.accept() + proxy = WebSocketProxyServerAgent( + "gateway", + bus=runner.bus, + websocket=websocket, + agent_name="worker", + remote_agent_name="voice", + ) + + @proxy.event_handler("on_client_connected") + async def on_client_connected(agent, websocket): + logger.info("Client connected") + + @proxy.event_handler("on_client_disconnected") + async def on_client_disconnected(agent, websocket): + logger.info("Client disconnected") + + await runner.add_task(proxy) + """ + + def __init__( + self, + name: str, + *, + bus: TaskBus, + websocket: WebSocket, + agent_name: str, + remote_agent_name: str, + forward_messages: tuple[type[BusMessage], ...] = (), + serializer: MessageSerializer | None = None, + ): + """Initialize the WebSocketProxyServerAgent. + + Args: + name: Unique name for this agent. + bus: The `TaskBus` for inter-agent communication. + websocket: An accepted FastAPI/Starlette WebSocket connection. + agent_name: Name of the local agent to route messages to/from. + Only messages from this agent are forwarded to the client. + remote_agent_name: Name of the agent on the remote client. + Only outbound messages targeted at this agent are sent. + Only inbound messages targeted at the local agent are accepted. + forward_messages: Additional message types to forward from + the local agent (e.g. ``(BusFrameMessage,)`` for frame + routing). These are forwarded based on source agent name + only, regardless of target. + serializer: Serializer for bus messages. Defaults to + `JSONMessageSerializer`. + """ + super().__init__(name, bus=bus) + self._ws = websocket + self._agent_name = agent_name + self._remote_agent_name = remote_agent_name + self._forward_messages = forward_messages + self._serializer = serializer or JSONMessageSerializer() + self._receive_task: asyncio.Task | None = None + + self._register_event_handler("on_client_connected") + self._register_event_handler("on_client_disconnected") + + async def cleanup(self): + """Cancel the receive loop task and release resources.""" + await super().cleanup() + + if self._receive_task: + await self.cancel_task(self._receive_task) + self._receive_task = None + + async def on_ready(self) -> None: + """Start receiving messages from the WebSocket and watch the local agent.""" + await super().on_ready() + + logger.debug(f"Agent '{self}': WebSocket proxy server ready") + + await self._call_event_handler("on_client_connected", self._ws) + + self._receive_task = self.create_task(self._receive_loop(), f"{self.name}::ws_receive") + + # Schedule task right away. + await asyncio.sleep(0) + + # Watch the local agent so we can notify the remote side when it's ready + await self.watch_task(self._agent_name) + + async def on_task_ready(self, data: TaskReadyData) -> None: + """Notify the remote client that the local agent is ready.""" + if not self._ws: + return + + if data.agent_name != self._agent_name: + return + + logger.debug(f"Agent '{self}': local agent '{self._agent_name}' ready, notifying remote") + + try: + msg = BusTaskRegistryMessage( + source=self.name, + runner=data.runner, + agents=[TaskRegistryEntry(name=self._agent_name)], + ) + + await self._send_ws(msg) + except Exception: + logger.exception(f"Agent '{self}': failed to send registry to remote") + + async def on_bus_message(self, message: BusMessage) -> None: + """Forward messages from the local agent to the remote client. + + Args: + message: The bus message to process. + """ + await super().on_bus_message(message) + + if not self._ws: + return + + if isinstance(message, BusLocalMessage): + return + + if message.source != self._agent_name: + return + + # Forward targeted messages from the local agent to the remote agent + if message.target == self._remote_agent_name: + await self._send_ws(message) + # Forward additional message types from the local agent + elif isinstance(message, self._forward_messages): + await self._send_ws(message) + + async def _stop(self) -> None: + """Close the WebSocket connection and stop.""" + if self._ws and self._ws.client_state == WebSocketState.CONNECTED: + await self._ws.close() + logger.debug(f"Agent '{self}': WebSocket connection closed") + await super()._stop() + + async def _send_ws(self, message: BusMessage) -> None: + """Serialize and send a message over the WebSocket.""" + if not self._ws: + return + try: + data = self._serializer.serialize(message) + await self._ws.send_bytes(data) + logger.trace(f"Agent '{self}': sent {message}") + except (WebSocketDisconnect, Exception): + logger.warning(f"Agent '{self}': connection closed, stopping forwarding") + ws = self._ws + self._ws = None + await self._call_event_handler("on_client_disconnected", ws) + + async def _receive_loop(self) -> None: + """Read messages from the WebSocket and put them on the local bus.""" + try: + while True: + data = await self._ws.receive_bytes() + try: + message = self._serializer.deserialize(data) + if not message: + continue + + # Accept additional message types (e.g. BusFrameMessage) + if self._forward_messages and isinstance(message, self._forward_messages): + logger.trace(f"Agent '{self}': received {message} from client") + await self.send_message(message) + continue + + # Only accept other messages targeted at the local agent + if message.target != self._agent_name: + logger.warning( + f"Agent '{self}': dropped inbound message with " + f"unexpected target '{message.target}'" + ) + continue + + logger.trace(f"Agent '{self}': received {message} from client") + await self.send_message(message) + except Exception: + logger.exception(f"Agent '{self}': failed to deserialize client message") + except WebSocketDisconnect: + logger.warning(f"Agent '{self}': client disconnected") + ws = self._ws + self._ws = None + await self._call_event_handler("on_client_disconnected", ws) + except asyncio.CancelledError: + pass