chore: import original project baseline
Import the pre-repair source tree as the history baseline. Runtime data (data/), virtualenvs, bytecode caches and logs are gitignored so local secrets and user state stay out of the repo.
This commit is contained in:
@@ -0,0 +1,204 @@
|
||||
"""
|
||||
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
|
||||
Reference in New Issue
Block a user