401 lines
16 KiB
Python
401 lines
16 KiB
Python
"""多智能体会话状态机。
|
||
|
||
一个 MultiAgentState 绑定到一个多智能体对话的 conversation_id，维护：
|
||
- 已创建的子智能体实例（agent_id ↔ role_id ↔ display_name ↔ task_id ↔ status）
|
||
- 待插入到主对话的待发 user 消息队列（pending_master_messages）
|
||
- 主智能体工具调用 answer_sub_agent_question / answer_other_agent 写回答案的 futomap
|
||
- 子智能体调用 ask_master / ask_other_agent 时挂起的 futomap
|
||
|
||
关键约定（来自 .astrion/memory/multi_agent_mode_design.md）：
|
||
- 消息格式：`来自 {显示名} 的{类型}\\nid: {消息id}\\n\\n<{显示名}>\\n<{标签}>\\n{内容}\\n</{标签}>\\n</{显示名}>`
|
||
- 接收方决定插入方式：
|
||
- 子智能体 ask 阻塞等待 → main 调 answer_* 返回到工具结果
|
||
- 子智能体 idle 状态 → 主对话的 pending_master_messages 直接插入新轮 user 消息
|
||
- 子智能体 running 中 → inline 插入到当前末尾（在下一轮 model 调用前合并 messages）
|
||
- 通信是「工具调用提问」+「回答返回到工具结果」；其他场景（输出/进度/完成/任务发布/消息/回答）
|
||
才以 user 消息格式插入对话。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import uuid
|
||
from asyncio import AbstractEventLoop
|
||
from dataclasses import dataclass, field
|
||
from datetime import datetime
|
||
from pathlib import Path
|
||
from typing import Any, Dict, List, Optional, TYPE_CHECKING
|
||
from modules.multi_agent.debug_logger import ma_debug
|
||
|
||
if TYPE_CHECKING:
|
||
from modules.sub_agent.task import SubAgentTask
|
||
|
||
# ---------- 消息类型常量 ----------
|
||
TYPE_TASK = "Task" # 主→子 任务发布
|
||
TYPE_OUTPUT = "Output" # 子→主 进度/完成输出（统一）
|
||
TYPE_ASK = "Ask" # 子→主 / 子→子 提问
|
||
TYPE_ANSWER = "Answer" # 主→子 / 子→子 回答（不插入对话，仅做工具结果）
|
||
TYPE_MESSAGE = "Message" # 任意方向 消息
|
||
# 内部枚举到此
|
||
|
||
QUESTION_PREFIX_ASK_MASTER = "ask_master"
|
||
QUESTION_PREFIX_ASK_OTHER = "ask_other"
|
||
|
||
|
||
def format_multi_agent_message(
|
||
*,
|
||
display_name: str,
|
||
msg_type: str,
|
||
content: str,
|
||
msg_id: Optional[str] = None,
|
||
target: Optional[str] = None,
|
||
extra_attrs: Optional[Dict[str, str]] = None,
|
||
msg_type_text: Optional[str] = None,
|
||
) -> str:
|
||
"""按统一格式构造 user 消息字符串。
|
||
|
||
Args:
|
||
display_name: 发出方显示名（如 UI Operator_1 / Team Leader）
|
||
msg_type: 消息类型，对应上方 TYPE_* 常量
|
||
content: 消息正文
|
||
msg_id: 消息 id；不传则自动生成
|
||
target: 接收方显示名（用于子→子 提问时标明对谁提问）
|
||
extra_attrs: 额外标签属性（如 question_id="ask_xxx"）
|
||
msg_type_text: 覆盖默认的中文消息类型文案（如"任务结束汇报"）
|
||
"""
|
||
if not msg_id:
|
||
msg_id = f"msg_{uuid.uuid4().hex[:10]}"
|
||
|
||
type_label = msg_type_text or msg_type_to_text(msg_type)
|
||
# 第一行：自然语言前缀（含 target 标识）
|
||
if target:
|
||
prefix = f"来自 {display_name} 向 {target} 的{type_label}"
|
||
else:
|
||
prefix = f"来自 {display_name} 的{type_label}"
|
||
|
||
# 第二行：id
|
||
id_line = f"id: {msg_id}"
|
||
|
||
# 属性 attr 字符串
|
||
attrs = ""
|
||
if target:
|
||
attrs += f' target="{target}"'
|
||
if extra_attrs:
|
||
for k, v in extra_attrs.items():
|
||
attrs += f' {k}="{v}"'
|
||
|
||
# XML 包裹
|
||
tag = msg_type
|
||
xml = (
|
||
f"<{display_name}>\n"
|
||
f"<{tag}{attrs}>\n"
|
||
f"{content}\n"
|
||
f"</{tag}>\n"
|
||
f"</{display_name}>"
|
||
)
|
||
|
||
return f"{prefix}\n{id_line}\n\n{xml}"
|
||
|
||
|
||
def msg_type_to_text(msg_type: str) -> str:
|
||
"""把 TYPE_* 转为中文短语，用于 prompt 前缀。"""
|
||
mapping = {
|
||
TYPE_TASK: "任务发布",
|
||
TYPE_OUTPUT: "任务进度输出",
|
||
TYPE_ASK: "提问",
|
||
TYPE_ANSWER: "回答",
|
||
TYPE_MESSAGE: "消息",
|
||
}
|
||
return mapping.get(msg_type, msg_type)
|
||
|
||
|
||
def build_master_dispatch_text(task: str, msg_id: Optional[str] = None) -> str:
|
||
"""主智能体发布任务时插入到子智能体对话的 user 消息文本。"""
|
||
return format_multi_agent_message(
|
||
display_name="Team Leader",
|
||
msg_type=TYPE_TASK,
|
||
content=task,
|
||
msg_id=msg_id,
|
||
)
|
||
|
||
|
||
def build_sub_agent_output_text(display_name: str, content: str, msg_id: Optional[str] = None, *, is_final: bool = False) -> str:
|
||
"""子智能体输出（进度或完成）插入到主对话的 user 消息文本。"""
|
||
return format_multi_agent_message(
|
||
display_name=display_name,
|
||
msg_type=TYPE_OUTPUT,
|
||
content=content,
|
||
msg_id=msg_id,
|
||
msg_type_text="任务结束汇报" if is_final else "任务进度输出",
|
||
)
|
||
|
||
|
||
def build_sub_agent_ask_master_text(display_name: str, question: str, question_id: str) -> str:
|
||
"""子智能体向主智能体提问时插入到主对话的 user 消息文本。"""
|
||
return format_multi_agent_message(
|
||
display_name=display_name,
|
||
msg_type=TYPE_ASK,
|
||
content=question,
|
||
msg_id=question_id,
|
||
)
|
||
|
||
|
||
def build_sub_agent_ask_other_text(
|
||
display_name: str,
|
||
target_display: str,
|
||
question: str,
|
||
question_id: str,
|
||
) -> str:
|
||
"""子智能体向另一个子智能体提问时插入到目标子智能体对话的文本。"""
|
||
return format_multi_agent_message(
|
||
display_name=display_name,
|
||
msg_type=TYPE_ASK,
|
||
content=question,
|
||
msg_id=question_id,
|
||
target=target_display,
|
||
)
|
||
|
||
|
||
def build_master_message_to_sub_agent(message: str, msg_id: Optional[str] = None) -> str:
|
||
"""主智能体 send_message_to_sub_agent 时插入子对话的 user 消息文本。"""
|
||
return format_multi_agent_message(
|
||
display_name="Team Leader",
|
||
msg_type=TYPE_MESSAGE,
|
||
content=message,
|
||
msg_id=msg_id,
|
||
)
|
||
|
||
|
||
def build_master_answer_to_sub_agent(
|
||
display_name: str,
|
||
target_display: str,
|
||
answer: str,
|
||
question_id: str,
|
||
) -> str:
|
||
"""主智能体回答插入到子对话（仅当子智能体 not waiting 或 idle 时走 user 消息路径）。"""
|
||
return format_multi_agent_message(
|
||
display_name=display_name,
|
||
msg_type=TYPE_ANSWER,
|
||
content=answer,
|
||
msg_id=question_id,
|
||
target=target_display,
|
||
extra_attrs={"question_id": question_id},
|
||
)
|
||
|
||
|
||
# ---------- 运行态状态机 ----------
|
||
@dataclass
|
||
class AgentInstance:
|
||
"""一个多智能体会话中已创建的子智能体实例。"""
|
||
|
||
agent_id: int
|
||
role_id: str
|
||
display_name: str
|
||
task_id: str
|
||
status: str = "running" # running / idle / terminated / failed / timeout
|
||
summary: str = ""
|
||
created_at: float = field(default_factory=lambda: datetime.now().timestamp())
|
||
last_output: str = ""
|
||
|
||
def to_dict(self) -> Dict[str, Any]:
|
||
return {
|
||
"agent_id": self.agent_id,
|
||
"role_id": self.role_id,
|
||
"display_name": self.display_name,
|
||
"task_id": self.task_id,
|
||
"status": self.status,
|
||
"summary": self.summary,
|
||
"created_at": self.created_at,
|
||
"last_output": self.last_output,
|
||
}
|
||
|
||
|
||
class MultiAgentState:
|
||
"""绑到一个 conversation_id 的多智能体运行态。
|
||
|
||
线程安全：所有 pubic 方法均假设在 SubAgentManager 的事件循环线程中调用，
|
||
或者由 chat task 主线程通过 manager 的 _run_coro 进入此循环。
|
||
跨线程访问通过 manager._run_coro 桥接，避免直接调用。
|
||
"""
|
||
|
||
def __init__(self, conversation_id: str):
|
||
self.conversation_id = conversation_id
|
||
# agent_id 映射；同一会话里 agent_id 唯一
|
||
self.agents: Dict[int, AgentInstance] = {}
|
||
# task_id -> agent_id（便于在 SubAgentTask 完成时回写）
|
||
self.task_id_to_agent_id: Dict[str, int] = {}
|
||
# 主智能体待插入消息队列（每条都是字符串，由 chat task 取走）
|
||
self.pending_master_messages: List[str] = []
|
||
# ask_master / ask_other_agent 的等待 future
|
||
# key = question_id, value = asyncio.Future (结果为 answer str 或 Exception)
|
||
self.pending_questions: Dict[str, asyncio.Future] = {}
|
||
# question_id -> 创建 future 时所在的事件循环，用于跨循环安全 set_result
|
||
self.pending_question_loops: Dict[str, AbstractEventLoop] = {}
|
||
# 回答早于 wait_for_answer 注册时先暂存
|
||
self.pending_answers: Dict[str, str] = {}
|
||
# 一个 agent 可能同时只阻塞在一个 ask 工具上（最简实现）
|
||
# key = agent_id, value = question_id（表示当前 agent 正阻塞等待）
|
||
self.agent_blocking_question: Dict[int, str] = {}
|
||
# 角色实例计数：role_id -> 已分配的最大 agent_id（数字）
|
||
# 用于创建新实例时自动递增编号，但允许调用方显式指定
|
||
self.role_counters: Dict[str, int] = {}
|
||
|
||
# ----- 创建/查询 -----
|
||
def next_agent_id_for_role(self, role_id: str) -> int:
|
||
"""为指定角色分配下一个 agent_id 编号。"""
|
||
n = self.role_counters.get(role_id, 0) + 1
|
||
self.role_counters[role_id] = n
|
||
return n
|
||
|
||
def register_instance(self, instance: AgentInstance) -> None:
|
||
if instance.agent_id in self.agents:
|
||
raise ValueError(f"agent_id {instance.agent_id} 已存在")
|
||
self.agents[instance.agent_id] = instance
|
||
self.task_id_to_agent_id[instance.task_id] = instance.agent_id
|
||
|
||
def get_instance(self, agent_id: int) -> Optional[AgentInstance]:
|
||
return self.agents.get(agent_id)
|
||
|
||
def get_instance_by_task_id(self, task_id: str) -> Optional[AgentInstance]:
|
||
aid = self.task_id_to_agent_id.get(task_id)
|
||
if aid is None:
|
||
return None
|
||
return self.agents.get(aid)
|
||
|
||
def list_active(self) -> List[AgentInstance]:
|
||
return [a for a in self.agents.values() if a.status == "running" or a.status == "idle"]
|
||
|
||
def list_all(self) -> List[AgentInstance]:
|
||
return list(self.agents.values())
|
||
|
||
def mark_status(self, agent_id: int, status: str, last_output: str = "") -> None:
|
||
a = self.agents.get(agent_id)
|
||
if a:
|
||
a.status = status
|
||
if last_output:
|
||
a.last_output = last_output
|
||
|
||
# ----- 主对话注入 -----
|
||
def push_master_message(self, message_text: str) -> None:
|
||
"""把一条 user 消息追加到主对话待插入队列。"""
|
||
self.pending_master_messages.append(message_text)
|
||
|
||
def drain_master_messages(self) -> List[str]:
|
||
"""取出（清空）所有待插入主对话的消息。"""
|
||
msgs = self.pending_master_messages
|
||
self.pending_master_messages = []
|
||
return msgs
|
||
|
||
def has_pending_master_messages(self) -> bool:
|
||
return len(self.pending_master_messages) > 0
|
||
|
||
# ----- 阻塞问答 -----
|
||
async def wait_for_answer(self, question_id: str, agent_id: int, timeout: float = 600.0) -> str:
|
||
"""子智能体 ask_* 工具调用后阻塞等待答案。
|
||
|
||
返回 answer 字符串；超时/取消抛 asyncio.TimeoutError 或 CancelledError。
|
||
"""
|
||
# 如果回答已经提前到达，直接返回
|
||
if question_id in self.pending_answers:
|
||
return self.pending_answers.pop(question_id)
|
||
if question_id in self.pending_questions:
|
||
old_fut = self.pending_questions[question_id]
|
||
try:
|
||
old_loop = old_fut.get_loop()
|
||
if old_loop.is_closed():
|
||
self.pending_questions.pop(question_id, None)
|
||
self.pending_question_loops.pop(question_id, None)
|
||
else:
|
||
raise RuntimeError(f"question_id 已存在: {question_id}")
|
||
except Exception:
|
||
self.pending_questions.pop(question_id, None)
|
||
self.pending_question_loops.pop(question_id, None)
|
||
loop = asyncio.get_running_loop()
|
||
fut: asyncio.Future = loop.create_future()
|
||
self.pending_questions[question_id] = fut
|
||
self.pending_question_loops[question_id] = loop
|
||
self.agent_blocking_question[agent_id] = question_id
|
||
try:
|
||
return await asyncio.wait_for(fut, timeout=timeout)
|
||
finally:
|
||
self.pending_questions.pop(question_id, None)
|
||
self.pending_question_loops.pop(question_id, None)
|
||
if self.agent_blocking_question.get(agent_id) == question_id:
|
||
self.agent_blocking_question.pop(agent_id, None)
|
||
|
||
async def _do_provide_answer(self, question_id: str, answer: str) -> bool:
|
||
"""在同 future 所属事件循环内设置结果。"""
|
||
fut = self.pending_questions.get(question_id)
|
||
if not fut or fut.done():
|
||
return False
|
||
try:
|
||
fut.set_result(answer)
|
||
except asyncio.InvalidStateError:
|
||
return False
|
||
return True
|
||
|
||
def provide_answer(self, question_id: str, answer: str) -> bool:
|
||
"""主/其他子智能体 answer_* 工具调用时回写答案。
|
||
|
||
返回 True 表示找到等待中的 future；False 表示无等待方或已超时。
|
||
支持跨事件循环调用（例如主对话循环回答子智能体循环里的提问）。
|
||
"""
|
||
ma_debug(
|
||
"state_provide_answer",
|
||
question_id=question_id,
|
||
has_pending=question_id in self.pending_questions,
|
||
answer_preview=str(answer)[:300],
|
||
)
|
||
# 如果 wait_for_answer 还没注册，先把答案暂存
|
||
if question_id not in self.pending_questions:
|
||
self.pending_answers[question_id] = answer
|
||
return True
|
||
fut = self.pending_questions.get(question_id)
|
||
if not fut or fut.done():
|
||
self.pending_answers[question_id] = answer
|
||
return False
|
||
loop = self.pending_question_loops.get(question_id)
|
||
if loop is None:
|
||
try:
|
||
loop = fut.get_loop()
|
||
except Exception:
|
||
pass
|
||
if loop is not None:
|
||
try:
|
||
asyncio.run_coroutine_threadsafe(self._do_provide_answer(question_id, answer), loop)
|
||
return True
|
||
except Exception:
|
||
pass
|
||
# 同循环回退（future 所属循环可能已关闭，失败时把答案暂存，避免阻塞方永远等不到）
|
||
try:
|
||
return asyncio.run_coroutine_threadsafe(self._do_provide_answer(question_id, answer), asyncio.get_event_loop()).result(timeout=5)
|
||
except Exception:
|
||
self.pending_questions.pop(question_id, None)
|
||
self.pending_question_loops.pop(question_id, None)
|
||
self.pending_answers[question_id] = answer
|
||
return True
|
||
|
||
def is_agent_blocking(self, agent_id: int) -> bool:
|
||
return agent_id in self.agent_blocking_question
|
||
|
||
def get_blocking_question_id(self, agent_id: int) -> Optional[str]:
|
||
return self.agent_blocking_question.get(agent_id)
|
||
|
||
# ----- 持久化（最简版） -----
|
||
def to_snapshot(self) -> Dict[str, Any]:
|
||
return {
|
||
"conversation_id": self.conversation_id,
|
||
"agents": [a.to_dict() for a in self.agents.values()],
|
||
"role_counters": self.role_counters,
|
||
}
|
||
|
||
@classmethod
|
||
def from_snapshot(cls, snapshot: Dict[str, Any]) -> "MultiAgentState":
|
||
state = cls(conversation_id=snapshot.get("conversation_id", ""))
|
||
state.role_counters = dict(snapshot.get("role_counters") or {})
|
||
for a_data in snapshot.get("agents") or []:
|
||
a = AgentInstance(**a_data)
|
||
state.agents[a.agent_id] = a
|
||
if a.task_id:
|
||
state.task_id_to_agent_id[a.task_id] = a.agent_id
|
||
return state |