mirror of
https://github.com/openswarm-ai/openswarm.git
synced 2026-08-20 19:52:23 +02:00
85 lines
3.0 KiB
Python
85 lines
3.0 KiB
Python
"""Bridge a broadcast `agent:message` into the typed `events.agent.message`.
|
|
|
|
Called from ws_manager.send_to_session, the single chokepoint every agent message
|
|
flows through. Best-effort: never raises into the broadcast path.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any, Optional
|
|
|
|
from pydantic import BaseModel, ConfigDict
|
|
from typeguard import typechecked
|
|
|
|
from backend.apps.agents.core.models import AgentSession
|
|
from backend.apps.service.analytics.client import track_agent_message
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class BroadcastMessage(BaseModel):
|
|
# An agent:message broadcast payload, validated at the WS boundary; extra fields ignored.
|
|
model_config = ConfigDict(validate_assignment=True, extra="ignore")
|
|
id: Optional[str] = None
|
|
role: Optional[str] = None
|
|
content: Any = None
|
|
parent_id: Optional[str] = None
|
|
branch_id: Optional[str] = None
|
|
|
|
|
|
@typechecked
|
|
def p_branch_version(session: AgentSession, message: BroadcastMessage) -> int:
|
|
# Edit marker for branch_id: only the message that CREATED a forked branch (the actual edit) scores non-zero; replies and new turns reset to 0.
|
|
branch_str = message.branch_id or "main"
|
|
branches = getattr(session, "branches", None) or {}
|
|
b = branches.get(branch_str)
|
|
fork_point = getattr(b, "fork_point_message_id", None) if b else None
|
|
if not fork_point:
|
|
return 0
|
|
branch_user_msgs = [
|
|
m for m in (getattr(session, "messages", None) or [])
|
|
if getattr(m, "branch_id", None) == branch_str and getattr(m, "role", None) == "user"
|
|
]
|
|
if not branch_user_msgs or getattr(branch_user_msgs[0], "id", None) != message.id:
|
|
return 0
|
|
siblings = sorted(
|
|
(x for x in branches.values()
|
|
if getattr(x, "fork_point_message_id", None) == fork_point),
|
|
key=lambda x: x.created_at,
|
|
)
|
|
for i, x in enumerate(siblings, start=1):
|
|
if x.id == branch_str:
|
|
return i
|
|
return 0
|
|
|
|
|
|
@typechecked
|
|
def bridge_agent_message(session_id: str, message: BroadcastMessage) -> None:
|
|
# seq is the message's stable index in the persisted history (survives close -> reopen -> restart); transient messages with no anchor are skipped.
|
|
if not message.id or not message.role:
|
|
return
|
|
try:
|
|
from backend.apps.agents.agent_manager import agent_manager
|
|
sess = agent_manager.sessions.get(session_id)
|
|
except Exception:
|
|
sess = None
|
|
if sess is None:
|
|
return
|
|
msgs = getattr(sess, "messages", None) or []
|
|
seq = next((i for i, m in enumerate(msgs) if getattr(m, "id", None) == message.id), None)
|
|
if seq is None:
|
|
return
|
|
track_agent_message(
|
|
agent_id=session_id,
|
|
seq=seq,
|
|
id=str(message.id),
|
|
role=str(message.role),
|
|
content=message.content,
|
|
parent_id=message.parent_id,
|
|
branch_id=p_branch_version(sess, message),
|
|
provider=getattr(sess, "provider", None),
|
|
model=getattr(sess, "model", None),
|
|
thinking_level=getattr(sess, "thinking_level", None),
|
|
)
|