document-management/backend/tests/legacy_ai/test.py

import json
import logging
import os

from src.extractor import InfoExtractor

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

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="glm-5.2",
        api_key=os.environ.get("INFO_EXTRACTOR_API_KEY", "your-api-key"),
        base_url=os.environ.get("INFO_EXTRACTOR_BASE_URL", "https://open.bigmodel.cn/api/coding/paas/v4"),
        temperature=0,
        max_retries=3,
    )

    # 执行提取
    result = extractor.extract(
        article=SAMPLE_ARTICLE,
        definition=SAMPLE_DEFINITION,
        mode="two_phase",
        verbose=True,
        annotate=True,
        debug=True,
    )

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


if __name__ == "__main__":
    main()