622f348cf1
- db.py: 新增 upsert_attachments/fetch_all_ids，实现 Common.Attachment 幂等写入 （先删受影响 NS 旧行再批量插入，依赖唯一索引，可安全重跑） - write_attachments.py: 写入入口，支持 --limit/--ids-file/--id（单个或逗号分隔多个） /--mode/--dry-run/--enable-other，运行结束打印 token 与缓存命中率汇总 - llm_client.py: LLMCallResult 捕获 usage/elapsed_ms/model/attempt - classifier.py: classify_batch 结果透传 meta（耗时分两种、上下文 token、缓存命中率）， 新增 summarize_results 聚合批统计 - main.py: 新增 --summary 把批汇总打到 stderr，stdout 保持干净 JSON Lines - order_logger.py: 每次 LLM 尝试补充 [调用统计] 段，便于排查耗时与缓存效果 - README.md: 对齐上述接口与统计说明 Co-Authored-By: WorkBuddy <workbuddy@tencent.com>
199 lines
9.0 KiB
Python
199 lines
9.0 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
|
||
# 以下字段为性能/用量统计，默认 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,
|
||
)
|