138 lines
6.5 KiB
Python
138 lines
6.5 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""LLM 调用层：OpenAI 兼容接口，要求模型输出固定的纯文本格式。
|
||
|
||
之所以不让模型直接输出 JSON，是因为 JSON 的括号/引号/转义更容易被模型写错；
|
||
"标签: 值"这种极简格式模型几乎不会出错，且比 json.loads 更容易做宽松解析。
|
||
本模块只负责"调用模型、拿到原始文本"，格式校验和清洗交给 parser.py。
|
||
|
||
提示词内容（system prompt + few-shot）全部来自 prompts.py，本文件不内嵌任何
|
||
提示词文字——修改分类边界/细分类目，只需要改 prompts.py。
|
||
|
||
为配合日志记录，classify_raw 返回时会连带一份完整的"本次调用消息列表"，
|
||
调用方（classifier.py）据此写日志，即使调用失败也能拿到"发出去的消息是什么"。
|
||
|
||
enable_other 原样透传给 prompts.get_prompt，决定这次调用的提示词里要不要
|
||
包含"其他"兜底类目；调用方（classifier.py）必须把同一个值也传给
|
||
parser.parse_llm_output，两边保持一致。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import time
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
from openai import APIError, APITimeoutError, OpenAI
|
||
|
||
from prompts import get_prompt
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class LLMClientError(Exception):
|
||
"""LLM 调用失败（网络/超时/接口错误），重试耗尽后抛出。"""
|
||
|
||
|
||
@dataclass
|
||
class LLMCallResult:
|
||
"""一次 LLM 调用的完整记录，无论成功与否都会产出，供日志模块使用。"""
|
||
|
||
messages: list[dict[str, str]] # 实际发送给模型的完整消息列表(含system+few-shot+用户)
|
||
raw_response: str | None # 模型原始回复；调用失败时为 None
|
||
error: str | None # 调用失败时的错误描述；成功时为 None
|
||
|
||
|
||
class LLMClient:
|
||
"""封装 OpenAI 兼容接口的调用，内置 few-shot 和重试。"""
|
||
|
||
def __init__(self, llm_cfg: dict[str, Any]):
|
||
self._cfg = llm_cfg
|
||
self._client = OpenAI(
|
||
base_url=llm_cfg["base_url"],
|
||
api_key=llm_cfg["api_key"],
|
||
timeout=llm_cfg["timeout"],
|
||
)
|
||
|
||
def _build_messages(
|
||
self, param_text: str, mode: str, enable_other: bool = False
|
||
) -> list[dict[str, str]]:
|
||
system_prompt, fewshot = get_prompt(mode, enable_other=enable_other)
|
||
messages: list[dict[str, str]] = [{"role": "system", "content": system_prompt}]
|
||
for user_text, assistant_text in fewshot:
|
||
messages.append(
|
||
{"role": "user", "content": f'请判断以下订单参数是否携带附件：\n"""\n{user_text}\n"""'}
|
||
)
|
||
messages.append({"role": "assistant", "content": assistant_text})
|
||
messages.append(
|
||
{"role": "user", "content": f'请判断以下订单参数是否携带附件：\n"""\n{param_text}\n"""'}
|
||
)
|
||
return messages
|
||
|
||
def classify_raw(
|
||
self, param_text: str, mode: str = "coarse", enable_other: bool = False
|
||
) -> LLMCallResult:
|
||
"""调用模型，返回完整调用记录（消息列表 + 原始回复/错误）。
|
||
|
||
与旧版不同：即使调用最终失败（重试耗尽），也不再抛异常中断调用方，
|
||
而是把失败信息装进 LLMCallResult 返回——这样 classifier.py 才能在
|
||
"调用彻底失败"的情况下依然拿到"发出去的消息是什么"，写进日志方便排查
|
||
（比如看是不是消息本身有问题导致接口一直拒绝）。
|
||
|
||
调用方如果需要区分"成功"还是"失败"，检查 result.error is None 即可。
|
||
"""
|
||
messages = self._build_messages(param_text, mode, enable_other=enable_other)
|
||
max_retry = self._cfg["max_retry"]
|
||
backoff = self._cfg["retry_backoff_seconds"]
|
||
last_err: Exception | None = None
|
||
|
||
# extra_body 用于透传标准 OpenAI SDK 不识别的厂商专属参数。
|
||
# disable_thinking=true 时关闭推理链——像 DeepSeek-V4 系列这类默认开启
|
||
# thinking 的模型，max_tokens 限制的是"推理token+正文token"的总量；
|
||
# 分类任务规则清晰、不需要推理，关闭后能避免推理阶段耗尽token预算导致
|
||
# 正文被截断为空（表现为调用"成功"但content是空字符串）。
|
||
# 若接口不支持该字段，通常会被直接忽略而非报错，但仍建议按需关闭本配置。
|
||
extra_body: dict[str, Any] = {}
|
||
if self._cfg.get("disable_thinking", False):
|
||
extra_body["thinking"] = {"type": "disabled"}
|
||
|
||
for attempt in range(1, max_retry + 1):
|
||
try:
|
||
resp = self._client.chat.completions.create(
|
||
model=self._cfg["model"],
|
||
temperature=self._cfg["temperature"],
|
||
max_tokens=self._cfg["max_tokens"],
|
||
messages=messages,
|
||
extra_body=extra_body or None,
|
||
)
|
||
content = resp.choices[0].message.content
|
||
# 空字符串和 None 同样视为"没拿到有效正文"——推理型模型在
|
||
# max_tokens 预算被推理阶段耗尽时，常表现为 content="" 而非
|
||
# None（HTTP层面调用是成功的），必须一并捕获，否则会把这种
|
||
# 情况误判为"调用成功、只是格式不对"，掩盖了真正的原因。
|
||
if not content:
|
||
raise LLMClientError(
|
||
f"模型返回内容为空 (content={content!r})，"
|
||
f"若模型支持思维链，可能是推理耗尽了max_tokens预算"
|
||
)
|
||
return LLMCallResult(messages=messages, raw_response=content, error=None)
|
||
except (APIError, APITimeoutError) as e:
|
||
last_err = e
|
||
logger.warning(
|
||
"LLM 调用失败 (第 %d/%d 次): %s", attempt, max_retry, e
|
||
)
|
||
if attempt < max_retry:
|
||
time.sleep(backoff * attempt)
|
||
except Exception as e: # noqa: BLE001 - 捕获SDK未明确分类的异常，统一包装
|
||
last_err = e
|
||
logger.warning(
|
||
"LLM 调用出现未预期异常 (第 %d/%d 次): %s", attempt, max_retry, e
|
||
)
|
||
if attempt < max_retry:
|
||
time.sleep(backoff * attempt)
|
||
|
||
return LLMCallResult(
|
||
messages=messages,
|
||
raw_response=None,
|
||
error=f"LLM 调用重试 {max_retry} 次后仍失败: {last_err}",
|
||
)
|