"""WebSocket publish helpers for Agency Channels.
Wire protocol and abstractions (Monitor, Messenger, Catalog) live in
``agency.channel.protocol``. Consumers live under ``agency.channel.consumer``.
This module exposes ``publish_event``, group helpers, and ``Publisher``.
"""
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
from agency.models import Run
from datetime import datetime
from typing import Any, Optional
from asgiref.sync import async_to_sync
from channels import DEFAULT_CHANNEL_LAYER
from channels.layers import BaseChannelLayer, get_channel_layer
from agency.utils import serializable
[docs]
def publish_event(
event_type: str,
groups: list[str],
**kwargs) -> None:
"""Publish an event to the task's channel group and to the general tasks-all group.
Args:
- event_type: The type of event (task_started, task_success, etc.)
- groups: The list of groups to publish the event to
- **kwargs: The event data to be passed to the consumer handler
"""
try:
channel_layer = get_channel_layer()
if channel_layer:
event_data = {
'type': event_type,
'ts': datetime.now().isoformat(),
**kwargs
}
if groups:
for group in groups:
async_to_sync(channel_layer.group_send)(group, event_data)
except Exception as e:
# Don't fail the task if event publishing fails
import logging
logging.getLogger(__name__).error(f"Failed to publish task event: {e}")
[docs]
def channel_log(message: str, data: Optional[Any] = None, obj: Optional[Any] = None, level: str = "info") -> None:
"""Publish a log event to the logs channel group.
Notes:
- `type="log"` is required for Django Channels dispatch to `LogConsumer.log(...)`.
- The browser-level "severity" should be sent as `level` (info|warning|error|debug).
"""
payload = data if data is not None else obj
publish_event(
"log",
["logs"],
level=level,
message=message,
data=serializable(payload) if payload is not None else None,
)
[docs]
def channel_run(event_type: str, run: "Run", **kwargs) -> None:
"""Publish a thread event to the thread channel groups.
Args:
event_type: The type of event (tool_call, tool_result, model_request, etc.)
run: Run instance to get thread_id
**kwargs: Event data to be passed to the consumer handler (will be serialized)
"""
thread_id = str(run.thread.id) if run.thread else None
groups = ["threads"] # Always send to all threads group
if thread_id:
# Note: Use hyphen (-), not colon (:) - Django Channels has issues with colons
# Using hyphen to distinguish from variables like thread_id (which use underscore)
groups.append(f"thread-{thread_id}") # Also send to specific thread group
# Serialize kwargs to ensure all data is JSON-serializable
serialized_kwargs = {}
for key, value in kwargs.items():
serialized_kwargs[key] = serializable(value) if value is not None else None
publish_event(
event_type,
groups,
run_id=str(run.id),
thread_id=thread_id,
**serialized_kwargs
)
[docs]
def channel_user(event_name: str, run: "Run", **kwargs) -> None:
"""Send a message to the user's channel group.
Args:
event_name: The event name (e.g., "Run.before_agent", "Run.after_agent")
run: Run instance to get user from thread
**kwargs: Event data to be passed (will be serialized)
"""
if not run.thread or not run.thread.user:
return # No user to send to
user_id = run.thread.user.id
groups = [f"user-{user_id}"]
run_data = serializable(run)
# Serialize kwargs
serialized_kwargs = {}
for key, value in kwargs.items():
serialized_kwargs[key] = serializable(value) if value is not None else None
# Send message to UserConsumer
try:
channel_layer = get_channel_layer()
if channel_layer:
event_data = {
'type': 'user_message',
'message': event_name,
'data': {
'Run': run_data,
**serialized_kwargs
}
}
for group in groups:
async_to_sync(channel_layer.group_send)(group, event_data)
except Exception as e:
import logging
logging.getLogger(__name__).error(f"Failed to send user message: {e}")
[docs]
class Publisher:
"""Centralized class for sending messages via Django Channels.
Django Channels provides a channel layer abstraction for sending messages between
different parts of your application, typically from backend code to WebSocket consumers.
How it works:
------------
1. **Channel Layer**: A messaging system (usually Redis or in-memory) that routes messages
between different parts of your application.
2. **Groups**: Named collections of WebSocket connections. When you send to a group,
all connected WebSocket clients in that group receive the message.
3. **Message Format**: Messages are dictionaries that must include a 'type' key.
The 'type' determines which handler method is called in the consumer.
4. **Consumers**: WebSocket handlers that receive messages. They define methods like
`async def event_type(self, event)` to handle different message types.
Note: All data is automatically serialized to ensure JSON compatibility.
"""
[docs]
def __init__(self, alias: str = DEFAULT_CHANNEL_LAYER):
"""Initialize the Channel with a specific channel layer.
Args:
alias: The channel layer alias from Django settings (default: 'default')
Raises:
ValueError: If the channel layer is not found
"""
self.alias = alias
self.layer: BaseChannelLayer = get_channel_layer(alias)
if not self.layer:
raise ValueError(f"Channel layer {alias} not found")
[docs]
def send(self, channel: str, message: dict) -> None:
"""Send a message directly to a specific channel.
This sends to a single channel name (not a group). Use this when you know
the exact channel name of a WebSocket connection.
Args:
channel: The specific channel name to send to
message: The message dict (must include 'type' key)
"""
async_to_sync(self.layer.send)(channel, message)
[docs]
def group_send(self, group: str, message: dict) -> None:
"""Send a message to a channel group.
All WebSocket connections in the group will receive this message.
Args:
group: The group name (e.g., "logs", "thread-123", "user-456")
message: The message dict (must include 'type' key)
"""
async_to_sync(self.layer.group_send)(group, message)
[docs]
def send_to_groups(self, groups: list[str], message: dict) -> None:
"""Send a message to multiple channel groups.
Args:
groups: List of group names to send to
message: The message dict (must include 'type' key)
"""
for group in groups:
self.group_send(group, message)
[docs]
def consumer_handle(
self,
event_type: str,
groups: list[str],
payload: Any | None = None,
) -> None:
"""Publish an event to channel groups with automatic timestamp and serialization.
This is a convenience method that:
- Adds a timestamp to the event
- Serializes all kwargs to ensure JSON compatibility
- Sends to multiple groups
Args:
event_type: The event type (determines consumer handler method)
groups: List of group names to send to
**kwargs: Event data (will be serialized)
"""
event_data = {
'type': event_type,
'ts': datetime.now().isoformat(),
**serializable(payload),
}
self.send_to_groups(groups, event_data)