97 lines
3.6 KiB
Python
97 lines
3.6 KiB
Python
"""AI 调用工具封装
|
|
|
|
每个协程独立重试，失败后 sleep 200ms 再试，直到成功。
|
|
不再有全局门闸，互不影响。
|
|
"""
|
|
|
|
import asyncio
|
|
import re
|
|
from datetime import datetime
|
|
from typing import Any, Optional
|
|
|
|
from json_repair import repair_json
|
|
from langchain_core.language_models import BaseChatModel
|
|
from langchain_core.messages import SystemMessage, HumanMessage
|
|
from snowflake import SnowflakeGenerator
|
|
|
|
from app.core.logger import log
|
|
|
|
# AI 日志专用雪花ID
|
|
_log_id_gen = SnowflakeGenerator(instance=2)
|
|
|
|
# 匹配 <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)
|
|
|
|
|
|
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, scene: Optional[str] = None) -> str:
|
|
"""异步调用 LLM，返回原始文本。失败后独立重试直到成功。"""
|
|
attempt = 0
|
|
|
|
while True:
|
|
attempt += 1
|
|
try:
|
|
messages = [
|
|
SystemMessage(content=system_prompt),
|
|
HumanMessage(content=user_message),
|
|
]
|
|
response = await llm.ainvoke(messages)
|
|
content = response.content
|
|
# 记录 AI 调用日志（失败不影响主流程）
|
|
if scene:
|
|
await _save_call_log(scene, system_prompt, user_message, content)
|
|
return content
|
|
except Exception as e:
|
|
log.warning("AI 调用失败(第{}次), 200ms 后重试: {}", attempt, e)
|
|
await asyncio.sleep(0.2)
|
|
|
|
|
|
async def ai_chat_json(llm: BaseChatModel, system_prompt: str, user_message: str, scene: Optional[str] = None) -> Any:
|
|
"""异步调用 LLM，返回解析后的 JSON 对象"""
|
|
raw = await ai_chat(llm, system_prompt, user_message, scene=scene)
|
|
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
|
|
|
|
|
|
async def _save_call_log(scene: str, system_prompt: str, user_message: str, response: Optional[str]) -> None:
|
|
"""异步写入 AI 调用日志到 PG，失败只 warning 不影响主流程"""
|
|
try:
|
|
from app.core.database import PgSession
|
|
from sqlalchemy import insert
|
|
from app.models.pg.ai_call_log import AiCallLog
|
|
|
|
async with PgSession() as pg:
|
|
await pg.execute(
|
|
insert(AiCallLog).values(
|
|
id=next(_log_id_gen),
|
|
scene=scene,
|
|
system_prompt=system_prompt,
|
|
user_message=user_message,
|
|
response=response,
|
|
created_at=datetime.now(),
|
|
)
|
|
)
|
|
await pg.commit()
|
|
except Exception as e:
|
|
log.warning("AI 调用日志写入失败: {}", e)
|