439 lines
19 KiB
Python
439 lines
19 KiB
Python
"""
|
||
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
|