240 lines
8.0 KiB
Python
240 lines
8.0 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 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
|
||
has_sent_thinking = 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_thinking and not has_sent_final:
|
||
try:
|
||
await ws_client.reply_stream(
|
||
frame, stream_id, "💭 正在思考...", False
|
||
)
|
||
has_sent_thinking = True
|
||
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 "抱歉，处理消息时出现错误，请稍后重试。"
|