110 lines
4.3 KiB
Python
110 lines
4.3 KiB
Python
"""AI 调用工具封装
|
||
|
||
核心机制：健康门闸（Health Gate）
|
||
- ainvoke 抛异常 → 标记故障，全部协程阻塞，单探针退避重试直到成功
|
||
- ainvoke 正常返回 → 放行，后续逻辑不变
|
||
"""
|
||
|
||
import asyncio
|
||
import re
|
||
from typing import Any
|
||
|
||
from json_repair import repair_json
|
||
from langchain_core.language_models import BaseChatModel
|
||
from langchain_core.messages import SystemMessage, HumanMessage
|
||
|
||
from app.core.logger import log
|
||
|
||
# 匹配 <think>任意内容</think>，用于剥离推理模型的思考过程
|
||
_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)
|
||
|
||
# 匹配 ```json ... ``` 代码块，提取中间的 JSON 内容
|
||
_CODE_BLOCK_RE = re.compile(r"```(?:json\w*)?\s*\n?(.*?)\n?\s*```", re.DOTALL | re.IGNORECASE)
|
||
|
||
# ──────────── 健康门闸 ────────────
|
||
_gate_event = asyncio.Event() # clear=故障阻塞中, set=健康放行
|
||
_gate_event.set() # 初始状态：健康
|
||
_gate_probing = False # 是否已有探针在重试
|
||
|
||
|
||
def parse_llm_json(text: str) -> Any:
|
||
"""解析 AI 输出的 JSON，自动去除思考标签、markdown 代码块，容错处理"""
|
||
# 1. 去掉 <think>...</think> 思考内容
|
||
cleaned = _THINK_RE.sub("", text).strip()
|
||
# 2. 如果有 ```json ... ``` 代码块，只取代码块里的内容
|
||
match = _CODE_BLOCK_RE.search(cleaned)
|
||
if match:
|
||
cleaned = match.group(1).strip()
|
||
# 3. repair_json 容错解析：修复不规范的 JSON（多余逗号、缺引号、非法转义等）
|
||
return repair_json(cleaned, return_objects=True)
|
||
|
||
|
||
async def ai_chat(llm: BaseChatModel, system_prompt: str, user_message: str) -> str:
|
||
"""异步调用 LLM，返回原始文本。接口异常时触发门闸阻塞重试直到成功。"""
|
||
global _gate_probing
|
||
|
||
while True:
|
||
# 门闸检查：如果当前处于故障状态，阻塞等待直到探针恢复门闸
|
||
await _gate_event.wait()
|
||
|
||
try:
|
||
messages = [
|
||
SystemMessage(content=system_prompt),
|
||
HumanMessage(content=user_message),
|
||
]
|
||
response = await llm.ainvoke(messages)
|
||
return response.content
|
||
except Exception as e:
|
||
# 接口失败 → 关闭门闸
|
||
if _gate_event.is_set():
|
||
_gate_event.clear()
|
||
log.error("AI 接口异常，门闸关闭，全部协程阻塞等待恢复: {}", e)
|
||
|
||
if not _gate_probing:
|
||
# 当前协程成为探针，负责退避重试直到成功
|
||
_gate_probing = True
|
||
result = await _probe_until_recover(llm, system_prompt, user_message)
|
||
# 探针成功，直接返回结果（不用再调一次）
|
||
return result
|
||
# 已有探针在工作，回到 while 顶部 await _gate_event.wait() 等待唤醒
|
||
|
||
|
||
async def _probe_until_recover(llm: BaseChatModel, system_prompt: str, user_message: str) -> str:
|
||
"""探针：退避重试直到接口恢复，返回成功的响应内容，并打开门闸唤醒所有等待协程"""
|
||
global _gate_probing
|
||
|
||
delay = 1.0 # 初始退避 1 秒
|
||
max_delay = 30.0 # 最大退避 30 秒
|
||
attempt = 0
|
||
|
||
while True:
|
||
await asyncio.sleep(delay)
|
||
attempt += 1
|
||
try:
|
||
messages = [
|
||
SystemMessage(content=system_prompt),
|
||
HumanMessage(content=user_message),
|
||
]
|
||
response = await llm.ainvoke(messages)
|
||
# 成功 → 打开门闸，唤醒所有等待协程
|
||
_gate_probing = False
|
||
_gate_event.set()
|
||
log.info("AI 接口恢复，门闸打开，第{}次探测成功", attempt)
|
||
return response.content
|
||
except Exception as e:
|
||
log.warning("AI 接口仍不可用，第{}次探测失败(退避{}s): {}", attempt, delay, e)
|
||
delay = min(delay * 2, max_delay)
|
||
|
||
|
||
async def ai_chat_json(llm: BaseChatModel, system_prompt: str, user_message: str) -> Any:
|
||
"""异步调用 LLM，返回解析后的 JSON 对象"""
|
||
raw = await ai_chat(llm, system_prompt, user_message)
|
||
if not raw or not raw.strip():
|
||
log.warning("AI 返回为空")
|
||
return None
|
||
try:
|
||
return parse_llm_json(raw)
|
||
except Exception as e:
|
||
log.warning("AI JSON 解析失败: {}, raw={}", e, raw[:200])
|
||
return None
|