189 lines
6.8 KiB
Python
189 lines
6.8 KiB
Python
"""上传后用 LLM 一次分析文件：一句话说明、图表推荐、列的含义
|
|
|
|
只把列概况和几行样例发给模型，不发全量数据。模型给出的推荐在返回前逐条校验并真实跑一遍汇总，
|
|
列不存在、聚合不合法、结果为空的推荐直接丢弃，所以前端拿到的推荐都能直接出图
|
|
"""
|
|
|
|
import asyncio
|
|
import json
|
|
from collections.abc import Awaitable, Callable
|
|
from pathlib import Path
|
|
|
|
from langchain_core.exceptions import OutputParserException
|
|
from langchain_core.messages import HumanMessage, SystemMessage
|
|
from loguru import logger
|
|
from pydantic import BaseModel, Field, ValidationError
|
|
|
|
from app.clients.llm import get_llm
|
|
from app.config import settings
|
|
from app.core.exceptions import BizException
|
|
from app.datafiles.aggregate import run_aggregate
|
|
from app.datafiles.query import QueryError
|
|
from app.datafiles.storage import table_path
|
|
from app.prompt.insights import INSIGHTS_SYSTEM_PROMPT
|
|
from app.schemas.data_file import (
|
|
MAX_SUGGESTIONS,
|
|
ChartSuggestion,
|
|
ColumnInsight,
|
|
DataFileRecord,
|
|
DataTable,
|
|
FileInsights,
|
|
)
|
|
|
|
# 单次分析的超时（秒），含模型自己的重试；超时记为失败，文件本身不受影响
|
|
INSIGHTS_TIMEOUT = 60
|
|
# 同时分析的文件数，其余排队，避免一次传很多文件时同时打满模型
|
|
INSIGHTS_CONCURRENCY = 4
|
|
# 发给模型的表数和每张表的列数上限，宽表只分析靠前的列
|
|
MAX_TABLES = 5
|
|
MAX_COLUMNS = 40
|
|
# 模型多给几条候补，校验丢弃后仍能凑够；再多的不处理
|
|
MAX_CANDIDATES = 6
|
|
# 模型输出格式不合法时重试的次数（包含第一次）
|
|
ASK_ATTEMPTS = 2
|
|
MAX_SUMMARY_CHARS = 100
|
|
MAX_TITLE_CHARS = 40
|
|
MAX_LABEL_CHARS = 40
|
|
MAX_UNIT_CHARS = 10
|
|
# 结果行数太少画不成图
|
|
MIN_RESULT_ROWS = 2
|
|
|
|
# 生成函数：参数是文件记录和数据文件根目录，返回完成的洞察；失败直接抛异常
|
|
InsightsGenerator = Callable[[DataFileRecord, Path], Awaitable[FileInsights]]
|
|
|
|
|
|
class InsightsDraft(BaseModel):
|
|
"""模型的输出，校验、清洗之后才成为 FileInsights"""
|
|
|
|
summary: str = Field(description="一句话说明这份数据是什么，40 字以内")
|
|
suggestions: list[ChartSuggestion] = Field(
|
|
description="图表推荐，最多 3 条，按推荐程度排序"
|
|
)
|
|
columns: list[ColumnInsight] = Field(description="输入里每一列的含义")
|
|
|
|
|
|
def insights_enabled() -> bool:
|
|
"""配置了模型且没有关闭开关才分析"""
|
|
return (
|
|
settings.data_file_insights
|
|
and bool(settings.llm_api_key.strip())
|
|
and bool(settings.llm_model.strip())
|
|
)
|
|
|
|
|
|
async def generate_insights(record: DataFileRecord, root: Path) -> FileInsights:
|
|
draft = await _ask(record)
|
|
return await asyncio.to_thread(check_draft, draft, record, root)
|
|
|
|
|
|
async def _ask(record: DataFileRecord) -> InsightsDraft:
|
|
messages = [
|
|
SystemMessage(INSIGHTS_SYSTEM_PROMPT),
|
|
HumanMessage(json.dumps(build_input(record), ensure_ascii=False)),
|
|
]
|
|
error: Exception | None = None
|
|
for _ in range(ASK_ATTEMPTS):
|
|
try:
|
|
draft = await get_llm(InsightsDraft).ainvoke(messages)
|
|
except (ValidationError, OutputParserException) as failure:
|
|
# 输出不符合格式时再问一次；网络类错误 SDK 已经重试过，直接抛出
|
|
error = failure
|
|
continue
|
|
if draft is not None:
|
|
return draft
|
|
error = ValueError("模型没有按格式输出")
|
|
assert error is not None
|
|
raise error
|
|
|
|
|
|
def build_input(record: DataFileRecord) -> dict:
|
|
"""发给模型的内容：每张表的列概况和样例行"""
|
|
return {
|
|
"file": record.name,
|
|
"tables": [_table_input(table) for table in record.tables[:MAX_TABLES]],
|
|
}
|
|
|
|
|
|
def _table_input(table: DataTable) -> dict:
|
|
return {
|
|
"table": table.table,
|
|
"label": table.label,
|
|
"rows": table.rows,
|
|
"columns": [
|
|
column.model_dump(mode="json", exclude_none=True)
|
|
for column in table.profile.columns[:MAX_COLUMNS]
|
|
],
|
|
# 样例行和 columns 都按表的列顺序排列，截掉同样的列
|
|
"sample": [row[:MAX_COLUMNS] for row in table.profile.sample],
|
|
}
|
|
|
|
|
|
def check_draft(
|
|
draft: InsightsDraft, record: DataFileRecord, root: Path
|
|
) -> FileInsights:
|
|
"""清洗模型的输出：不合法的推荐和列直接丢弃；同步函数，里面会真实跑汇总查询"""
|
|
tables = {table.table: table for table in record.tables}
|
|
suggestions: list[ChartSuggestion] = []
|
|
seen: set[str] = set()
|
|
for item in draft.suggestions[:MAX_CANDIDATES]:
|
|
if len(suggestions) == MAX_SUGGESTIONS:
|
|
break
|
|
title = item.title.strip()[:MAX_TITLE_CHARS]
|
|
table = tables.get(item.table)
|
|
key = item.model_dump_json(include={"table", "request"})
|
|
if not title or table is None or key in seen:
|
|
continue
|
|
if not _fits_chart(item) or not _runs(
|
|
item, table, table_path(root, record.file_id, table)
|
|
):
|
|
logger.info("丢弃不可用的图表推荐：{} {}", record.file_id, item.title)
|
|
continue
|
|
seen.add(key)
|
|
suggestions.append(item.model_copy(update={"title": title}))
|
|
return FileInsights(
|
|
status="ready",
|
|
summary=draft.summary.strip()[:MAX_SUMMARY_CHARS],
|
|
suggestions=suggestions,
|
|
columns=_clean_columns(draft.columns, tables),
|
|
)
|
|
|
|
|
|
def _fits_chart(item: ChartSuggestion) -> bool:
|
|
"""前端图表只认“一个分类列 + 若干指标”；饼图、漏斗只取一个指标，多写会和画面对不上"""
|
|
request = item.request
|
|
if len(request.group_by) != 1:
|
|
return False
|
|
return len(request.metrics) == 1 or item.chart not in ("pie", "funnel")
|
|
|
|
|
|
def _runs(item: ChartSuggestion, table: DataTable, path: Path) -> bool:
|
|
"""真实跑一遍汇总：能编译、能执行（文件已被删除也算不能），并且有足够的行数画图"""
|
|
try:
|
|
result = run_aggregate(path, item.request, table.columns)
|
|
except (BizException, QueryError):
|
|
return False
|
|
return len(result.rows) >= MIN_RESULT_ROWS
|
|
|
|
|
|
def _clean_columns(
|
|
columns: list[ColumnInsight], tables: dict[str, DataTable]
|
|
) -> list[ColumnInsight]:
|
|
known = {
|
|
(name, column.name)
|
|
for name, table in tables.items()
|
|
for column in table.columns
|
|
}
|
|
cleaned: dict[tuple[str, str], ColumnInsight] = {}
|
|
for item in columns:
|
|
key = (item.table, item.name)
|
|
if key not in known or key in cleaned:
|
|
continue
|
|
unit = (item.unit or "").strip()[:MAX_UNIT_CHARS]
|
|
cleaned[key] = item.model_copy(
|
|
update={
|
|
"label": item.label.strip()[:MAX_LABEL_CHARS] or item.name,
|
|
"unit": unit or None,
|
|
}
|
|
)
|
|
return list(cleaned.values())
|