more_dots/scripts/demo_chat.py

#!/usr/bin/env python3
"""简洁演示问数脚本：只保留问题输入与结果输出。"""

from __future__ import annotations

import argparse
import sys
import time
import traceback
import uuid
from pathlib import Path
from typing import Any, Dict, Iterable, Optional

# 允许直接使用 `python scripts/demo_chat.py` 运行
PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from agent.agents.conversation import ConversationAgent
from config import Config
from scripts.console_chat import extract_table_rows, parse_sr_api_result, render_text_table


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="演示版控制台问数：仅展示耗时、SQL、结果表和数据行数。"
    )
    parser.add_argument("--model-section", default=None, help="可选：指定 config.ini 中的模型配置段")
    parser.add_argument("--conversation-id", default=None, help="可选：会话 ID")
    parser.add_argument("--user", default="demo-user", help="可选：用户标识")
    parser.add_argument("--query", "-q", default=None, help="单次执行模式：直接执行一条问题后退出")
    return parser


def _last_answer_text(result: Dict[str, Any]) -> str:
    messages = (result or {}).get("messages") or []
    if not messages:
        return ""
    last = messages[-1]
    return getattr(last, "content", "") or str(last)


def format_demo_result(result: Dict[str, Any], elapsed_seconds: float) -> str:
    context = (result or {}).get("context") or {}
    final_sql = str(context.get("final_sql") or "")
    parsed_result = parse_sr_api_result(context.get("sr_api_result"))
    headers, rows = extract_table_rows(parsed_result)
    row_count = len(rows)
    has_structured_result = context.get("sr_api_result") is not None and isinstance(parsed_result, (dict, list))
    is_empty_result = bool(context.get("is_empty_result"))
    response_source = str(context.get("response_source") or "")
    answer = _last_answer_text(result)

    blocks = [f"耗时: {elapsed_seconds:.2f}s"]

    if final_sql:
        blocks.append(f"SQL:\n{final_sql}")
    else:
        blocks.append("SQL:\n<未生成 SQL>")

    blocks.append(f"数据行数: {row_count}")

    if is_empty_result and response_source in {"model_empty_result_fallback", "empty_result_fixed_fallback"}:
        blocks.append(f"结果说明:\n{answer or '未查询到符合条件的数据，请尝试调整筛选条件后再查询。'}")
    elif has_structured_result:
        table_text = render_text_table(headers or ['result'], rows)
        blocks.append(f"SQL执行结果表:\n{table_text}")
    else:
        blocks.append(f"SQL执行结果:\n{answer or str(parsed_result or '<无结果>')}")

    return "\n\n".join(blocks)


def run_turn(
    agent: ConversationAgent,
    query: str,
    *,
    user: str,
    conversation_id: str,
) -> Dict[str, Any]:
    started_at = time.perf_counter()
    result = agent.run(
        query,
        user=user,
        conversation_id=conversation_id,
        skip_sr_api=False,
        debug_node_trace=False,
    )
    elapsed_seconds = time.perf_counter() - started_at
    print(format_demo_result(result, elapsed_seconds))
    return result


def interactive_loop(args: argparse.Namespace) -> int:
    conversation_id = args.conversation_id or f"demo_{uuid.uuid4().hex[:8]}"
    agent = ConversationAgent(model_section=args.model_section)

    while True:
        try:
            user_input = input("\n问题> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye.")
            return 0

        if not user_input:
            continue

        if user_input.lower() in {"/quit", "/exit", "quit", "exit"}:
            print("Bye.")
            return 0

        try:
            run_turn(
                agent,
                user_input,
                user=args.user,
                conversation_id=conversation_id,
            )
        except Exception as exc:
            print(f"\n[ERROR] {exc}")
            traceback.print_exc()

    return 0


def one_shot(args: argparse.Namespace) -> int:
    conversation_id = args.conversation_id or f"demo_{uuid.uuid4().hex[:8]}"
    agent = ConversationAgent(model_section=args.model_section)
    query = str(args.query or "")
    try:
        run_turn(
            agent,
            query,
            user=args.user,
            conversation_id=conversation_id,
        )
        return 0
    except Exception as exc:
        print(f"[ERROR] {exc}", file=sys.stderr)
        traceback.print_exc()
        return 1


def main(argv: Optional[Iterable[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(list(argv) if argv is not None else None)

    try:
        Config.validate_config()
    except Exception as exc:
        print(f"Configuration error: {exc}", file=sys.stderr)
        print("Please check `config/config.ini` and your model/API settings.", file=sys.stderr)
        return 1

    if args.query:
        return one_shot(args)
    return interactive_loop(args)


if __name__ == "__main__":
    raise SystemExit(main())