proposa-ai/scripts/test_step1.py

"""
Step 1 测试：从招标文件、采购需求、澄清公告中提取表/函/表格。

运行: uv run python scripts/test_step1.py

本脚本会调用当前配置的 LLM，并把结果保存为:
  output/171-上海群众艺术馆/step1_extracted_items.pkl
后续 test_step2.py 会读取该文件，不再重复提取。
"""

import logging
import os
import pickle
import sys

from _bootstrap import PROJECT_ROOT

# ---------- 测试数据（环境变量可覆盖） ----------
TENDER_FILE = os.environ.get(
    "PROPOSA_TENDER_FILE",
    "test_data/171-上海群众艺术馆/上海市群众艺术馆物业管理服务采购项目招标文件.pdf",
)
PROCUREMENT_FILE = os.environ.get(
    "PROPOSA_PROCUREMENT_FILE",
    "test_data/171-上海群众艺术馆/采购需求.docx",
)
CLARIFICATION_FILE = os.environ.get("PROPOSA_CLARIFICATION_FILE", "")
OUTPUT_DIR = os.environ.get(
    "PROPOSA_STEP1_OUTPUT_DIR",
    os.environ.get("PROPOSA_WORK_DIR", "output/171-上海群众艺术馆"),
)
STEP1_ITEMS_FILE = os.environ.get(
    "PROPOSA_STEP1_ITEMS_FILE",
    os.path.join(OUTPUT_DIR, "step1_extracted_items.pkl"),
)
VISUAL_DOCX_DIR = os.environ.get(
    "PROPOSA_STEP1_VISUAL_DIR",
    os.path.join(OUTPUT_DIR, "内容提取"),
)
os.environ.setdefault("BID_LLM_CACHE_DIR", os.path.join(OUTPUT_DIR, ".llm_cache"))

# ---------- 日志 ----------
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%H:%M:%S",
    stream=sys.stderr,
)

from step1_parsing import extract_items_from_tender

print("=" * 60)
print("Step 1: 提取招标文件中的表/函/表格（提示词+代码统一入口）")
print("=" * 60)

items = extract_items_from_tender(
    tender_file=TENDER_FILE,
    procurement_file=PROCUREMENT_FILE,
    clarification_file=CLARIFICATION_FILE,
    output_dir=VISUAL_DOCX_DIR,
)

os.makedirs(OUTPUT_DIR, exist_ok=True)
with open(STEP1_ITEMS_FILE, "wb") as file:
    pickle.dump(items, file)

table_count = sum(1 for i in items if i.item_type in ("表", "表格"))
letter_count = sum(1 for i in items if i.item_type == "函")
print(f"\n提取项数: {len(items)}（表/表格 {table_count} 项，函 {letter_count} 项）")
print(f"可视化 DOCX 输出目录: {VISUAL_DOCX_DIR}")
print(f"Step1 中间产物: {STEP1_ITEMS_FILE}")

for it in items:
    src = f"pages={it.pages}" if it.pages else ""
    print(f"  [{it.item_type}] {it.name} {src}")

print()
print("Step 1 测试完成 [OK]")