cli.py

"""
Interactive CLI for the 飞享IM Q&A chatbot (no FastAPI required).
"""

import asyncio
import os
import time

from config import DOCS_PERSIST_PATH
from graph import build_graph, QAState
from ingest import build_retriever, load_retriever

MAX_HISTORY_TURNS = 10  # 本地 session 保留最近 10 轮


async def stream_answer(graph, question: str, history: list[dict]) -> str:
    """Stream answer tokens to stdout and return the full collected answer."""
    print("助手：", end="", flush=True)
    collected: list[str] = []

    async for event in graph.astream_events(
        QAState(question=question, history=history), version="v2"
    ):
        kind = event["event"]
        node = event.get("metadata", {}).get("langgraph_node", "")

        if kind == "on_chat_model_stream" and node in ("generate", "reject"):
            content = event["data"]["chunk"].content
            if isinstance(content, list):
                for part in content:
                    if isinstance(part, dict) and part.get("type") == "text" and part["text"]:
                        print(part["text"], end="", flush=True)
                        collected.append(part["text"])
            elif isinstance(content, str) and content:
                print(content, end="", flush=True)
                collected.append(content)

        elif kind == "on_chain_end" and node == "fallback":
            output = event["data"].get("output", {})
            answer = output.get("answer", "") if isinstance(output, dict) else ""
            if answer:
                print(answer, end="", flush=True)
                collected.append(answer)

    print()  # 输出换行
    return "".join(collected)


def main():
    if not os.path.exists(DOCS_PERSIST_PATH):
        print("首次运行，正在构建知识库（约需 1-2 分钟）...")
        retriever = build_retriever()
    else:
        print("加载知识库...")
        retriever = load_retriever()

    graph = build_graph(retriever)
    history: list[dict] = []  # session 级别会话历史

    print("\n" + "=" * 60)
    print("  飞享IM 智能问答助手（输入 q 退出）")
    print("=" * 60 + "\n")

    while True:
        try:
            question = input("您好，请问有什么问题？\n> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n再见！")
            break

        if question.lower() in ("q", "quit", "exit", "退出"):
            print("再见！")
            break

        if not question:
            continue

        print("\n思考中...\n")
        t0 = time.perf_counter()
        answer = asyncio.run(stream_answer(graph, question, history))
        print(f"\n[timing] total: {time.perf_counter() - t0:.2f}s")
        print("-" * 60 + "\n")

        # 保存本轮到历史，超出限制时丢弃最早的一轮
        history.append({"role": "user", "content": question})
        history.append({"role": "assistant", "content": answer})
        if len(history) > MAX_HISTORY_TURNS * 2:
            history = history[-(MAX_HISTORY_TURNS * 2):]


if __name__ == "__main__":
    main()