168 lines
5.1 KiB
Python
168 lines
5.1 KiB
Python
#!/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())
|
|
|
|
|
|
|
|
|
|
|
|
|