init
This commit is contained in:
237
wechat-bot/bot/agent_client.py
Normal file
237
wechat-bot/bot/agent_client.py
Normal file
@@ -0,0 +1,237 @@
|
||||
"""
|
||||
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 extract_final_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)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# 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 = ""
|
||||
has_sent_final = False
|
||||
|
||||
async for line in response.aiter_lines():
|
||||
if not line.startswith("data: "):
|
||||
continue
|
||||
|
||||
data_str = line[6:]
|
||||
try:
|
||||
data = json.loads(data_str)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
|
||||
# 错误处理
|
||||
if "error" in data:
|
||||
logger.error("Agent SSE 错误: %s", data["error"])
|
||||
if "not found" in data["error"].lower():
|
||||
clear_session(user_id)
|
||||
return "抱歉，处理消息时出现错误，请稍后重试。"
|
||||
|
||||
# 收集所有文本（含思考过程）
|
||||
parts = data.get("content", {}).get("parts", [])
|
||||
for part in parts:
|
||||
text = part.get("text", "")
|
||||
if not text:
|
||||
continue
|
||||
if part.get("thought"):
|
||||
continue
|
||||
full_text += text
|
||||
|
||||
# 实时推送"正在思考"状态
|
||||
if full_text and not has_sent_final:
|
||||
try:
|
||||
await ws_client.reply_stream(
|
||||
frame, stream_id, "💭 正在思考...", False
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# 最终完成时提取正式回复并发送
|
||||
is_final = data.get("finishReason") == "STOP"
|
||||
if is_final and full_text and not has_sent_final:
|
||||
clean_reply = extract_final_reply(full_text)
|
||||
logger.info("Agent 回复完成，长度: %d", len(clean_reply))
|
||||
try:
|
||||
await ws_client.reply_stream(
|
||||
frame, stream_id, clean_reply, True
|
||||
)
|
||||
has_sent_final = True
|
||||
except Exception as e:
|
||||
logger.warning("推送最终回复失败: %s", e)
|
||||
|
||||
# 兜底：如果流结束但未发送最终回复
|
||||
if full_text and not has_sent_final:
|
||||
clean_reply = extract_final_reply(full_text)
|
||||
logger.info("Agent SSE 兜底发送，长度: %d", len(clean_reply))
|
||||
try:
|
||||
await ws_client.reply_stream(
|
||||
frame, stream_id, clean_reply, True
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning("推送兜底回复失败: %s", e)
|
||||
|
||||
return extract_final_reply(full_text) if full_text 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 "抱歉，处理消息时出现错误，请稍后重试。"
|
||||
Reference in New Issue
Block a user