205 lines
8.7 KiB
Python
205 lines
8.7 KiB
Python
"""
|
||
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<cb>
|
||
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
|