"""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_RE = re.compile(r".*?", 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. 去掉 ... 思考内容 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