kg-agent/scripts/ask.py

"""CLI 交互式 Agent 问答（含 human-in-the-loop 确认）。

用法：
  uv run python scripts/ask.py                     # 交互模式（多轮，可确认/缩小范围）
  uv run python scripts/ask.py "问题" --auto       # 单次自动确认（批处理/评测用）
"""

from __future__ import annotations

import argparse
import json
import sys
import time

from langgraph.checkpoint.memory import InMemorySaver
from langgraph.types import Command


def _make_checkpointer():
    """上下文记忆：当前存内存（InMemorySaver，进程内多轮有效）。
    后续持久化/图数据库化时只改这里（如 SqliteSaver / PostgresSaver / 自研存储）。"""
    return InMemorySaver()


def _print_result(st: dict, debug: bool = False,
                  total_sec: float | None = None,
                  process_sec: float | None = None) -> None:
    print("=" * 50)
    print("分类:", st.get("category"))
    if st.get("chat_answer"):
        print("闲聊回复:", st["chat_answer"])
        _print_elapsed(total_sec, process_sec)
        return
    print("槽位:", json.dumps(st.get("slots", {}), ensure_ascii=False))
    print("接地:", json.dumps(st.get("grounded", {}), ensure_ascii=False))
    if debug and st.get("capability"):
        print("能力拓展:", json.dumps(st.get("capability"), ensure_ascii=False))
    cap_check = st.get("cap_check") or {}
    if cap_check:
        tail = cap_check.get("note") or ("；".join(cap_check.get("issues", [])) if not cap_check.get("ok") else "")
        print(f"能力拓展检查: {'✅ 通过' if cap_check.get('ok') else '⚠️ 未通过'} {tail}")
    cap_llm = st.get("cap_llm_check") or {}
    if cap_llm:
        tail = cap_llm.get("reason") or cap_llm.get("note") or ""
        print(f"能力拓展LLM审查: {'✅ 通过' if cap_llm.get('ok') else '⚠️ 未通过'} {tail}")
    print("计划:", json.dumps(st.get("plan", {}), ensure_ascii=False))
    print("执行详情:")
    for s in st.get("subgraph", {}).get("steps", []):
        tool = s.get("tool", "")
        row_count = s.get("row_count", len(s.get("rows", [])))
        note = s.get("note") or ""
        dropped = s.get("dropped_params") or []
        print(f"  {s.get('step_id')} [{tool}] 行数={row_count} "
              f"实参={json.dumps(s.get('executed_params', {}), ensure_ascii=False)}"
              + (f" 丢弃参数={json.dumps(dropped, ensure_ascii=False)}" if dropped else "")
              + (f" 备注={note}" if note else ""))
        if debug:
            if s.get("query"):
                print(f"      Cypher: {s['query']}")
            if s.get("query_params"):
                print(f"      参数: {json.dumps(s['query_params'], ensure_ascii=False)}")
            first = s.get("rows", [])[:2]
            if first:
                print(f"      首行: {json.dumps(first, ensure_ascii=False, default=str)}")
    print("子图行数:", st.get("subgraph", {}).get("total_rows"))
    run_check = st.get("run_check") or {}
    if run_check:
        tail = "；".join(run_check.get("issues", []) + run_check.get("warns", []))
        print(f"执行检查: {'✅ 通过' if run_check.get('ok') else '⚠️ 未通过'} {tail}")
    run_llm = st.get("run_llm_check") or {}
    if run_llm:
        tail = run_llm.get("reason") or ""
        print(f"执行LLM审查: {'✅ 通过' if run_llm.get('ok') else '⚠️ 未通过'} {tail}")
    print("迭代次数:", st.get("iterations"))
    print("-" * 50)
    if st.get("reuse") and not st.get("need_full_query", True):
        print("回答:（基于前序对话推导）", st.get("answer"))
    else:
        print("回答:", st.get("answer"))
    _print_elapsed(total_sec, process_sec)
    print("=" * 50)


def _print_elapsed(total_sec: float | None, process_sec: float | None) -> None:
    if total_sec is not None:
        wait_sec = max(total_sec - (process_sec or 0.0), 0.0)
        print(f"耗时: 总 {total_sec:.1f}s（处理 {process_sec or 0.0:.1f}s / 确认等待 {wait_sec:.1f}s）")


def _run_once(graph, question: str, thread_id: str, auto: bool, user_id: str, debug: bool = False) -> None:
    t0 = time.monotonic()
    process_sec = 0.0
    config = {"configurable": {"thread_id": thread_id}}
    t1 = time.monotonic()
    result = graph.invoke({"question": question, "user_id": user_id}, config)
    process_sec += time.monotonic() - t1
    while "__interrupt__" in result:
        intr = result["__interrupt__"][0]
        payload = intr.value
        print("\n[需要确认]", payload.get("message"))
        if payload.get("options"):
            for i, o in enumerate(payload["options"], 1):
                print(f"  {i}. {o}")
        if auto:
            default = "确认" if payload.get("type") in ("confirm_slots", "confirm_plan") else "继续返回全部"
            print(f"（自动确认: {default}）")
            reply = default
        else:
            reply = input("你的回复 > ").strip()
        t1 = time.monotonic()
        result = graph.invoke(Command(resume=reply), config)
        process_sec += time.monotonic() - t1
    _print_result(result, debug=debug, total_sec=time.monotonic() - t0, process_sec=process_sec)


def main() -> None:
    sys.stdout.reconfigure(encoding="utf-8")
    ap = argparse.ArgumentParser()
    ap.add_argument("question", nargs="?", default=None, help="问题（不填则进入交互模式）")
    ap.add_argument("--auto", action="store_true", help="自动确认，跳过人工确认")
    ap.add_argument("--debug", action="store_true", help="显示每步实际执行的 Cypher/参数/首行（调试用）")
    ap.add_argument("--user", default="admin01", help="提问人身份（工号或用户名，默认 admin01）")
    args = ap.parse_args()

    from knowledge_agent.agent.graph import build_agent_graph

    graph = build_agent_graph(checkpointer=_make_checkpointer())
    thread = "default"

    if args.question:
        _run_once(graph, args.question, thread, args.auto, args.user, args.debug)
        return

    print("申勤物业知识助手（交互模式；输入 退出/quit 结束，输入 清空/新话题 重置上下文）")
    while True:
        try:
            q = input("\n问题 > ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if not q:
            continue
        if q in ("退出", "quit", "exit"):
            break
        if q in ("清空", "新话题", "新会话", "重置", "reset", "clear", "new"):
            thread = f"t{time.monotonic():.3f}"
            print("（已清空上下文，开始新话题）")
            continue
        _run_once(graph, q, thread, args.auto, args.user, args.debug)


if __name__ == "__main__":
    main()