334 lines
12 KiB
Python
334 lines
12 KiB
Python
"""
|
||
Agent HTTP 客户端模块
|
||
负责与远程 ADK Agent 服务通信，包括：
|
||
- HTTP 客户端生命周期管理
|
||
- 用户会话管理（创建、缓存、清除）
|
||
- SSE 流式调用（主方案）
|
||
- 同步调用（降级方案）
|
||
"""
|
||
import json
|
||
import logging
|
||
from typing import Optional
|
||
|
||
import httpx
|
||
|
||
from config import Config
|
||
from bot.reply_parser import parse_reply
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 用户会话映射: user_id -> session_id
|
||
_user_sessions: dict[str, str] = {}
|
||
|
||
# Agent HTTP 客户端
|
||
_http_client: Optional[httpx.AsyncClient] = None
|
||
|
||
|
||
def get_http_client() -> httpx.AsyncClient:
|
||
"""获取或创建 HTTP 客户端（全局单例）"""
|
||
global _http_client
|
||
if _http_client is None:
|
||
_http_client = httpx.AsyncClient(
|
||
base_url=Config.AGENT_BASE_URL,
|
||
timeout=Config.AGENT_TIMEOUT,
|
||
)
|
||
return _http_client
|
||
|
||
|
||
async def close_http_client():
|
||
"""关闭 HTTP 客户端"""
|
||
global _http_client
|
||
if _http_client:
|
||
await _http_client.aclose()
|
||
_http_client = None
|
||
|
||
|
||
# ============================================================
|
||
# 会话管理
|
||
# ============================================================
|
||
|
||
async def get_or_create_session(user_id: str) -> str:
|
||
"""获取或创建用户的 Agent 会话"""
|
||
if user_id in _user_sessions:
|
||
return _user_sessions[user_id]
|
||
|
||
client = get_http_client()
|
||
try:
|
||
url = f"/apps/{Config.AGENT_APP_NAME}/users/{user_id}/sessions"
|
||
resp = await client.post(url, json={})
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
session_id = data["id"]
|
||
_user_sessions[user_id] = session_id
|
||
logger.info("为用户 %s 创建 Agent 会话: %s", user_id, session_id)
|
||
return session_id
|
||
except Exception as e:
|
||
logger.error("创建 Agent 会话失败: %s", e)
|
||
raise
|
||
|
||
|
||
def clear_session(user_id: str):
|
||
"""清除用户的会话缓存（会话失效时调用）"""
|
||
_user_sessions.pop(user_id, None)
|
||
|
||
|
||
# ============================================================
|
||
# 显示格式化
|
||
# ============================================================
|
||
|
||
def _format_turn_for_display(text: str, use_emoji: bool = True) -> str:
|
||
"""
|
||
将一轮的原始文本（ thinking... response...）格式化为微信展示文本。
|
||
use_emoji=True 用 💭 前缀（流式阶段，企微会渲染为可折叠的"已深度思考"）
|
||
use_emoji=False 用纯文本标题（最终回复，避免触发"思考中"指示器）
|
||
"""
|
||
if not text:
|
||
return ""
|
||
|
||
prefix = "💭 思考过程:" if use_emoji else "【思考过程】"
|
||
|
||
# 查找 response 分隔标记
|
||
parts = text.split('\n response\n', 1)
|
||
if len(parts) == 2:
|
||
think_raw, body_raw = parts
|
||
think_clean = think_raw.strip()
|
||
if think_clean.startswith(' thinking'):
|
||
think_clean = think_clean[len(' thinking'):].lstrip('\n')
|
||
think_clean = think_clean.strip()
|
||
body_clean = body_raw.strip()
|
||
|
||
result = ""
|
||
if think_clean:
|
||
result = f"{prefix}\n{think_clean}\n\n"
|
||
if body_clean:
|
||
result += body_clean
|
||
return result
|
||
else:
|
||
# 还没有 response，全是思考内容
|
||
think_clean = text.strip()
|
||
if think_clean.startswith(' thinking'):
|
||
think_clean = think_clean[len(' thinking'):].lstrip('\n')
|
||
think_clean = think_clean.strip()
|
||
return f"{prefix}\n{think_clean}" if think_clean else ""
|
||
|
||
|
||
# ============================================================
|
||
# Agent 调用
|
||
# ============================================================
|
||
|
||
async def call_agent_stream(
|
||
ws_client,
|
||
frame: dict,
|
||
user_id: str,
|
||
message: str,
|
||
) -> str:
|
||
"""
|
||
通过 SSE 流式调用 Agent，实现打字机效果。
|
||
失败时降级为同步调用。
|
||
返回完整的回复文本。
|
||
"""
|
||
from aibot import generate_req_id
|
||
|
||
session_id = await get_or_create_session(user_id)
|
||
client = get_http_client()
|
||
|
||
try:
|
||
req_body = {
|
||
"app_name": Config.AGENT_APP_NAME,
|
||
"user_id": user_id,
|
||
"session_id": session_id,
|
||
"new_message": {
|
||
"role": "user",
|
||
"parts": [{"text": message}],
|
||
},
|
||
"streaming": True,
|
||
}
|
||
|
||
async with client.stream(
|
||
"POST",
|
||
"/run_sse",
|
||
json=req_body,
|
||
timeout=Config.AGENT_TIMEOUT,
|
||
) as response:
|
||
if response.status_code != 200:
|
||
logger.warning(
|
||
"SSE 流式调用返回 %d，降级为同步调用",
|
||
response.status_code,
|
||
)
|
||
return await call_agent_sync(user_id, message)
|
||
|
||
stream_id = generate_req_id("stream")
|
||
full_text = "" # 累积非 STOP 的增量文本
|
||
turn_start = 0 # 当前轮在 full_text 中的起始位置
|
||
stop_texts: list[str] = [] # 每轮 STOP 时的完整文本快照
|
||
event_count = 0
|
||
last_sent_display = "" # 上次已发送到微信的展示文本
|
||
has_sent_final = False
|
||
first_event_sent = False # 是否已发送过流式更新（用于5秒超时保护）
|
||
|
||
async for line in response.aiter_lines():
|
||
if not line.startswith("data: "):
|
||
continue
|
||
|
||
data_str = line[6:]
|
||
event_count += 1
|
||
try:
|
||
data = json.loads(data_str)
|
||
except json.JSONDecodeError:
|
||
logger.warning("SSE #%d JSON 解析失败: %s", event_count, data_str[:200])
|
||
continue
|
||
|
||
finish_reason = data.get("finishReason", "")
|
||
content = data.get("content", {})
|
||
parts = content.get("parts", [])
|
||
|
||
# 日志：每个 SSE 事件的摘要
|
||
part_summary = []
|
||
for p in parts:
|
||
t = p.get("text", "")
|
||
part_summary.append(
|
||
f"text_len={len(t)}, thought={p.get('thought', False)}, "
|
||
f"preview={t[:50]!r}"
|
||
)
|
||
logger.info(
|
||
"SSE #%d | finish=%s | parts=%d | %s",
|
||
event_count, finish_reason, len(parts),
|
||
" | ".join(part_summary) if part_summary else "(无parts)",
|
||
)
|
||
|
||
# 错误处理
|
||
if "error" in data:
|
||
logger.error("Agent SSE 错误: %s", data["error"])
|
||
if "not found" in data["error"].lower():
|
||
clear_session(user_id)
|
||
return "抱歉，处理消息时出现错误，请稍后重试。"
|
||
|
||
is_final = finish_reason == "STOP"
|
||
|
||
if is_final:
|
||
# STOP 事件：ADK 会重复发送完整文本，不追加到 full_text
|
||
# 保存当前轮的文本快照
|
||
turn_text = full_text[turn_start:]
|
||
if turn_text:
|
||
stop_texts.append(turn_text)
|
||
logger.info(
|
||
">>> STOP #%d，当前轮文本长度: %d，前200字: %s",
|
||
len(stop_texts), len(turn_text), turn_text[:200],
|
||
)
|
||
turn_start = len(full_text)
|
||
last_sent_display = "" # 新轮次重置展示缓存
|
||
else:
|
||
# 非 STOP：增量文本
|
||
for part in parts:
|
||
text = part.get("text", "")
|
||
if text:
|
||
full_text += text
|
||
|
||
# 流式更新微信消息
|
||
turn_text = full_text[turn_start:]
|
||
display = _format_turn_for_display(turn_text)
|
||
if display and display != last_sent_display:
|
||
try:
|
||
await ws_client.reply_stream(
|
||
frame, stream_id, display, False
|
||
)
|
||
last_sent_display = display
|
||
first_event_sent = True
|
||
except Exception:
|
||
pass
|
||
|
||
# === SSE 流结束，发送最终回复 ===
|
||
logger.info(
|
||
"=== SSE 流结束 === %d 个事件, %d 个 STOP, full_text 总长: %d",
|
||
event_count, len(stop_texts), len(full_text),
|
||
)
|
||
for idx, st in enumerate(stop_texts):
|
||
tc, bc = parse_reply(st)
|
||
logger.info(
|
||
"STOP[%d] 解析: 思考=%d, 正文=%d, 正文前100字: %s",
|
||
idx, len(tc), len(bc), bc[:100] if bc else "(空)",
|
||
)
|
||
|
||
# 从最后一个有正文的 STOP 构建最终内容
|
||
final_body = ""
|
||
for stop_text in reversed(stop_texts):
|
||
_, body_content = parse_reply(stop_text)
|
||
if body_content:
|
||
final_body = body_content
|
||
break
|
||
|
||
# 兜底
|
||
if not final_body and stop_texts:
|
||
_, final_body = parse_reply(stop_texts[-1])
|
||
if not final_body and full_text:
|
||
_, final_body = parse_reply(full_text[turn_start:])
|
||
if not final_body and full_text:
|
||
final_body = full_text[turn_start:]
|
||
|
||
if final_body:
|
||
try:
|
||
# 用纯正文 + finish=True 正常关闭流
|
||
# 思考过程在流式阶段已通过 💭 展示，finish=True 替换为正文
|
||
await ws_client.reply_stream(
|
||
frame, stream_id, final_body, True
|
||
)
|
||
has_sent_final = True
|
||
logger.info(">>> 最终回复发送成功，正文长度: %d", len(final_body))
|
||
except Exception as e:
|
||
logger.warning("推送最终回复失败: %s", e)
|
||
|
||
_, body = parse_reply(full_text[turn_start:] if turn_start else full_text)
|
||
return body if body else ""
|
||
|
||
except httpx.ConnectError:
|
||
logger.error("无法连接到 Agent 服务: %s", Config.AGENT_BASE_URL)
|
||
return "抱歉，Agent 服务连接失败，请稍后重试。"
|
||
except Exception as e:
|
||
logger.error("SSE 流式调用异常，降级为同步: %s", e)
|
||
return await call_agent_sync(user_id, message)
|
||
|
||
|
||
async def call_agent_sync(user_id: str, message: str) -> str:
|
||
"""
|
||
同步调用 Agent（/run 接口），作为流式调用的降级方案。
|
||
"""
|
||
session_id = await get_or_create_session(user_id)
|
||
client = get_http_client()
|
||
|
||
def _build_body(sid: str) -> dict:
|
||
return {
|
||
"app_name": Config.AGENT_APP_NAME,
|
||
"user_id": user_id,
|
||
"session_id": sid,
|
||
"new_message": {
|
||
"role": "user",
|
||
"parts": [{"text": message}],
|
||
},
|
||
}
|
||
|
||
def _extract_reply(events: list) -> str:
|
||
"""从 Event[] 中提取最终回复（跳过 thought）"""
|
||
for event in events:
|
||
parts = event.get("content", {}).get("parts", [])
|
||
for part in parts:
|
||
if part.get("text") and not part.get("thought"):
|
||
return part["text"]
|
||
return ""
|
||
|
||
try:
|
||
resp = await client.post("/run", json=_build_body(session_id))
|
||
resp.raise_for_status()
|
||
return _extract_reply(resp.json())
|
||
|
||
except httpx.HTTPStatusError as e:
|
||
if e.response.status_code == 404:
|
||
clear_session(user_id)
|
||
new_session = await get_or_create_session(user_id)
|
||
resp = await client.post("/run", json=_build_body(new_session))
|
||
resp.raise_for_status()
|
||
return _extract_reply(resp.json())
|
||
logger.error("Agent HTTP 错误 [%d]: %s", e.response.status_code, e)
|
||
return "抱歉，服务暂时不可用，请稍后重试。"
|
||
except Exception as e:
|
||
logger.error("同步调用 Agent 失败: %s", e)
|
||
return "抱歉，处理消息时出现错误，请稍后重试。" |