337 lines
12 KiB
Python
337 lines
12 KiB
Python
#!/usr/bin/env python3
|
|
"""控制台问数脚本：支持交互式多轮问答和单次执行。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import sys
|
|
import traceback
|
|
import uuid
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Iterable, Optional, Sequence
|
|
|
|
# 允许直接使用 `python scripts/console_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
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(
|
|
description="在控制台中直接输入自然语言问题，调用 ConversationAgent 进行问数。"
|
|
)
|
|
parser.add_argument("--model-section", default=None, help="可选：指定 config.ini 中的模型配置段")
|
|
parser.add_argument("--conversation-id", default=None, help="可选：会话 ID，仅用于展示和上下文标识")
|
|
parser.add_argument("--user", default="console-user", help="可选：用户标识")
|
|
parser.add_argument("--query", "-q", default=None, help="单次执行模式：直接执行一条问题后退出")
|
|
parser.add_argument("--skip-sr-api", action="store_true", help="仅生成 SQL，不执行 SR API")
|
|
parser.add_argument("--show-sql", action="store_true", help="额外打印生成的 SQL")
|
|
parser.add_argument("--show-context", action="store_true", help="额外打印完整上下文 JSON")
|
|
parser.add_argument("--show-plan", action="store_true", help="额外打印 SQL 规划 JSON")
|
|
parser.add_argument("--no-banner", action="store_true", help="不显示启动横幅")
|
|
return parser
|
|
|
|
|
|
def print_banner(model_section: Optional[str], conversation_id: str, skip_sr_api: bool) -> None:
|
|
print("=" * 72)
|
|
print("APBO Console Chat")
|
|
print(f"Model Section : {model_section or 'default'}")
|
|
print(f"Conversation : {conversation_id}")
|
|
print(f"Execution : {'SQL only' if skip_sr_api else 'SQL + SR API'}")
|
|
print("Commands : /quit /exit /sql /context /plan /exec")
|
|
print("=" * 72)
|
|
|
|
|
|
def _safe_json(data: Any) -> str:
|
|
try:
|
|
return json.dumps(data, ensure_ascii=False, indent=2, default=str)
|
|
except Exception:
|
|
return str(data)
|
|
|
|
|
|
def _looks_like_json(text: str) -> bool:
|
|
stripped = (text or "").strip()
|
|
return stripped.startswith("{") or stripped.startswith("[")
|
|
|
|
|
|
def _try_json_loads(value: Any) -> Any:
|
|
if isinstance(value, (dict, list)):
|
|
return value
|
|
if isinstance(value, str) and _looks_like_json(value):
|
|
try:
|
|
return json.loads(value)
|
|
except Exception:
|
|
return value
|
|
return value
|
|
|
|
|
|
def parse_sr_api_result(raw_result: Any) -> Any:
|
|
"""解析 SR API 返回值，兼容外层 envelope 和内层 text JSON。"""
|
|
parsed = _try_json_loads(raw_result)
|
|
if isinstance(parsed, dict) and "text" in parsed:
|
|
text_payload = _try_json_loads(parsed.get("text"))
|
|
parsed = {**parsed, "text": text_payload}
|
|
return parsed
|
|
|
|
|
|
def _find_table_candidate(value: Any) -> Any:
|
|
if isinstance(value, list):
|
|
return value
|
|
if isinstance(value, dict):
|
|
if "columns" in value and any(key in value for key in ("rows", "data", "values")):
|
|
return value
|
|
for key in ("data", "rows", "records", "record", "items", "list", "result", "text"):
|
|
nested = value.get(key)
|
|
if isinstance(nested, (list, dict)):
|
|
found = _find_table_candidate(nested)
|
|
if found is not None:
|
|
return found
|
|
return None
|
|
|
|
|
|
def extract_table_rows(parsed_result: Any) -> tuple[list[str], list[list[str]]]:
|
|
"""从常见查询结果结构中提取表头和二维行数据。"""
|
|
candidate = _find_table_candidate(parsed_result)
|
|
if candidate is None:
|
|
return [], []
|
|
|
|
if isinstance(candidate, dict) and isinstance(candidate.get("columns"), list):
|
|
headers = [str(col) for col in candidate.get("columns") or []]
|
|
raw_rows = candidate.get("rows") or candidate.get("data") or candidate.get("values") or []
|
|
if raw_rows and all(isinstance(row, dict) for row in raw_rows):
|
|
return headers or list(raw_rows[0].keys()), [
|
|
[str((row or {}).get(header, "")) for header in (headers or list(raw_rows[0].keys()))]
|
|
for row in raw_rows
|
|
]
|
|
return headers, [[str(cell) for cell in row] for row in raw_rows if isinstance(row, (list, tuple))]
|
|
|
|
if isinstance(candidate, list) and candidate:
|
|
if all(isinstance(row, dict) for row in candidate):
|
|
headers: list[str] = []
|
|
for row in candidate:
|
|
for key in row.keys():
|
|
if key not in headers:
|
|
headers.append(str(key))
|
|
return headers, [[str((row or {}).get(header, "")) for header in headers] for row in candidate]
|
|
if all(isinstance(row, (list, tuple)) for row in candidate):
|
|
width = max(len(row) for row in candidate)
|
|
headers = [f"col_{idx + 1}" for idx in range(width)]
|
|
return headers, [[str(row[idx]) if idx < len(row) else "" for idx in range(width)] for row in candidate]
|
|
|
|
return [], []
|
|
|
|
|
|
def render_text_table(headers: Sequence[str], rows: Sequence[Sequence[str]], *, max_width: int = 28, max_rows: int = 20) -> str:
|
|
"""将二维数据渲染成纯文本表格。"""
|
|
if not headers or not rows:
|
|
return "<empty table>"
|
|
|
|
def clip(value: Any) -> str:
|
|
text = str(value).replace("\r", " ").replace("\n", " ")
|
|
return text if len(text) <= max_width else text[: max_width - 3] + "..."
|
|
|
|
display_rows = list(rows[:max_rows])
|
|
str_rows = [[clip(cell) for cell in row] for row in display_rows]
|
|
clipped_headers = [clip(header) for header in headers]
|
|
|
|
widths = []
|
|
for idx, header in enumerate(clipped_headers):
|
|
col_values = [row[idx] if idx < len(row) else "" for row in str_rows]
|
|
widths.append(max(len(header), *(len(value) for value in col_values)) if col_values else len(header))
|
|
|
|
def render_row(values: Sequence[str]) -> str:
|
|
padded = []
|
|
for idx, width in enumerate(widths):
|
|
value = values[idx] if idx < len(values) else ""
|
|
padded.append(value.ljust(width))
|
|
return "| " + " | ".join(padded) + " |"
|
|
|
|
separator = "+-" + "-+-".join("-" * width for width in widths) + "-+"
|
|
lines = [separator, render_row(clipped_headers), separator]
|
|
lines.extend(render_row(row) for row in str_rows)
|
|
lines.append(separator)
|
|
if len(rows) > max_rows:
|
|
lines.append(f"... showing first {max_rows} of {len(rows)} rows")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _format_sr_api_table(sr_api_result: Any) -> Optional[str]:
|
|
parsed = parse_sr_api_result(sr_api_result)
|
|
headers, rows = extract_table_rows(parsed)
|
|
if not headers or not rows:
|
|
return None
|
|
|
|
summary = f"Query Result: {len(rows)} row(s)"
|
|
if isinstance(parsed, dict) and parsed.get("status_code") is not None:
|
|
summary += f" | status={parsed.get('status_code')}"
|
|
return summary + "\n" + render_text_table(headers, rows)
|
|
|
|
|
|
def format_result(
|
|
result: Dict[str, Any],
|
|
*,
|
|
show_sql: bool = False,
|
|
show_context: bool = False,
|
|
show_plan: bool = False,
|
|
) -> str:
|
|
context = (result or {}).get("context") or {}
|
|
messages = (result or {}).get("messages") or []
|
|
sr_api_result = context.get("sr_api_result")
|
|
|
|
answer = ""
|
|
if messages:
|
|
last = messages[-1]
|
|
answer = getattr(last, "content", "") or str(last)
|
|
if not answer:
|
|
answer = str(sr_api_result or context.get("final_sql") or "<empty response>")
|
|
|
|
table_block = _format_sr_api_table(sr_api_result) if sr_api_result else None
|
|
if table_block and (_looks_like_json(answer) or answer.strip().lower() in {"ok", "success"}):
|
|
answer = "查询成功，结果已按二维表格展示如下。"
|
|
|
|
blocks = [f"Answer:\n{answer}"]
|
|
|
|
if table_block:
|
|
blocks.append(table_block)
|
|
|
|
if show_sql and context.get("final_sql"):
|
|
blocks.append(f"SQL:\n{context['final_sql']}")
|
|
|
|
if show_plan and context.get("sql_plan"):
|
|
blocks.append(f"SQL Plan:\n{_safe_json(context['sql_plan'])}")
|
|
|
|
if show_context:
|
|
blocks.append(f"Context:\n{_safe_json(context)}")
|
|
|
|
return "\n\n".join(blocks)
|
|
|
|
|
|
def run_turn(
|
|
agent: ConversationAgent,
|
|
query: str,
|
|
*,
|
|
user: str,
|
|
conversation_id: str,
|
|
skip_sr_api: bool,
|
|
show_sql: bool,
|
|
show_context: bool,
|
|
show_plan: bool,
|
|
) -> Dict[str, Any]:
|
|
result = agent.run(
|
|
query,
|
|
user=user,
|
|
conversation_id=conversation_id,
|
|
skip_sr_api=skip_sr_api,
|
|
debug_node_trace=True,
|
|
)
|
|
print(format_result(result, show_sql=show_sql, show_context=show_context, show_plan=show_plan))
|
|
return result
|
|
|
|
|
|
def interactive_loop(args: argparse.Namespace) -> int:
|
|
conversation_id = args.conversation_id or f"console_{uuid.uuid4().hex[:8]}"
|
|
agent = ConversationAgent(model_section=args.model_section)
|
|
|
|
show_sql = bool(args.show_sql)
|
|
show_context = bool(args.show_context)
|
|
show_plan = bool(args.show_plan)
|
|
skip_sr_api = bool(args.skip_sr_api)
|
|
|
|
if not args.no_banner:
|
|
print_banner(args.model_section, conversation_id, skip_sr_api)
|
|
|
|
while True:
|
|
try:
|
|
user_input = input("\n问数> ").strip()
|
|
except (EOFError, KeyboardInterrupt):
|
|
print("\nBye.")
|
|
return 0
|
|
|
|
if not user_input:
|
|
continue
|
|
|
|
lowered = user_input.lower()
|
|
if lowered in {"/quit", "/exit", "quit", "exit"}:
|
|
print("Bye.")
|
|
return 0
|
|
if lowered == "/sql":
|
|
show_sql = not show_sql
|
|
print(f"show_sql = {show_sql}")
|
|
continue
|
|
if lowered == "/context":
|
|
show_context = not show_context
|
|
print(f"show_context = {show_context}")
|
|
continue
|
|
if lowered == "/plan":
|
|
show_plan = not show_plan
|
|
print(f"show_plan = {show_plan}")
|
|
continue
|
|
if lowered == "/exec":
|
|
skip_sr_api = not skip_sr_api
|
|
print(f"skip_sr_api = {skip_sr_api}")
|
|
continue
|
|
|
|
try:
|
|
run_turn(
|
|
agent,
|
|
user_input,
|
|
user=args.user,
|
|
conversation_id=conversation_id,
|
|
skip_sr_api=skip_sr_api,
|
|
show_sql=show_sql,
|
|
show_context=show_context,
|
|
show_plan=show_plan,
|
|
)
|
|
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"console_{uuid.uuid4().hex[:8]}"
|
|
agent = ConversationAgent(model_section=args.model_section)
|
|
try:
|
|
run_turn(
|
|
agent,
|
|
args.query,
|
|
user=args.user,
|
|
conversation_id=conversation_id,
|
|
skip_sr_api=bool(args.skip_sr_api),
|
|
show_sql=bool(args.show_sql),
|
|
show_context=bool(args.show_context),
|
|
show_plan=bool(args.show_plan),
|
|
)
|
|
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())
|
|
|
|
|