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,58 @@
|
||||
"""
|
||||
core/agent —— pi agent 核心框架的 Python 1:1 重构
|
||||
==================================================
|
||||
模块对照（pi-main → 本包）：
|
||||
packages/agent/src/types.ts → types.py 数据模型（消息/事件/工具/配置）
|
||||
packages/agent/src/agent-loop.ts → loop.py 核心循环（runLoop 1:1）
|
||||
packages/agent/src/agent.ts → agent.py Agent 状态机（prompt/steer/followUp/abort）
|
||||
packages/ai/src/api/*.ts → stream_fn.py OpenAI 兼容流式（vLLM）+ 输出预算钳制
|
||||
packages/ai/src/api/simple-options→ context.py token 估算 + clampMaxTokensToContext
|
||||
packages/agent/src/compaction.ts → compaction.py 上下文压缩（切分/摘要/替换）
|
||||
agent-session.ts 后置恢复逻辑 → recovery.py 重试退避 + 溢出/截断压缩恢复
|
||||
packages/agent/src/tools/*.ts → tools.py 工具管线 + 内置 read/bash/write/edit
|
||||
|
||||
使用示例（最小闭环）：
|
||||
from core.agent import Agent, AgentConfig, ModelConfig
|
||||
from core.agent.stream_fn import openai_stream
|
||||
from core.agent.recovery import AgentRunner
|
||||
from core.agent.tools import default_tools
|
||||
|
||||
cfg = AgentConfig(model=ModelConfig(...), tools=default_tools())
|
||||
agent = Agent(cfg)
|
||||
agent.set_stream_fn(openai_stream)
|
||||
runner = AgentRunner(agent)
|
||||
result = runner.run("你好")
|
||||
"""
|
||||
from .types import (AgentConfig, AgentError, AgentEvent, AgentMessage,
|
||||
AgentState, AgentTool, AgentToolResult, AbortSignal,
|
||||
AssistantMessageEvent, ModelConfig, RetryConfig, RunResult,
|
||||
ToolCall, new_id)
|
||||
from .agent import Agent
|
||||
from .context import (CONTEXT_SAFETY_TOKENS, calculate_context_tokens,
|
||||
clamp_max_tokens_to_context,
|
||||
estimate_context_tokens, estimate_message_tokens,
|
||||
should_compact)
|
||||
from .compaction import (CompactionSettings,
|
||||
DEFAULT_COMPACTION_SETTINGS, compact_context,
|
||||
find_cut_point, prepare_compaction)
|
||||
from .recovery import (AgentRunner, compute_retry_delay_ms, find_last_assistant,
|
||||
is_context_overflow, is_recoverable_length,
|
||||
is_retryable_assistant_error)
|
||||
from .stream_fn import (classify_error, from_openai_messages, openai_stream,
|
||||
to_openai_messages)
|
||||
from .tools import default_tools
|
||||
|
||||
# 🆕 版本记录：核心框架为 pi (badlogic/pi-mono) 0.81.x 时代的 Python 移植，
|
||||
# 2026-09 对齐 pi 0.85.1 的估算/锚定/钳制/重试语义（P0/P1/P2/M1/M2/M3）。
|
||||
__version__ = "1.0.0 (pi 0.81.x port, aligned 0.85.1 semantics)"
|
||||
__all__ = [
|
||||
"Agent", "AgentConfig", "AgentError", "AgentEvent", "AgentMessage",
|
||||
"AgentRunner", "AgentState", "AgentTool", "AgentToolResult",
|
||||
"AbortSignal", "AssistantMessageEvent", "CONTEXT_SAFETY_TOKENS",
|
||||
"ModelConfig", "RetryConfig", "RunResult", "ToolCall",
|
||||
"clamp_max_tokens_to_context", "calculate_context_tokens", "classify_error", "compact_context",
|
||||
"compute_retry_delay_ms", "default_tools", "estimate_context_tokens",
|
||||
"find_last_assistant", "find_cut_point", "is_context_overflow",
|
||||
"is_recoverable_length", "is_retryable_assistant_error", "new_id",
|
||||
"openai_stream", "should_compact", "to_openai_messages",
|
||||
]
|
||||
@@ -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
|
||||
@@ -0,0 +1,637 @@
|
||||
"""
|
||||
core/agent/compaction.py
|
||||
========================
|
||||
🌟 pi 上下文压缩算法的 Python 1:1 移植（harness 原版算法）
|
||||
|
||||
对照 pi-main 源码（逐函数对应）：
|
||||
packages/agent/src/harness/compaction/compaction.ts
|
||||
DEFAULT_COMPACTION_SETTINGS -> CompactionSettings / DEFAULT_COMPACTION_SETTINGS
|
||||
shouldCompact -> 见 context.py: should_compact（同一公式）
|
||||
estimateTokens / estimateContextTokens -> context.py（usage 锚定，CJK 感知见说明）
|
||||
findValidCutPoints / findCutPoint -> find_valid_cut_points / find_cut_point
|
||||
findTurnStartIndex -> find_turn_start
|
||||
prepareCompaction -> prepare_compaction
|
||||
SUMMARIZATION_SYSTEM_PROMPT -> 同名（逐字移植）
|
||||
SUMMARIZATION_PROMPT -> 同名（逐字移植）
|
||||
UPDATE_SUMMARIZATION_PROMPT -> 同名（逐字移植，迭代更新用）
|
||||
TURN_PREFIX_SUMMARIZATION_PROMPT -> 同名（逐字移植，断轮前缀用）
|
||||
generateSummaryWithUsage -> generate_summary（maxTokens = 0.8×reserve）
|
||||
generateTurnPrefixSummary -> generate_turn_prefix_summary（0.5×reserve）
|
||||
compact -> compact_context（断轮双摘要 + 拼接格式 1:1）
|
||||
packages/agent/src/harness/compaction/utils.ts
|
||||
serializeConversation -> serialize_conversation（1:1，含 2000 字符截断）
|
||||
extractFileOpsFromMessage -> extract_file_ops_from_message
|
||||
computeFileLists -> compute_file_lists
|
||||
formatFileOperations -> format_file_operations
|
||||
TOOL_RESULT_MAX_CHARS = 2000 -> 同名常量
|
||||
|
||||
摘要 LLM 调用由上层注入：
|
||||
summarize_fn(prompt_text: str, system_prompt: str, max_tokens: int) -> str
|
||||
（pi 里是 models.completeSimple + retry；haocode 用 OpenAI 客户端非流式调用，
|
||||
由 llm_engine.AgentWorker 实现并注入。）
|
||||
|
||||
已声明的偏差（仅 2 处，见 context.py 头注）：
|
||||
1. 单条消息 token 估算用 CJK 感知启发式（pi 是 chars/4）——对中文会话更安全
|
||||
2. 摘要 LLM 调用失败时降级为机械摘录（pi 返回 CompactionError）——桌面应用优先不丢上下文
|
||||
其余全部 1:1：触发公式、usage 锚定、token 预算切点、有效切点规则、断轮双摘要、
|
||||
迭代式 previousSummary 更新、摘要提示词逐字、文件操作附录、拼接格式。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Callable, List, Optional, Set, Tuple
|
||||
|
||||
from .types import AgentMessage, ModelConfig
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 压缩设置 —— 1:1 对照 DEFAULT_COMPACTION_SETTINGS
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class CompactionSettings:
|
||||
"""pi: interface CompactionSettings { enabled; reserveTokens; keepRecentTokens }"""
|
||||
enabled: bool = True
|
||||
reserve_tokens: int = 16384 # pi: 16384（摘要提示词与输出预留）
|
||||
keep_recent_tokens: int = 20000 # pi: 20000（压缩后保留的近期上下文预算）
|
||||
|
||||
|
||||
DEFAULT_COMPACTION_SETTINGS = CompactionSettings()
|
||||
|
||||
# pi utils.ts: const TOOL_RESULT_MAX_CHARS = 2000
|
||||
TOOL_RESULT_MAX_CHARS = 2000
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 摘要提示词 —— 从 pi compaction.ts 逐字移植（不得改写，摘要质量依赖它）
|
||||
# ======================================================================
|
||||
SUMMARIZATION_SYSTEM_PROMPT = \
|
||||
"You are a context summarization assistant. Your task is to read a conversation " \
|
||||
"between a user and an AI assistant, then produce a structured summary following " \
|
||||
"the exact format specified.\n\n" \
|
||||
"Do NOT continue the conversation. Do NOT respond to any questions in the " \
|
||||
"conversation. ONLY output the structured summary."
|
||||
|
||||
SUMMARIZATION_PROMPT = """The messages above are a conversation to summarize. Create a structured context checkpoint summary that another LLM will use to continue the work.
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[What is the user trying to accomplish? Can be multiple items if the session covers different tasks.]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Any constraints, preferences, or requirements mentioned by user]
|
||||
- [Or "(none)" if none were mentioned]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Completed tasks/changes]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Current work]
|
||||
|
||||
### Blocked
|
||||
- [Issues preventing progress, if any]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale]
|
||||
|
||||
## Next Steps
|
||||
1. [Ordered list of what should happen next]
|
||||
|
||||
## Critical Context
|
||||
- [Any data, examples, or references needed to continue]
|
||||
- [Or "(none)" if not applicable]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages."""
|
||||
|
||||
UPDATE_SUMMARIZATION_PROMPT = """The messages above are NEW conversation messages to incorporate into the existing summary provided in <previous-summary> tags.
|
||||
|
||||
Update the existing structured summary with new information. RULES:
|
||||
- PRESERVE all existing information from the previous summary
|
||||
- ADD new progress, decisions, and context from the new messages
|
||||
- UPDATE the Progress section: move items from "In Progress" to "Done" when completed
|
||||
- UPDATE "Next Steps" based on what was accomplished
|
||||
- PRESERVE exact file paths, function names, and error messages
|
||||
- If something is no longer relevant, you may remove it
|
||||
|
||||
Use this EXACT format:
|
||||
|
||||
## Goal
|
||||
[Preserve existing goals, add new ones if the task expanded]
|
||||
|
||||
## Constraints & Preferences
|
||||
- [Preserve existing, add new ones discovered]
|
||||
|
||||
## Progress
|
||||
### Done
|
||||
- [x] [Include previously done items AND newly completed items]
|
||||
|
||||
### In Progress
|
||||
- [ ] [Current work - update based on progress]
|
||||
|
||||
### Blocked
|
||||
- [Current blockers - remove if resolved]
|
||||
|
||||
## Key Decisions
|
||||
- **[Decision]**: [Brief rationale] (preserve all previous, add new)
|
||||
|
||||
## Next Steps
|
||||
1. [Update based on current state]
|
||||
|
||||
## Critical Context
|
||||
- [Preserve important context, add new if needed]
|
||||
|
||||
Keep each section concise. Preserve exact file paths, function names, and error messages."""
|
||||
|
||||
TURN_PREFIX_SUMMARIZATION_PROMPT = """This is the PREFIX of a turn that was too large to keep. The SUFFIX (recent work) is retained.
|
||||
|
||||
Summarize the prefix to provide context for the retained suffix:
|
||||
|
||||
## Original Request
|
||||
[What did the user ask for in this turn?]
|
||||
|
||||
## Early Progress
|
||||
- [Key decisions and work done in the prefix]
|
||||
|
||||
## Context for Suffix
|
||||
- [Information needed to understand the retained recent work]
|
||||
|
||||
Be concise. Focus on what's needed to understand the kept suffix."""
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 对话序列化 —— 1:1 对照 utils.ts serializeConversation
|
||||
# ======================================================================
|
||||
def _safe_json(value) -> str:
|
||||
try:
|
||||
s = json.dumps(value, ensure_ascii=False)
|
||||
return s if s is not None else "undefined"
|
||||
except Exception:
|
||||
return "[unserializable]"
|
||||
|
||||
|
||||
def _content_text(content) -> str:
|
||||
"""pi contentText: str 或 [{type:"text",text}] 列表取文本拼接"""
|
||||
if content is None:
|
||||
return ""
|
||||
if isinstance(content, str):
|
||||
return content
|
||||
if isinstance(content, list):
|
||||
parts = []
|
||||
for b in content:
|
||||
if isinstance(b, dict) and b.get("type") == "text":
|
||||
parts.append(str(b.get("text", "")))
|
||||
return "\n".join(p for p in parts if p)
|
||||
return str(content)
|
||||
|
||||
|
||||
def _truncate_for_summary(text: str, max_chars: int) -> str:
|
||||
"""pi utils.ts truncateForSummary（逐字逻辑）"""
|
||||
if len(text) <= max_chars:
|
||||
return text
|
||||
truncated = len(text) - max_chars
|
||||
return f"{text[:max_chars]}\n\n[... {truncated} more characters truncated]"
|
||||
|
||||
|
||||
def serialize_conversation(messages: List[AgentMessage]) -> str:
|
||||
"""
|
||||
1:1 对照 utils.ts serializeConversation（输出格式逐字一致）：
|
||||
[User]: ...
|
||||
[Assistant thinking]: ...
|
||||
[Assistant]: ...
|
||||
[Assistant tool calls]: name(k=v, k2=v2); name2(...)
|
||||
[Tool result]: ...（超 2000 字符截断）
|
||||
"""
|
||||
parts: List[str] = []
|
||||
for msg in messages:
|
||||
if msg.role == "user":
|
||||
content = _content_text(msg.content)
|
||||
if content:
|
||||
parts.append(f"[User]: {content}")
|
||||
elif msg.role == "assistant":
|
||||
thinking_parts = []
|
||||
tool_calls = []
|
||||
if msg.reasoning:
|
||||
thinking_parts.append(msg.reasoning)
|
||||
text = _content_text(msg.content)
|
||||
for tc in (msg.tool_calls or []):
|
||||
args_str = ", ".join(f"{k}={_safe_json(v)}"
|
||||
for k, v in (tc.arguments or {}).items())
|
||||
tool_calls.append(f"{tc.name}({args_str})")
|
||||
if thinking_parts:
|
||||
parts.append(f"[Assistant thinking]: {chr(10).join(thinking_parts)}")
|
||||
if text:
|
||||
parts.append(f"[Assistant]: {text}")
|
||||
if tool_calls:
|
||||
parts.append(f"[Assistant tool calls]: {'; '.join(tool_calls)}")
|
||||
elif msg.role == "toolResult":
|
||||
content = _content_text(msg.content)
|
||||
if content:
|
||||
parts.append(f"[Tool result]: "
|
||||
f"{_truncate_for_summary(content, TOOL_RESULT_MAX_CHARS)}")
|
||||
return "\n\n".join(parts)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 文件操作提取 —— 1:1 对照 utils.ts extractFileOps*/computeFileLists/formatFileOperations
|
||||
# ======================================================================
|
||||
class FileOperations:
|
||||
def __init__(self):
|
||||
self.read: Set[str] = set()
|
||||
self.written: Set[str] = set()
|
||||
self.edited: Set[str] = set()
|
||||
|
||||
|
||||
def extract_file_ops_from_message(message: AgentMessage, file_ops: FileOperations):
|
||||
"""pi: assistant 的 toolCall 参数里 path 字段 → read/write/edit 归类"""
|
||||
if message.role != "assistant":
|
||||
return
|
||||
for tc in (message.tool_calls or []):
|
||||
args = tc.arguments or {}
|
||||
path = args.get("path")
|
||||
if not isinstance(path, str) or not path:
|
||||
continue
|
||||
if tc.name == "read":
|
||||
file_ops.read.add(path)
|
||||
elif tc.name == "write":
|
||||
file_ops.written.add(path)
|
||||
elif tc.name == "edit":
|
||||
file_ops.edited.add(path)
|
||||
|
||||
|
||||
def compute_file_lists(file_ops: FileOperations) -> Tuple[List[str], List[str]]:
|
||||
"""pi computeFileLists: modified = edited|written；readOnly = read-modified；均排序"""
|
||||
modified = file_ops.edited | file_ops.written
|
||||
read_only = sorted(f for f in file_ops.read if f not in modified)
|
||||
return read_only, sorted(modified)
|
||||
|
||||
|
||||
def format_file_operations(read_files: List[str], modified_files: List[str]) -> str:
|
||||
"""pi formatFileOperations: <read-files>/<modified-files> 标签拼接"""
|
||||
sections = []
|
||||
if read_files:
|
||||
sections.append("<read-files>\n" + "\n".join(read_files) + "\n</read-files>")
|
||||
if modified_files:
|
||||
sections.append("<modified-files>\n" + "\n".join(modified_files) + "\n</modified-files>")
|
||||
if not sections:
|
||||
return ""
|
||||
return "\n\n" + "\n\n".join(sections)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 切分点 —— 1:1 对照 findValidCutPoints / findTurnStartIndex / findCutPoint
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class CutPointResult:
|
||||
"""pi: interface CutPointResult"""
|
||||
first_kept_index: int # 保留段首条在 compactable 列表里的下标
|
||||
turn_start_index: int = -1 # 断轮时：该轮起点（user 消息）下标；否则 -1
|
||||
is_split_turn: bool = False
|
||||
|
||||
|
||||
def find_valid_cut_points(messages: List[AgentMessage],
|
||||
start_index: int, end_index: int) -> List[int]:
|
||||
"""
|
||||
pi findValidCutPoints：消息角色为 user/assistant 的位置是有效切点
|
||||
（toolResult 不能做切点——它会与前面的 toolCall 分离）。
|
||||
pi 里的 bashExecution/branchSummary/compactionSummary 等角色
|
||||
在 haocode 消息模型中不存在，等价规则即 role in (user, assistant)。
|
||||
"""
|
||||
cut_points = []
|
||||
for i in range(start_index, end_index):
|
||||
if messages[i].role in ("user", "assistant"):
|
||||
cut_points.append(i)
|
||||
return cut_points
|
||||
|
||||
|
||||
def find_turn_start(messages: List[AgentMessage], entry_index: int,
|
||||
start_index: int) -> int:
|
||||
"""pi findTurnStartIndex：向前找本轮起点（user 消息 / branch_summary）"""
|
||||
for i in range(entry_index, start_index - 1, -1):
|
||||
if messages[i].role == "user":
|
||||
return i
|
||||
return -1
|
||||
|
||||
|
||||
def _estimate(msg: AgentMessage) -> int:
|
||||
from .context import estimate_message_tokens
|
||||
return estimate_message_tokens(msg)
|
||||
|
||||
|
||||
def find_cut_point(messages: List[AgentMessage], start_index: int,
|
||||
end_index: int, keep_recent_tokens: int) -> CutPointResult:
|
||||
"""
|
||||
1:1 对照 pi findCutPoint：
|
||||
1. 从尾部向前累计 token，直到累计 >= keep_recent_tokens
|
||||
2. 取该位置（含）之后的第一个有效切点
|
||||
3. 切点不是 user 消息 → 断轮：找本轮起点，前缀单独摘要
|
||||
（pi 的「回退跳过状态条目」循环针对 session 状态条目；
|
||||
haocode 消息列表没有状态条目，等价省略。）
|
||||
"""
|
||||
cut_points = find_valid_cut_points(messages, start_index, end_index)
|
||||
if not cut_points:
|
||||
return CutPointResult(first_kept_index=start_index)
|
||||
|
||||
accumulated = 0
|
||||
cut_index = cut_points[0]
|
||||
for i in range(end_index - 1, start_index - 1, -1):
|
||||
accumulated += _estimate(messages[i])
|
||||
if accumulated >= keep_recent_tokens:
|
||||
for c in cut_points:
|
||||
if c >= i:
|
||||
cut_index = c
|
||||
break
|
||||
break
|
||||
|
||||
is_user = messages[cut_index].role == "user"
|
||||
turn_start = -1 if is_user else find_turn_start(messages, cut_index, start_index)
|
||||
is_split = (not is_user) and turn_start != -1
|
||||
return CutPointResult(first_kept_index=cut_index,
|
||||
turn_start_index=turn_start,
|
||||
is_split_turn=is_split)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 压缩准备 —— 1:1 对照 prepareCompaction
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class CompactionPreparation:
|
||||
"""pi: interface CompactionPreparation"""
|
||||
messages_to_summarize: List[AgentMessage] = field(default_factory=list)
|
||||
turn_prefix_messages: List[AgentMessage] = field(default_factory=list)
|
||||
retained_tail: List[AgentMessage] = field(default_factory=list)
|
||||
is_split_turn: bool = False
|
||||
tokens_before: int = 0
|
||||
previous_summary: Optional[str] = None
|
||||
file_ops: FileOperations = field(default_factory=FileOperations)
|
||||
settings: CompactionSettings = DEFAULT_COMPACTION_SETTINGS
|
||||
|
||||
|
||||
def _compaction_diag(line: str) -> None:
|
||||
"""压缩自诊断日志（与 recovery.compact_diag_log 同一文件/规范：
|
||||
print + 追加 compaction_diag.log + 失败静默）。"""
|
||||
import os as _os
|
||||
import time as _time
|
||||
t = _time.time()
|
||||
stamp = (f"[{_time.strftime('%H:%M:%S', _time.localtime(t))}"
|
||||
f".{int(t * 1000) % 1000:03d}]")
|
||||
print(f"{stamp} {line}", flush=True)
|
||||
try:
|
||||
path = _os.path.join(
|
||||
_os.path.dirname(_os.path.dirname(_os.path.dirname(_os.path.abspath(__file__)))),
|
||||
"compaction_diag.log")
|
||||
with open(path, "a", encoding="utf-8") as f:
|
||||
f.write(f"{stamp} {line}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _nothing_to_summarize_diag(messages: List[AgentMessage],
|
||||
compactable: List[AgentMessage],
|
||||
cut: CutPointResult,
|
||||
settings: CompactionSettings) -> str:
|
||||
"""prepare_compaction 返回 None（无可摘要内容）时的自诊断行：
|
||||
记录分支子类型 + 关键量，下次失败可直接从日志定位原因。
|
||||
|
||||
子类型：
|
||||
no_valid_cut_points —— 可压缩范围内没有任何 user/assistant 条目
|
||||
total_below_keep_recent —— 可压缩范围总量 < keep_recent（通常意味着
|
||||
上下文主体是旧摘要本身，无新内容）
|
||||
cut_pinned_at_zero —— 预算从尾部累加只在 i=0 才达标，即首条
|
||||
条目独占 ≥ (总量-keep_recent) 的 token
|
||||
（单条超长消息/巨型工具输出主导上下文）
|
||||
"""
|
||||
from .context import estimate_context_tokens, estimate_message_tokens
|
||||
try:
|
||||
total = estimate_context_tokens(messages).tokens
|
||||
n_user = sum(1 for m in compactable if m.role == "user")
|
||||
n_asst = sum(1 for m in compactable if m.role == "assistant")
|
||||
ests = [(estimate_message_tokens(m), i) for i, m in enumerate(compactable)]
|
||||
e0 = ests[0][0] if ests else 0
|
||||
rest = sum(e for e, _ in ests[1:])
|
||||
top3 = sorted(ests, reverse=True)[:3]
|
||||
roles = " ".join(("T" if m.role == "toolResult" else m.role[0].upper())
|
||||
for m in compactable[:10])
|
||||
if not find_valid_cut_points(compactable, 0, len(compactable)):
|
||||
sub = "no_valid_cut_points"
|
||||
elif sum(e for e, _ in ests) < settings.keep_recent_tokens:
|
||||
sub = "total_below_keep_recent"
|
||||
else:
|
||||
sub = "cut_pinned_at_zero"
|
||||
big = "; ".join(f"idx{i}={e}" for e, i in top3)
|
||||
return (f"[COMPACT_NONE] branch=nothing_to_summarize sub={sub} "
|
||||
f"msgs={len(messages)} est_total={total} compactable={len(compactable)} "
|
||||
f"keep_recent={settings.keep_recent_tokens} "
|
||||
f"first_kept={cut.first_kept_index} turn_start={cut.turn_start_index} "
|
||||
f"split={cut.is_split_turn} user={n_user} asst={n_asst} "
|
||||
f"est_first={e0} est_rest={rest} top3=[{big}] head_roles=[{roles}]")
|
||||
except Exception as ex: # 诊断本身失败不能影响主流程
|
||||
return f"[COMPACT_NONE] branch=nothing_to_summarize (diag failed: {ex})"
|
||||
|
||||
|
||||
def prepare_compaction(messages: List[AgentMessage],
|
||||
settings: Optional[CompactionSettings] = None
|
||||
) -> Optional[CompactionPreparation]:
|
||||
"""
|
||||
1:1 对照 pi prepareCompaction：
|
||||
- 上一条压缩摘要（messages[0].kind == "compaction_summary"）不重复摘要，
|
||||
其内容作为 previousSummary 走迭代更新提示词
|
||||
- 可压缩范围 = 摘要之后的全部消息（pi 里等价于「上次保留尾 + 新消息」）
|
||||
- 切点在可压缩范围内选；tokens_before 按完整上下文（含摘要消息）估算
|
||||
不可压缩（空/无摘要对象）时返回 None（对照 pi 返回 ok(undefined)）。
|
||||
"""
|
||||
from .context import estimate_context_tokens
|
||||
settings = settings or DEFAULT_COMPACTION_SETTINGS
|
||||
if not messages:
|
||||
_compaction_diag("[COMPACT_NONE] branch=empty_messages")
|
||||
return None
|
||||
if messages[-1].kind == "compaction_summary":
|
||||
_compaction_diag(
|
||||
f"[COMPACT_NONE] branch=tail_is_summary msgs={len(messages)} "
|
||||
f"tail_head={_content_text(messages[-1].content)[:60]!r}")
|
||||
return None
|
||||
|
||||
previous_summary = None
|
||||
if messages[0].kind == "compaction_summary":
|
||||
previous_summary = _content_text(messages[0].content)
|
||||
compactable = messages[1:]
|
||||
else:
|
||||
compactable = messages
|
||||
if not compactable:
|
||||
_compaction_diag("[COMPACT_NONE] branch=compactable_empty (上下文只剩旧摘要)")
|
||||
return None
|
||||
|
||||
tokens_before = estimate_context_tokens(messages).tokens
|
||||
cut = find_cut_point(compactable, 0, len(compactable),
|
||||
settings.keep_recent_tokens)
|
||||
|
||||
history_end = cut.turn_start_index if cut.is_split_turn else cut.first_kept_index
|
||||
messages_to_summarize = compactable[:history_end]
|
||||
turn_prefix_messages = []
|
||||
if cut.is_split_turn:
|
||||
turn_prefix_messages = compactable[cut.turn_start_index:cut.first_kept_index]
|
||||
retained_tail = compactable[cut.first_kept_index:]
|
||||
|
||||
if not messages_to_summarize and not turn_prefix_messages:
|
||||
# 🆕 自诊断：记录是哪种子条件导致无东西可摘要（见 _nothing_to_summarize_diag）
|
||||
_compaction_diag(_nothing_to_summarize_diag(messages, compactable, cut,
|
||||
settings))
|
||||
return None # 没有可摘要内容
|
||||
|
||||
file_ops = FileOperations()
|
||||
for m in messages_to_summarize:
|
||||
extract_file_ops_from_message(m, file_ops)
|
||||
if cut.is_split_turn:
|
||||
for m in turn_prefix_messages:
|
||||
extract_file_ops_from_message(m, file_ops)
|
||||
|
||||
return CompactionPreparation(
|
||||
messages_to_summarize=messages_to_summarize,
|
||||
turn_prefix_messages=turn_prefix_messages,
|
||||
retained_tail=retained_tail,
|
||||
is_split_turn=cut.is_split_turn,
|
||||
tokens_before=tokens_before,
|
||||
previous_summary=previous_summary,
|
||||
file_ops=file_ops,
|
||||
settings=settings,
|
||||
)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 摘要生成 —— 1:1 对照 generateSummaryWithUsage / generateTurnPrefixSummary
|
||||
# ======================================================================
|
||||
def _build_summary_prompt(conversation_text: str, previous_summary: Optional[str]) -> str:
|
||||
"""pi generateSummaryWithUsage 的 prompt 组装（逐字结构）"""
|
||||
base = UPDATE_SUMMARIZATION_PROMPT if previous_summary else SUMMARIZATION_PROMPT
|
||||
prompt = f"<conversation>\n{conversation_text}\n</conversation>\n\n"
|
||||
if previous_summary:
|
||||
prompt += f"<previous-summary>\n{previous_summary}\n</previous-summary>\n\n"
|
||||
prompt += base
|
||||
return prompt
|
||||
|
||||
|
||||
def generate_summary(messages: List[AgentMessage],
|
||||
summarize_fn: Callable[[str, str, int], str],
|
||||
reserve_tokens: int,
|
||||
model_max_tokens: int,
|
||||
previous_summary: Optional[str] = None
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""
|
||||
返回 (summary_text, error)。
|
||||
maxTokens = min(0.8 × reserveTokens, model.maxTokens) —— 1:1 对照。
|
||||
"""
|
||||
max_tokens = min(
|
||||
int(0.8 * reserve_tokens),
|
||||
model_max_tokens if model_max_tokens > 0 else (1 << 30),
|
||||
)
|
||||
conversation = serialize_conversation(messages)
|
||||
prompt = _build_summary_prompt(conversation, previous_summary)
|
||||
try:
|
||||
text = summarize_fn(prompt, SUMMARIZATION_SYSTEM_PROMPT, max_tokens)
|
||||
except Exception as e:
|
||||
return None, f"Summarization failed: {e}"
|
||||
if not text or not text.strip():
|
||||
return None, "Summarization failed: empty response"
|
||||
return text.strip(), None
|
||||
|
||||
|
||||
def generate_turn_prefix_summary(messages: List[AgentMessage],
|
||||
summarize_fn: Callable[[str, str, int], str],
|
||||
reserve_tokens: int,
|
||||
model_max_tokens: int
|
||||
) -> Tuple[Optional[str], Optional[str]]:
|
||||
"""maxTokens = min(0.5 × reserveTokens, model.maxTokens) —— 1:1 对照"""
|
||||
max_tokens = min(
|
||||
int(0.5 * reserve_tokens),
|
||||
model_max_tokens if model_max_tokens > 0 else (1 << 30),
|
||||
)
|
||||
conversation = serialize_conversation(messages)
|
||||
prompt = f"<conversation>\n{conversation}\n</conversation>\n\n{TURN_PREFIX_SUMMARIZATION_PROMPT}"
|
||||
try:
|
||||
text = summarize_fn(prompt, SUMMARIZATION_SYSTEM_PROMPT, max_tokens)
|
||||
except Exception as e:
|
||||
return None, f"Turn prefix summarization failed: {e}"
|
||||
if not text or not text.strip():
|
||||
return None, "Turn prefix summarization failed: empty response"
|
||||
return text.strip(), None
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 主入口 —— 1:1 对照 compact()
|
||||
# ======================================================================
|
||||
def compact_context(messages: List[AgentMessage],
|
||||
model: ModelConfig,
|
||||
summarize_fn: Callable[[str, str, int], str],
|
||||
settings: Optional[CompactionSettings] = None
|
||||
) -> Optional[List[AgentMessage]]:
|
||||
"""
|
||||
执行压缩。返回 [压缩摘要消息] + 保留尾巴；不可压缩时返回 None。
|
||||
|
||||
摘要消息: role="user", kind="compaction_summary"，content 为纯摘要文本
|
||||
（pi 的 compaction 条目；下次压缩时自动走迭代更新提示词）。
|
||||
|
||||
断轮（切点落在某轮中间）时 1:1 对照 pi compact()：
|
||||
历史摘要 与 轮前缀摘要 分两次 LLM 调用，拼接为
|
||||
{history}\n\n---\n\n**Turn Context (split turn):**\n\n{prefix}
|
||||
最后追加文件操作附录（<read-files>/<modified-files>）。
|
||||
"""
|
||||
settings = settings or DEFAULT_COMPACTION_SETTINGS
|
||||
prep = prepare_compaction(messages, settings)
|
||||
if prep is None:
|
||||
return None
|
||||
|
||||
history_text: Optional[str] = None
|
||||
prefix_error: Optional[str] = None
|
||||
history_error: Optional[str] = None
|
||||
|
||||
if prep.is_split_turn and prep.turn_prefix_messages:
|
||||
if prep.messages_to_summarize:
|
||||
history_text, history_error = generate_summary(
|
||||
prep.messages_to_summarize, summarize_fn,
|
||||
prep.settings.reserve_tokens, model.max_tokens,
|
||||
prep.previous_summary)
|
||||
if history_error:
|
||||
return _degraded_compact(prep, history_error)
|
||||
else:
|
||||
history_text = "No prior history."
|
||||
prefix_text, prefix_error = generate_turn_prefix_summary(
|
||||
prep.turn_prefix_messages, summarize_fn,
|
||||
prep.settings.reserve_tokens, model.max_tokens)
|
||||
if prefix_error:
|
||||
return _degraded_compact(prep, prefix_error)
|
||||
summary = (f"{history_text}\n\n---\n\n"
|
||||
f"**Turn Context (split turn):**\n\n{prefix_text}")
|
||||
else:
|
||||
if not prep.messages_to_summarize:
|
||||
return None
|
||||
summary, history_error = generate_summary(
|
||||
prep.messages_to_summarize, summarize_fn,
|
||||
prep.settings.reserve_tokens, model.max_tokens,
|
||||
prep.previous_summary)
|
||||
if history_error:
|
||||
return _degraded_compact(prep, history_error)
|
||||
|
||||
read_files, modified_files = compute_file_lists(prep.file_ops)
|
||||
summary += format_file_operations(read_files, modified_files)
|
||||
|
||||
summary_msg = AgentMessage(role="user", content=summary,
|
||||
kind="compaction_summary")
|
||||
return [summary_msg] + prep.retained_tail
|
||||
|
||||
|
||||
def _degraded_compact(prep: CompactionPreparation,
|
||||
error: str) -> Optional[List[AgentMessage]]:
|
||||
"""
|
||||
已声明偏差（对照 pi: 直接返回 CompactionError）：
|
||||
桌面应用优先「不丢上下文」——摘要失败时降级为机械摘录，
|
||||
保留尾巴原样不动。
|
||||
"""
|
||||
old = prep.messages_to_summarize + prep.turn_prefix_messages
|
||||
if not old:
|
||||
return None
|
||||
excerpt = serialize_conversation(old)[-500:]
|
||||
summary = (f"(自动压缩：摘要生成失败 [{error}]，以下为旧对话尾部摘录)\n\n{excerpt}")
|
||||
summary_msg = AgentMessage(role="user", content=summary,
|
||||
kind="compaction_summary")
|
||||
return [summary_msg] + prep.retained_tail
|
||||
@@ -0,0 +1,341 @@
|
||||
"""
|
||||
core/agent/context.py
|
||||
=====================
|
||||
🌟 pi 上下文管理的 Python 1:1 移植 —— 令牌估算 + 输出预算钳制
|
||||
|
||||
对照 pi-main 源码：
|
||||
packages/ai/src/api/simple-options.ts
|
||||
- CONTEXT_SAFETY_TOKENS = 4096 (行 12)
|
||||
- clampMaxTokensToContext() (行 12-34)
|
||||
maxOutput = contextWindow - inputTokens - 4096
|
||||
return min(model.maxTokens, maxOutput) (下限 1)
|
||||
- clampOutputsToContext() (行 36-60)
|
||||
多模型共享窗口时按比例分配剩余输出预算
|
||||
packages/ai/src/api/token-utils.ts
|
||||
- 每条消息估算: ceil(chars/4) + 4
|
||||
|
||||
🌟 已声明偏差（仅 1 处，其余全部 1:1）：
|
||||
pi 按 chars/4 估算（针对英文）。中文 1 字 ≈ 1 token，chars/4 会严重低估。
|
||||
这里采用 CJK 感知估算：CJK 字符按 1 token/字，其余按 4 字符/token。
|
||||
这直接影响「压缩触发时机」和「输出预算钳制」，必须更准才不会撞 vLLM 上限。
|
||||
|
||||
pi 的 usage 锚定机制 1:1 保留：若历史里存在有效 assistant usage，
|
||||
总估算 = 该 usage 的 totalTokens（服务商精确值）+ 其后消息的逐条估算。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
from .types import AgentMessage, ModelConfig
|
||||
|
||||
# ======================================================================
|
||||
# 对照 simple-options.ts: export const CONTEXT_SAFETY_TOKENS = 4096
|
||||
# 用途：上下文钳制时预留的安全余量（tokenizer 估算误差 + vLLM 精确计数的偏差）
|
||||
# ======================================================================
|
||||
CONTEXT_SAFETY_TOKENS = 4096
|
||||
|
||||
# 图片 token 估算（对照 pi 对 image part 的固定估算）
|
||||
IMAGE_TOKENS = 1600
|
||||
|
||||
_CJK_RE = re.compile(
|
||||
r"[\u2e80-\u2eff\u3040-\u30ff\u3400-\u4dbf\u4e00-\u9fff"
|
||||
r"\uf900-\ufaff\uff00-\uffef]"
|
||||
)
|
||||
|
||||
|
||||
def _count_text_chars(value: Any) -> int:
|
||||
"""统计内容里的字符数（对照 pi token-utils 的字符统计口径）"""
|
||||
if isinstance(value, str):
|
||||
return len(value)
|
||||
if isinstance(value, list): # OpenAI 多模态 content 数组
|
||||
n = 0
|
||||
for part in value:
|
||||
if isinstance(part, dict):
|
||||
if part.get("type") == "text":
|
||||
n += len(part.get("text", "") or "")
|
||||
elif part.get("type") == "image_url":
|
||||
n += 0 # 图片单独按 IMAGE_TOKENS 计
|
||||
else:
|
||||
n += len(str(part))
|
||||
return n
|
||||
return len(str(value or ""))
|
||||
|
||||
|
||||
def _count_image_parts(content: Any) -> int:
|
||||
if isinstance(content, list):
|
||||
return sum(1 for p in content
|
||||
if isinstance(p, dict) and p.get("type") == "image_url")
|
||||
return 0
|
||||
|
||||
|
||||
def estimate_message_tokens(message: AgentMessage) -> int:
|
||||
"""
|
||||
单条消息 token 估算（CJK 感知版）。
|
||||
pi 口径: ceil(chars/4) + 4。本项目: CJK 字 ×1 + 其余 /4 + 4 + 图片×1600。
|
||||
"""
|
||||
raw = message.content
|
||||
if isinstance(raw, str):
|
||||
chars = len(raw)
|
||||
cjk = len(_CJK_RE.findall(raw))
|
||||
else:
|
||||
# 数组 content：拼接所有文本部分
|
||||
joined = []
|
||||
for part in raw if isinstance(raw, list) else []:
|
||||
if isinstance(part, dict) and part.get("type") == "text":
|
||||
joined.append(part.get("text", "") or "")
|
||||
text = "".join(joined)
|
||||
chars = len(text)
|
||||
cjk = len(_CJK_RE.findall(text))
|
||||
tokens = cjk + math.ceil((chars - cjk) / 4) + 4
|
||||
tokens += _count_image_parts(raw) * IMAGE_TOKENS
|
||||
# 思考内容与工具参数也占上下文（assistant 的 reasoning / tool_calls）
|
||||
if message.reasoning:
|
||||
rcjk = len(_CJK_RE.findall(message.reasoning))
|
||||
tokens += rcjk + math.ceil((len(message.reasoning) - rcjk) / 4)
|
||||
for tc in message.tool_calls:
|
||||
arg_text = _json_str(tc.arguments)
|
||||
tokens += math.ceil(len(arg_text) / 4)
|
||||
# toolResult 的 content 走 content 字段（str 或 list）
|
||||
return tokens
|
||||
|
||||
|
||||
def _json_str(obj: Any) -> str:
|
||||
import json
|
||||
try:
|
||||
return json.dumps(obj, ensure_ascii=False)
|
||||
except Exception:
|
||||
return str(obj)
|
||||
|
||||
|
||||
def _estimate_text_tokens(text: str) -> int:
|
||||
"""CJK 感知文本 token 估算（与 estimate_message_tokens 同公式，不含每条 +4 开销）"""
|
||||
if not text:
|
||||
return 0
|
||||
cjk = len(_CJK_RE.findall(text))
|
||||
return cjk + math.ceil((len(text) - cjk) / 4)
|
||||
|
||||
|
||||
def _estimate_tools_tokens(tools) -> int:
|
||||
"""
|
||||
🆕 P2: 工具 schema token 估算 —— name/description/parameters 的 JSON 字符数（CJK 感知）。
|
||||
对照 pi Context 分支: all tools = Σ estimate_text(json(tool))。
|
||||
"""
|
||||
total = 0
|
||||
for t in tools or []:
|
||||
schema = {
|
||||
"name": getattr(t, "name", "") or "",
|
||||
"description": getattr(t, "description", "") or "",
|
||||
"parameters": getattr(t, "parameters", None),
|
||||
}
|
||||
total += _estimate_text_tokens(_json_str(schema))
|
||||
return total
|
||||
|
||||
|
||||
def estimate_context_tokens(messages: List[AgentMessage],
|
||||
system_prompt: str = "",
|
||||
tools: Optional[List[Any]] = None
|
||||
) -> "ContextUsageEstimate":
|
||||
"""
|
||||
🌟 1:1 对照 pi estimateContextTokens（usage 锚定）：
|
||||
- 找到最后一条带有效 usage 的 assistant 消息（🆕 P0: 含 timestamp 时效校验）
|
||||
- 总 tokens = 该 usage 的精确值 + 其后消息的逐条估算
|
||||
- 无 usage 时 = 全部消息逐条估算
|
||||
🆕 P2: 可选 system_prompt / tools 参数（对照 pi Context 分支）：
|
||||
- 仅在无锚点分支计入 system 提示词 + 工具 schema 的估算
|
||||
（有锚点时 usage 已是含 system+tools 的精确总值，不重复计）
|
||||
逐条估算用 CJK 感知版 estimate_message_tokens（已声明偏差）。
|
||||
"""
|
||||
last_idx, last_usage = _find_last_usage(messages)
|
||||
|
||||
sys_tokens = _estimate_text_tokens(system_prompt) if system_prompt else 0
|
||||
tool_tokens = _estimate_tools_tokens(tools) if tools else 0
|
||||
|
||||
if last_idx is None:
|
||||
estimated = (sum(estimate_message_tokens(m) for m in messages)
|
||||
+ sys_tokens + tool_tokens)
|
||||
return ContextUsageEstimate(tokens=estimated, usage_tokens=0,
|
||||
trailing_tokens=estimated, last_usage_index=None)
|
||||
|
||||
usage_tokens = calculate_context_tokens(last_usage)
|
||||
trailing = sum(estimate_message_tokens(m) for m in messages[last_idx + 1:])
|
||||
return ContextUsageEstimate(tokens=usage_tokens + trailing,
|
||||
usage_tokens=usage_tokens,
|
||||
trailing_tokens=trailing,
|
||||
last_usage_index=last_idx)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ContextUsageEstimate:
|
||||
"""
|
||||
🌟 1:1 对照 pi estimateContextTokens 的返回值
|
||||
packages/agent/src/harness/compaction/compaction.ts ContextUsageEstimate:
|
||||
tokens 总估算
|
||||
usageTokens 最后一条有效 assistant usage 报告的精确 token 数
|
||||
trailingTokens 该 usage 之后消息的估算值
|
||||
lastUsageIndex 提供 usage 的消息下标（无则 None）
|
||||
"""
|
||||
tokens: int = 0
|
||||
usage_tokens: int = 0
|
||||
trailing_tokens: int = 0
|
||||
last_usage_index: Optional[int] = None
|
||||
|
||||
|
||||
def calculate_context_tokens(usage: Dict[str, Any]) -> int:
|
||||
"""pi calculateContextTokens: totalTokens || input+output+cacheRead+cacheWrite"""
|
||||
if not usage:
|
||||
return 0
|
||||
total = int(usage.get("totalTokens", 0) or 0)
|
||||
if total:
|
||||
return total
|
||||
return (int(usage.get("input", 0) or 0) + int(usage.get("output", 0) or 0)
|
||||
+ int(usage.get("cacheRead", 0) or 0)
|
||||
+ int(usage.get("cacheWrite", 0) or 0))
|
||||
|
||||
|
||||
def _get_assistant_usage(msg: AgentMessage) -> Dict[str, Any]:
|
||||
"""
|
||||
pi getAssistantUsage：只认「有效」的 assistant usage——
|
||||
stopReason 不是 aborted/error，且 calculateContextTokens > 0。
|
||||
（🆕 P0: 锚点扫描还需通过 timestamp 时效校验，见 _find_last_usage）
|
||||
"""
|
||||
if msg.role != "assistant":
|
||||
return {}
|
||||
if msg.stop_reason in ("aborted", "error"):
|
||||
return {}
|
||||
usage = msg.usage or {}
|
||||
if usage and calculate_context_tokens(usage) > 0:
|
||||
return usage
|
||||
return {}
|
||||
|
||||
|
||||
def _row_block_tool_flags(messages: List["AgentMessage"]) -> List[bool]:
|
||||
"""
|
||||
🆕 G1: 逐条标记「其 DB 行（同非空 db_msg_id 的连续块）是否含工具活动」。
|
||||
含工具输出的行，其入库 usage 是「本轮最后一个子请求」的内存快照
|
||||
（工具结果未受 4000 字回放截断），不等于下一请求（截断回放）→ 锚点失效。
|
||||
轮中内存条目（db_msg_id=""）一律 False：保持原规则
|
||||
（轮中条目的内存载荷 == 下一子请求载荷，锚点有效）。
|
||||
"""
|
||||
n = len(messages)
|
||||
flags = [False] * n
|
||||
i = 0
|
||||
while i < n:
|
||||
rid = getattr(messages[i], "db_msg_id", "") or ""
|
||||
if not rid:
|
||||
i += 1
|
||||
continue
|
||||
j = i
|
||||
has_tool = False
|
||||
while j < n and (getattr(messages[j], "db_msg_id", "") or "") == rid:
|
||||
mj = messages[j]
|
||||
if mj.role == "toolResult" or (mj.role == "assistant" and mj.tool_calls):
|
||||
has_tool = True
|
||||
j += 1
|
||||
for k in range(i, j):
|
||||
flags[k] = has_tool
|
||||
i = j
|
||||
return flags
|
||||
|
||||
|
||||
def _find_last_usage(messages: List["AgentMessage"]):
|
||||
"""
|
||||
🆕 P0: 1:1 对照 pi getLastAssistantUsageInfo（含锚点失效校验）：
|
||||
assistant 的 usage 要成为锚点，必须满足
|
||||
① timestamp >= 它前面所有消息的最大 timestamp（即它不比任何前消息「更旧」）。
|
||||
场景：压缩后 state = [新摘要消息(T_new)] + [保留的旧消息(T_old)]，
|
||||
保留窗里 assistant 的 usage 来自压缩前（如 60K 精确值）——不校验的话，
|
||||
下一次 should_compact 会锚定到过期值 → 误触发二次压缩（摘要套摘要）。
|
||||
摘要消息 timestamp 晚于全部保留消息 → 前缀游标抬高 → 旧 usage 自动失效；
|
||||
压缩后第一个新 assistant 回复到达 → 新 usage 锚点自动恢复。
|
||||
② 🆕 G1: 该 usage 所在 DB 行不含工具活动（toolResult / assistant.tool_calls）。
|
||||
含工具输出的行，入库 usage = 本轮内存快照（工具结果未截断）≠ 下一请求
|
||||
（4000 字截断回放）→ 系统性高估 → 误触发压缩（实测 113904 锚点 vs
|
||||
真实下一请求 ~90k）。纯文本行 / 轮中内存条目的 usage 仍是精确值，保留锚定。
|
||||
（消息无 timestamp 信息（全 0）时行为等同旧版「取最后一条有效」，完全向后兼容。）
|
||||
"""
|
||||
latest_prefix_ts = -1 # 对照 pi: Number.NEGATIVE_INFINITY
|
||||
last_idx = None
|
||||
last_usage: Dict[str, Any] = {}
|
||||
row_tool_flags = _row_block_tool_flags(messages)
|
||||
for i, msg in enumerate(messages):
|
||||
if msg.role == "assistant":
|
||||
applies = (msg.timestamp or 0) >= latest_prefix_ts
|
||||
if (applies
|
||||
and not row_tool_flags[i]
|
||||
and msg.stop_reason not in ("aborted", "error")
|
||||
and msg.usage
|
||||
and calculate_context_tokens(msg.usage) > 0):
|
||||
last_idx, last_usage = i, msg.usage
|
||||
ts = msg.timestamp or 0
|
||||
if ts > latest_prefix_ts:
|
||||
latest_prefix_ts = ts
|
||||
return last_idx, last_usage
|
||||
|
||||
|
||||
def clamp_max_tokens_to_context(model: ModelConfig,
|
||||
context: List[AgentMessage],
|
||||
system_prompt: str = "",
|
||||
tools: Optional[List[Any]] = None
|
||||
) -> Optional[Tuple[int, int]]:
|
||||
"""
|
||||
🌟 对照 simple-options.ts:12-34 clampMaxTokensToContext（1:1 公式）：
|
||||
|
||||
input_estimate = estimateContextTokens(context) # 🆕 P2: 含 system+tools（pi Context 实参）
|
||||
if (input_estimate > model.contextWindow) return null // 上下文已溢出
|
||||
maxOutput = contextWindow - inputTokens - CONTEXT_SAFETY_TOKENS
|
||||
return max(1, min(model.maxTokens, maxOutput))
|
||||
|
||||
这就是 pi 约束公式的本体：
|
||||
输入 + 输出 ≤ contextWindow − 4096 ≤ vLLM max_model_len
|
||||
🆕 P2: system_prompt / tools 仅在无 usage 锚点分支计入（首轮流式精确记账，
|
||||
与 pi 传完整 Context 的行为一致）；有锚点时 usage 已精确覆盖，不重复计。
|
||||
返回 (max_tokens, input_tokens)；输入已溢出窗口时返回 None（由恢复逻辑接管）。
|
||||
"""
|
||||
input_tokens = estimate_context_tokens(
|
||||
context, system_prompt=system_prompt, tools=tools).tokens
|
||||
if input_tokens > model.context_window:
|
||||
return None
|
||||
max_output = model.context_window - input_tokens - CONTEXT_SAFETY_TOKENS
|
||||
return max(1, min(model.max_tokens, max_output)), input_tokens
|
||||
|
||||
|
||||
def clamp_outputs_to_context(context: List[AgentMessage],
|
||||
models: List[ModelConfig],
|
||||
reserve_tokens: int = CONTEXT_SAFETY_TOKENS
|
||||
) -> List[int]:
|
||||
"""
|
||||
对照 simple-options.ts:36-60 clampOutputsToContext：
|
||||
多个模型共享同一窗口时，把剩余输出预算按比例分配给各模型。
|
||||
（haocode 单模型场景用不到，保留以求框架完整）
|
||||
"""
|
||||
input_tokens = estimate_context_tokens(context).tokens
|
||||
available = max(0, sum(m.context_window for m in models) - input_tokens - reserve_tokens)
|
||||
if not models:
|
||||
return []
|
||||
weights = [m.max_tokens for m in models]
|
||||
total_w = sum(weights) or 1
|
||||
alloc = [max(1, int(available * w / total_w)) for w in weights]
|
||||
# 各自不超过自身 max_tokens
|
||||
alloc = [min(a, m.max_tokens) for a, m in zip(alloc, models)]
|
||||
return alloc
|
||||
|
||||
|
||||
def should_compact(messages: List[AgentMessage], model: ModelConfig,
|
||||
reserve_tokens: int = 16384,
|
||||
system_prompt: str = "",
|
||||
tools: Optional[List[Any]] = None) -> Tuple[bool, int]:
|
||||
"""
|
||||
🌟 1:1 对照 pi compaction.ts shouldCompact（harness 原版公式）：
|
||||
contextTokens > contextWindow - reserveTokens
|
||||
reserveTokens 默认 16384（DEFAULT_COMPACTION_SETTINGS.reserveTokens），
|
||||
即「为摘要提示词与输出预留 16K」。估算值取 usage 锚定估算。
|
||||
🆕 G2: system_prompt/tools 透传 —— 无锚点分支（工具行快照锚点被 G1 失效后
|
||||
更常走到）也要计入 system + 工具 schema，度量「下一请求真实载荷」，
|
||||
与显示端（update_context_display）同口径。
|
||||
"""
|
||||
tokens = estimate_context_tokens(messages, system_prompt=system_prompt,
|
||||
tools=tools).tokens
|
||||
return (tokens > model.context_window - reserve_tokens), tokens
|
||||
@@ -0,0 +1,482 @@
|
||||
"""
|
||||
core/agent/loop.py
|
||||
==================
|
||||
🌟 pi 核心循环的 Python 1:1 移植 —— 整个框架的心脏
|
||||
|
||||
对照 pi-main 源码（packages/agent/src/agent-loop.ts, 796 行）：
|
||||
runAgentLoop() → run_loop(agent, new_message, ...)
|
||||
runLoop() → _run_loop() (行 163-278)
|
||||
streamAssistantResponse() → _stream_turn() (行 281-380)
|
||||
executeToolCalls() → execute_tool_calls() (行 413-427)
|
||||
executeToolCallsParallel/Sequential → _execute_parallel/_execute_sequential
|
||||
prepareToolCall() → tools.prepare_tool_calls
|
||||
failToolCallsFromTruncatedMessage → tools.fail_tool_calls_from_truncated_message
|
||||
shouldTerminateToolBatch → _should_terminate_batch (行 561-563)
|
||||
|
||||
pi runLoop 的结构（本文件逐行对应）：
|
||||
emit agent_start
|
||||
pending = getSteeringMessages() # 循环开始前先取一次
|
||||
outer: while True:
|
||||
inner: while hasMoreToolCalls or pending:
|
||||
turn_start（首轮不发，首轮 turn_start 在 runAgentLoop 入口发）
|
||||
注入 pending（message_start/end → 入 context + newMessages）
|
||||
assistant = streamAssistantResponse(...)
|
||||
if stopReason in (error, aborted): turn_end + agent_end + return
|
||||
toolCalls = assistant 的 toolCall 块
|
||||
if toolCalls:
|
||||
length → failToolCallsFromTruncatedMessage（不执行！）
|
||||
否则 → executeToolCalls（sequential/parallel 二选一）
|
||||
hasMoreToolCalls = !batch.terminate
|
||||
toolResults 入 context + newMessages
|
||||
turn_end
|
||||
prepareNextTurn 钩子（可换 context/model）
|
||||
shouldStopAfterTurn 钩子 → agent_end + return
|
||||
pending = getSteeringMessages() # 每轮结束取一次
|
||||
followUps = getFollowUpMessages()
|
||||
if followUps: pending = followUps; continue # 外层续跑
|
||||
break
|
||||
agent_end
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple
|
||||
|
||||
from .context import clamp_max_tokens_to_context
|
||||
from .tools import (PreparedToolCall, execute_tool_call, fail_tool_calls_from_truncated_message,
|
||||
prepare_tool_calls)
|
||||
from .types import (AgentConfig, AgentError, AgentEvent, AgentMessage,
|
||||
AgentTool, AgentToolResult, AbortSignal, AssistantMessageEvent,
|
||||
RunResult, ToolCall, new_id)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 事件发射辅助
|
||||
# ======================================================================
|
||||
def _emit(agent, event: AgentEvent):
|
||||
agent._emit(event)
|
||||
|
||||
|
||||
def _message_events(agent, msg: AgentMessage):
|
||||
_emit(agent, AgentEvent(type="message_start", message=msg))
|
||||
_emit(agent, AgentEvent(type="message_end", message=msg))
|
||||
|
||||
|
||||
def _tool_result_message(finalized: Dict[str, Any]) -> AgentMessage:
|
||||
"""对照 createToolResultMessage"""
|
||||
tc: ToolCall = finalized["tool_call"]
|
||||
result: AgentToolResult = finalized["result"]
|
||||
return AgentMessage(
|
||||
role="toolResult",
|
||||
tool_call_id=tc.id,
|
||||
tool_name=tc.name,
|
||||
content=result.content,
|
||||
is_error=finalized.get("is_error", False) or result.is_error,
|
||||
)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 流式生成一轮助手消息 —— 对照 streamAssistantResponse (行 281-380)
|
||||
# ======================================================================
|
||||
def _stream_turn(agent, current_context: List[AgentMessage],
|
||||
config: AgentConfig, signal: AbortSignal,
|
||||
stream_fn: Callable) -> Tuple[AgentMessage, Optional[AgentError]]:
|
||||
"""
|
||||
返回 (assistant_message, error)。
|
||||
error 非 None 时 message.stop_reason == "error"。
|
||||
"""
|
||||
# 🌟 pi 同款：流开始前检查中止 → 立即产出 aborted 消息
|
||||
if signal.aborted:
|
||||
msg = AgentMessage(role="assistant", stop_reason="aborted")
|
||||
agent.state.messages.append(msg)
|
||||
_message_events(agent, msg)
|
||||
return msg, None
|
||||
|
||||
# 🌟 系统提示词注入（对照 pi：systemPrompt 放在每次 API 请求头部，
|
||||
# 不进 state.messages、不占压缩/历史）
|
||||
if config.system_prompt:
|
||||
api_context: List[AgentMessage] = [
|
||||
AgentMessage(role="system", content=config.system_prompt)
|
||||
] + current_context
|
||||
else:
|
||||
api_context = current_context
|
||||
|
||||
# 🌟 输出预算钳制（对照 simple-options clampMaxTokensToContext）
|
||||
# 🆕 P2: 把 system + 工具 schema 传给估算器（对照 pi 传完整 Context），
|
||||
# 仅在无 usage 锚点分支生效；否则新会话首轮会多预留 ~system+tools 的预算。
|
||||
clamped = clamp_max_tokens_to_context(
|
||||
config.model, current_context,
|
||||
system_prompt=config.system_prompt or "", tools=config.tools)
|
||||
if clamped is None:
|
||||
err = AgentError(
|
||||
message="上下文溢出：估算输入 token 已超过模型窗口（需要压缩）",
|
||||
kind="overflow", recoverable=True)
|
||||
msg = AgentMessage(role="assistant", stop_reason="error",
|
||||
error_message=err.message)
|
||||
agent.state.messages.append(msg)
|
||||
_message_events(agent, msg)
|
||||
return msg, err
|
||||
max_tokens, _input_tokens = clamped
|
||||
|
||||
msg = AgentMessage(role="assistant")
|
||||
agent.state.streaming_message = msg
|
||||
agent.state.streaming_delta = {}
|
||||
_emit(agent, AgentEvent(type="message_start", message=msg))
|
||||
|
||||
raw_tc: Dict[int, Dict[str, str]] = {} # 中止时用于保留部分工具调用（仅展示）
|
||||
got_final = False
|
||||
error: Optional[AgentError] = None
|
||||
|
||||
try:
|
||||
for kind, payload in stream_fn(api_context, config.model, signal,
|
||||
max_tokens, config.tools):
|
||||
if kind == "event":
|
||||
ev = payload
|
||||
if ev.type == "text_delta":
|
||||
msg.content = (msg.content or "") + ev.text
|
||||
agent.state.streaming_delta["text"] = ev.text
|
||||
_emit(agent, AgentEvent(type="message_update", message=msg,
|
||||
assistant_message_event=ev))
|
||||
elif ev.type == "thinking_delta":
|
||||
msg.reasoning += ev.text
|
||||
agent.state.streaming_delta["thinking"] = ev.text
|
||||
_emit(agent, AgentEvent(type="message_update", message=msg,
|
||||
assistant_message_event=ev))
|
||||
elif ev.type == "toolcall_delta":
|
||||
slot = raw_tc.setdefault(ev.tool_call_index,
|
||||
{"id": "", "name": "", "args": ""})
|
||||
if ev.tool_call_field == "id":
|
||||
slot["id"] = ev.tool_call_delta
|
||||
elif ev.tool_call_field == "name":
|
||||
slot["name"] += ev.tool_call_delta
|
||||
else:
|
||||
slot["args"] += ev.tool_call_delta
|
||||
_emit(agent, AgentEvent(type="message_update", message=msg,
|
||||
assistant_message_event=ev))
|
||||
else: # final
|
||||
final: AgentMessage = payload
|
||||
got_final = True
|
||||
# 采用权威 final：usage / stop_reason / 解析好的 tool_calls
|
||||
msg.usage = final.usage
|
||||
msg.stop_reason = final.stop_reason
|
||||
msg.tool_calls = final.tool_calls
|
||||
if not final.content and msg.content:
|
||||
pass # 以循环侧累积为准（二者一致）
|
||||
else:
|
||||
msg.content = final.content if final.content else msg.content
|
||||
msg.reasoning = final.reasoning or msg.reasoning
|
||||
except AgentError as e:
|
||||
error = e
|
||||
msg.stop_reason = "error"
|
||||
msg.error_message = e.message
|
||||
_emit(agent, AgentEvent(type="message_update", message=msg,
|
||||
assistant_message_event=
|
||||
AssistantMessageEvent.error(e.message)))
|
||||
except Exception as e:
|
||||
from .stream_fn import classify_error
|
||||
error = classify_error(e)
|
||||
msg.stop_reason = "error"
|
||||
msg.error_message = error.message
|
||||
|
||||
if not got_final and error is None:
|
||||
# 生成器中途结束且无异常 = 中止（流被关闭，pi 同款语义）
|
||||
msg.stop_reason = "aborted"
|
||||
# 保留部分工具调用（仅用于 UI 展示；中止后不会执行，也不会进入下次上下文）
|
||||
for idx in sorted(raw_tc.keys()):
|
||||
slot = raw_tc[idx]
|
||||
if slot.get("name"):
|
||||
msg.tool_calls.append(ToolCall(id=slot.get("id") or new_id("call"),
|
||||
name=slot["name"],
|
||||
raw_arguments=slot.get("args", "")))
|
||||
if error is None and msg.stop_reason not in ("stop", "length", "aborted"):
|
||||
msg.stop_reason = "stop"
|
||||
|
||||
# 🌟 兜底（haocode 扩展，pi 无此层）：不支持 tools API 的供应商/模型会把
|
||||
# 工具调用用文字"演"出来（如 <bash>ls</bash>）。识别单参工具 bash/read
|
||||
# 转成真 tool_call 继续执行；write/edit 多参歧义大不兜底。
|
||||
if error is None and not msg.tool_calls and msg.content:
|
||||
from .tools import parse_text_tool_calls
|
||||
cleaned, txt_calls = parse_text_tool_calls(msg.content)
|
||||
if txt_calls:
|
||||
msg.tool_calls = txt_calls
|
||||
|
||||
agent.state.streaming_message = None
|
||||
agent.state.streaming_delta = {}
|
||||
agent.state.messages.append(msg)
|
||||
_emit(agent, AgentEvent(type="message_end", message=msg))
|
||||
return msg, error
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 工具批量执行 —— 对照 executeToolCalls (行 413-427)
|
||||
# ======================================================================
|
||||
def _should_terminate_batch(finalized: List[Dict[str, Any]]) -> bool:
|
||||
"""对照 shouldTerminateToolBatch: 全部结果都请求 terminate 才终止"""
|
||||
return len(finalized) > 0 and all(
|
||||
f["result"].terminate for f in finalized)
|
||||
|
||||
|
||||
def _run_prepared_in_pool(prep: PreparedToolCall, assistant: AgentMessage,
|
||||
config: AgentConfig, signal: AbortSignal,
|
||||
agent) -> Dict[str, Any]:
|
||||
"""线程池内执行单个工具（对照 parallel 版的 async 闭包）"""
|
||||
on_update = None
|
||||
on_timer = None
|
||||
if agent is not None:
|
||||
def on_update(partial: str):
|
||||
_emit(agent, AgentEvent(type="tool_execution_update",
|
||||
tool_call=prep.tool_call, arg=partial))
|
||||
def on_timer(elapsed_i: int, timeout_i: int):
|
||||
# 🆕 bash 运行中每秒滴一次 → 前端气泡读秒 N/Ts
|
||||
_emit(agent, AgentEvent(type="tool_execution_timer",
|
||||
tool_call=prep.tool_call,
|
||||
arg=(int(elapsed_i), int(timeout_i))))
|
||||
result = execute_tool_call(prep, assistant, config, signal, on_update,
|
||||
on_timer)
|
||||
finalized = {"tool_call": prep.tool_call, "result": result,
|
||||
"is_error": result.is_error}
|
||||
_emit(agent, AgentEvent(type="tool_execution_end",
|
||||
tool_call=prep.tool_call, result=result,
|
||||
is_error=result.is_error))
|
||||
return finalized
|
||||
|
||||
|
||||
def _execute_parallel(current_context, assistant, config, signal, agent,
|
||||
prepared: List[PreparedToolCall]) -> Dict[str, Any]:
|
||||
"""对照 executeToolCallsParallel：start 事件串行发、准备串行、执行并发、
|
||||
结果消息按原始顺序产出"""
|
||||
finalized: List[Dict[str, Any]] = []
|
||||
pending_futures: List = []
|
||||
max_workers = max(1, min(8, len(prepared)))
|
||||
with ThreadPoolExecutor(max_workers=max_workers,
|
||||
thread_name_prefix="tool") as pool:
|
||||
for prep in prepared:
|
||||
_emit(agent, AgentEvent(type="tool_execution_start",
|
||||
tool_call=prep.tool_call,
|
||||
arg=str(prep.args) if prep.args else ""))
|
||||
if prep.error:
|
||||
# 准备失败（未知工具/参数非法）→ 立即结果（对照 kind:"immediate"）
|
||||
result = AgentToolResult.text(prep.error, is_error=True)
|
||||
fin = {"tool_call": prep.tool_call, "result": result,
|
||||
"is_error": True}
|
||||
_emit(agent, AgentEvent(type="tool_execution_end",
|
||||
tool_call=prep.tool_call, result=result,
|
||||
is_error=True))
|
||||
finalized.append(fin)
|
||||
if signal.aborted:
|
||||
break
|
||||
continue
|
||||
fut = pool.submit(_run_prepared_in_pool, prep, assistant, config,
|
||||
signal, agent)
|
||||
pending_futures.append(fut)
|
||||
if signal.aborted:
|
||||
# 已提交的任务仍会完成（它们内部检查 signal），不再追加
|
||||
pass
|
||||
for fut in pending_futures:
|
||||
finalized.append(fut.result())
|
||||
|
||||
messages: List[AgentMessage] = []
|
||||
for fin in finalized:
|
||||
tr = _tool_result_message(fin)
|
||||
_message_events(agent, tr)
|
||||
messages.append(tr)
|
||||
return {"messages": messages, "terminate": _should_terminate_batch(finalized)}
|
||||
|
||||
|
||||
def _execute_sequential(current_context, assistant, config, signal, agent,
|
||||
prepared: List[PreparedToolCall]) -> Dict[str, Any]:
|
||||
"""对照 executeToolCallsSequential：一次一个，完成一个再下一个"""
|
||||
finalized: List[Dict[str, Any]] = []
|
||||
messages: List[AgentMessage] = []
|
||||
for prep in prepared:
|
||||
_emit(agent, AgentEvent(type="tool_execution_start",
|
||||
tool_call=prep.tool_call,
|
||||
arg=str(prep.args) if prep.args else ""))
|
||||
if not prep.error:
|
||||
fin = _run_prepared_in_pool(prep, assistant, config, signal, agent)
|
||||
else:
|
||||
result = AgentToolResult.text(prep.error, is_error=True)
|
||||
fin = {"tool_call": prep.tool_call, "result": result, "is_error": True}
|
||||
_emit(agent, AgentEvent(type="tool_execution_end",
|
||||
tool_call=prep.tool_call, result=result,
|
||||
is_error=True))
|
||||
finalized.append(fin)
|
||||
tr = _tool_result_message(fin)
|
||||
_message_events(agent, tr)
|
||||
messages.append(tr)
|
||||
if signal.aborted:
|
||||
break
|
||||
return {"messages": messages, "terminate": _should_terminate_batch(finalized)}
|
||||
|
||||
|
||||
def execute_tool_calls(current_context, assistant: AgentMessage,
|
||||
config: AgentConfig, signal: AbortSignal, agent) -> Dict[str, Any]:
|
||||
"""对照 executeToolCalls: 批次里有 sequential 工具 → 整批串行"""
|
||||
prepared = prepare_tool_calls(assistant, config.tools)
|
||||
has_sequential = any(
|
||||
(p.tool and p.tool.execution_mode == "sequential") for p in prepared
|
||||
if not p.error)
|
||||
if config.tool_execution == "sequential" or has_sequential:
|
||||
return _execute_sequential(current_context, assistant, config, signal,
|
||||
agent, prepared)
|
||||
return _execute_parallel(current_context, assistant, config, signal, agent,
|
||||
prepared)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 主循环 —— 对照 runLoop (行 163-278)
|
||||
# ======================================================================
|
||||
def run_loop(agent, new_message: Optional[AgentMessage], signal: AbortSignal,
|
||||
stream_fn: Callable) -> RunResult:
|
||||
config = agent.config
|
||||
new_messages: List[AgentMessage] = []
|
||||
|
||||
# 入口校验（对照 runAgentLoopContinue 的前置检查由 Agent 层负责）
|
||||
_emit(agent, AgentEvent(type="agent_start"))
|
||||
_emit(agent, AgentEvent(type="turn_start")) # 首轮 turn_start（对照行 142/170）
|
||||
|
||||
# context 准备 + transformContext 钩子
|
||||
base = list(agent.state.messages)
|
||||
if config.transform_context:
|
||||
try:
|
||||
base = config.transform_context(base) or base
|
||||
except Exception:
|
||||
pass
|
||||
current_context = list(base)
|
||||
|
||||
first_turn = True
|
||||
pending: List[AgentMessage] = agent._take_steering()
|
||||
|
||||
while True: # outer loop
|
||||
has_more_tool_calls = True
|
||||
|
||||
while has_more_tool_calls or pending: # inner loop
|
||||
if not first_turn:
|
||||
_emit(agent, AgentEvent(type="turn_start"))
|
||||
else:
|
||||
first_turn = False
|
||||
|
||||
# ---- 注入 pending 消息（steering / followUp）----
|
||||
if pending:
|
||||
for m in pending:
|
||||
_message_events(agent, m)
|
||||
current_context.append(m)
|
||||
new_messages.append(m)
|
||||
agent.state.messages.append(m)
|
||||
pending = []
|
||||
|
||||
# ---- 🆕 轮中主动压缩检查（haocode 增强，偏离 pi 1:1）----
|
||||
# 单条工具输出可能把上下文顶出窗口；发下一次请求前主动检查
|
||||
#（与轮首 should_compact 同公式）。compact_fn 返回新列表
|
||||
#（发生了压缩）→ 同步循环局部上下文。
|
||||
if config.compact_fn is not None and agent.state.messages:
|
||||
try:
|
||||
compacted = config.compact_fn(list(agent.state.messages))
|
||||
if compacted is not None:
|
||||
agent.state.messages = compacted
|
||||
current_context = list(compacted)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---- 流式生成助手消息 ----
|
||||
assistant, error = _stream_turn(agent, current_context, config,
|
||||
signal, stream_fn)
|
||||
new_messages.append(assistant)
|
||||
current_context.append(assistant)
|
||||
|
||||
if assistant.stop_reason in ("error", "aborted"):
|
||||
agent.state.error = error
|
||||
_emit(agent, AgentEvent(type="turn_end", message=assistant))
|
||||
_emit(agent, AgentEvent(type="agent_end",
|
||||
stop_reason=assistant.stop_reason,
|
||||
error=error,
|
||||
messages=new_messages))
|
||||
agent._finish_run(new_messages, assistant.stop_reason, error)
|
||||
return RunResult(stop_reason=assistant.stop_reason, error=error,
|
||||
message_count=len(new_messages))
|
||||
|
||||
# ---- 工具调用 ----
|
||||
tool_results: List[AgentMessage] = []
|
||||
has_more_tool_calls = False
|
||||
if assistant.tool_calls:
|
||||
if assistant.stop_reason == "length":
|
||||
# 🌟 截断保护（对照行 212-213）：参数可能残缺，一律失败不执行
|
||||
fail_msgs = fail_tool_calls_from_truncated_message(
|
||||
assistant, "length")
|
||||
# 事件流与正常执行对齐
|
||||
for tc, tr in zip(assistant.tool_calls, fail_msgs):
|
||||
_emit(agent, AgentEvent(type="tool_execution_start",
|
||||
tool_call=tc))
|
||||
res = AgentToolResult.text(
|
||||
f'工具调用 "{tc.name}" 未执行：输出达到 token 上限，'
|
||||
f'参数可能被截断。请用完整参数重新发起。', is_error=True)
|
||||
_emit(agent, AgentEvent(type="tool_execution_end",
|
||||
tool_call=tc, result=res,
|
||||
is_error=True))
|
||||
_message_events(agent, tr)
|
||||
batch = {"messages": fail_msgs, "terminate": False}
|
||||
else:
|
||||
batch = execute_tool_calls(current_context, assistant, config,
|
||||
signal, agent)
|
||||
tool_results = batch["messages"]
|
||||
has_more_tool_calls = not batch["terminate"]
|
||||
for r in tool_results:
|
||||
current_context.append(r)
|
||||
new_messages.append(r)
|
||||
agent.state.messages.append(r)
|
||||
|
||||
_emit(agent, AgentEvent(type="turn_end", message=assistant))
|
||||
|
||||
# ---- prepareNextTurn 钩子 ----
|
||||
if config.prepare_next_turn:
|
||||
try:
|
||||
snap = config.prepare_next_turn({
|
||||
"message": assistant, "tool_results": tool_results,
|
||||
"context": current_context, "new_messages": new_messages,
|
||||
})
|
||||
if snap:
|
||||
current_context = snap.get("context") or current_context
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---- shouldStopAfterTurn 钩子 ----
|
||||
if config.should_stop_after_turn:
|
||||
try:
|
||||
if config.should_stop_after_turn({
|
||||
"message": assistant, "tool_results": tool_results,
|
||||
"context": current_context,
|
||||
"new_messages": new_messages}):
|
||||
stop = _last_assistant_stop(new_messages)
|
||||
_emit(agent, AgentEvent(type="agent_end",
|
||||
stop_reason=stop,
|
||||
messages=new_messages))
|
||||
agent._finish_run(new_messages, stop, None)
|
||||
return RunResult(stop_reason=stop,
|
||||
message_count=len(new_messages))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---- 每轮结束取 steering ----
|
||||
pending = agent._take_steering()
|
||||
|
||||
# ---- 外层：followUp ----
|
||||
follow_ups = agent._take_follow_ups()
|
||||
if follow_ups:
|
||||
pending = follow_ups
|
||||
continue
|
||||
break
|
||||
|
||||
stop = _last_assistant_stop(new_messages)
|
||||
_emit(agent, AgentEvent(type="agent_end", stop_reason=stop,
|
||||
messages=new_messages))
|
||||
agent._finish_run(new_messages, stop, None)
|
||||
return RunResult(stop_reason=stop, message_count=len(new_messages))
|
||||
|
||||
|
||||
def _last_assistant_stop(new_messages: List[AgentMessage]) -> str:
|
||||
"""对照 pi: agent_end 不携带 stopReason，会话层从最后一条 assistant 消息读取"""
|
||||
for m in reversed(new_messages):
|
||||
if m.role == "assistant":
|
||||
return m.stop_reason or "stop"
|
||||
return "stop"
|
||||
@@ -0,0 +1,519 @@
|
||||
"""
|
||||
core/agent/recovery.py
|
||||
======================
|
||||
🌟 pi 会话层恢复逻辑的 Python 1:1 移植
|
||||
|
||||
对照 pi-main 源码：
|
||||
packages/coding-agent/src/core/agent-session.ts
|
||||
_runAgentPrompt (行 1074): prompt → while _handlePostAgentRun(): continue()
|
||||
_handlePostAgentRun (行 1088): 基于「最后一条 assistant 消息」的三路决策
|
||||
_prepareRetry (行 2811): 计数 → 移除坏消息 → 退避睡眠
|
||||
_isRetryableError (行 2770): 溢出交给压缩；其余走重试判定
|
||||
_checkCompaction (行 ~2034): 溢出压缩恢复（只试一次）
|
||||
packages/ai/src/utils/retry.ts
|
||||
isRetryableAssistantError (行 223) + 可重试/不可重试正则模式表
|
||||
退避: baseDelayMs × 2^(attempt-1)（默认 2s, 4s, 8s；maxRetries 默认 3）
|
||||
packages/ai/src/utils/overflow.ts
|
||||
isContextOverflow (行 134)：三种溢出识别
|
||||
isRecoverableLength：length 且 output < 原始输出上限
|
||||
|
||||
🌟 决策表（与 pi 完全一致，基于最后一条 assistant 消息）：
|
||||
(1) stop_reason=="error"：
|
||||
a. 是上下文溢出 → 删除坏消息 + 压缩 + continue（只试一次）
|
||||
b. 命中可重试模式（429/5xx/超时/断连…且非配额耗尽）
|
||||
且未超 maxRetries → 删除坏消息 + 退避(2s×2^(n-1)) + continue
|
||||
c. 其他 → 结束
|
||||
(2) isRecoverableLength（length 且实际输出 < 原始 maxTokens，
|
||||
说明是被上下文窗口挤断而非输出上限截断）且未尝试过
|
||||
→ 删除该助手消息 + 压缩 + continue（只试一次）
|
||||
(3) 队列里还有 followUp/steering → continue（续跑）
|
||||
(4) 都不满足 → settle
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Callable, Optional
|
||||
|
||||
from .compaction import CompactionSettings, _content_text, compact_context
|
||||
from .context import estimate_context_tokens, should_compact
|
||||
from .types import (AgentMessage, AgentConfig, RetryConfig, RunResult)
|
||||
|
||||
|
||||
# ========== 压缩诊断日志（与 main_window.diag_log 同一规范：
|
||||
# [HH:MM:SS.mmm] 时间戳 + 项目根目录追加 + UTF-8 + 失败静默） ==========
|
||||
_COMPACT_LOG_PATH = os.path.join(
|
||||
os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))),
|
||||
"compaction_diag.log")
|
||||
|
||||
|
||||
def compact_diag_log(line: str):
|
||||
"""压缩日志只记录：同一行内容 print（控制台）+ 追加写 compaction_diag.log"""
|
||||
t = time.time()
|
||||
stamp = (f"[{time.strftime('%H:%M:%S', time.localtime(t))}"
|
||||
f".{int(t * 1000) % 1000:03d}]")
|
||||
print(f"{stamp} {line}", flush=True)
|
||||
try:
|
||||
with open(_COMPACT_LOG_PATH, "a", encoding="utf-8") as f:
|
||||
f.write(f"{stamp} {line}\n")
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
|
||||
def _settings_of(cfg: AgentConfig) -> CompactionSettings:
|
||||
"""AgentConfig 压缩参数 → CompactionSettings（1:1 映射 pi settings 字段）"""
|
||||
return CompactionSettings(reserve_tokens=cfg.compaction_reserve,
|
||||
keep_recent_tokens=cfg.compaction_keep_recent)
|
||||
|
||||
|
||||
def _cut_ids_of(msgs: list, new_msgs: list) -> tuple:
|
||||
"""🆕 压缩持久化切点：new_msgs = [摘要] + 保留尾巴（旧 msgs 的连续后缀）。
|
||||
按 identity 定位尾巴起点 → (cut_before_id, first_retained_id)（DB 行 id）。
|
||||
|
||||
⚠️ 行粒度：同一 DB 行（assistant timeline）会被 build_api_context 回放成
|
||||
多条 API 消息（asst/tool 条目共享同一 _db_msg_id）。若尾巴起点落在某行
|
||||
回放序列中间，必须回退到该行第一条，否则 cut_before 与 first_retained
|
||||
会是同一行 → DB 自环（链死循环卡死）。
|
||||
回退后下一轮会从该行完整重放（保守方向，宁多勿漏）。
|
||||
任一为空 = 无法在 DB 链上切（如尾巴全是轮中内存消息）→ UI 不插标记。"""
|
||||
if not new_msgs or getattr(new_msgs[0], "kind", None) != "compaction_summary":
|
||||
return "", ""
|
||||
tail = new_msgs[1:]
|
||||
if not tail:
|
||||
return "", ""
|
||||
first = tail[0]
|
||||
for i, m in enumerate(msgs):
|
||||
if m is first:
|
||||
# 回退到同一 DB 行回放序列的头部
|
||||
fid = getattr(first, "db_msg_id", "") or ""
|
||||
j = i
|
||||
while j > 0 and fid and getattr(msgs[j - 1], "db_msg_id", "") == fid:
|
||||
j -= 1
|
||||
if j == 0:
|
||||
return "", "" # 没有“之前”的行 → 无切点
|
||||
return (getattr(msgs[j - 1], "db_msg_id", "") or "",
|
||||
getattr(first, "db_msg_id", "") or "")
|
||||
return "", "" # 尾巴首条不在旧列表（不应发生）→ 不插标记
|
||||
|
||||
# ======================================================================
|
||||
# 对照 retry.ts 的模式表（原样移植）
|
||||
# ======================================================================
|
||||
_NON_RETRYABLE_LIMIT_RE = re.compile(
|
||||
r"GoUsageLimitError|FreeUsageLimitError|"
|
||||
r"Monthly usage limit reached|available balance|"
|
||||
r"insufficient_quota|out of budget|quota exceeded|billing",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
_RETRYABLE_RE = re.compile(
|
||||
r"overloaded|rate.?limit|too many requests|429|500|502|503|504|524|"
|
||||
r"service.?unavailable|server.?error|internal.?error|"
|
||||
r"provider.?returned.?error|exceeded request buffer limit while retrying upstream|"
|
||||
r"network.?error|connection.?error|connection.?refused|connection.?lost|"
|
||||
r"other side closed|fetch failed|getaddrinfo|ENOTFOUND|EAI_AGAIN|"
|
||||
r"upstream.?connect|reset before headers|socket hang up|timeout|timed?\s?out",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
# 对照 overflow.ts OVERFLOW_PATTERNS（各服务商话术）
|
||||
_OVERFLOW_RE = re.compile(
|
||||
r"context\s*window|prompt is too long|maximum context|too many tokens|"
|
||||
r"context length|exceeds the (model|maximum) context|"
|
||||
r"maximum context length is \d+ tokens|exceeds the maximum allowed input length|"
|
||||
r"longer than the model's context length|exceeds the available context size|"
|
||||
r"greater than the context length|exceeded model token limit|"
|
||||
r"Range of input length should be|configured context size|"
|
||||
r"长于|超出.*上下文|上下文.*(超出|超过|不足)",
|
||||
re.IGNORECASE,
|
||||
)
|
||||
|
||||
# 限速类不是溢出（对照 overflow.ts NON_OVERFLOW_PATTERNS 的判定意图）
|
||||
_NON_OVERFLOW_RE = re.compile(
|
||||
r"rate.?limit|too many requests|429|overloaded", re.IGNORECASE)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 判定函数（对照 pi overflow.ts / retry.ts，基于 assistant 消息）
|
||||
# ======================================================================
|
||||
def is_context_overflow(message: AgentMessage,
|
||||
context_window: int = 0) -> bool:
|
||||
"""
|
||||
对照 isContextOverflow 的三种情形：
|
||||
1. error + 消息文本命中溢出模式（且不是限速类）
|
||||
2. 静默溢出：stop 但 usage.input > contextWindow（z.ai 风格）
|
||||
3. 服务端截断式溢出：length 且 output==0 且 input ≥ 0.99×窗口
|
||||
"""
|
||||
if message.role != "assistant":
|
||||
return False
|
||||
if message.stop_reason == "error" and message.error_message:
|
||||
if not _NON_OVERFLOW_RE.search(message.error_message) \
|
||||
and _OVERFLOW_RE.search(message.error_message):
|
||||
return True
|
||||
usage = message.usage or {}
|
||||
input_tokens = int(usage.get("input", 0) or 0) + int(usage.get("cacheRead", 0) or 0)
|
||||
if context_window and message.stop_reason == "stop":
|
||||
if input_tokens > context_window:
|
||||
return True
|
||||
if context_window and message.stop_reason == "length" \
|
||||
and int(usage.get("output", 0) or 0) == 0:
|
||||
if input_tokens >= context_window * 0.99:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def is_retryable_assistant_error(message: AgentMessage) -> bool:
|
||||
"""对照 isRetryableAssistantError：基于 errorMessage 文本模式"""
|
||||
if message.role != "assistant" or message.stop_reason != "error":
|
||||
return False
|
||||
if not message.error_message:
|
||||
return False
|
||||
if _NON_RETRYABLE_LIMIT_RE.search(message.error_message):
|
||||
return False
|
||||
return bool(_RETRYABLE_RE.search(message.error_message))
|
||||
|
||||
|
||||
def is_recoverable_length(message: AgentMessage,
|
||||
desired_max_output: int) -> bool:
|
||||
"""
|
||||
对照 pi isRecoverableLength：
|
||||
length 停止 且 实际输出 < 「原始输出上限」（钳制前的 maxTokens）。
|
||||
"""
|
||||
if message.role != "assistant" or message.stop_reason != "length":
|
||||
return False
|
||||
if desired_max_output <= 0:
|
||||
return False
|
||||
out = int((message.usage or {}).get("output", 0) or 0)
|
||||
return out < desired_max_output
|
||||
|
||||
|
||||
def compute_retry_delay_ms(attempt: int, cfg: RetryConfig) -> float:
|
||||
"""对照 pi _prepareRetry: baseDelayMs × 2^(attempt-1)（无抖动、默认无封顶）"""
|
||||
delay = cfg.base_delay_ms * (cfg.factor ** (attempt - 1))
|
||||
if cfg.max_delay_ms and delay > cfg.max_delay_ms:
|
||||
delay = cfg.max_delay_ms
|
||||
return float(delay)
|
||||
|
||||
|
||||
def find_last_assistant(messages) -> Optional[AgentMessage]:
|
||||
"""对照 agent-session _findLastAssistantMessage"""
|
||||
for m in reversed(messages):
|
||||
if m.role == "assistant":
|
||||
return m
|
||||
return None
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 编排器 —— 对照 agent-session 的 run 循环
|
||||
# ======================================================================
|
||||
class AgentRunner:
|
||||
"""
|
||||
对照 agent-session.ts 的 run 编排（agent 之外的「会话层」）。
|
||||
haocode 里由 AgentWorker(QThread) 持有一个 AgentRunner。
|
||||
|
||||
summarize_fn(transcript) -> str：压缩用的非流式 LLM 调用（缺省 → 禁压缩）。
|
||||
"""
|
||||
|
||||
def __init__(self, agent: Agent,
|
||||
summarize_fn: Optional[Callable[[str], str]] = None,
|
||||
on_retry_scheduled: Optional[Callable] = None,
|
||||
on_retry_finished: Optional[Callable] = None,
|
||||
on_compaction_started: Optional[Callable] = None,
|
||||
on_compaction_finished: Optional[Callable] = None):
|
||||
self.agent = agent
|
||||
self.summarize_fn = summarize_fn
|
||||
# 🆕 M3: 重试回调（对照 pi retryAssistantCall 的
|
||||
# onRetryScheduled / onRetryFinished）——UI 可提示「正在重试」
|
||||
self.on_retry_scheduled = on_retry_scheduled
|
||||
self.on_retry_finished = on_retry_finished
|
||||
# 🆕 压缩开始回调（摘要 LLM 调用阻塞前触发）——UI 显示「执行中」动态气泡
|
||||
self.on_compaction_started = on_compaction_started
|
||||
# 🆕 压缩完成回调（压缩结束瞬间触发，不等整轮 run 结束）
|
||||
# ——前端气泡即时定格「已完成上下文压缩」（与思考块定格同款时机）
|
||||
self.on_compaction_finished = on_compaction_finished
|
||||
self._overflow_recovered = False # 对照 _overflowRecoveryAttempted
|
||||
self._retry_attempt = 0 # 对照 _retryAttempt
|
||||
self.last_action = "none" # 诊断: none/retry/overflow_compact/length_compact/queued
|
||||
self.compactions_performed = 0 # 成功压缩次数（UI 提示用）
|
||||
# 🆕 压缩事件队列（UI 气泡用：摘要全文/前后 token/耗时/路径）
|
||||
self.compaction_events: list = []
|
||||
# 🆕 轮中主动压缩护栏（同一 run 内；AgentRunner 每轮新建 → 自动复位）
|
||||
self._mid_turn_fail_streak = 0
|
||||
self._mid_turn_exhausted = False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 发送前压缩（🌟 1:1 对照 pi shouldCompact：tokens > window - reserveTokens）
|
||||
# ------------------------------------------------------------------
|
||||
def pre_prompt_compaction(self) -> bool:
|
||||
if self.summarize_fn is None:
|
||||
return False
|
||||
cfg = self.agent.config
|
||||
msgs = self.agent.state.messages
|
||||
m = cfg.model
|
||||
# 🆕 G2: 透传 system_prompt/tools —— 无锚点时度量「下一请求真实载荷」
|
||||
should, tokens = should_compact(msgs, m, cfg.compaction_reserve,
|
||||
system_prompt=cfg.system_prompt,
|
||||
tools=cfg.tools)
|
||||
# 🆕 每次请求前的校验：print + 写 compaction_diag.log（触发与否都记）
|
||||
compact_diag_log(
|
||||
f"[PRE_CHECK] model={m.name} window={m.context_window} "
|
||||
f"reserve={cfg.compaction_reserve} "
|
||||
f"threshold={m.context_window - cfg.compaction_reserve} "
|
||||
f"est={tokens} triggered={str(should).lower()}")
|
||||
if not should:
|
||||
return False
|
||||
self._notify_compaction_started("pre_prompt", tokens)
|
||||
t0 = time.time()
|
||||
new_msgs = compact_context(msgs, cfg.model, self.summarize_fn,
|
||||
settings=_settings_of(cfg))
|
||||
if new_msgs is not None:
|
||||
cut_before, first_retained = _cut_ids_of(msgs, new_msgs)
|
||||
self.agent.state.messages = new_msgs
|
||||
self.compactions_performed += 1
|
||||
self._record_compaction("pre_prompt", tokens, new_msgs, t0,
|
||||
cut_before_id=cut_before,
|
||||
first_retained_id=first_retained)
|
||||
return True
|
||||
self._record_compaction_failed("pre_prompt", tokens, t0)
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 🆕 轮中主动压缩（haocode 增强，偏离 pi 1:1）
|
||||
# ------------------------------------------------------------------
|
||||
def compact_if_needed(self, msgs: list) -> Optional[list]:
|
||||
"""接线到 AgentConfig.compact_fn：内层循环每次 LLM 请求前调用。
|
||||
|
||||
与轮首 pre_prompt_compaction 同公式（tokens > window − reserve）。
|
||||
轮中单条工具输出可把上下文顶出窗口（轮中内存工具结果没有
|
||||
4000 字回放上限），发请求前主动压缩，避免「一次失败往返 +
|
||||
响应式兜底」。返回新消息列表（发生了压缩）或 None（不需要/
|
||||
不可压缩 → 原样发请求，响应式 overflow 安全网仍在）。
|
||||
|
||||
防刷屏护栏：同一 run 内连续失败 2 次后不再尝试（形状无有效
|
||||
切点，如单条巨型条目，硬试无意义）；成功则复位计数。
|
||||
"""
|
||||
if self.summarize_fn is None or self._mid_turn_exhausted:
|
||||
return None
|
||||
cfg = self.agent.config
|
||||
m = cfg.model
|
||||
# 🆕 G2: 透传 system_prompt/tools（与轮首同口径）
|
||||
should, tokens = should_compact(msgs, m, cfg.compaction_reserve,
|
||||
system_prompt=cfg.system_prompt,
|
||||
tools=cfg.tools)
|
||||
compact_diag_log(
|
||||
f"[PRE_CHECK] path=mid_turn est={tokens} "
|
||||
f"threshold={m.context_window - cfg.compaction_reserve} "
|
||||
f"triggered={str(should).lower()}")
|
||||
if not should:
|
||||
return None
|
||||
if self._do_compaction("mid_turn", msgs=msgs):
|
||||
self._mid_turn_fail_streak = 0
|
||||
return list(self.agent.state.messages)
|
||||
self._mid_turn_fail_streak += 1
|
||||
if self._mid_turn_fail_streak >= 2:
|
||||
self._mid_turn_exhausted = True
|
||||
compact_diag_log(
|
||||
"[PRE_CHECK] path=mid_turn exhausted "
|
||||
"（连续 2 次失败，本轮不再尝试）")
|
||||
return None
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 主入口 —— 对照 _runAgentPrompt
|
||||
# ------------------------------------------------------------------
|
||||
def run(self, message) -> RunResult:
|
||||
self.pre_prompt_compaction()
|
||||
self._overflow_recovered = False
|
||||
self._retry_attempt = 0
|
||||
result = self.agent.prompt(message)
|
||||
return self._post_loop(result)
|
||||
|
||||
def _post_loop(self, result: RunResult) -> RunResult:
|
||||
# 对照: while (await this._handlePostAgentRun()) await this.agent.continue()
|
||||
while True:
|
||||
action = self._handle_post_agent_run(result)
|
||||
if action == "stop":
|
||||
return result
|
||||
if action == "retry":
|
||||
result = self.agent.continue_()
|
||||
# 🆕 M3: 重试结果回调（成功 = 下一轮未以 error/aborted 收尾）
|
||||
if self.on_retry_finished is not None:
|
||||
la = find_last_assistant(self.agent.state.messages)
|
||||
ok = (la is not None
|
||||
and la.stop_reason not in ("error", "aborted"))
|
||||
try:
|
||||
self.on_retry_finished(bool(ok))
|
||||
except Exception:
|
||||
pass
|
||||
continue
|
||||
if action == "compact_retry":
|
||||
if not (self._remove_last_bad_assistant()
|
||||
and self._do_compaction(self.last_action)):
|
||||
return result # 压缩不可用 → 放弃，交 UI 报错
|
||||
result = self.agent.continue_()
|
||||
continue
|
||||
if action == "queued":
|
||||
result = self.agent.continue_()
|
||||
continue
|
||||
return result
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 三路决策 —— 对照 _handlePostAgentRun
|
||||
# ------------------------------------------------------------------
|
||||
def _handle_post_agent_run(self, result: RunResult) -> str:
|
||||
msg = find_last_assistant(self.agent.state.messages)
|
||||
if msg is None:
|
||||
return "stop"
|
||||
cfg = self.agent.config
|
||||
|
||||
# (1) 错误路径
|
||||
if msg.stop_reason == "error":
|
||||
# a. 溢出 → 压缩恢复（只试一次）
|
||||
if is_context_overflow(msg, cfg.model.context_window) \
|
||||
and not self._overflow_recovered \
|
||||
and self.summarize_fn is not None:
|
||||
self._overflow_recovered = True
|
||||
self.last_action = "overflow_compact"
|
||||
return "compact_retry"
|
||||
# b. 可重试错误 → 移除坏消息 + 退避 + continue
|
||||
if is_retryable_assistant_error(msg) \
|
||||
and self._prepare_retry(cfg.retry,
|
||||
reason=msg.error_message or ""):
|
||||
self.last_action = "retry"
|
||||
return "retry"
|
||||
return "stop"
|
||||
|
||||
# (2) 可恢复的 length 截断
|
||||
if is_recoverable_length(msg, cfg.model.max_tokens) \
|
||||
and not self._overflow_recovered \
|
||||
and self.summarize_fn is not None:
|
||||
self._overflow_recovered = True
|
||||
self.last_action = "length_compact"
|
||||
return "compact_retry"
|
||||
|
||||
# (3) 队列里还有消息 → 续跑
|
||||
if self.agent.has_queued():
|
||||
self.last_action = "queued"
|
||||
return "queued"
|
||||
|
||||
# (4) settle
|
||||
return "stop"
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 对照 _prepareRetry：计数 → 移除坏消息 → 退避睡眠
|
||||
# ------------------------------------------------------------------
|
||||
def _prepare_retry(self, cfg: RetryConfig, reason: str = "") -> bool:
|
||||
if self._retry_attempt >= cfg.max_attempts:
|
||||
return False
|
||||
self._retry_attempt += 1
|
||||
delay_ms = compute_retry_delay_ms(self._retry_attempt, cfg)
|
||||
# 移除错误助手消息（对照: messages.slice(0,-1)，会话历史保留由 UI 层负责）
|
||||
messages = self.agent.state.messages
|
||||
if messages and messages[-1].role == "assistant":
|
||||
self.agent.state.messages = messages[:-1]
|
||||
# 🆕 M3: 通知重试已调度（参数: attempt, max_attempts, delay_ms, reason）
|
||||
if self.on_retry_scheduled is not None:
|
||||
try:
|
||||
self.on_retry_scheduled(self._retry_attempt, cfg.max_attempts,
|
||||
delay_ms, reason)
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(delay_ms / 1000.0)
|
||||
return True
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 压缩恢复辅助
|
||||
# ------------------------------------------------------------------
|
||||
def _remove_last_bad_assistant(self) -> bool:
|
||||
msgs = self.agent.state.messages
|
||||
if msgs and msgs[-1].role == "assistant" \
|
||||
and msgs[-1].stop_reason in ("length", "error"):
|
||||
self.agent.state.messages = msgs[:-1]
|
||||
return True
|
||||
return False
|
||||
|
||||
def _do_compaction(self, path: str = "reactive",
|
||||
msgs: Optional[list] = None) -> bool:
|
||||
if self.summarize_fn is None:
|
||||
return False
|
||||
cfg = self.agent.config
|
||||
if msgs is None:
|
||||
msgs = self.agent.state.messages
|
||||
tokens_before = estimate_context_tokens(msgs).tokens
|
||||
self._notify_compaction_started(path, tokens_before)
|
||||
t0 = time.time()
|
||||
new_msgs = compact_context(msgs, cfg.model, self.summarize_fn,
|
||||
settings=_settings_of(cfg))
|
||||
if new_msgs is not None:
|
||||
cut_before, first_retained = _cut_ids_of(msgs, new_msgs)
|
||||
self.agent.state.messages = new_msgs
|
||||
self.compactions_performed += 1
|
||||
self._record_compaction(path, tokens_before, new_msgs, t0,
|
||||
cut_before_id=cut_before,
|
||||
first_retained_id=first_retained)
|
||||
return True
|
||||
self._record_compaction_failed(path, tokens_before, t0)
|
||||
return False
|
||||
|
||||
# ------------------------------------------------------------------
|
||||
# 🆕 压缩可视化辅助：开始/完成 两阶段（print + 日志 + 事件供 UI 气泡）
|
||||
# ------------------------------------------------------------------
|
||||
def _notify_compaction_started(self, path: str, tokens_before: int) -> None:
|
||||
compact_diag_log(f"[COMPACT_START] path={path} before={tokens_before}")
|
||||
if self.on_compaction_started is not None:
|
||||
try:
|
||||
self.on_compaction_started(path)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _record_compaction_failed(self, path: str, tokens_before: int,
|
||||
t0: float) -> None:
|
||||
"""压缩不可行（如历史无法切分）→ 气泡显示失败原因，不能永远卡在「执行中」"""
|
||||
duration_ms = int((time.time() - t0) * 1000)
|
||||
compact_diag_log(
|
||||
f"[COMPACT_FAIL] path={path} before={tokens_before} "
|
||||
f"duration_ms={duration_ms}")
|
||||
ev = {
|
||||
"summary": "本次压缩未能执行（历史无法切分），继续以原有历史对话。",
|
||||
"before": tokens_before,
|
||||
"after": tokens_before,
|
||||
"duration_ms": duration_ms,
|
||||
"path": path,
|
||||
"failed": True,
|
||||
}
|
||||
self.compaction_events.append(ev)
|
||||
if self.on_compaction_finished is not None:
|
||||
try:
|
||||
self.on_compaction_finished(dict(ev))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _record_compaction(self, path: str, tokens_before: int,
|
||||
new_msgs: list, t0: float,
|
||||
cut_before_id: str = "",
|
||||
first_retained_id: str = "") -> None:
|
||||
summary = ""
|
||||
if new_msgs and getattr(new_msgs[0], "kind", None) == "compaction_summary":
|
||||
summary = _content_text(new_msgs[0].content) or ""
|
||||
tokens_after = estimate_context_tokens(new_msgs).tokens
|
||||
duration_ms = int((time.time() - t0) * 1000)
|
||||
compact_diag_log(
|
||||
f"[COMPACT_DONE] path={path} before={tokens_before} "
|
||||
f"after={tokens_after} summary_chars={len(summary)} "
|
||||
f"duration_ms={duration_ms} "
|
||||
f"cut_before={cut_before_id or '-'} first_retained={first_retained_id or '-'}")
|
||||
ev = {
|
||||
"summary": summary,
|
||||
"before": tokens_before,
|
||||
"after": tokens_after,
|
||||
"duration_ms": duration_ms,
|
||||
"path": path,
|
||||
# 🆕 压缩持久化切点：UI 收到后在 DB 链上插标记行（两者齐备才插）
|
||||
"cut_before_id": cut_before_id,
|
||||
"first_retained_id": first_retained_id,
|
||||
}
|
||||
self.compaction_events.append(ev)
|
||||
if self.on_compaction_finished is not None:
|
||||
try:
|
||||
self.on_compaction_finished(dict(ev))
|
||||
except Exception:
|
||||
pass
|
||||
@@ -0,0 +1,438 @@
|
||||
"""
|
||||
core/agent/stream_fn.py
|
||||
=======================
|
||||
🌟 pi-ai 流式接口的 Python 1:1 移植 —— OpenAI 兼容（含 vLLM）
|
||||
|
||||
对照 pi-main 源码：
|
||||
packages/ai/src/api/openai-completions.ts
|
||||
- streamSimple() (行 469 起): 消费 SSE chunk → 产出 AssistantMessage
|
||||
- buildParams() (行 536 起): 组装请求体
|
||||
* vLLM/OpenAI 默认字段名 = max_completion_tokens
|
||||
* 少数服务商(chutes/deepseek/moonshot 等)用 max_tokens
|
||||
- mapStopReason(): finish_reason "length" → stopReason "length"
|
||||
packages/ai/src/api/simple-options.ts
|
||||
- buildBaseOptions(): maxTokens = options?.maxTokens ?? model.maxTokens
|
||||
（即模型定义里的 maxTokens 一定会进入请求体，这就是它被发给 vLLM 的原因）
|
||||
|
||||
本文件职责：
|
||||
1. to_openai_messages() —— pi 消息格式 → OpenAI API 格式（每轮调用前转换）
|
||||
2. openai_stream() —— 流式请求，逐 chunk 产出 AssistantMessageEvent，
|
||||
结束产出最终 AgentMessage（含 usage / stop_reason / tool_calls）
|
||||
3. classify_error() —— 异常分类（对照 pi isRetryable 的输入）
|
||||
|
||||
🌟 修复 haocode 原有 P1 bug：openai 客户端带 timeout（原来无超时，
|
||||
挂死的流会让线程永久阻塞）。max_retries=0（重试统一交给 recovery 层，
|
||||
与 pi 一致：网络层不重试，会话层按 1600ms×1.6^n 退避重试）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, Iterator, List, Optional, Tuple
|
||||
|
||||
from openai import OpenAI
|
||||
from openai import APIConnectionError, APIStatusError, APITimeoutError, RateLimitError
|
||||
|
||||
from .types import (AgentError, AgentMessage, AssistantMessageEvent,
|
||||
AbortSignal, ModelConfig, ToolCall, new_id)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# pi 消息 → OpenAI 消息（每轮发送前调用）
|
||||
# ======================================================================
|
||||
def to_openai_messages(context: List[AgentMessage],
|
||||
pass_reasoning: bool = True) -> List[Dict[str, Any]]:
|
||||
"""pi 消息 → OpenAI chat 格式。
|
||||
🆕 thinking 模式（DeepSeek 等）：assistant 消息必须把 reasoning_content
|
||||
原样回传，否则服务端 400
|
||||
（'The reasoning_content in the thinking mode must be passed back to the API'）。
|
||||
默认开启：轮中续问/工具回灌场景，内存里本轮刚流式产出的 assistant 带
|
||||
reasoning → 回传；DB 回放的旧行 reasoning=""（不入库）→ 无该字段，
|
||||
跨轮全新请求不受影响。pass_reasoning=False 显式关闭（特殊网关逃生阀）。
|
||||
"""
|
||||
out: List[Dict[str, Any]] = []
|
||||
for m in context:
|
||||
if m.role == "system":
|
||||
# 只放行请求头部的系统提示词（loop 注入）；
|
||||
# 历史中间出现的 system（DB 防御）仍然跳过
|
||||
if not out:
|
||||
out.append({"role": "system", "content": m.content})
|
||||
continue
|
||||
if m.role == "toolResult":
|
||||
out.append({
|
||||
"role": "tool",
|
||||
"tool_call_id": m.tool_call_id,
|
||||
"content": m.content if isinstance(m.content, str)
|
||||
else "".join(c.get("text", "") for c in m.content
|
||||
if isinstance(c, dict)),
|
||||
})
|
||||
continue
|
||||
if m.role == "assistant":
|
||||
entry: Dict[str, Any] = {"role": "assistant"}
|
||||
# 纯工具调用轮：content 可能为空 → 用 ""（vLLM 接受，避免 null 报错）
|
||||
entry["content"] = m.content if m.content else ""
|
||||
if m.tool_calls:
|
||||
entry["tool_calls"] = [
|
||||
{
|
||||
"id": tc.id,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": tc.name,
|
||||
"arguments": json.dumps(tc.arguments, ensure_ascii=False)
|
||||
if tc.arguments else tc.raw_arguments or "{}",
|
||||
},
|
||||
}
|
||||
for tc in m.tool_calls
|
||||
]
|
||||
if pass_reasoning and m.reasoning:
|
||||
entry["reasoning_content"] = m.reasoning
|
||||
out.append(entry)
|
||||
continue
|
||||
# user（content 可以是 str 或 OpenAI 多模态 list，原样透传）
|
||||
out.append({"role": "user", "content": m.content})
|
||||
return out
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 异常分类 —— 对照 pi isRetryable / retryable error 判定
|
||||
# ======================================================================
|
||||
def classify_error(e: Exception) -> AgentError:
|
||||
"""把 openai SDK 异常分类，恢复层据此决定重试/放弃。"""
|
||||
if isinstance(e, RateLimitError):
|
||||
return AgentError(message=str(e), kind="rate_limit",
|
||||
status_code=429, recoverable=True)
|
||||
if isinstance(e, APITimeoutError):
|
||||
return AgentError(message=f"请求超时: {e}", kind="timeout", recoverable=True)
|
||||
if isinstance(e, APIConnectionError):
|
||||
return AgentError(message=f"连接失败: {e}", kind="connection", recoverable=True)
|
||||
if isinstance(e, APIStatusError):
|
||||
code = e.status_code
|
||||
msg = str(e)
|
||||
if code in (502, 503, 504):
|
||||
return AgentError(message=msg, kind="server_error",
|
||||
status_code=code, recoverable=True)
|
||||
if code == 529:
|
||||
return AgentError(message=msg, kind="overload",
|
||||
status_code=code, recoverable=True)
|
||||
if code == 408:
|
||||
return AgentError(message=msg, kind="timeout",
|
||||
status_code=code, recoverable=True)
|
||||
if code in (401, 403):
|
||||
return AgentError(message=msg, kind="auth",
|
||||
status_code=code, recoverable=False)
|
||||
# 4xx（参数错误/模型不存在/上下文超长等）不可重试，交恢复层细分
|
||||
return AgentError(message=msg, kind="unknown",
|
||||
status_code=code, recoverable=False)
|
||||
if isinstance(e, (ConnectionError, TimeoutError)):
|
||||
return AgentError(message=str(e), kind="connection", recoverable=True)
|
||||
return AgentError(message=str(e), kind="unknown", recoverable=False)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 🆕 M1/M2: 流式字段提取（对照 pi openai-completions.ts 的字段优先级）
|
||||
# ======================================================================
|
||||
# 🆕 M1: 1:1 对照 pi chunk 级 reasoning 字段优先级（llama.cpp/vLLM/OpenRouter 各不同）
|
||||
_REASONING_FIELDS = ("reasoning_content", "reasoning", "reasoning_text")
|
||||
|
||||
|
||||
def _pick_reasoning(delta) -> str:
|
||||
"""从 delta 上取第一个非空思考字段（只取一个，防同内容多字段重复计）"""
|
||||
if delta is None:
|
||||
return ""
|
||||
dump = None
|
||||
for name in _REASONING_FIELDS:
|
||||
v = getattr(delta, name, None)
|
||||
if v is None and dump is None and hasattr(delta, "model_dump"):
|
||||
try:
|
||||
dump = delta.model_dump()
|
||||
except Exception:
|
||||
dump = {}
|
||||
if v is None and isinstance(dump, dict):
|
||||
v = dump.get(name)
|
||||
if v:
|
||||
return str(v)
|
||||
return ""
|
||||
|
||||
|
||||
def _pick_usage(chunk):
|
||||
"""🆕 M2: 对照 pi：先 chunk.usage，再 choices[0].usage（Moonshot 系只放 choice 里）"""
|
||||
u = getattr(chunk, "usage", None)
|
||||
if u is not None:
|
||||
return u
|
||||
try:
|
||||
choices = getattr(chunk, "choices", None)
|
||||
if choices:
|
||||
return getattr(choices[0], "usage", None)
|
||||
except Exception:
|
||||
pass
|
||||
return None
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 流式请求核心 —— 对照 openai-completions.ts streamSimple
|
||||
# ======================================================================
|
||||
def openai_stream(context: List[AgentMessage],
|
||||
model: ModelConfig,
|
||||
signal: AbortSignal,
|
||||
max_tokens: int,
|
||||
tools: Optional[List[Any]] = None,
|
||||
) -> Iterator[Tuple[str, Any]]:
|
||||
"""
|
||||
生成器协议（对照 pi streamSimple 的 for-await 消费方式）：
|
||||
yield ("event", AssistantMessageEvent) # 增量事件（text/thinking/toolcall/done）
|
||||
yield ("final", AgentMessage) # 最终助手消息（含 stop_reason/usage）
|
||||
异常统一转 AgentError 抛出（由循环层捕获 → stop_reason="error"）。
|
||||
|
||||
🌟 中止语义与 pi 一致：每个 chunk 边界检查 signal.aborted，
|
||||
命中即关闭流、以 stop_reason="aborted" 收尾（不是 error）。
|
||||
"""
|
||||
client = OpenAI(
|
||||
api_key=model.api_key or "EMPTY",
|
||||
base_url=model.base_url,
|
||||
# 🌟 P1 修复：显式超时（pi 侧由 fetch timeout 保证）
|
||||
timeout=model.timeout_seconds,
|
||||
max_retries=0, # 重试统一由 recovery 层负责
|
||||
)
|
||||
|
||||
params: Dict[str, Any] = {
|
||||
"model": model.name,
|
||||
"messages": to_openai_messages(context, model.pass_reasoning),
|
||||
"stream": True,
|
||||
"stream_options": {"include_usage": True}, # 拿真实 usage（pi 同款）
|
||||
"temperature": model.temperature,
|
||||
}
|
||||
# 对照 pi buildParams：默认 max_completion_tokens（vLLM 兼容）
|
||||
if getattr(model, "use_max_tokens_field", False):
|
||||
params["max_tokens"] = max_tokens
|
||||
else:
|
||||
params["max_completion_tokens"] = max_tokens
|
||||
|
||||
# 🌟 工具定义（对照 pi buildParams 的 tools 段）——不发给模型，
|
||||
# 模型就不可能发出真正的 tool_call，只会用文字"演"工具调用！
|
||||
if tools:
|
||||
params["tools"] = [
|
||||
{
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": t.name,
|
||||
"description": t.description,
|
||||
"parameters": t.parameters,
|
||||
},
|
||||
}
|
||||
for t in tools
|
||||
]
|
||||
|
||||
# 累积状态
|
||||
content_parts: List[str] = []
|
||||
reasoning_parts: List[str] = []
|
||||
tc_acc: Dict[int, Dict[str, Any]] = {} # index -> {id, name, args}
|
||||
finish_reason: Optional[str] = None
|
||||
usage: Dict[str, Any] = {}
|
||||
stream = None
|
||||
|
||||
try:
|
||||
stream = client.chat.completions.create(**params)
|
||||
for chunk in stream:
|
||||
# 🌟 pi 同款：chunk 边界检查中止
|
||||
if signal.aborted:
|
||||
finish_reason = "aborted"
|
||||
try:
|
||||
stream.close()
|
||||
except Exception:
|
||||
pass
|
||||
break
|
||||
|
||||
# 🆕 M2: usage 读取顺序对照 pi —— chunk.usage 优先，choice.usage 兜底
|
||||
usage_obj = _pick_usage(chunk)
|
||||
if usage_obj is not None:
|
||||
usage = {
|
||||
"input": getattr(usage_obj, "prompt_tokens", 0) or 0,
|
||||
"output": getattr(usage_obj, "completion_tokens", 0) or 0,
|
||||
"cacheRead": 0, "cacheWrite": 0,
|
||||
}
|
||||
if not getattr(chunk, "choices", None):
|
||||
continue
|
||||
|
||||
choice = chunk.choices[0]
|
||||
delta = choice.delta
|
||||
|
||||
# 🆕 M1: 思考增量（对照 pi reasoningFields 优先级，含多字段去重）
|
||||
reasoning = _pick_reasoning(delta)
|
||||
if reasoning:
|
||||
yield ("event", AssistantMessageEvent.thinking_delta(reasoning))
|
||||
reasoning_parts.append(reasoning)
|
||||
|
||||
if getattr(delta, "content", None):
|
||||
yield ("event", AssistantMessageEvent.text_delta(delta.content))
|
||||
content_parts.append(delta.content)
|
||||
|
||||
# 工具调用增量（按 index 累积 id/name/arguments 三段）
|
||||
for tc_delta in (getattr(delta, "tool_calls", None) or []):
|
||||
idx = tc_delta.index or 0
|
||||
slot = tc_acc.setdefault(idx, {"id": "", "name": "", "args": ""})
|
||||
if tc_delta.id:
|
||||
slot["id"] = tc_delta.id
|
||||
yield ("event", AssistantMessageEvent.toolcall_delta(
|
||||
idx, tc_delta.id, field="id"))
|
||||
fn = tc_delta.function
|
||||
if fn is not None:
|
||||
if fn.name:
|
||||
slot["name"] += fn.name
|
||||
yield ("event", AssistantMessageEvent.toolcall_delta(
|
||||
idx, fn.name, field="name"))
|
||||
if fn.arguments:
|
||||
slot["args"] += fn.arguments
|
||||
yield ("event", AssistantMessageEvent.toolcall_delta(
|
||||
idx, fn.arguments, field="arguments"))
|
||||
|
||||
if choice.finish_reason:
|
||||
finish_reason = choice.finish_reason
|
||||
except Exception as e:
|
||||
# 中止过程中断网/断流不当错误处理（pi 同款语义）
|
||||
if signal.aborted:
|
||||
finish_reason = "aborted"
|
||||
else:
|
||||
raise classify_error(e)
|
||||
finally:
|
||||
try:
|
||||
if stream is not None:
|
||||
stream.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# ---- 组装最终助手消息（对照 mapStopReason + AssistantMessage 构造）----
|
||||
# pi mapStopReason: "length"→"length", "stop"/"tool_calls"→"stop", 其余→"stop"
|
||||
if finish_reason == "aborted":
|
||||
stop_reason = "aborted"
|
||||
elif finish_reason == "length":
|
||||
stop_reason = "length"
|
||||
else:
|
||||
stop_reason = "stop"
|
||||
|
||||
tool_calls: List[ToolCall] = []
|
||||
for idx in sorted(tc_acc.keys()):
|
||||
slot = tc_acc[idx]
|
||||
try:
|
||||
args = json.loads(slot["args"]) if slot["args"] else {}
|
||||
except json.JSONDecodeError:
|
||||
# 参数 JSON 被截断/损坏：保留 raw，arguments 置空，
|
||||
# 由循环层按「截断保护」路径处理（不执行残缺调用）
|
||||
args = {}
|
||||
tool_calls.append(ToolCall(
|
||||
id=slot["id"] or new_id("call"),
|
||||
name=slot["name"],
|
||||
arguments=args,
|
||||
raw_arguments=slot["args"],
|
||||
))
|
||||
|
||||
final = AgentMessage(
|
||||
role="assistant",
|
||||
content="".join(content_parts),
|
||||
reasoning="".join(reasoning_parts),
|
||||
tool_calls=tool_calls,
|
||||
stop_reason=stop_reason,
|
||||
usage=usage,
|
||||
)
|
||||
yield ("event", AssistantMessageEvent.done())
|
||||
yield ("final", final)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# OpenAI 格式 → pi 格式（from_openai_messages 的逆转换，用于从 DB 装载历史）
|
||||
# ======================================================================
|
||||
def _parse_tool_call(tc: Dict[str, Any]) -> Optional[ToolCall]:
|
||||
"""
|
||||
解析一条 OpenAI tool_call → ToolCall。
|
||||
🌟 参数 JSON 可能残缺（历史里存了被截断的调用）：
|
||||
解析失败 → arguments={} 且 raw_arguments 兜底成合法 "{}"，
|
||||
保证 to_openai_messages 回发 API 时参数始终是合法 JSON
|
||||
（API 只校验 JSON 合法性 + tool_call_id 配对，不校验参数内容）。
|
||||
"""
|
||||
if not isinstance(tc, dict):
|
||||
return None
|
||||
fn = tc.get("function") or {}
|
||||
raw_args = fn.get("arguments")
|
||||
if not isinstance(raw_args, str):
|
||||
raw_args = "" if raw_args is None else str(raw_args)
|
||||
try:
|
||||
args = json.loads(raw_args) if raw_args else {}
|
||||
if not isinstance(args, dict):
|
||||
args = {"_": args}
|
||||
except (json.JSONDecodeError, TypeError, ValueError):
|
||||
args = {}
|
||||
raw_args = "{}" # 残缺 JSON 兜底，防回发 API 时 400
|
||||
return ToolCall(
|
||||
id=tc.get("id") or new_id("call"),
|
||||
name=fn.get("name") or "",
|
||||
arguments=args,
|
||||
raw_arguments=raw_args or "{}",
|
||||
)
|
||||
|
||||
|
||||
def from_openai_messages(messages: List[Dict[str, Any]]) -> List[AgentMessage]:
|
||||
"""
|
||||
把 OpenAI chat 消息数组转成 pi 的 AgentMessage 列表（全保真）。
|
||||
|
||||
🌟 修复：worker 模式此前在此丢失全部工具历史（tool 消息被跳过、
|
||||
assistant.tool_calls 不解析）→ 任务完成/中断后再次提问，模型看不到
|
||||
之前读过什么文件、执行过什么命令。现改为无损转换：
|
||||
- user: content 保持 str 或 OpenAI 多模态 list（原样透传）
|
||||
- assistant: content + reasoning + tool_calls（纯工具轮保留 tool_calls，
|
||||
content 置 ""，不再变成空 assistant 污染上下文）
|
||||
- tool: → AgentMessage(role="toolResult", tool_call_id, content)
|
||||
- system: 防御性跳过（DB 链表里不会出现）
|
||||
|
||||
下游 to_openai_messages 原生支持 toolResult / assistant.tool_calls，
|
||||
与 pi 语义一致，往返（to(from(x))）保真。
|
||||
"""
|
||||
out: List[AgentMessage] = []
|
||||
for m in messages:
|
||||
if not isinstance(m, dict):
|
||||
continue
|
||||
role = m.get("role")
|
||||
content = m.get("content")
|
||||
|
||||
if role == "user":
|
||||
if content is None:
|
||||
content = ""
|
||||
out.append(AgentMessage(role="user", content=content,
|
||||
# 🆕 压缩持久化：跨轮携带 DB 行 id + 摘要 kind
|
||||
# （kind=compaction_summary → 二次压缩走迭代更新）
|
||||
kind=m.get("_kind") or "",
|
||||
db_msg_id=m.get("_db_msg_id") or "",
|
||||
timestamp=int(m.get("timestamp") or 0)))
|
||||
|
||||
elif role == "assistant":
|
||||
# tool_calls 解析（容忍缺失/损坏）
|
||||
tool_calls: List[ToolCall] = []
|
||||
for tc in (m.get("tool_calls") or []):
|
||||
parsed = _parse_tool_call(tc)
|
||||
if parsed is not None:
|
||||
tool_calls.append(parsed)
|
||||
out.append(AgentMessage(
|
||||
role="assistant",
|
||||
content=content if isinstance(content, str) else "",
|
||||
reasoning=m.get("reasoning") or "",
|
||||
tool_calls=tool_calls,
|
||||
stop_reason="stop",
|
||||
# 🆕 P1: 回放入库的 usage（锚点）+ timestamp（P0 时效校验）
|
||||
usage=m.get("usage") or {},
|
||||
timestamp=int(m.get("timestamp") or 0),
|
||||
db_msg_id=m.get("_db_msg_id") or "",
|
||||
))
|
||||
|
||||
elif role == "tool":
|
||||
# tool 结果 → toolResult（保留 tool_call_id 供 API 配对）
|
||||
if content is None:
|
||||
content = ""
|
||||
elif not isinstance(content, str):
|
||||
content = str(content)
|
||||
out.append(AgentMessage(
|
||||
role="toolResult",
|
||||
content=content,
|
||||
tool_call_id=m.get("tool_call_id") or "",
|
||||
timestamp=int(m.get("timestamp") or 0),
|
||||
db_msg_id=m.get("_db_msg_id") or "",
|
||||
))
|
||||
|
||||
# 其他 role（system 等）防御性跳过
|
||||
return out
|
||||
@@ -0,0 +1,873 @@
|
||||
"""
|
||||
core/agent/tools.py
|
||||
===================
|
||||
🌟 pi 工具执行管线的 Python 1:1 移植 + 内置工具
|
||||
|
||||
对照 pi-main 源码：
|
||||
packages/agent/src/agent-loop.ts
|
||||
- prepareToolCalls() (行 ~470): 校验参数 → 标记错误（不执行）
|
||||
- executeTool() (行 ~520): before 钩子 → 执行 → after 钩子 → 结果定型
|
||||
- 并行语义: prepare 串行 → 执行并发(Promise.all) → 结果按原始顺序回写
|
||||
packages/agent/src/tools/*.ts (coding-agent 内置工具 read/bash/edit/write)
|
||||
|
||||
JSON Schema 校验：pi 用 validate-json-schema + ai/src/utils/validation.ts
|
||||
（structuredClone → normalizeOptionalNulls → Value.Convert → coerceWithJsonSchema →
|
||||
全量错误上报 + 回显收到的参数）；这里实现核心子集
|
||||
（type/required/properties/enum/items + 可选字段 null 归一化 + 数字/布尔轻量转换），
|
||||
零外部依赖。工具自定义预处理对照 pi 的 tool.prepareArguments（见 edit 的 legacy 兼容）。
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import threading
|
||||
import time
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
from .types import (AgentMessage, AgentTool, AgentToolResult, AbortSignal,
|
||||
ToolCall, new_id)
|
||||
|
||||
try:
|
||||
from core.debug_log import debug_log as _dbg_log # 🆕 计时观察日志（线程安全/静默）
|
||||
except Exception: # 导入失败也不影响工具执行
|
||||
def _dbg_log(msg, tag="APP"):
|
||||
pass
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# JSON Schema 校验（子集）
|
||||
# ======================================================================
|
||||
# ======================================================================
|
||||
# 参数归一化 + JSON Schema 校验（子集）
|
||||
# 对照 pi ai/src/utils/validation.ts:317-347 validateToolArguments：
|
||||
# structuredClone → normalizeOptionalNulls → Value.Convert
|
||||
# → coerceWithJsonSchema → Check → 【报全部错误 + 回显收到的参数】
|
||||
# 这里实现核心子集（type/required/properties/enum/items），零外部依赖。
|
||||
# ======================================================================
|
||||
_TYPE_MAP = {
|
||||
"string": str,
|
||||
"integer": int,
|
||||
"number": (int, float),
|
||||
"boolean": bool,
|
||||
"array": list,
|
||||
"object": dict,
|
||||
}
|
||||
|
||||
|
||||
def _norm_types(schema: Dict[str, Any]) -> List[str]:
|
||||
"""schema 声明的类型（兼容 TypeBox 的 type 数组形式）"""
|
||||
t = schema.get("type")
|
||||
if isinstance(t, list):
|
||||
return [x for x in t if isinstance(x, str)]
|
||||
return [t] if isinstance(t, str) else []
|
||||
|
||||
|
||||
def _type_ok(value: Any, t: str) -> bool:
|
||||
"""类型匹配。🌟 bool 不算 integer/number（python 里 bool 是 int 子类）"""
|
||||
if t == "integer":
|
||||
return isinstance(value, int) and not isinstance(value, bool)
|
||||
if t == "number":
|
||||
return isinstance(value, (int, float)) and not isinstance(value, bool)
|
||||
if t == "boolean":
|
||||
return isinstance(value, bool)
|
||||
py = _TYPE_MAP.get(t)
|
||||
return isinstance(value, py) if py else True
|
||||
|
||||
|
||||
def _join_path(path: str, key: str) -> str:
|
||||
return f"{path}.{key}" if path else key
|
||||
|
||||
|
||||
def normalize_and_coerce(args: Any, schema: Dict[str, Any]) -> Any:
|
||||
"""归一化 + 轻量强制转换（对照 pi normalizeOptionalNulls + coerceWithJsonSchema）
|
||||
|
||||
- 可选字段的显式 null → 删除该键（模型常对「未用到的参数」发 null）
|
||||
- 数字字段收到纯数字字符串 → integer/number（"30" → 30）
|
||||
- boolean 字段收到 "true"/"false" → 布尔
|
||||
返回新对象，不修改入参。
|
||||
"""
|
||||
if not isinstance(schema, dict) or not isinstance(args, dict):
|
||||
return args
|
||||
props = schema.get("properties") or {}
|
||||
required = set(schema.get("required") or [])
|
||||
out: Dict[str, Any] = dict(args)
|
||||
for key, val in list(out.items()):
|
||||
sub = props.get(key)
|
||||
if not isinstance(sub, dict):
|
||||
continue
|
||||
if val is None:
|
||||
if key not in required:
|
||||
del out[key]
|
||||
continue
|
||||
out[key] = _coerce_value(val, sub)
|
||||
return out
|
||||
|
||||
|
||||
def _coerce_value(val: Any, sub: Dict[str, Any]) -> Any:
|
||||
types = _norm_types(sub)
|
||||
if isinstance(val, bool) or val is None:
|
||||
return val
|
||||
if isinstance(val, str):
|
||||
s = val.strip()
|
||||
if s and ("integer" in types or "number" in types):
|
||||
try:
|
||||
f = float(s)
|
||||
if f.is_integer():
|
||||
return int(f)
|
||||
if "number" in types:
|
||||
return f
|
||||
return val # integer 字段收到 1.5 → 保持原值（交由校验报错）
|
||||
except ValueError:
|
||||
return val
|
||||
if "boolean" in types and s.lower() in ("true", "false"):
|
||||
return s.lower() == "true"
|
||||
return val
|
||||
if isinstance(val, dict):
|
||||
return normalize_and_coerce(val, sub)
|
||||
if isinstance(val, list):
|
||||
items = sub.get("items")
|
||||
if isinstance(items, dict):
|
||||
return [_coerce_value(v, items) if not isinstance(v, dict)
|
||||
else normalize_and_coerce(v, items) for v in val]
|
||||
return val
|
||||
|
||||
|
||||
def _type_error(path: str, types: List[str], value: Any) -> str:
|
||||
t = types[0] if types else "object"
|
||||
if t == "string":
|
||||
return f"参数 {path} 应为 string"
|
||||
if t in ("integer", "number") and isinstance(value, bool):
|
||||
return f"参数 {path} 应为 {t}"
|
||||
return f"参数 {path} 类型错误: 期望 {t}"
|
||||
|
||||
|
||||
def _check_value(value: Any, schema: Dict[str, Any], path: str,
|
||||
errs: List[str], root: bool = False) -> None:
|
||||
"""递归收集【全部】校验错误（对照 pi Errors() 全量上报）"""
|
||||
if not isinstance(schema, dict):
|
||||
return
|
||||
if root and not isinstance(value, dict):
|
||||
errs.append(f"参数必须是对象，实际是 {type(value).__name__}")
|
||||
return
|
||||
types = _norm_types(schema) or (["object"] if root else [])
|
||||
if types and not any(_type_ok(value, t) for t in types):
|
||||
errs.append(_type_error(path, types, value))
|
||||
return # 类型不符 → 后续检查无意义
|
||||
if "enum" in schema and value not in schema["enum"]:
|
||||
errs.append(f"参数 {path} 取值必须是 {schema['enum']} 之一")
|
||||
if isinstance(value, dict):
|
||||
props = schema.get("properties") or {}
|
||||
for req in schema.get("required") or []:
|
||||
if req not in value:
|
||||
errs.append(f"缺少必填参数: {req}" if root
|
||||
else f"参数 {path} 缺少必填字段 {req}")
|
||||
for k, v in value.items():
|
||||
if k in props:
|
||||
_check_value(v, props[k], _join_path(path, k), errs)
|
||||
elif isinstance(value, list):
|
||||
items = schema.get("items")
|
||||
if isinstance(items, dict):
|
||||
for i, item in enumerate(value):
|
||||
_check_value(item, items, f"{path}[{i}]", errs)
|
||||
|
||||
|
||||
def validate_json_schema(args: Any, schema: Dict[str, Any]) -> Optional[str]:
|
||||
"""返回错误描述（多条以「; 」连接）；合法返回 None。"""
|
||||
errs: List[str] = []
|
||||
_check_value(args, schema, "", errs, root=True)
|
||||
return "; ".join(errs) if errs else None
|
||||
|
||||
|
||||
def _validate_value(value: Any, schema: Dict[str, Any], path: str) -> Optional[str]:
|
||||
"""（保留旧签名：返回该节点的首个错误）"""
|
||||
errs: List[str] = []
|
||||
_check_value(value, schema, path, errs)
|
||||
return errs[0] if errs else None
|
||||
|
||||
|
||||
def _brief_json(obj: Any, limit: int = 600) -> str:
|
||||
"""参数回显（长内容截断，避免 write 的大 content 撑爆错误消息）"""
|
||||
try:
|
||||
s = json.dumps(obj, ensure_ascii=False)
|
||||
except Exception:
|
||||
s = repr(obj)
|
||||
return s if len(s) <= limit else s[:limit] + f"…(共 {len(s)} 字符)"
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 准备阶段 —— 对照 agent-loop.ts prepareToolCalls
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class PreparedToolCall:
|
||||
"""一次工具调用的完整准备结果（执行前定型）"""
|
||||
tool_call: ToolCall
|
||||
tool: Optional[AgentTool]
|
||||
args: Dict[str, Any]
|
||||
error: str = "" # 准备阶段失败原因（未知工具/参数非法）→ 直接返回错误 toolResult
|
||||
|
||||
|
||||
def prepare_tool_calls(assistant: AgentMessage,
|
||||
tools: List[AgentTool]) -> List[PreparedToolCall]:
|
||||
tool_map = {t.name: t for t in tools}
|
||||
prepared: List[PreparedToolCall] = []
|
||||
for tc in assistant.tool_calls:
|
||||
tool = tool_map.get(tc.name)
|
||||
if tool is None:
|
||||
prepared.append(PreparedToolCall(
|
||||
tool_call=tc, tool=None, args={},
|
||||
error=f"未知工具: {tc.name}",
|
||||
))
|
||||
continue
|
||||
if not isinstance(tc.arguments, dict):
|
||||
prepared.append(PreparedToolCall(
|
||||
tool_call=tc, tool=tool, args={},
|
||||
error="工具参数解析失败（JSON 不完整）",
|
||||
))
|
||||
continue
|
||||
# 🌟 工具自定义参数预处理（对照 pi tool.prepareArguments）
|
||||
raw_args = tc.arguments
|
||||
if tool.prepare_arguments is not None:
|
||||
try:
|
||||
raw_args = tool.prepare_arguments(raw_args)
|
||||
except Exception as e:
|
||||
prepared.append(PreparedToolCall(
|
||||
tool_call=tc, tool=tool, args=tc.arguments,
|
||||
error=f"参数预处理失败: {e}",
|
||||
))
|
||||
continue
|
||||
if not isinstance(raw_args, dict):
|
||||
prepared.append(PreparedToolCall(
|
||||
tool_call=tc, tool=tool, args=tc.arguments,
|
||||
error="参数预处理返回的不是对象",
|
||||
))
|
||||
continue
|
||||
# 🌟 归一化 + 轻量强制转换（null 可选字段删除 / "30" → 30）
|
||||
args = normalize_and_coerce(raw_args, tool.parameters)
|
||||
err = validate_json_schema(args, tool.parameters)
|
||||
if err:
|
||||
prepared.append(PreparedToolCall(
|
||||
tool_call=tc, tool=tool, args=args,
|
||||
error=f"参数校验失败: {err}；收到的参数: {_brief_json(raw_args)}",
|
||||
))
|
||||
continue
|
||||
prepared.append(PreparedToolCall(tool_call=tc, tool=tool,
|
||||
args=args))
|
||||
return prepared
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 执行阶段 —— 对照 agent-loop.ts executeTool
|
||||
# before 钩子 → 执行 → after 钩子 → 异常兜底
|
||||
# ======================================================================
|
||||
def execute_tool_call(prepared: PreparedToolCall,
|
||||
assistant: AgentMessage,
|
||||
config,
|
||||
signal: AbortSignal,
|
||||
on_update: Optional[Callable[[str], None]],
|
||||
on_timer: Optional[Callable[[int, int], None]] = None,
|
||||
) -> AgentToolResult:
|
||||
"""对照 agent-loop.ts executePreparedToolCall + prepareToolCall 的钩子语义
|
||||
🆕 on_timer(elapsed_s, timeout_s)：bash 运行中每秒滴一次（气泡读秒）"""
|
||||
tc = prepared.tool_call
|
||||
|
||||
# 准备阶段已失败 → 直接错误结果（pi 同款：不进入执行）
|
||||
if prepared.error:
|
||||
return AgentToolResult.text(prepared.error, is_error=True)
|
||||
if signal.aborted:
|
||||
return AgentToolResult.text("操作已中止 (Operation aborted)", is_error=True)
|
||||
|
||||
# before 钩子（对照 pi: 可修改 args / 拒绝 block / 请求 terminate）
|
||||
if config.before_tool_call:
|
||||
try:
|
||||
decision = config.before_tool_call(
|
||||
{"assistant_message": assistant, "tool_call": tc,
|
||||
"args": prepared.args, "context": config.tool_context},
|
||||
signal)
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"before 钩子异常: {e}", is_error=True)
|
||||
if signal.aborted:
|
||||
return AgentToolResult.text("操作已中止 (Operation aborted)",
|
||||
is_error=True)
|
||||
if decision:
|
||||
if decision.get("block"):
|
||||
# 🐛 修复：旧代码向 AgentToolResult.text() 传了不存在的 terminate 参数
|
||||
# （钩子一旦 block 就抛 TypeError）—— 改为直接构造，保留 terminate 语义
|
||||
return AgentToolResult(
|
||||
content=[{"type": "text",
|
||||
"text": decision.get("reason") or "工具执行被拦截"}],
|
||||
is_error=True,
|
||||
terminate=bool(decision.get("terminate")))
|
||||
if decision.get("args") is not None:
|
||||
# 🌟 对照 pi applyBeforeToolDecision：钩子改参后【重新校验】
|
||||
new_args = decision["args"]
|
||||
if not isinstance(new_args, dict):
|
||||
return AgentToolResult.text(
|
||||
"before 钩子返回的参数不是对象", is_error=True)
|
||||
new_args = normalize_and_coerce(new_args, prepared.tool.parameters)
|
||||
_verr = validate_json_schema(new_args, prepared.tool.parameters)
|
||||
if _verr:
|
||||
return AgentToolResult.text(
|
||||
f"before 钩子修改后的参数校验失败: {_verr}；"
|
||||
f"收到的参数: {_brief_json(decision['args'])}",
|
||||
is_error=True)
|
||||
prepared.args = new_args
|
||||
|
||||
t0 = time.time()
|
||||
try:
|
||||
# 🆕 秒级滴答回调注入 ctx（不污染调用方的 tool_context 字典）
|
||||
_ctx = dict(config.tool_context or {})
|
||||
if on_timer is not None:
|
||||
_ctx["on_timer"] = on_timer
|
||||
result = prepared.tool.execute(
|
||||
tc.id, prepared.args, signal, on_update, _ctx,
|
||||
)
|
||||
if not isinstance(result, AgentToolResult):
|
||||
# 宽容处理：工具返回 str 也接受
|
||||
result = AgentToolResult.text(str(result))
|
||||
except Exception as e:
|
||||
result = AgentToolResult.text(f"工具执行异常: {e}", is_error=True)
|
||||
|
||||
# after 钩子
|
||||
if config.after_tool_call:
|
||||
try:
|
||||
config.after_tool_call(tc, result, result.is_error)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 结果定型（对照 pi finalizeToolResult：确保 content 结构合法）
|
||||
if not result.content:
|
||||
result = AgentToolResult.text("(无输出)")
|
||||
result.details = {"duration_ms": int((time.time() - t0) * 1000),
|
||||
**(result.details or {} if isinstance(result.details, dict) else {})}
|
||||
return result
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 截断保护 —— 对照 agent-loop.ts failToolCallsFromTruncatedMessage
|
||||
# 助手消息被 length/aborted 截断时，其工具调用参数可能残缺：
|
||||
# 一律替换为错误 toolResult，绝不执行残缺调用。
|
||||
# ======================================================================
|
||||
def fail_tool_calls_from_truncated_message(assistant: AgentMessage,
|
||||
reason: str = "aborted") -> List[AgentMessage]:
|
||||
"""为每个工具调用生成错误 toolResult 消息（对照 pi 返回值的组装）"""
|
||||
msgs: List[AgentMessage] = []
|
||||
for tc in assistant.tool_calls:
|
||||
# 对照 pi: "Tool call {name} was not executed: the response hit the
|
||||
# output token limit before the arguments were complete..."
|
||||
text = (f"工具调用 {tc.name} 未执行（{reason}）："
|
||||
f"响应在参数完整前达到输出长度上限，部分参数已被丢弃。"
|
||||
f"请用完整参数重试该操作。")
|
||||
msgs.append(AgentMessage(
|
||||
role="toolResult",
|
||||
tool_call_id=tc.id,
|
||||
tool_name=tc.name,
|
||||
content=text,
|
||||
is_error=True,
|
||||
))
|
||||
return msgs
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 内置工具 —— 对照 pi coding-agent 的 read / bash / write / edit
|
||||
# ======================================================================
|
||||
_MAX_READ_LINES = 2000
|
||||
_MAX_OUTPUT_BYTES = 50 * 1024
|
||||
|
||||
|
||||
def _resolve_path(path: str, ctx: Dict[str, Any]) -> str:
|
||||
"""相对路径基于 tool_context 的 cwd（默认项目根）"""
|
||||
if os.path.isabs(path):
|
||||
return os.path.abspath(path)
|
||||
cwd = ctx.get("cwd", os.getcwd())
|
||||
return os.path.abspath(os.path.join(cwd, path))
|
||||
|
||||
|
||||
def tool_read(tool_call_id: str, args: Dict[str, Any], signal: AbortSignal,
|
||||
on_update, ctx: Dict[str, Any]) -> AgentToolResult:
|
||||
"""read: 读取文本文件（支持 offset/limit 行窗口），带行号输出"""
|
||||
if not args.get("path"):
|
||||
return AgentToolResult.text("path 不能为空", is_error=True)
|
||||
path = _resolve_path(args["path"], ctx)
|
||||
if not os.path.isfile(path):
|
||||
return AgentToolResult.text(f"文件不存在: {path}", is_error=True)
|
||||
try:
|
||||
offset = max(1, int(args.get("offset", 1)))
|
||||
except (TypeError, ValueError):
|
||||
return AgentToolResult.text("offset 必须是整数", is_error=True)
|
||||
try:
|
||||
# 🌟 limit 必须有下限：负数会被 python 负索引语义吃掉
|
||||
# （旧版 limit=-5 → lines[0:-5] → 除末尾 5 行外全部返回，与直觉完全相反）
|
||||
limit = max(1, min(int(args.get("limit", 2000)), _MAX_READ_LINES))
|
||||
except (TypeError, ValueError):
|
||||
return AgentToolResult.text("limit 必须是整数", is_error=True)
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8", errors="replace") as f:
|
||||
lines = f.readlines()
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"读取失败: {e}", is_error=True)
|
||||
total = len(lines)
|
||||
if total == 0:
|
||||
return AgentToolResult.text("[文件为空（0 行）]")
|
||||
if offset > total:
|
||||
return AgentToolResult.text(
|
||||
f"[起始行 offset={offset} 超出文件范围，该文件共 {total} 行]")
|
||||
chunk = lines[offset - 1: offset - 1 + limit]
|
||||
numbered = "".join(
|
||||
f"{i + offset:6d}\t{line}" for i, line in enumerate(chunk)
|
||||
)
|
||||
shown_hi = offset - 1 + len(chunk)
|
||||
footer = f"\n[已显示 {offset}–{shown_hi} 行，共 {total} 行]"
|
||||
if shown_hi < total:
|
||||
footer += f"（还有 {total - shown_hi} 行未显示，用 offset={shown_hi + 1} 继续）"
|
||||
out = numbered + footer
|
||||
if len(out.encode("utf-8")) > _MAX_OUTPUT_BYTES:
|
||||
out = out.encode("utf-8")[:_MAX_OUTPUT_BYTES].decode("utf-8", "ignore")
|
||||
out += "\n[输出超过 50KB 已截断]"
|
||||
return AgentToolResult.text(out)
|
||||
|
||||
|
||||
def tool_bash(tool_call_id: str, args: Dict[str, Any], signal: AbortSignal,
|
||||
on_update, ctx: Dict[str, Any]) -> AgentToolResult:
|
||||
"""bash: 执行 shell 命令（默认 120s 超时）。对照 pi: bash 是 sequential 工具
|
||||
|
||||
🆕 秒级滴答：Popen + communicate(timeout=1) 循环，每秒：
|
||||
① 日志 [timer] bash 读秒 N/Ts
|
||||
② on_timer(N, T) 推前端气泡读秒
|
||||
③ 到期杀进程树（Windows taskkill /T，连孤儿子进程一起杀）→ 真超时
|
||||
"""
|
||||
if signal.aborted:
|
||||
return AgentToolResult.text("操作已中止 (Operation aborted)", is_error=True)
|
||||
command = args.get("command", "")
|
||||
if not command.strip():
|
||||
return AgentToolResult.text("command 不能为空", is_error=True)
|
||||
timeout = min(float(args.get("timeout", 120)), 600)
|
||||
cwd = ctx.get("cwd") or os.getcwd()
|
||||
# 对照 pi bash 工具：命令始终经 shell 解释（支持管道/别名/内置命令）
|
||||
use_shell = ctx.get("shell", True)
|
||||
on_timer = ctx.get("on_timer")
|
||||
t0 = time.time()
|
||||
# 🆕 计时观察①：计时器启动时刻 + 模型实际传的 timeout 值
|
||||
_dbg_log(f"[timer] bash 开始 timeout={timeout:.0f}s "
|
||||
f"(显式={args.get('timeout')}) cmd={command[:80]!r}")
|
||||
|
||||
def _tick(elapsed_i: int):
|
||||
"""每秒一次：日志 + 推前端（静默吞异常，绝不影响执行）"""
|
||||
_dbg_log(f"[timer] bash 读秒 {elapsed_i}/{int(timeout)}s")
|
||||
try:
|
||||
if on_timer:
|
||||
on_timer(elapsed_i, int(timeout))
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _kill_tree(proc):
|
||||
"""杀整个进程树：Windows 上默认 kill 只杀 cmd 壳，孤儿子进程继续
|
||||
持管道 → 假超时（设定10s 实际20s）。taskkill /T 整树杀。"""
|
||||
try:
|
||||
if os.name == "nt":
|
||||
subprocess.run(["taskkill", "/F", "/T", "/PID", str(proc.pid)],
|
||||
capture_output=True, timeout=10)
|
||||
else:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
try:
|
||||
proc.kill()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 🐛 修复：text=True 不带 encoding 时按系统码页（中文 Windows=GBK）解码，
|
||||
# 子进程输出 UTF-8（python/git/中文 echo）→ _readerthread UnicodeDecodeError。
|
||||
# 强制 UTF-8 + 容错替换；PYTHONIOENCODING 让 python 子进程也按 UTF-8 输出。
|
||||
_env = dict(os.environ, PYTHONIOENCODING="utf-8")
|
||||
try:
|
||||
proc = subprocess.Popen(
|
||||
command, shell=use_shell,
|
||||
stdout=subprocess.PIPE, stderr=subprocess.PIPE,
|
||||
text=True, encoding="utf-8", errors="replace",
|
||||
cwd=cwd, env=_env,
|
||||
)
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"执行失败: {e}", is_error=True)
|
||||
|
||||
timed_out = False
|
||||
out_parts: List[str] = []
|
||||
err_parts: List[str] = []
|
||||
err_header_fed = False
|
||||
|
||||
# 🆕 实时输出：stdout/stderr 各起一个 reader 线程 → 队列，
|
||||
# 主循环（仍在 worker 线程内）抽干队列并回调 on_update。
|
||||
# 绝不从 reader 线程直接回调 —— UI 侧 _on_tool_updated 会改 timeline，跨线程不安全。
|
||||
_q: "queue.Queue" = queue.Queue()
|
||||
|
||||
def _reader(stream, tag):
|
||||
try:
|
||||
for line in iter(stream.readline, ""):
|
||||
_q.put((tag, line))
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
try:
|
||||
stream.close()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
_t_out = threading.Thread(target=_reader, args=(proc.stdout, "out"), daemon=True)
|
||||
_t_err = threading.Thread(target=_reader, args=(proc.stderr, "err"), daemon=True)
|
||||
_t_out.start()
|
||||
_t_err.start()
|
||||
|
||||
def _feed(tag, chunk):
|
||||
"""抽到一块输出：累积 + 推实时流（只在 worker 线程内调用）"""
|
||||
nonlocal err_header_fed
|
||||
piece = chunk
|
||||
if tag == "err":
|
||||
err_parts.append(chunk)
|
||||
if not err_header_fed:
|
||||
err_header_fed = True
|
||||
piece = "[stderr]\n" + chunk
|
||||
else:
|
||||
out_parts.append(chunk)
|
||||
try:
|
||||
if on_update and piece:
|
||||
on_update(piece)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
def _drain():
|
||||
"""抽干队列里已到达的输出（单轮上限防极端刷屏卡死）"""
|
||||
for _ in range(500):
|
||||
try:
|
||||
tag, chunk = _q.get_nowait()
|
||||
except queue.Empty:
|
||||
return
|
||||
_feed(tag, chunk)
|
||||
|
||||
last_tick = 0 # 首次滴答仍在 elapsed=1（与旧 communicate 实现一致，不在 0 秒多滴一次）
|
||||
try:
|
||||
while True:
|
||||
# ① 最多 0.2s 醒一次 → 把新输出实时推给前端
|
||||
try:
|
||||
tag, chunk = _q.get(timeout=0.2)
|
||||
_feed(tag, chunk)
|
||||
except queue.Empty:
|
||||
pass
|
||||
_drain()
|
||||
|
||||
# ② 进程已退出：等 reader 读完管道残余 → 收尾
|
||||
if proc.poll() is not None:
|
||||
_t_out.join(timeout=1.0)
|
||||
_t_err.join(timeout=1.0)
|
||||
_drain()
|
||||
break
|
||||
|
||||
# ③ 每秒读秒 + 中止/超时判定（原语义不变）
|
||||
elapsed_i = int(time.time() - t0)
|
||||
if elapsed_i > last_tick:
|
||||
last_tick = elapsed_i
|
||||
if signal.aborted:
|
||||
_kill_tree(proc)
|
||||
_t_out.join(timeout=1.0)
|
||||
_t_err.join(timeout=1.0)
|
||||
_drain()
|
||||
return AgentToolResult.text(
|
||||
"操作已中止 (Operation aborted)", is_error=True)
|
||||
_tick(elapsed_i)
|
||||
if time.time() - t0 >= timeout:
|
||||
# 到期 → 真杀（进程树）
|
||||
timed_out = True
|
||||
_kill_tree(proc)
|
||||
try:
|
||||
proc.wait(timeout=15)
|
||||
except Exception:
|
||||
pass
|
||||
_t_out.join(timeout=5.0)
|
||||
_t_err.join(timeout=5.0)
|
||||
_drain()
|
||||
break
|
||||
except Exception as e:
|
||||
try:
|
||||
_kill_tree(proc)
|
||||
except Exception:
|
||||
pass
|
||||
return AgentToolResult.text(f"执行失败: {e}", is_error=True)
|
||||
|
||||
out = "".join(out_parts)
|
||||
err = "".join(err_parts)
|
||||
code = proc.returncode if proc.returncode is not None else -1
|
||||
|
||||
dur = time.time() - t0
|
||||
if timed_out:
|
||||
# 🆕 计时观察②：计时器到期
|
||||
_dbg_log(f"[timer] bash 超时触发 设定={timeout:.0f}s "
|
||||
f"实际={dur:.1f}s cmd={command[:80]!r}")
|
||||
return AgentToolResult.text(f"命令超时（>{timeout:.0f}s）已终止", is_error=True)
|
||||
# 🆕 计时观察③：正常结束 + 实际耗时
|
||||
_dbg_log(f"[timer] bash 正常结束 dur={dur:.1f}s "
|
||||
f"exit={code} 设定timeout={timeout:.0f}s")
|
||||
result = f"$ {command}\n"
|
||||
if out:
|
||||
result += out if out.endswith("\n") else out + "\n"
|
||||
if err:
|
||||
result += f"[stderr]\n{err}"
|
||||
result += f"\n[exit {code}] ({dur:.1f}s)"
|
||||
if len(result.encode("utf-8")) > _MAX_OUTPUT_BYTES:
|
||||
result = result.encode("utf-8")[:_MAX_OUTPUT_BYTES].decode("utf-8", "ignore")
|
||||
result += "\n[输出超过 50KB 已截断]"
|
||||
return AgentToolResult.text(result, is_error=code != 0,
|
||||
details={"exit_code": code})
|
||||
|
||||
|
||||
def _atomic_write(path: str, content: str) -> None:
|
||||
"""原子写：同目录临时文件 + os.replace（避免半截文件）
|
||||
|
||||
保留原有 newline 语义（默认 None → 平台换行翻译），仅增加原子性。
|
||||
"""
|
||||
d = os.path.dirname(os.path.abspath(path)) or "."
|
||||
fd, tmp = tempfile.mkstemp(dir=d, prefix=".hocode_w_", suffix=".tmp")
|
||||
try:
|
||||
with os.fdopen(fd, "w", encoding="utf-8") as f:
|
||||
f.write(content)
|
||||
os.replace(tmp, path) # Windows/POSIX 均为原子替换
|
||||
except BaseException:
|
||||
try:
|
||||
os.remove(tmp)
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
|
||||
|
||||
def tool_write(tool_call_id: str, args: Dict[str, Any], signal: AbortSignal,
|
||||
on_update, ctx: Dict[str, Any]) -> AgentToolResult:
|
||||
"""write: 创建/覆盖文件（自动建父目录）"""
|
||||
if not args.get("path"):
|
||||
return AgentToolResult.text("path 不能为空", is_error=True)
|
||||
path = _resolve_path(args["path"], ctx)
|
||||
content = args.get("content", "")
|
||||
if content is None:
|
||||
content = ""
|
||||
if not isinstance(content, str):
|
||||
content = str(content)
|
||||
try:
|
||||
parent = os.path.dirname(path)
|
||||
if parent:
|
||||
os.makedirs(parent, exist_ok=True)
|
||||
_atomic_write(path, content)
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"写入失败: {e}", is_error=True)
|
||||
return AgentToolResult.text(f"已写入 {len(content)} 字符 → {path}")
|
||||
|
||||
|
||||
def _prepare_edit_args(args: Any) -> Any:
|
||||
"""归一化 edit 参数（对照 pi edit.ts:56-70 prepareEditArguments）
|
||||
|
||||
- edits 为 JSON 字符串 → 解析为数组
|
||||
- edits 内条目为 JSON 字符串 → 解析为对象
|
||||
- 兼容 legacy 单条形式 {path, oldText, newText} → 包装为 edits:[{...}]
|
||||
"""
|
||||
if not isinstance(args, dict):
|
||||
return args
|
||||
out = dict(args)
|
||||
edits = out.get("edits")
|
||||
if isinstance(edits, str):
|
||||
try:
|
||||
parsed = json.loads(edits)
|
||||
if isinstance(parsed, list):
|
||||
out["edits"] = parsed
|
||||
elif isinstance(parsed, dict):
|
||||
out["edits"] = [parsed]
|
||||
except Exception:
|
||||
pass # 解析不了就交给校验层报错
|
||||
if out.get("edits") is None and ("oldText" in out or "newText" in out):
|
||||
ed: Dict[str, Any] = {}
|
||||
if "oldText" in out:
|
||||
ed["oldText"] = out["oldText"]
|
||||
if "newText" in out:
|
||||
ed["newText"] = out["newText"]
|
||||
out = {"path": out.get("path"), "edits": [ed]}
|
||||
if isinstance(out.get("edits"), list):
|
||||
norm = []
|
||||
for it in out["edits"]:
|
||||
if isinstance(it, str):
|
||||
try:
|
||||
it = json.loads(it)
|
||||
except Exception:
|
||||
pass
|
||||
norm.append(it)
|
||||
out["edits"] = norm
|
||||
return out
|
||||
|
||||
|
||||
def _plan_edits(content: str, edits: List[Any]):
|
||||
"""对【原始内容】定位每条 oldText（要求恰好 1 次）+ 区间重叠检测
|
||||
|
||||
对照 pi edit-diff.ts:348「edits[i] and edits[j] overlap … Merge them into one edit」。
|
||||
返回 (错误信息, [(start, end, index, new_text), ...] 已按 start 排序)
|
||||
"""
|
||||
spans: List[tuple] = []
|
||||
for i, ed in enumerate(edits):
|
||||
if not isinstance(ed, dict):
|
||||
return f"第 {i + 1} 条 edit 不是对象", None
|
||||
old = ed.get("oldText", "")
|
||||
if not isinstance(old, str):
|
||||
return f"第 {i + 1} 条 edit 的 oldText 必须是字符串", None
|
||||
if old == "":
|
||||
return f"第 {i + 1} 条 edit 的 oldText 不能为空", None
|
||||
c = content.count(old)
|
||||
if c == 0:
|
||||
return f"第 {i + 1} 条 edit 未找到匹配文本（oldText 不存在或已变化）", None
|
||||
if c > 1:
|
||||
return f"第 {i + 1} 条 edit 匹配到 {c} 处（要求唯一），请提供更长的上下文", None
|
||||
start = content.index(old)
|
||||
spans.append((start, start + len(old), i, ed.get("newText", "")))
|
||||
ordered = sorted(spans)
|
||||
for a, b in zip(ordered, ordered[1:]):
|
||||
if b[0] < a[1]:
|
||||
return (f"edits[{a[2]}] 与 edits[{b[2]}] 区域重叠，"
|
||||
f"请合并为一条 edit 或改为互不相交的修改"), None
|
||||
return None, ordered
|
||||
|
||||
|
||||
def tool_edit(tool_call_id: str, args: Dict[str, Any], signal: AbortSignal,
|
||||
on_update, ctx: Dict[str, Any]) -> AgentToolResult:
|
||||
"""edit: 精确文本替换（edits: [{oldText, newText}]）
|
||||
|
||||
🌟 语义对齐 pi：全部 edit 都对【原始文件内容】定位，要求各自唯一且区间互不重叠，
|
||||
然后按偏移一次性重建（不是逐条 replace 的增量语义）。
|
||||
"""
|
||||
if not args.get("path"):
|
||||
return AgentToolResult.text("path 不能为空", is_error=True)
|
||||
path = _resolve_path(args["path"], ctx)
|
||||
if not os.path.isfile(path):
|
||||
return AgentToolResult.text(f"文件不存在: {path}", is_error=True)
|
||||
edits = args.get("edits", [])
|
||||
if not isinstance(edits, list) or not edits:
|
||||
return AgentToolResult.text("edits 不能为空", is_error=True)
|
||||
try:
|
||||
with open(path, "r", encoding="utf-8") as f:
|
||||
content = f.read()
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"读取失败: {e}", is_error=True)
|
||||
err, ordered = _plan_edits(content, edits)
|
||||
if err:
|
||||
return AgentToolResult.text(err, is_error=True)
|
||||
# 按偏移从后往前替换（前面的偏移不受影响）
|
||||
out = content
|
||||
for start, end, _i, new_text in reversed(ordered):
|
||||
out = out[:start] + (new_text if isinstance(new_text, str) else str(new_text)) + out[end:]
|
||||
try:
|
||||
_atomic_write(path, out)
|
||||
except Exception as e:
|
||||
return AgentToolResult.text(f"编辑失败: {e}", is_error=True)
|
||||
return AgentToolResult.text(f"已应用 {len(ordered)} 处编辑 → {path}")
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 默认工具集 —— 对照 pi 默认启用 read/bash/edit/write
|
||||
# ======================================================================
|
||||
|
||||
_TEXT_TOOL_RE = re.compile(r"<(bash|read)>(.*?)</\1>", re.DOTALL)
|
||||
|
||||
|
||||
def parse_text_tool_calls(content: str):
|
||||
"""
|
||||
兜底解析（haocode 扩展，pi 无此层）：
|
||||
对不支持 tools API 的供应商/模型——它们会把工具调用用纯文字"演"出来
|
||||
（例如 <bash>ls</bash>、<read>路径</read>），本函数识别单参数工具
|
||||
bash / read 并转成真 ToolCall 供循环执行。
|
||||
write / edit 参数多、文本歧义大，不做兜底（保持安全）。
|
||||
返回 (原文, [ToolCall])；未命中时 calls 为空列表。
|
||||
"""
|
||||
if not content or "<" not in content:
|
||||
return content, []
|
||||
calls: List[ToolCall] = []
|
||||
|
||||
def _sub(m):
|
||||
name, payload = m.group(1), m.group(2).strip()
|
||||
if not payload:
|
||||
return m.group(0)
|
||||
if name == "bash":
|
||||
calls.append(ToolCall(id=new_id("txtcall"), name="bash",
|
||||
arguments={"command": payload}))
|
||||
elif name == "read":
|
||||
calls.append(ToolCall(id=new_id("txtcall"), name="read",
|
||||
arguments={"path": payload}))
|
||||
return m.group(0) # 保留原文（UI 已渲染，不回改；执行由 tool_calls 驱动）
|
||||
|
||||
cleaned = _TEXT_TOOL_RE.sub(_sub, content)
|
||||
return cleaned, calls
|
||||
|
||||
def default_tools() -> List[AgentTool]:
|
||||
return [
|
||||
AgentTool(
|
||||
name="read", label="读取文件",
|
||||
description="读取文本文件内容（带行号）。支持 offset/limit 按行窗口读取大文件。",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "文件路径"},
|
||||
"offset": {"type": "integer", "description": "起始行号（从 1 开始）"},
|
||||
"limit": {"type": "integer", "description": "最多读取行数（默认 2000）"},
|
||||
},
|
||||
"required": ["path"],
|
||||
},
|
||||
execute=tool_read,
|
||||
),
|
||||
AgentTool(
|
||||
name="bash", label="执行命令",
|
||||
description="执行 shell 命令并返回 stdout/stderr/退出码。默认 120 秒超时。",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"command": {"type": "string", "description": "要执行的命令"},
|
||||
"timeout": {"type": "number", "description": "超时秒数（默认 120，最大 600）"},
|
||||
},
|
||||
"required": ["command"],
|
||||
},
|
||||
execute=tool_bash,
|
||||
execution_mode="sequential", # 对照 pi: bash 整批串行
|
||||
),
|
||||
AgentTool(
|
||||
name="write", label="写入文件",
|
||||
description="创建或覆盖写入文件（自动创建父目录）。",
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "文件路径"},
|
||||
"content": {"type": "string", "description": "文件内容"},
|
||||
},
|
||||
"required": ["path", "content"],
|
||||
},
|
||||
execute=tool_write,
|
||||
),
|
||||
AgentTool(
|
||||
name="edit", label="编辑文件",
|
||||
description=("对文件做精确文本替换。edits 中每条 oldText 必须在原文件中唯一，"
|
||||
"且各条区间不得重叠（重叠请合并为一条）。"),
|
||||
parameters={
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"path": {"type": "string", "description": "文件路径"},
|
||||
"edits": {
|
||||
"type": "array",
|
||||
"items": {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"oldText": {"type": "string"},
|
||||
"newText": {"type": "string"},
|
||||
},
|
||||
"required": ["oldText"],
|
||||
},
|
||||
"description": "替换操作列表",
|
||||
},
|
||||
},
|
||||
"required": ["path", "edits"],
|
||||
},
|
||||
execute=tool_edit,
|
||||
prepare_arguments=_prepare_edit_args,
|
||||
),
|
||||
]
|
||||
@@ -0,0 +1,331 @@
|
||||
"""
|
||||
core/agent/types.py
|
||||
====================
|
||||
🌟 pi agent-core 的 Python 1:1 移植 —— 数据模型层
|
||||
|
||||
对照 pi-main 源码：
|
||||
packages/agent/src/types.ts (443 行)
|
||||
- AgentMessage 联合类型 (user / assistant / toolResult)
|
||||
- AgentEvent 11 种事件
|
||||
- AgentTool / AgentToolCall / AgentToolResult
|
||||
- AgentState / AgentConfig / AgentError
|
||||
|
||||
设计原则（与 pi 完全一致）：
|
||||
1. 消息是不可变事实（append-only），循环只追加、不修改
|
||||
2. 事件是唯一对外输出（TUI/Qt 都只是事件订阅者）
|
||||
3. 工具 = (name, description, JSON-Schema 参数, execute 函数)
|
||||
4. 中止用 AbortSignal 标志位，不用异常控制流
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Callable, Dict, List, Optional
|
||||
|
||||
|
||||
def _now_ms() -> int:
|
||||
return int(time.time() * 1000)
|
||||
|
||||
|
||||
def new_id(prefix: str) -> str:
|
||||
return f"{prefix}_{uuid.uuid4().hex[:16]}"
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 中止信号 —— 对照 pi: AbortController/AbortSignal（Web 标准）
|
||||
# pi 里 streamSimple 每消费一个 chunk 检查 signal.aborted
|
||||
# ======================================================================
|
||||
class AbortSignal:
|
||||
def __init__(self):
|
||||
self.aborted: bool = False
|
||||
self.reason: str = ""
|
||||
|
||||
def abort(self, reason: str = "aborted"):
|
||||
self.aborted = True
|
||||
self.reason = reason
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 消息模型 —— 对照 types.ts 的 AgentMessage 联合类型
|
||||
# pi: type AgentMessage = AgentUserMessage | AgentAssistantMessage | AgentToolResultMessage
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class ToolCall:
|
||||
"""pi: AssistantMessage 内容块里的 {type:"toolCall", id, name, arguments}"""
|
||||
id: str
|
||||
name: str
|
||||
arguments: Dict[str, Any] = field(default_factory=dict)
|
||||
# 流式阶段累积的原始 JSON 字符串（参数被截断时可能是残缺 JSON）
|
||||
raw_arguments: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentMessage:
|
||||
"""
|
||||
统一消息容器。role 决定哪些字段有效：
|
||||
user -> content: str 或 OpenAI 多模态 list
|
||||
assistant -> content(str 正文) + reasoning(思考) + tool_calls + stop_reason + usage
|
||||
toolResult -> tool_call_id + tool_name + content + is_error
|
||||
与 pi 的三种 Message 类型字段一一对应。
|
||||
"""
|
||||
role: str # "user" | "assistant" | "toolResult"
|
||||
content: Any = "" # str | List[dict]
|
||||
reasoning: str = "" # pi: thinking block（vLLM: reasoning_content）
|
||||
tool_calls: List[ToolCall] = field(default_factory=list)
|
||||
tool_call_id: str = "" # toolResult 专用
|
||||
tool_name: str = "" # toolResult 专用
|
||||
is_error: bool = False # toolResult 专用
|
||||
stop_reason: str = "" # assistant 专用: stop|length|aborted|error
|
||||
error_message: str = "" # assistant 出错时的说明
|
||||
usage: Dict[str, Any] = field(default_factory=dict) # {input, output, cacheRead...}
|
||||
kind: str = "" # "" | "compaction_summary"（对照 pi compaction 条目）
|
||||
db_msg_id: str = "" # 🆕 压缩持久化：对应 DB 行的 id（build_api_context 注入，切点计算用）
|
||||
timestamp: int = field(default_factory=_now_ms)
|
||||
id: str = field(default_factory=lambda: new_id("msg"))
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"role": self.role, "content": self.content, "reasoning": self.reasoning,
|
||||
"tool_calls": [
|
||||
{"id": tc.id, "name": tc.name, "arguments": tc.arguments,
|
||||
"raw_arguments": tc.raw_arguments}
|
||||
for tc in self.tool_calls
|
||||
],
|
||||
"tool_call_id": self.tool_call_id, "tool_name": self.tool_name,
|
||||
"is_error": self.is_error, "stop_reason": self.stop_reason,
|
||||
"error_message": self.error_message, "usage": self.usage,
|
||||
"kind": self.kind,
|
||||
"db_msg_id": self.db_msg_id,
|
||||
"timestamp": self.timestamp, "id": self.id,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def from_dict(d: Dict[str, Any]) -> "AgentMessage":
|
||||
return AgentMessage(
|
||||
role=d.get("role", "user"), content=d.get("content", ""),
|
||||
reasoning=d.get("reasoning", ""),
|
||||
tool_calls=[ToolCall(id=t.get("id", ""), name=t.get("name", ""),
|
||||
arguments=t.get("arguments", {}),
|
||||
raw_arguments=t.get("raw_arguments", ""))
|
||||
for t in d.get("tool_calls", [])],
|
||||
tool_call_id=d.get("tool_call_id", ""), tool_name=d.get("tool_name", ""),
|
||||
is_error=d.get("is_error", False), stop_reason=d.get("stop_reason", ""),
|
||||
error_message=d.get("error_message", ""), usage=d.get("usage", {}),
|
||||
kind=d.get("kind", ""),
|
||||
db_msg_id=d.get("db_msg_id", ""),
|
||||
timestamp=d.get("timestamp", 0), id=d.get("id") or new_id("msg"),
|
||||
)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 助手消息的流式增量事件 —— 对照 pi: AssistantMessageEvent
|
||||
# {type:"text_delta", textDelta} | {type:"thinking_delta", ...}
|
||||
# | {type:"toolcall_delta", ...} | {type:"done"} | {type:"error"}
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class AssistantMessageEvent:
|
||||
type: str # "text_delta" | "thinking_delta" | "toolcall_delta" | "done" | "error"
|
||||
text: str = "" # text_delta / thinking_delta 的增量
|
||||
tool_call_index: int = -1 # toolcall_delta 属于第几个 toolCall
|
||||
tool_call_delta: str = "" # toolcall_delta 的原始增量片段
|
||||
tool_call_field: str = "arguments" # 片段属于 id | name | arguments 哪一段
|
||||
error_message: str = "" # error 时的说明
|
||||
|
||||
@staticmethod
|
||||
def text_delta(text: str) -> "AssistantMessageEvent":
|
||||
return AssistantMessageEvent(type="text_delta", text=text)
|
||||
|
||||
@staticmethod
|
||||
def thinking_delta(text: str) -> "AssistantMessageEvent":
|
||||
return AssistantMessageEvent(type="thinking_delta", text=text)
|
||||
|
||||
@staticmethod
|
||||
def toolcall_delta(index: int, delta: str,
|
||||
field: str = "arguments") -> "AssistantMessageEvent":
|
||||
return AssistantMessageEvent(type="toolcall_delta", tool_call_index=index,
|
||||
tool_call_delta=delta, tool_call_field=field)
|
||||
|
||||
@staticmethod
|
||||
def done() -> "AssistantMessageEvent":
|
||||
return AssistantMessageEvent(type="done")
|
||||
|
||||
@staticmethod
|
||||
def error(message: str) -> "AssistantMessageEvent":
|
||||
return AssistantMessageEvent(type="error", error_message=message)
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# Agent 事件（11 种）—— 对照 types.ts 的 AgentEvent 联合类型
|
||||
# ======================================================================
|
||||
EVENT_TYPES = (
|
||||
"agent_start", "agent_end",
|
||||
"turn_start", "turn_end",
|
||||
"message_start", "message_update", "message_end",
|
||||
"tool_execution_start", "tool_execution_update", "tool_execution_end",
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentEvent:
|
||||
"""
|
||||
事件容器。字段随 type 变化：
|
||||
agent_end -> stop_reason, error
|
||||
message_start/end -> message (AgentMessage)
|
||||
message_update -> message(当前累积快照), assistant_message_event
|
||||
tool_execution_start -> tool_call, tool_name
|
||||
tool_execution_update -> tool_call, arg(增量)
|
||||
tool_execution_timer -> tool_call, arg=(已耗秒, 超时秒) 🆕 bash 读秒
|
||||
tool_execution_end -> tool_call, result, is_error
|
||||
"""
|
||||
type: str
|
||||
message: Optional[AgentMessage] = None
|
||||
assistant_message_event: Optional[AssistantMessageEvent] = None
|
||||
tool_call: Optional[ToolCall] = None
|
||||
arg: str = ""
|
||||
result: Optional["AgentToolResult"] = None
|
||||
is_error: bool = False
|
||||
stop_reason: str = ""
|
||||
error: Optional["AgentError"] = None
|
||||
messages: List[AgentMessage] = field(default_factory=list) # agent_end 携带的本轮新消息
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 工具 —— 对照 types.ts: AgentTool / AgentToolCall / AgentToolResult
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class AgentToolResult:
|
||||
"""pi: {content: OutputContent[], details?, isError, terminate?}"""
|
||||
content: List[Dict[str, Any]] = field(default_factory=list) # [{type:"text", text}]
|
||||
details: Any = None
|
||||
is_error: bool = False
|
||||
terminate: bool = False # 对照 pi: 工具可请求终止整个批次/循环
|
||||
|
||||
@staticmethod
|
||||
def text(text: str, is_error: bool = False, details: Any = None) -> "AgentToolResult":
|
||||
return AgentToolResult(content=[{"type": "text", "text": text}],
|
||||
details=details, is_error=is_error)
|
||||
|
||||
def as_text(self) -> str:
|
||||
return "".join(c.get("text", "") for c in self.content if c.get("type") == "text")
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentTool:
|
||||
"""
|
||||
pi: interface AgentTool<T> { name; description; parameters(JSON Schema);
|
||||
execute(toolCallId, args, signal, onUpdate, context) => Promise<ToolResult> }
|
||||
Python 版 execute 签名完全一致（同步执行，循环里用线程池并发）。
|
||||
"""
|
||||
name: str
|
||||
description: str
|
||||
parameters: Dict[str, Any] # JSON Schema
|
||||
execute: Callable # (tool_call_id, args, signal, on_update, context) -> AgentToolResult
|
||||
label: str = "" # UI 展示用短标签
|
||||
# 对照 pi types.ts: executionMode —— "sequential" 工具会把整个批次拉回串行执行
|
||||
execution_mode: str = "parallel"
|
||||
# 🌟 对照 pi AgentHarnessTool.prepareArguments：校验前的确定性参数预处理
|
||||
# （如 edit 兼容 legacy 单条形式 / edits 为 JSON 字符串），签名 (args) -> args
|
||||
prepare_arguments: Optional[Callable] = None
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 运行时状态 / 配置 —— 对照 types.ts: AgentState / AgentConfig
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class AgentState:
|
||||
messages: List[AgentMessage] = field(default_factory=list)
|
||||
is_streaming: bool = False
|
||||
streaming_message: Optional[AgentMessage] = None # 正在流式生成的助手消息
|
||||
streaming_delta: Dict[str, str] = field(default_factory=dict) # 各通道当前增量缓冲
|
||||
error: Optional["AgentError"] = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RetryConfig:
|
||||
"""对照 pi settings-manager.ts retry 默认值（行 32-33）：
|
||||
maxRetries=3, baseDelayMs=2000，指数退避 2s → 4s → 8s（base × 2^(n-1)）
|
||||
max_delay_ms=0 表示不设上限（pi _prepareRetry 无封顶）"""
|
||||
max_attempts: int = 3
|
||||
base_delay_ms: int = 2000
|
||||
factor: float = 2.0
|
||||
max_delay_ms: int = 0 # 0 = 不封顶
|
||||
|
||||
|
||||
@dataclass
|
||||
class ModelConfig:
|
||||
"""
|
||||
对照 pi models.json 里的模型定义。
|
||||
context_window = pi 的 contextWindow；max_tokens = pi 的 maxTokens。
|
||||
"""
|
||||
provider: str = ""
|
||||
name: str = ""
|
||||
context_window: int = 128000
|
||||
max_tokens: int = 8192
|
||||
# 请求参数
|
||||
temperature: float = 0.7
|
||||
timeout_seconds: float = 180.0
|
||||
api_key: str = ""
|
||||
base_url: str = "https://api.openai.com/v1"
|
||||
# 🆕 thinking 模式（DeepSeek 等）：assistant 消息必须把 reasoning_content
|
||||
# 原样回传，否则 400。默认开启（只影响有 reasoning 的轮中消息；
|
||||
# DB 回放行 reasoning="" 不带字段 → 全新请求体不变）。
|
||||
pass_reasoning: bool = True
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentConfig:
|
||||
"""
|
||||
对照 types.ts: AgentConfig
|
||||
model / tools / systemPrompt / maxTurns / steeringMode / followUpMode
|
||||
+ 6 个循环钩子（transformContext / beforeToolCall / afterToolCall /
|
||||
shouldStopAfterTurn / prepareNextTurn）
|
||||
"""
|
||||
model: ModelConfig
|
||||
system_prompt: str = ""
|
||||
tools: List[AgentTool] = field(default_factory=list)
|
||||
max_turns: int = 50 # pi 默认 50
|
||||
steering_mode: str = "one-at-a-time" # pi 默认 "one-at-a-time"（"all" 一次性全注入）
|
||||
follow_up_mode: str = "one-at-a-time"
|
||||
# 对照 pi types.ts ToolExecutionMode: 批次级执行模式（有 sequential 工具自动转串行）
|
||||
tool_execution: str = "parallel" # "sequential" | "parallel"
|
||||
retry: RetryConfig = field(default_factory=RetryConfig)
|
||||
# 压缩（🌟 1:1 对照 pi DEFAULT_COMPACTION_SETTINGS）
|
||||
compaction_reserve: int = 16384 # reserveTokens: 摘要提示词与输出预留
|
||||
compaction_keep_recent: int = 20000 # keepRecentTokens: 压缩后保留的近期上下文预算
|
||||
# 钩子（None = 无钩子，对照 pi 的可选字段）
|
||||
transform_context: Optional[Callable] = None # (messages) -> messages
|
||||
before_tool_call: Optional[Callable] = None # (tool_call, args) -> 可修改/拒绝
|
||||
after_tool_call: Optional[Callable] = None # (tool_call, result, is_error) -> None
|
||||
should_stop_after_turn: Optional[Callable] = None # (messages) -> bool
|
||||
prepare_next_turn: Optional[Callable] = None # (context) -> context
|
||||
# 🆕 haocode 增强（偏离 pi 1:1）：轮中请求前的主动压缩检查
|
||||
# (messages) -> messages|None：返回新列表 = 发生了压缩（循环需同步上下文）；
|
||||
# None = 不需要/不可压缩（原样发请求，响应式安全网仍在）
|
||||
compact_fn: Optional[Callable] = None # (messages) -> messages|None
|
||||
# 运行上下文（透传给 tool.execute，如工作目录）
|
||||
tool_context: Dict[str, Any] = field(default_factory=dict)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentError(Exception):
|
||||
"""对照 pi 的 AgentError —— 携带分类信息供恢复逻辑判定"""
|
||||
message: str = ""
|
||||
# 分类: "rate_limit" | "overload" | "server_error" | "timeout" |
|
||||
# "connection" | "auth" | "unknown"
|
||||
kind: str = "unknown"
|
||||
status_code: Optional[int] = None
|
||||
recoverable: bool = False # 恢复逻辑判定后的标记
|
||||
|
||||
def __str__(self):
|
||||
return self.message or "AgentError"
|
||||
|
||||
|
||||
# ======================================================================
|
||||
# 运行结果 —— 对照 agent-session.ts 里对 agent_end 事件的读取
|
||||
# ======================================================================
|
||||
@dataclass
|
||||
class RunResult:
|
||||
stop_reason: str = "stop" # stop | length | aborted | error
|
||||
error: Optional[AgentError] = None
|
||||
message_count: int = 0
|
||||
Reference in New Issue
Block a user