53 lines
1.8 KiB
Python
53 lines
1.8 KiB
Python
"""NL2SQL 结果摘要和基础图表配置。"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from typing import Any
|
|
|
|
|
|
def build_chart_config(columns: list[str], rows: list[dict[str, Any]]) -> dict[str, Any] | None:
|
|
"""识别简单的分类数值结果，返回安全的柱状图配置。"""
|
|
if not rows or len(columns) < 2:
|
|
return None
|
|
numeric_column = next(
|
|
(
|
|
column
|
|
for column in columns
|
|
if all(isinstance(row.get(column), (int, float)) and not isinstance(row.get(column), bool) for row in rows)
|
|
),
|
|
None,
|
|
)
|
|
category_column = next((column for column in columns if column != numeric_column), None)
|
|
if numeric_column is None or category_column is None:
|
|
return None
|
|
return {"type": "bar", "category": category_column, "value": numeric_column}
|
|
|
|
|
|
async def summarize_result(
|
|
question: str,
|
|
columns: list[str],
|
|
rows: list[dict[str, Any]],
|
|
*,
|
|
llm_client,
|
|
max_prompt_rows: int = 20,
|
|
) -> str:
|
|
"""使用脱敏后的受控结果生成摘要，模型失败时返回固定话术。"""
|
|
fallback = f"查询完成，共返回 {len(rows)} 条记录。"
|
|
prompt = (
|
|
"请用简洁中文总结查询结果，只基于提供的数据，不要猜测。\n"
|
|
f"问题：{question}\n"
|
|
f"列名：{json.dumps(columns, ensure_ascii=False)}\n"
|
|
f"结果：{json.dumps(rows[:max_prompt_rows], ensure_ascii=False, default=str)}"
|
|
)
|
|
try:
|
|
answer = await llm_client.chat(
|
|
[
|
|
{"role": "system", "content": "你是数据查询结果摘要助手。"},
|
|
{"role": "user", "content": prompt},
|
|
],
|
|
temperature=0,
|
|
)
|
|
except Exception: # noqa: BLE001 摘要失败回退固定文本
|
|
return fallback
|
|
return answer.strip() or fallback
|