""" core/agent/agent.py =================== 🌟 pi Agent 状态机的 Python 1:1 移植 对照 pi-main 源码：packages/agent/src/agent.ts (592 行) class Agent { state: AgentState # messages / isStreaming / error config: AgentConfig private listeners: Set subscribe(cb) => unsubscribe # 事件订阅（TUI/Qt 都靠它） prompt(message | text) => Promise # 追加 user 消息并跑一轮 run-to-settle continue() => Promise # 从当前上下文继续（不能以 assistant 结尾） steer(text) # 中途注入（turn 边界消费） followUp(text) # 停止后注入（agent 停止时消费 → 续跑） abort() # 中止当前 run（下一边界生效） } 关键语义（与 pi 完全一致）： 1. steering 在「循环开始前」和「每个 turn 结束后」被消费 （one-at-a-time 每次取 1 条；"all" 一次取光） 2. followUp 只在「run 即将结束」时被外层循环消费 → 触发续跑 3. abort 不抛异常：信号置位 → 当前流/工具在下一检查点收尾 → 助手消息 stop_reason="aborted" → agent_end 4. 同一个 Agent 实例可反复 prompt（state.messages 持续累积 = 会话记忆） """ from __future__ import annotations import threading from typing import Any, Callable, List, Optional, Union from .loop import run_loop from .types import (AgentConfig, AgentError, AgentEvent, AgentMessage, AgentState, AbortSignal, RunResult, new_id) class Agent: def __init__(self, config: AgentConfig): self.config = config self.state = AgentState() self._listeners: List[Callable[[AgentEvent], None]] = [] self._steering_queue: List[AgentMessage] = [] self._follow_up_queue: List[AgentMessage] = [] self._lock = threading.Lock() self._active_signal: Optional[AbortSignal] = None # ------------------------------------------------------------------ # 事件订阅 —— 对照 agent.ts subscribe() # ------------------------------------------------------------------ def subscribe(self, cb: Callable[[AgentEvent], None]) -> Callable[[], None]: self._listeners.append(cb) def unsubscribe(): try: self._listeners.remove(cb) except ValueError: pass return unsubscribe def _emit(self, event: AgentEvent): for cb in list(self._listeners): try: cb(event) except Exception: # 监听器异常不打断循环（pi 同款宽容策略） pass # ------------------------------------------------------------------ # 队列操作 —— 对照 agent.ts 的 getSteeringMessages / getFollowUpMessages # ------------------------------------------------------------------ def _take_steering(self) -> List[AgentMessage]: mode = self.config.steering_mode with self._lock: if not self._steering_queue: return [] if mode == "all": out = self._steering_queue[:] self._steering_queue.clear() else: # one-at-a-time（pi 默认） out = [self._steering_queue.pop(0)] return out def _take_follow_ups(self) -> List[AgentMessage]: mode = self.config.follow_up_mode with self._lock: if not self._follow_up_queue: return [] if mode == "all": out = self._follow_up_queue[:] self._follow_up_queue.clear() else: out = [self._follow_up_queue.pop(0)] return out def has_queued(self) -> bool: with self._lock: return bool(self._steering_queue or self._follow_up_queue) # ------------------------------------------------------------------ # 用户入口 —— 对照 agent.ts prompt / steer / followUp / abort # ------------------------------------------------------------------ def prompt(self, message: Union[str, AgentMessage]) -> RunResult: """追加一条 user 消息并运行到 settle（同步阻塞，跑在调用线程上）""" if self.state.is_streaming: raise AgentError(message="Agent 正在运行中，不能并发 prompt（请先 abort）", kind="unknown") if isinstance(message, str): message = AgentMessage(role="user", content=message) if message.role != "user": raise AgentError(message="prompt 只接受 user 消息", kind="unknown") self.state.messages.append(message) return self._run(new_message=message) def continue_(self) -> RunResult: """ 对照 agent.ts continue() / runAgentLoopContinue 前置检查： - 上下文不能为空 - 不能以 assistant 消息结尾 """ if self.state.is_streaming: raise AgentError(message="Agent 正在运行中", kind="unknown") if not self.state.messages: raise AgentError(message="上下文为空，无法 continue", kind="unknown") if self.state.messages[-1].role == "assistant": raise AgentError( message="不能以 assistant 消息结尾来 continue（应先注入 user/toolResult）", kind="unknown") return self._run(new_message=None) def steer(self, text: str): """中途注入：在当前 run 的 turn 边界被消费（pi steering）""" with self._lock: self._steering_queue.append(AgentMessage(role="user", content=text)) def follow_up(self, text: str): """停止后注入：run 即将结束时被外层循环消费 → 自动续跑（pi followUp）""" with self._lock: self._follow_up_queue.append(AgentMessage(role="user", content=text)) def abort(self): """ 中止当前 run。 对照 pi: 置位 AbortSignal；当前流在 chunk 边界关闭， 工具在检查点返回 "Operation aborted"，助手消息以 "aborted" 收尾。 """ sig = self._active_signal if sig is not None: sig.abort("aborted") with self._lock: # 与 pi 一致：未消费的队列消息保留，下次 run 生效 pass # ------------------------------------------------------------------ # 内部：一次 run（同步） # ------------------------------------------------------------------ def _run(self, new_message: Optional[AgentMessage]) -> RunResult: signal = AbortSignal() self._active_signal = signal self.state.is_streaming = True self.state.error = None try: result = run_loop(self, new_message, signal, self._make_stream_fn()) except AgentError: raise except Exception as e: self.state.error = AgentError(message=f"循环异常: {e}", kind="unknown") self._emit(AgentEvent(type="agent_end", stop_reason="error", error=self.state.error)) result = RunResult(stop_reason="error", error=self.state.error) finally: self.state.is_streaming = False self.state.streaming_message = None self.state.streaming_delta = {} self._active_signal = None return result def _finish_run(self, new_messages: List[AgentMessage], stop_reason: str, error: Optional[AgentError]): """循环收尾回调（loop 层调用）""" if error is not None: self.state.error = error # 本轮新消息已在循环中逐条 append 进 state.messages；此处仅记录收尾 self._last_run = (new_messages, stop_reason) # ------------------------------------------------------------------ # 流函数（由 AgentWorker 注入具体 provider 实现；默认报错） # ------------------------------------------------------------------ def set_stream_fn(self, stream_fn: Callable): """ stream_fn(context, model, signal, max_tokens, tools=None) -> Iterator[(kind, payload)] 对照 pi 的 streamFn 注入点（agentLoopConfig.streamFn）。 tools 参数 = 本轮可用的 AgentTool 列表，需序列化进 provider 请求。 """ self._stream_fn = stream_fn def _make_stream_fn(self) -> Callable: fn = getattr(self, "_stream_fn", None) if fn is None: def _missing(context, model, signal, max_tokens, tools=None): raise AgentError(message="未配置 stream_fn", kind="unknown") yield # pragma: no cover return _missing return fn