attachment_classifier/main.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""布莱迪压力表 - 订单附件识别 命令行入口。

用法:
    python main.py --sn 26B742                      # 单个总排号(粗分类,默认)
    python main.py --sn 26B742,26B743,26B744        # 批量，逗号分隔
    python main.py --id 802                         # 单个数据库真实 ID
    python main.py --id 802,803,804                 # 批量真实 ID，逗号分隔
    python main.py --ids-file ids.txt                # 批量，文件每行一个总排号(键类型 sn)
    python main.py --sn 26B742 --mode fine           # 精分类(输出"大类:细分"格式)
    python main.py --id 802 --config other.yaml     # 指定其他配置文件
    python main.py --sn 26B742 --pretty              # 格式化输出JSON(默认单行紧凑)
    python main.py --sn 26B742 --log-dir /tmp/logs   # 覆盖配置文件里的日志目录
    python main.py --sn 26B742 --enable-other        # 允许模型使用"其他"兜底类目
    python main.py --sn 26B742 --summary             # 额外在 stderr 打印本批运行汇总统计

指定方式（两种键类型可混用）:
    --id  数据库真实 ID 列(id_field)，如 802
    --sn  总排号列(id_column)，如 26B742（即原先 --id 的语义）
    --ids-file 每行一个总排号(键类型 sn)
无论用哪种方式指定，输出 JSON 中的 zong_pai_hao 一律为回查到的总排号。

分类模式(--mode):
    coarse (默认) - 只判断大类: 资料/配件/耗材(启用 --enable-other 时还有"其他")
    fine          - 输出"大类:细分类目"，如"配件:针型阀""资料:说明书"

"其他"兜底类目(--enable-other 或配置文件 business.enable_other_category):
    默认关闭。关闭时，模型和格式校验都不知道"其他"这个选项存在，行为与
    未支持"其他"之前完全一致。开启后，遇到确实是随货附件、但不属于任何
    已列出细分类目的情况，模型可以归为"其他"，避免这类附件被漏判、或被
    模型勉强套进不准确的类目里。命令行只能把它从"关"打开成"开"，无法
    反过来用命令行强制关闭配置文件里已经打开的设置。

输出:
    每个总排号一行 JSON（JSON Lines 格式），例如:
    {"zong_pai_hao": "26B742", "status": "ok", "has_attachment": true, "types": ["资料"]}
    {"zong_pai_hao": "26B744", "status": "ok", "has_attachment": true, "types": ["配件:针型阀", "配件:表弯管"]}
    {"zong_pai_hao": "XXXXXX", "status": "not_found", "has_attachment": null, "types": []}

    types 是唯一的类型字段，不再有单独的 fine_types：
        coarse模式下 types 是大类列表；fine模式下每一项都是"大类:细分类目"。

    status 字段含义:
        ok                - 识别成功
        not_found         - 数据库中查不到该总排号
        empty_param       - 该总排号存在，但新参数字段为空
        llm_call_error    - LLM 调用失败(网络/接口错误,重试耗尽)
        llm_format_error  - LLM 返回内容格式不符合约定(重试耗尽)
        db_error          - 数据库连接/查询失败(影响整批)

日志:
    每次实际发起LLM调用的总排号，都会在日志目录下生成一个独立的日志文件
    (文件名含总排号和时间戳)，完整记录发给模型的对话内容、模型的每一次
    原始回复(含被判定格式错误、触发重试的)、以及最终解析结果。查不到
    (not_found)或参数为空(empty_param)的总排号不产生日志文件。
    日志目录默认读取配置文件 business.log_dir，可用 --log-dir 临时覆盖。
