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.
This commit is contained in:
Aleix Conchillo Flaqué
18
src/pipecat/tasks/llm/__init__.py
Normal file
18
src/pipecat/tasks/llm/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
119
src/pipecat/tasks/llm/llm_context_task.py
Normal file
119
src/pipecat/tasks/llm/llm_context_task.py
Normal file
@@ -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()
|
||||
338
src/pipecat/tasks/llm/llm_task.py
Normal file
338
src/pipecat/tasks/llm/llm_task.py
Normal file
@@ -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()
|
||||
62
src/pipecat/tasks/llm/tool_decorator.py
Normal file
62
src/pipecat/tasks/llm/tool_decorator.py
Normal file
@@ -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
|
||||
17
src/pipecat/tasks/proxy/__init__.py
Normal file
17
src/pipecat/tasks/proxy/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
15
src/pipecat/tasks/proxy/websocket/__init__.py
Normal file
15
src/pipecat/tasks/proxy/websocket/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
214
src/pipecat/tasks/proxy/websocket/client.py
Normal file
214
src/pipecat/tasks/proxy/websocket/client.py
Normal file
@@ -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)
|
||||
231
src/pipecat/tasks/proxy/websocket/server.py
Normal file
231
src/pipecat/tasks/proxy/websocket/server.py
Normal file
@@ -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
|
||||
Reference in New Issue
Block a user