document-management/backend/main.py

"""Info Extractor - 使用示例.

演示如何使用 InfoExtractor 从文章中提取结构化信息。

用法:
    # 设置 API 密钥（支持 OpenAI 格式的任何模型）
    set OPENAI_API_KEY=sk-xxx
    # 或使用第三方服务
    set OPENAI_API_KEY=your-key
    set OPENAI_BASE_URL=https://api.deepseek.com/v1

    # 运行
    python main.py                           # 默认: single_phase 模式
    python main.py two_phase                 # 两阶段模式
    python main.py --debug                   # 开启 debug 日志 (打印每次 LLM 交互的 prompt)
    python main.py two_phase --debug         # 组合使用
    set DEBUG=true && python main.py         # 通过环境变量开启 debug
"""

from __future__ import annotations

import json
import logging
import os
import sys

from src.extractor import InfoExtractor

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
)
logger = logging.getLogger(__name__)

# ── 示例数据 ──────────────────────────────────────────────────────────────────

SAMPLE_ARTICLE = (
    "2026年6月1日晚8点，好友们决定为王若飞（男）组织一场生日派对。"
    "这次小范围聚会共有5人参加，刘阳（女）作为总策划，"
    "除寿星和总策划外，还有李四（女）、王五（男）和赵六（男）参加。"
    "活动总预算为3000元，且已完成备案，"
    "现场安排了掷骰子、扔飞镖和打桌球等破冰游戏，确保大家玩得尽兴。"
)

SAMPLE_DEFINITION = {
    "metadata": {
        "node_label": "PDJH",
        "node_name": "派对计划",
        "node_type": "concept",
    },
    "subnodes": {
        "JBXX": {
            "metadata": {
                "node_label": "JBXX",
                "node_name": "基本信息",
                "node_type": "concept",
            },
            "subnodes": {
                "PDMC": {
                    "metadata": {
                        "node_label": "PDMC",
                        "node_name": "派对名称",
                        "node_type": "property",
                        "value_type": "string",
                        "is_list": False,
                    }
                },
                "PDKSSJ": {
                    "metadata": {
                        "node_label": "PDKSSJ",
                        "node_name": "派对开始时间",
                        "node_type": "property",
                        "value_type": "date_time",
                        "is_list": False,
                    }
                },
                "PDZRS": {
                    "metadata": {
                        "node_label": "PDZRS",
                        "node_name": "派对总人数",
                        "node_type": "property",
                        "value_type": "integer",
                        "is_list": False,
                    }
                },
                "PDZYS": {
                    "metadata": {
                        "node_label": "PDZYS",
                        "node_name": "派对总预算",
                        "node_type": "property",
                        "value_type": "float",
                        "is_list": False,
                    }
                },
                "SFBA": {
                    "metadata": {
                        "node_label": "SFBA",
                        "node_name": "是否备案",
                        "node_type": "property",
                        "value_type": "boolean",
                        "is_list": False,
                    }
                },
            },
        },
        "RYXX": {
            "metadata": {
                "node_label": "RYXX",
                "node_name": "人员信息",
                "node_type": "concept",
            },
            "subnodes": {
                "ZCH": {
                    "metadata": {
                        "node_label": "ZCH",
                        "node_name": "总策划",
                        "node_type": "property",
                        "value_type": "class",
                        "is_list": False,
                    },
                    "subnodes": {
                        "XM": {
                            "metadata": {
                                "node_label": "XM",
                                "node_name": "姓名",
                                "node_type": "property",
                                "value_type": "string",
                                "is_list": False,
                            }
                        },
                        "XB": {
                            "metadata": {
                                "node_label": "XB",
                                "node_name": "性别",
                                "node_type": "property",
                                "value_type": "string",
                                "is_list": False,
                            }
                        },
                    },
                },
                "RYQD": {
                    "metadata": {
                        "node_label": "RYQD",
                        "node_name": "人员清单",
                        "node_type": "property",
                        "value_type": "class",
                        "is_list": True,
                    },
                    "subnodes": {
                        "XM": {
                            "metadata": {
                                "node_label": "XM",
                                "node_name": "姓名",
                                "node_type": "property",
                                "value_type": "string",
                                "is_list": False,
                            }
                        },
                        "XB": {
                            "metadata": {
                                "node_label": "XB",
                                "node_name": "性别",
                                "node_type": "property",
                                "value_type": "string",
                                "is_list": False,
                            }
                        },
                    },
                },
            },
        },
        "HDXX": {
            "metadata": {
                "node_label": "HDXX",
                "node_name": "活动信息",
                "node_type": "concept",
            },
            "subnodes": {
                "PBYXQD": {
                    "metadata": {
                        "node_label": "PBYXQD",
                        "node_name": "破冰游戏清单",
                        "node_type": "property",
                        "value_type": "string",
                        "is_list": True,
                    }
                }
            },
        },
    },
}


def main() -> None:
    # 初始化提取器
    # 方式 1: 使用环境变量（OPENAI_API_KEY, OPENAI_BASE_URL）
    # 方式 2: 直接传入参数
    extractor = InfoExtractor(
        model_name=os.getenv("MODEL_NAME", "gpt-4o"),
        api_key=os.getenv("OPENAI_API_KEY"),
        base_url=os.getenv("OPENAI_BASE_URL"),
        temperature=0,
        max_retries=3,
    )

    # 解析命令行参数
    args = sys.argv[1:]
    mode = "single_phase"
    debug = os.getenv("DEBUG", "").lower() in ("1", "true", "yes")

    for arg in args:
        if arg == "--debug":
            debug = True
        elif arg in ("single_phase", "two_phase"):
            mode = arg

    logger.info(f"使用 {mode} 模式提取信息 (debug={debug})")

    # 执行提取（含标注）
    result = extractor.extract(
        article=SAMPLE_ARTICLE,
        definition=SAMPLE_DEFINITION,
        mode=mode,
        annotate=True,
        debug=debug,
    )

    # 输出提取结果
    print("\n" + "=" * 60)
    print("提取结果:")
    print("=" * 60)
    print(json.dumps(result["instance"], ensure_ascii=False, indent=2))

    # 输出标注后的文章
    print("\n" + "=" * 60)
    print("标注后的文章:")
    print("=" * 60)
    print(result["annotated_text"])


if __name__ == "__main__":
    main()