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>
291 lines
12 KiB
Python
291 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""编排层：串联 数据库查询 -> LLM调用 -> 格式校验解析 -> 组装最终JSON。
|
||
|
||
单条记录的处理逻辑（classify_param_text）:
|
||
1. 若新参数为空/None -> 直接判定无附件，不调用LLM(省成本，也没有可判断内容)。
|
||
2. 调用 LLM 拿到调用结果(LLMCallResult)。
|
||
- 若调用本身失败(网络/接口错误,重试耗尽) -> 记日志,标记 llm_call_error。
|
||
3. 用 parser.parse_llm_output(mode=..., enable_other=...) 做严格格式校验。
|
||
- 校验通过 -> 组装结果,记日志,结束。
|
||
- 校验失败(FormatError) -> 记日志,重新调用LLM重试。
|
||
4. 格式校验重试全部耗尽 -> 标记 status="llm_format_error"，不让脏数据进入最终结果。
|
||
|
||
无论上述哪一步，只要 OrderLogger 存在，每一次尝试都会被完整记录——包括发给
|
||
模型的完整对话、模型的原始回复（如果有）、格式校验是否通过。这是本模块与
|
||
日志模块的核心约定：不因为"这次调用失败了"就跳过记录，失败的调用恰恰最
|
||
需要被记下来供排查。
|
||
|
||
批量处理(classify_batch)在此基础上先查数据库拿到 {总排号: 新参数} 映射，
|
||
对查不到的总排号直接标记 status="not_found"，不发起LLM调用（也不生成日志，
|
||
因为根本没有"新参数"可供判断，没有对话可记）。
|
||
|
||
enable_other 控制是否允许模型使用"其他"兜底类目（默认 False），从
|
||
classify_batch/classify_single 一路透传到 llm_client.classify_raw（决定
|
||
提示词要不要包含"其他"）和 parser.parse_llm_output（决定校验要不要放行
|
||
"其他"）。本模块是唯一同时持有这两个调用点的地方，因此由本模块负责把
|
||
同一个值传给两边，不暴露给更上层去分别设置从而导致不一致。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import time
|
||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||
from pathlib import Path
|
||
from typing import Any
|
||
|
||
from db import DatabaseError, fetch_params_by_ids
|
||
from llm_client import LLMCallResult, LLMClient
|
||
from order_logger import OrderLogger
|
||
from parser import FormatError, ParsedResult, parse_llm_output
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def _empty_result(zong_pai_hao: str, status: str) -> dict[str, Any]:
|
||
return {
|
||
"zong_pai_hao": zong_pai_hao,
|
||
"status": status,
|
||
"has_attachment": None,
|
||
"types": [],
|
||
# 未发起 LLM 调用的结果（not_found/db_error），meta 给出一致的空壳，
|
||
# 便于下游统一读取（elapsed_ms 为 None，统计时按 0 处理）。
|
||
"meta": {"elapsed_ms": None, "attempts": 0, "llm": None},
|
||
}
|
||
|
||
|
||
def _build_usage(u: dict[str, Any]) -> dict[str, Any]:
|
||
"""把原始 usage dict 整理成对外结构，并算出缓存率。
|
||
|
||
缓存率 = prompt_cache_hit_tokens / prompt_tokens（prompt_tokens 为 0 时按
|
||
0.0 处理，避免除零）。保留 4 位小数。
|
||
"""
|
||
prompt = u.get("prompt_tokens", 0) or 0
|
||
completion = u.get("completion_tokens", 0) or 0
|
||
total = u.get("total_tokens", 0) or 0
|
||
hit = u.get("prompt_cache_hit_tokens", 0) or 0
|
||
miss = u.get("prompt_cache_miss_tokens", 0) or 0
|
||
cache_hit_rate = round(hit / prompt, 4) if prompt else 0.0
|
||
return {
|
||
"prompt_tokens": prompt,
|
||
"completion_tokens": completion,
|
||
"total_tokens": total,
|
||
"prompt_cache_hit_tokens": hit,
|
||
"prompt_cache_miss_tokens": miss,
|
||
"cache_hit_rate": cache_hit_rate,
|
||
}
|
||
|
||
|
||
def _make_meta(
|
||
call_result: LLMCallResult | None, start: float, attempt: int
|
||
) -> dict[str, Any]:
|
||
"""组装单条结果的 meta 对象。
|
||
|
||
elapsed_ms 为从 classify_param_text 入口到此处的总耗时（含重试+解析）；
|
||
若本次有成功调用，则 llm 块记录模型名、单次调用耗时与整理后的 usage；
|
||
否则（调用失败/未调用）llm 为 None。
|
||
"""
|
||
elapsed_ms = (time.perf_counter() - start) * 1000
|
||
if call_result is not None and call_result.error is None and call_result.usage is not None:
|
||
llm_block = {
|
||
"model": call_result.model,
|
||
"elapsed_ms": call_result.elapsed_ms,
|
||
"usage": _build_usage(call_result.usage),
|
||
}
|
||
else:
|
||
llm_block = None
|
||
return {"elapsed_ms": elapsed_ms, "attempts": attempt, "llm": llm_block}
|
||
|
||
|
||
def classify_param_text(
|
||
llm_client: LLMClient,
|
||
param_text: str | None,
|
||
mode: str = "coarse",
|
||
format_retry: int = 2,
|
||
order_logger: OrderLogger | None = None,
|
||
enable_other: bool = False,
|
||
) -> tuple[str, bool | None, list[str], dict[str, Any]]:
|
||
"""对单条"新参数"文本做分类，返回 (status, has_attachment, types)。
|
||
|
||
mode 决定分类粒度："coarse"只输出大类(资料/配件/耗材，启用 enable_other
|
||
时还有"其他")；"fine"输出"大类:细分类目"（如"配件:针型阀"）的列表。
|
||
不再有单独的 fine_types 返回值——fine 模式下细分类目已经内嵌在 types
|
||
的每个元素里。
|
||
|
||
enable_other 控制是否允许模型使用"其他"兜底那些不在既有枚举范围内的
|
||
附件，默认关闭；关闭时的行为与不支持"其他"之前完全一致。
|
||
|
||
status 取值:
|
||
"ok" - 正常识别成功
|
||
"empty_param" - 新参数字段为空，直接判无附件，未调用LLM
|
||
"llm_call_error" - LLM 网络/接口调用失败（重试耗尽）
|
||
"llm_format_error" - LLM 返回内容格式校验多次失败（重试耗尽）
|
||
|
||
format_retry 控制"格式校验失败后重新请求LLM"的次数，与 LLMClient 内部的
|
||
网络层重试是两回事：网络重试解决的是"请求没成功"，这里解决的是
|
||
"请求成功了但模型没按格式输出"。两层重试互不影响，各自独立计数。
|
||
|
||
order_logger 若提供，会记录本次分类过程中的每一次LLM往返（无论成败）。
|
||
"""
|
||
start = time.perf_counter()
|
||
|
||
if param_text is None or not str(param_text).strip():
|
||
if order_logger is not None:
|
||
order_logger.log_param_text(param_text)
|
||
order_logger.log_final_result("empty_param", None)
|
||
return "empty_param", False, [], _make_meta(None, start, 0)
|
||
|
||
if order_logger is not None:
|
||
order_logger.log_param_text(param_text)
|
||
|
||
last_call_result: LLMCallResult | None = None
|
||
last_format_err: Exception | None = None
|
||
for attempt in range(1, format_retry + 1):
|
||
call_result = llm_client.classify_raw(param_text, mode=mode, enable_other=enable_other)
|
||
last_call_result = call_result
|
||
|
||
if call_result.error is not None:
|
||
logger.error("LLM 调用失败: %s", call_result.error)
|
||
if order_logger is not None:
|
||
order_logger.log_llm_attempt(attempt, call_result, None, None)
|
||
order_logger.log_final_result("llm_call_error", None)
|
||
return "llm_call_error", None, [], _make_meta(call_result, start, attempt)
|
||
|
||
try:
|
||
parsed = parse_llm_output(call_result.raw_response, mode=mode, enable_other=enable_other)
|
||
if order_logger is not None:
|
||
order_logger.log_llm_attempt(attempt, call_result, parsed, None)
|
||
order_logger.log_final_result("ok", parsed)
|
||
return "ok", parsed.has_attachment, parsed.types, _make_meta(call_result, start, attempt)
|
||
except FormatError as e:
|
||
last_format_err = e
|
||
logger.warning(
|
||
"LLM 输出格式校验失败 (第 %d/%d 次): %s | 原始输出: %r",
|
||
attempt, format_retry, e, call_result.raw_response,
|
||
)
|
||
if order_logger is not None:
|
||
order_logger.log_llm_attempt(attempt, call_result, None, str(e))
|
||
|
||
logger.error("格式校验重试 %d 次后仍失败，放弃: %s", format_retry, last_format_err)
|
||
if order_logger is not None:
|
||
order_logger.log_final_result("llm_format_error", None)
|
||
return "llm_format_error", None, [], _make_meta(last_call_result, start, format_retry)
|
||
|
||
|
||
def classify_batch(
|
||
db_cfg: dict[str, Any],
|
||
llm_cfg: dict[str, Any],
|
||
zong_pai_hao_list: list[str],
|
||
max_workers: int = 8,
|
||
mode: str = "coarse",
|
||
log_dir: str | Path | None = None,
|
||
enable_other: bool = False,
|
||
) -> list[dict[str, Any]]:
|
||
"""批量分类入口：查库 -> 并发调用LLM -> 组装结果列表。
|
||
|
||
log_dir 若提供，会为每个实际发起LLM调用的总排号在该目录下生成一个日志
|
||
文件；查库失败(not_found)或参数为空(empty_param)的总排号不生成日志文件，
|
||
因为它们本就没有与LLM的对话内容可记。
|
||
|
||
返回结果顺序与输入 zong_pai_hao_list 一致，即使某些查询失败或格式错误也会
|
||
补全为对应 status 的占位结果，保证"输入N个总排号，输出N条结果"。
|
||
"""
|
||
# 去重但保留顺序，避免用户传入重复总排号导致重复查询/调用
|
||
seen: set[str] = set()
|
||
unique_ids: list[str] = []
|
||
for zph in zong_pai_hao_list:
|
||
if zph not in seen:
|
||
seen.add(zph)
|
||
unique_ids.append(zph)
|
||
|
||
try:
|
||
param_map = fetch_params_by_ids(db_cfg, unique_ids)
|
||
except DatabaseError as e:
|
||
logger.error("数据库查询失败，本批次全部标记为 db_error: %s", e)
|
||
return [_empty_result(zph, "db_error") for zph in zong_pai_hao_list]
|
||
|
||
llm_client = LLMClient(llm_cfg)
|
||
results_by_id: dict[str, dict[str, Any]] = {}
|
||
|
||
ids_to_classify = [zph for zph in unique_ids if param_map.get(zph) is not None]
|
||
for zph in unique_ids:
|
||
if param_map.get(zph) is None:
|
||
results_by_id[zph] = _empty_result(zph, "not_found")
|
||
|
||
def _worker(zph: str) -> tuple[str, dict[str, Any]]:
|
||
order_logger = OrderLogger(log_dir, zph, mode) if log_dir is not None else None
|
||
status, has_attachment, types, meta = classify_param_text(
|
||
llm_client, param_map[zph], mode=mode, order_logger=order_logger,
|
||
enable_other=enable_other,
|
||
)
|
||
return zph, {
|
||
"zong_pai_hao": zph,
|
||
"status": status,
|
||
"has_attachment": has_attachment,
|
||
"types": types,
|
||
"meta": meta,
|
||
}
|
||
|
||
if ids_to_classify:
|
||
with ThreadPoolExecutor(max_workers=max_workers) as pool:
|
||
futures = [pool.submit(_worker, zph) for zph in ids_to_classify]
|
||
for fut in as_completed(futures):
|
||
zph, result = fut.result()
|
||
results_by_id[zph] = result
|
||
|
||
# 按原始输入顺序（含重复项）展开最终结果
|
||
return [results_by_id[zph] for zph in zong_pai_hao_list]
|
||
|
||
|
||
def classify_single(
|
||
db_cfg: dict[str, Any],
|
||
llm_cfg: dict[str, Any],
|
||
zong_pai_hao: str,
|
||
mode: str = "coarse",
|
||
log_dir: str | Path | None = None,
|
||
enable_other: bool = False,
|
||
) -> dict[str, Any]:
|
||
"""单个总排号分类的便捷封装。"""
|
||
return classify_batch(
|
||
db_cfg, llm_cfg, [zong_pai_hao], max_workers=1, mode=mode, log_dir=log_dir,
|
||
enable_other=enable_other,
|
||
)[0]
|
||
|
||
|
||
def summarize_results(results: list[dict[str, Any]]) -> dict[str, Any]:
|
||
"""汇总一批 classify_batch 结果的运行统计。
|
||
|
||
从每条结果的 meta 里累加：总耗时、实际 LLM 调用次数、各类 token 总量、
|
||
以及整体缓存命中率（按 prompt_cache_hit_tokens / prompt_tokens 加权）。
|
||
供 CLI 打印批统计用，不改变 classify_batch 的返回结构。
|
||
"""
|
||
total_elapsed = 0.0
|
||
llm_calls = 0
|
||
sum_prompt = sum_completion = sum_total = sum_hit = sum_miss = 0
|
||
for r in results:
|
||
meta = r.get("meta") or {}
|
||
e = meta.get("elapsed_ms")
|
||
if e:
|
||
total_elapsed += e
|
||
llm = meta.get("llm")
|
||
if isinstance(llm, dict):
|
||
u = llm.get("usage")
|
||
if isinstance(u, dict):
|
||
llm_calls += 1
|
||
sum_prompt += u.get("prompt_tokens", 0) or 0
|
||
sum_completion += u.get("completion_tokens", 0) or 0
|
||
sum_total += u.get("total_tokens", 0) or 0
|
||
sum_hit += u.get("prompt_cache_hit_tokens", 0) or 0
|
||
sum_miss += u.get("prompt_cache_miss_tokens", 0) or 0
|
||
avg_cache_hit_rate = round(sum_hit / sum_prompt, 4) if sum_prompt else 0.0
|
||
return {
|
||
"count": len(results),
|
||
"total_elapsed_ms": round(total_elapsed, 1),
|
||
"llm_calls": llm_calls,
|
||
"total_prompt_tokens": sum_prompt,
|
||
"total_completion_tokens": sum_completion,
|
||
"total_tokens": sum_total,
|
||
"total_cache_hit_tokens": sum_hit,
|
||
"total_cache_miss_tokens": sum_miss,
|
||
"avg_cache_hit_rate": avg_cache_hit_rate,
|
||
}
|