kg-agent/scripts/test_retrieval.py

"""检索层冒烟测试：实体接地 + 查询模板。"""

from __future__ import annotations

import sys

from step3_qa_agent.retrieval import ground_entities, load_entity_index
from step3_qa_agent.retrieval.templates import (query_attendance, query_finance,
                                                 query_project_children,
                                                 query_project_renewals,
                                                 query_service_period)


def main() -> None:
    sys.stdout.reconfigure(encoding="utf-8")
    idx = load_entity_index()
    print(f"索引实体数: {len(idx.entities)}")

    print("\n== 实体接地 ==")
    for q in ["图书馆", "陆雷雷", "青东片区", "公安局", "董一青"]:
        r = ground_entities(q, idx)
        best = r.best
        if best:
            print(f"  {q!r} -> {best.entity.kind}:{best.entity.label} (key={best.entity.key}, 置信度={best.score:.2f})")
        else:
            print(f"  {q!r} -> 无候选")

    print("\n== 查询模板 ==")
    for label, res in [
        ("区图书馆 2026-03 考勤", query_attendance("XMSQ0102", "2026-03")),
        ("区图书馆 2026-03 项目财务", query_finance("XMSQ0102", "2026-03")),
        ("董一青 服务期", query_service_period(emp_id="KWL2307024")),
        ("税务局 子项目", query_project_children("XMSQ0108")),
        ("区图书馆 续签链", query_project_renewals("XMSQ0102")),
    ]:
        print(f"  [{label}] 来源={res.source}, {len(res.rows)} 行")
        for row in res.rows[:3]:
            print("    ", row)


if __name__ == "__main__":
    main()