attachment_classifier/llm_client.py

# -*- 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
    # 以下字段为性能/用量统计，默认 None（调用失败时无有效值）：
    usage: dict[str, Any] | None = None  # resp.usage 转成的 dict（含 prompt/completion/total
    #   tokens 及 DeepSeek 的 prompt_cache_hit/miss_tokens），失败时为 None
    elapsed_ms: float | None = None  # 单次成功调用的耗时（毫秒，含网络+推理），失败时为 None
    model: str | None = None  # 实际服务的模型名（来自 resp.model），失败时为 None
    attempt: int | None = None  # 本次返回对应的是第几次尝试（成功那次；全失败则为 max_retry）


def _extract_usage(resp: Any) -> dict[str, Any] | None:
    """从 chat.completions 响应里稳健地取出 usage 字典。

    DeepSeek 在标准 OpenAI usage 之上，额外返回 `prompt_cache_hit_tokens` /
    `prompt_cache_miss_tokens`（与 prompt_tokens 同级）。openai SDK 未必为这些
    厂商字段建模，所以先按属性访问，取不到再回退到 pydantic 的 model_dump()
    （extra 字段通常会被带出来），两层都试，避免漏掉缓存统计。
    """
    if resp is None or getattr(resp, "usage", None) is None:
        return None
    u = resp.usage
    out: dict[str, Any] = {}
    for key in (
        "prompt_tokens",
        "completion_tokens",
        "total_tokens",
        "prompt_cache_hit_tokens",
        "prompt_cache_miss_tokens",
    ):
        val = getattr(u, key, None)
        if val is not None:
            out[key] = val
    if not out:
        try:
            dumped = u.model_dump()  # type: ignore[attr-defined]
        except Exception:
            dumped = {}
        for key in (
            "prompt_tokens",
            "completion_tokens",
            "total_tokens",
            "prompt_cache_hit_tokens",
            "prompt_cache_miss_tokens",
        ):
            if key in dumped and dumped[key] is not None:
                out[key] = dumped[key]
    return out or 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):
            call_start = time.perf_counter()
            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预算"
                    )
                # 统计：单次调用耗时 + token 用量 + 实际服务模型
                call_elapsed_ms = (time.perf_counter() - call_start) * 1000
                usage = _extract_usage(resp)
                return LLMCallResult(
                    messages=messages,
                    raw_response=content,
                    error=None,
                    usage=usage,
                    elapsed_ms=call_elapsed_ms,
                    model=getattr(resp, "model", None),
                    attempt=attempt,
                )
            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}",
            usage=None,
            elapsed_ms=None,
            model=None,
            attempt=max_retry,
        )