Mutual_Fund/nl2sql/result.py

"""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:
    """兼容旧调用方：返回面向用户的自然语言答案。"""
    return await render_answer(
        question,
        columns,
        rows,
        llm_client=llm_client,
        max_prompt_rows=max_prompt_rows,
    )


async def render_answer(
    question: str,
    columns: list[str],
    rows: list[dict[str, Any]],
    *,
    llm_client,
    max_prompt_rows: int = 20,
) -> str:
    """将受控查询结果渲染为用户回答，不输出原始 JSON 或无关字段。"""
    fallback = f"查询完成，共返回 {len(rows)} 条记录。"
    prompt = (
        "请根据用户问题生成简洁、准确的中文业务回答。\n"
        "只使用查询结果中的必要字段回答问题，不要输出 JSON、SQL、Markdown 表格或内部字段名。\n"
        "如果有多条记录，优先做合计、数量或关键项概括；只有用户明确需要明细时才列出明细。\n"
        "不得补充查询结果中没有的数据，不要解释你的处理过程。\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": "你是基金平台的数据回答助手。不要输出 JSON。"},
                {"role": "user", "content": prompt},
            ],
            temperature=0,
            max_tokens=256,
        )
    except Exception:  # noqa: BLE001 摘要失败回退固定文本
        return fallback
    return answer.strip() or fallback