"""
from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent / "src"))

from classifier import classify_batch, summarize_results  # noqa: E402
from config_loader import ConfigError, load_config  # noqa: E402


def _parse_ids(args: argparse.Namespace) -> list[tuple[str, str]]:
    """从 --id / --sn / --ids-file 解析出 (标识符, 键类型) 列表。

    --id     -> 键类型 "id"（数据库真实 ID 列）
    --sn     -> 键类型 "sn"（总排号列，即原先 --id 的语义）
    --ids-file -> 每行一个总排号，键类型 "sn"
    三者可同时提供、合并后返回（顺序：--id, --sn, --ids-file）。
    """
    items: list[tuple[str, str]] = []
    if args.id:
        items.extend((s.strip(), "id") for s in args.id.split(",") if s.strip())
    if args.sn:
        items.extend((s.strip(), "sn") for s in args.sn.split(",") if s.strip())
    if args.ids_file:
        p = Path(args.ids_file)
        if not p.exists():
            print(f"[错误] 总排号文件不存在: {p.resolve()}", file=sys.stderr)
            sys.exit(1)
        with p.open("r", encoding="utf-8") as f:
            items.extend((line.strip(), "sn") for line in f if line.strip())
    return items


def main() -> None:
    parser = argparse.ArgumentParser(
        description="布莱迪压力表订单附件识别工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument(
        "--id", type=str, default=None,
        help="数据库真实 ID（如 802），单个或逗号分隔的多个，例如: 802 或 802,803",
    )
    parser.add_argument(
        "--sn", type=str, default=None,
        help="总排号（如 26B742），单个或逗号分隔的多个，例如: 26B742 或 26B742,26B743；"
             "与 --id 同义但键类型不同（--sn 查总排号列，--id 查真实 ID 列）",
    )
    parser.add_argument(
        "--ids-file", type=str, default=None,
        help="包含总排号的文本文件路径，每行一个总排号（键类型 sn）",
    )
    parser.add_argument(
        "--config", type=str, default="config.yaml",
        help="配置文件路径 (默认: config.yaml)",
    )
    parser.add_argument(
        "--mode", type=str, default=None, choices=["coarse", "fine"],
        help="分类粒度: coarse(粗分类,默认) 或 fine(精分类)。不指定则使用配置文件 business.default_mode",
    )
    parser.add_argument(
        "--log-dir", type=str, default=None,
        help="LLM对话日志存放目录。不指定则使用配置文件 business.log_dir",
    )
    parser.add_argument(
        "--pretty", action="store_true",
        help="以缩进格式输出JSON数组，而非默认的紧凑JSON Lines格式",
    )
    parser.add_argument(
        "--enable-other", action="store_true",
        help="允许模型使用'其他'兜底类目。不指定则使用配置文件 business.enable_other_category"
             "（缺省为关闭）；本开关只能从命令行打开，无法用命令行强制关闭配置文件里已打开的设置",
    )
    parser.add_argument(
        "--summary", action="store_true",
        help="在 stderr 额外打印本批运行的汇总统计（耗时/token/缓存命中率）；"
             "stdout 仍只输出干净的 JSON Lines 结果，便于管道/重定向",
    )
    args = parser.parse_args()

    if not args.id and not args.sn and not args.ids_file:
        parser.error("必须提供 --id / --sn / --ids-file 其中之一")

    ids = _parse_ids(args)
    if not ids:
        print("[错误] 未解析到任何有效的标识符（--id / --sn / --ids-file 均为空）", file=sys.stderr)
        sys.exit(1)

    try:
        cfg = load_config(args.config)
    except ConfigError as e:
        print(f"[配置错误] {e}", file=sys.stderr)
        sys.exit(1)

    logging.basicConfig(
        level=getattr(logging, cfg["business"]["log_level"], logging.INFO),
        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
        stream=sys.stderr,  # 日志走stderr，保证stdout只有干净的JSON结果，便于管道/重定向
    )

    mode = args.mode or cfg["business"]["default_mode"]
    log_dir = args.log_dir or cfg["business"]["log_dir"]
    # enable_other_category 是后加的可选配置项，缺省按"关闭"处理；--enable-other
    # 只能从"关"打开成"开"，不支持反过来用命令行强制关闭配置文件里已打开的设置。
    enable_other = bool(args.enable_other) or bool(cfg["business"].get("enable_other_category", False))

    # 整批分类的真实墙钟：在并发 classify_batch 外层计时（perf_counter）。
    # summarize_results 累加的 total_elapsed_ms 是各任务自身耗时的总和（并发下的
    # 累计工作量，会明显大于墙钟），二者之比 ≈ 并发增益；wall_clock_ms 才是真实墙钟。
    t0 = time.perf_counter()
    results = classify_batch(
        db_cfg=cfg["database"],
        llm_cfg=cfg["llm"],
        id_list=ids,
        max_workers=cfg["business"]["max_workers"],
        mode=mode,
        log_dir=log_dir,
        enable_other=enable_other,
    )
    wall_ms = (time.perf_counter() - t0) * 1000

    if args.pretty:
        print(json.dumps(results, ensure_ascii=False, indent=2))
    else:
        for r in results:
            print(json.dumps(r, ensure_ascii=False))

    if args.summary:
        summary = summarize_results(results, wall_clock_ms=wall_ms)
        print("--- 运行汇总 ---", file=sys.stderr)
        print(json.dumps(summary, ensure_ascii=False, indent=2), file=sys.stderr)


if __name__ == "__main__":
    main()