feat：LLM 分析文件并生成结构化建议
This commit is contained in:
@@ -0,0 +1,111 @@
|
||||
"""把结构化的聚合请求编译成 DuckDB 的 SELECT：前端按列分组、汇总，不用自己写 SQL，也不用关心引号和类型
|
||||
|
||||
列名、别名一律通过 quote_identifier 引用，请求里的字符串不会拼进 SQL 的其他位置
|
||||
"""
|
||||
|
||||
from app.core.exceptions import BizException
|
||||
from app.datafiles.query import quote_identifier
|
||||
from app.schemas.data_file import AggregateRequest, DataColumn, GroupBy, Metric
|
||||
|
||||
# 只能用于数值列的聚合；其余聚合对列类型没有要求
|
||||
NUMBER_ONLY = {"sum", "avg"}
|
||||
TEMPLATES = {
|
||||
"sum": "sum({})",
|
||||
"avg": "avg({})",
|
||||
"min": "min({})",
|
||||
"max": "max({})",
|
||||
"count": "count({})",
|
||||
"count_distinct": "count(DISTINCT {})",
|
||||
}
|
||||
TIME_TYPES = {"date", "datetime"}
|
||||
|
||||
|
||||
def build_aggregate_sql(
|
||||
request: AggregateRequest, columns: list[DataColumn]
|
||||
) -> tuple[str, list[str]]:
|
||||
"""返回 SQL（表名固定为 t）和结果列名；请求不合法时抛 422
|
||||
|
||||
SQL 的 LIMIT 比请求多 1，调用方据此判断结果是否被截断
|
||||
"""
|
||||
by_name = {column.name: column for column in columns}
|
||||
names: list[str] = []
|
||||
selects: list[str] = []
|
||||
|
||||
for item in request.group_by:
|
||||
column = _column(by_name, item.column)
|
||||
selects.append(
|
||||
f"{_group_expression(item, column)} AS {quote_identifier(item.column)}"
|
||||
)
|
||||
names.append(item.column)
|
||||
|
||||
for metric in request.metrics:
|
||||
alias = metric.alias or _default_alias(metric)
|
||||
selects.append(
|
||||
f"{_metric_expression(metric, by_name)} AS {quote_identifier(alias)}"
|
||||
)
|
||||
names.append(alias)
|
||||
|
||||
duplicated = sorted({name for name in names if names.count(name) > 1})
|
||||
if duplicated:
|
||||
raise BizException(
|
||||
422, f"结果列名重复：{'、'.join(duplicated)}，请给指标指定不同的 alias"
|
||||
)
|
||||
|
||||
sql = f"SELECT {', '.join(selects)} FROM t"
|
||||
if request.group_by:
|
||||
positions = ", ".join(str(i) for i in range(1, len(request.group_by) + 1))
|
||||
sql += f" GROUP BY {positions}"
|
||||
sql += f" ORDER BY {_order(request, names)} LIMIT {request.limit + 1}"
|
||||
return sql, names
|
||||
|
||||
|
||||
def _column(by_name: dict[str, DataColumn], name: str) -> DataColumn:
|
||||
if name not in by_name:
|
||||
raise BizException(422, f"列不存在：{name}")
|
||||
return by_name[name]
|
||||
|
||||
|
||||
def _group_expression(item: GroupBy, column: DataColumn) -> str:
|
||||
identifier = quote_identifier(item.column)
|
||||
if item.grain is None:
|
||||
return identifier
|
||||
if column.type not in TIME_TYPES:
|
||||
raise BizException(422, f"只有日期列才能按粒度分组：{item.column}")
|
||||
# 都转成字符串，方便直接做图表的分类轴；字符串顺序和时间顺序一致
|
||||
if item.grain == "quarter":
|
||||
return (
|
||||
f"strftime({identifier}, '%Y') || '-Q' || "
|
||||
f"CAST(quarter({identifier}) AS VARCHAR)"
|
||||
)
|
||||
pattern = {"year": "%Y", "month": "%Y-%m", "day": "%Y-%m-%d"}[item.grain]
|
||||
return f"strftime({identifier}, '{pattern}')"
|
||||
|
||||
|
||||
def _metric_expression(metric: Metric, by_name: dict[str, DataColumn]) -> str:
|
||||
if metric.column is None:
|
||||
if metric.agg != "count":
|
||||
raise BizException(422, f"{metric.agg} 需要指定 column")
|
||||
return "count(*)"
|
||||
column = _column(by_name, metric.column)
|
||||
if metric.agg in NUMBER_ONLY and column.type != "number":
|
||||
raise BizException(422, f"{metric.agg} 只能用于数值列：{metric.column}")
|
||||
return TEMPLATES[metric.agg].format(quote_identifier(metric.column))
|
||||
|
||||
|
||||
def _default_alias(metric: Metric) -> str:
|
||||
"""汇总列默认沿用列名，直接对应图表物料按列名引用的写法；计数类加后缀避免和原列混淆"""
|
||||
if metric.agg == "count":
|
||||
return f"{metric.column}_count" if metric.column else "count"
|
||||
if metric.agg == "count_distinct":
|
||||
return f"{metric.column}_distinct"
|
||||
return metric.column or metric.agg
|
||||
|
||||
|
||||
def _order(request: AggregateRequest, names: list[str]) -> str:
|
||||
if request.order_by is None:
|
||||
# 有分组时按第一个分组列升序，结果顺序稳定；只有汇总时只有一行
|
||||
return "1 ASC"
|
||||
if request.order_by.by not in names:
|
||||
raise BizException(422, f"排序列必须是结果里的列：{request.order_by.by}")
|
||||
position = names.index(request.order_by.by) + 1
|
||||
return f"{position} {request.order_by.direction.upper()} NULLS LAST"
|
||||
@@ -0,0 +1,203 @@
|
||||
"""上传后用 LLM 一次分析文件：一句话说明、图表推荐、列的含义
|
||||
|
||||
只把列概况和几行样例发给模型，不发全量数据。模型给出的推荐在返回前逐条校验并真实跑一遍汇总，
|
||||
列不存在、聚合不合法、结果为空的推荐直接丢弃，所以前端拿到的推荐都能直接出图
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable
|
||||
from pathlib import Path
|
||||
|
||||
import duckdb
|
||||
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 build_aggregate_sql
|
||||
from app.datafiles.query import QueryError, run_query
|
||||
from app.datafiles.storage import file_dir
|
||||
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:
|
||||
columns = table.profile.get("columns", [])[:MAX_COLUMNS]
|
||||
keep = {column["name"] for column in columns}
|
||||
positions = [
|
||||
position for position, column in enumerate(table.columns) if column.name in keep
|
||||
]
|
||||
return {
|
||||
"table": table.table,
|
||||
"label": table.label,
|
||||
"rows": table.rows,
|
||||
"columns": columns,
|
||||
# 样例行只保留上面那些列，顺序与 columns 一致
|
||||
"sample": [
|
||||
[row[position] for position in positions]
|
||||
for row in table.profile.get("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, file_dir(root, record.file_id)
|
||||
):
|
||||
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, directory: Path) -> bool:
|
||||
"""真实跑一遍汇总：能编译、能执行，并且有足够的行数画图"""
|
||||
path = directory / table.parquet
|
||||
if not path.is_file():
|
||||
return False
|
||||
try:
|
||||
sql, _ = build_aggregate_sql(item.request, table.columns)
|
||||
result = run_query(
|
||||
sql,
|
||||
{"t": path},
|
||||
timeout=settings.data_query_timeout,
|
||||
max_rows=item.request.limit,
|
||||
)
|
||||
except (BizException, QueryError, duckdb.Error):
|
||||
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())
|
||||
@@ -62,6 +62,29 @@ def profile_table(
|
||||
return profile, preview
|
||||
|
||||
|
||||
def read_rows(
|
||||
path: Path,
|
||||
offset: int,
|
||||
limit: int,
|
||||
sort_by: str | None = None,
|
||||
descending: bool = False,
|
||||
) -> list[list[Any]]:
|
||||
"""读一页数据，单元格截断规则与预览相同；sort_by 必须是已校验过的列名，空值排最后"""
|
||||
order = (
|
||||
f" ORDER BY {quote_identifier(sort_by)} {'DESC' if descending else 'ASC'} NULLS LAST"
|
||||
if sort_by
|
||||
else ""
|
||||
)
|
||||
connection = connect({"t": path})
|
||||
try:
|
||||
rows = connection.execute(
|
||||
f"FROM t{order} LIMIT {int(limit)} OFFSET {int(offset)}"
|
||||
).fetchall()
|
||||
finally:
|
||||
connection.close()
|
||||
return [[json_value(value, PREVIEW_CELL_CHARS) for value in row] for row in rows]
|
||||
|
||||
|
||||
def _profile_column(
|
||||
connection: duckdb.DuckDBPyConnection,
|
||||
column: DataColumn,
|
||||
|
||||
+158
-2
@@ -1,4 +1,4 @@
|
||||
"""上传流程：保存原文件 → 子进程解析（超时终止）→ 分配表名 → 保存记录；以及查看、删除和过期清理"""
|
||||
"""上传流程：保存原文件 → 子进程解析（超时终止）→ 分配表名 → 保存记录 → 后台用 LLM 分析；以及查看、删除和过期清理"""
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
@@ -14,16 +14,29 @@ from pathlib import Path
|
||||
from typing import BinaryIO
|
||||
from uuid import uuid4
|
||||
|
||||
import duckdb
|
||||
from loguru import logger
|
||||
|
||||
from app.config import PROJECT_ROOT, settings
|
||||
from app.core.exceptions import BizException
|
||||
from app.datafiles.aggregate import build_aggregate_sql
|
||||
from app.datafiles.insights import (
|
||||
INSIGHTS_CONCURRENCY,
|
||||
INSIGHTS_TIMEOUT,
|
||||
InsightsGenerator,
|
||||
)
|
||||
from app.datafiles.parse import ParsedFile, ParsedTable
|
||||
from app.datafiles.profile import read_rows
|
||||
from app.datafiles.query import QueryError, run_query
|
||||
from app.datafiles.storage import SOURCE_NAME, DataFileStore, file_dir
|
||||
from app.schemas.base import JsonObject
|
||||
from app.schemas.data_file import (
|
||||
AggregateRequest,
|
||||
AggregateResult,
|
||||
DataFileRecord,
|
||||
DataTable,
|
||||
DataTableRows,
|
||||
FileInsights,
|
||||
RequestDataFile,
|
||||
RequestDataTable,
|
||||
)
|
||||
@@ -54,6 +67,7 @@ class DataFileService:
|
||||
*,
|
||||
max_mb: int | None = None,
|
||||
ttl: int | None = None,
|
||||
insights: InsightsGenerator | None = None,
|
||||
):
|
||||
self.store = store
|
||||
self.root = root or settings.data_file_dir
|
||||
@@ -63,6 +77,12 @@ class DataFileService:
|
||||
self._parsing = asyncio.Semaphore(PARSE_CONCURRENCY)
|
||||
# 文件数检查、分配表名和写入记录要一起完成，并发上传时才不会超限或重名；单实例部署，进程内的锁就够了
|
||||
self._saving = asyncio.Lock()
|
||||
# 上传后的 LLM 分析；为空表示不分析，文件的 insights 为空
|
||||
self._insights = insights
|
||||
self._insight_slots = asyncio.Semaphore(INSIGHTS_CONCURRENCY)
|
||||
# 正在本进程里分析的文件。记录里是 pending 却不在这里，说明进程重启过、任务已丢失
|
||||
self._analysing: set[str] = set()
|
||||
self._tasks: set[asyncio.Task] = set()
|
||||
|
||||
async def upload(
|
||||
self, user_id: int, conversation_id: str, name: str | None, source: BinaryIO
|
||||
@@ -89,11 +109,20 @@ class DataFileService:
|
||||
parsed,
|
||||
_next_table_number(existing),
|
||||
)
|
||||
if self._insights is not None:
|
||||
# 先登记再写入，别的请求读到 pending 时不会把它当成丢失的任务
|
||||
self._analysing.add(file_id)
|
||||
record = record.model_copy(
|
||||
update={"insights": FileInsights(status="pending")}
|
||||
)
|
||||
await self.store.insert(record)
|
||||
except BaseException:
|
||||
# 失败的上传不留文件，会话内的文件数不变
|
||||
self._analysing.discard(file_id)
|
||||
await asyncio.to_thread(shutil.rmtree, directory, ignore_errors=True)
|
||||
raise
|
||||
if self._insights is not None:
|
||||
self._start_insights(record)
|
||||
logger.info(
|
||||
"数据文件已上传：{} 格式 {}，{} 张表",
|
||||
file_id,
|
||||
@@ -107,15 +136,92 @@ class DataFileService:
|
||||
) -> list[DataFileRecord]:
|
||||
await self._check_conversation(conversation_id, user_id)
|
||||
records = await self.store.find(conversation_id)
|
||||
return [record for record in records if record.owner_id == user_id]
|
||||
return [
|
||||
await self._settle(record)
|
||||
for record in records
|
||||
if record.owner_id == user_id
|
||||
]
|
||||
|
||||
async def get_file(self, user_id: int, file_id: str) -> DataFileRecord:
|
||||
record = await self.store.get(file_id)
|
||||
# 他人的文件与不存在同样返回 404
|
||||
if record is None or record.owner_id != user_id:
|
||||
raise BizException(404, NOT_FOUND)
|
||||
return await self._settle(record)
|
||||
|
||||
async def regenerate_insights(self, user_id: int, file_id: str) -> DataFileRecord:
|
||||
"""重新分析，用于上次失败之后；正在分析时 409"""
|
||||
if self._insights is None:
|
||||
raise BizException(422, "未启用文件分析")
|
||||
record = await self.get_file(user_id, file_id)
|
||||
if file_id in self._analysing:
|
||||
raise BizException(409, "正在分析，请稍后")
|
||||
self._analysing.add(file_id)
|
||||
pending = FileInsights(status="pending")
|
||||
try:
|
||||
await self.store.update_insights(file_id, pending)
|
||||
except BaseException:
|
||||
self._analysing.discard(file_id)
|
||||
raise
|
||||
record = record.model_copy(update={"insights": pending})
|
||||
self._start_insights(record)
|
||||
return record
|
||||
|
||||
async def get_rows(
|
||||
self,
|
||||
user_id: int,
|
||||
file_id: str,
|
||||
table_name: str,
|
||||
offset: int,
|
||||
limit: int,
|
||||
sort_by: str | None = None,
|
||||
descending: bool = False,
|
||||
) -> DataTableRows:
|
||||
"""表的一页数据；表或排序列不存在时 404 / 422，文件已被清理时同样 404"""
|
||||
record = await self.get_file(user_id, file_id)
|
||||
table = next((item for item in record.tables if item.table == table_name), None)
|
||||
if table is None:
|
||||
raise BizException(404, "数据表不存在")
|
||||
if sort_by is not None and sort_by not in {c.name for c in table.columns}:
|
||||
raise BizException(422, f"排序列不存在：{sort_by}")
|
||||
path = file_dir(self.root, file_id) / table.parquet
|
||||
try:
|
||||
rows = await asyncio.to_thread(
|
||||
read_rows, path, offset, limit, sort_by, descending
|
||||
)
|
||||
except duckdb.IOException:
|
||||
# 记录还在但 Parquet 已被删除（过期清理、删除时目录没删净）
|
||||
raise BizException(404, NOT_FOUND) from None
|
||||
return DataTableRows(total=table.rows, offset=offset, rows=rows)
|
||||
|
||||
async def aggregate(
|
||||
self, user_id: int, file_id: str, table_name: str, request: AggregateRequest
|
||||
) -> AggregateResult:
|
||||
"""对一张表分组汇总，供前端做图表数据；和 agent 的 query_data 同一套只读沙箱"""
|
||||
record = await self.get_file(user_id, file_id)
|
||||
table = next((item for item in record.tables if item.table == table_name), None)
|
||||
if table is None:
|
||||
raise BizException(404, "数据表不存在")
|
||||
sql, names = build_aggregate_sql(request, table.columns)
|
||||
path = file_dir(self.root, file_id) / table.parquet
|
||||
if not path.is_file():
|
||||
raise BizException(404, NOT_FOUND)
|
||||
try:
|
||||
result = await asyncio.to_thread(
|
||||
run_query,
|
||||
sql,
|
||||
{"t": path},
|
||||
timeout=settings.data_query_timeout,
|
||||
max_rows=request.limit,
|
||||
)
|
||||
except QueryError as error:
|
||||
raise BizException(422, str(error)) from None
|
||||
return AggregateResult(
|
||||
columns=names,
|
||||
rows=[dict(zip(names, row)) for row in result.rows],
|
||||
truncated=result.truncated,
|
||||
)
|
||||
|
||||
async def delete_file(self, user_id: int, file_id: str) -> None:
|
||||
await self.get_file(user_id, file_id)
|
||||
await self.store.mark_deleted(file_id)
|
||||
@@ -166,6 +272,12 @@ class DataFileService:
|
||||
for record in records
|
||||
]
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""应用退出时取消还在进行的分析，记录里留下的 pending 下次读取时会记为失败"""
|
||||
for task in list(self._tasks):
|
||||
task.cancel()
|
||||
await asyncio.gather(*self._tasks, return_exceptions=True)
|
||||
|
||||
async def sweep(self) -> int:
|
||||
"""删除没有对应记录（过期、已删除、内存模式重启后丢失）的目录，返回删除的个数"""
|
||||
live = await self.store.live_ids()
|
||||
@@ -180,6 +292,50 @@ class DataFileService:
|
||||
logger.opt(exception=error).error("清理数据文件目录失败")
|
||||
await asyncio.sleep(SWEEP_INTERVAL)
|
||||
|
||||
def _start_insights(self, record: DataFileRecord) -> None:
|
||||
task = asyncio.create_task(self._analyse(record))
|
||||
self._tasks.add(task)
|
||||
task.add_done_callback(self._tasks.discard)
|
||||
|
||||
async def _analyse(self, record: DataFileRecord) -> None:
|
||||
"""后台分析并写回记录；任何失败都只记为 failed，不影响文件使用"""
|
||||
assert self._insights is not None
|
||||
try:
|
||||
try:
|
||||
async with self._insight_slots:
|
||||
insights = await asyncio.wait_for(
|
||||
self._insights(record, self.root), INSIGHTS_TIMEOUT
|
||||
)
|
||||
except Exception as error:
|
||||
logger.opt(exception=error).warning(
|
||||
"数据文件分析失败：{}", record.file_id
|
||||
)
|
||||
insights = FileInsights(status="failed")
|
||||
await self.store.update_insights(record.file_id, insights)
|
||||
logger.info(
|
||||
"数据文件分析完成：{} {}，{} 条推荐",
|
||||
record.file_id,
|
||||
insights.status,
|
||||
len(insights.suggestions),
|
||||
)
|
||||
except Exception as error:
|
||||
logger.opt(exception=error).error("保存数据文件分析结果失败")
|
||||
finally:
|
||||
self._analysing.discard(record.file_id)
|
||||
|
||||
async def _settle(self, record: DataFileRecord) -> DataFileRecord:
|
||||
"""pending 但本进程没有在分析（重启丢了任务），记为失败，免得前端一直等"""
|
||||
insights = record.insights
|
||||
if (
|
||||
insights is None
|
||||
or insights.status != "pending"
|
||||
or record.file_id in self._analysing
|
||||
):
|
||||
return record
|
||||
failed = FileInsights(status="failed")
|
||||
await self.store.update_insights(record.file_id, failed)
|
||||
return record.model_copy(update={"insights": failed})
|
||||
|
||||
async def _check_conversation(self, conversation_id: str, user_id: int) -> None:
|
||||
if not await self._conversations.exists(conversation_id, user_id):
|
||||
raise BizException(404, "会话不存在或已过期")
|
||||
|
||||
@@ -11,7 +11,7 @@ from typing import Any, Protocol
|
||||
from bson.codec_options import CodecOptions
|
||||
from pymongo import ASCENDING
|
||||
|
||||
from app.schemas.data_file import DataFileRecord
|
||||
from app.schemas.data_file import DataFileRecord, FileInsights
|
||||
|
||||
# 原文件不带扩展名：calamine 见到 .xls 等扩展名会按扩展名打开，用户的扩展名和内容不符时会失败
|
||||
SOURCE_NAME = "source"
|
||||
@@ -42,6 +42,10 @@ class DataFileStore(Protocol):
|
||||
|
||||
async def mark_deleted(self, file_id: str) -> None: ...
|
||||
|
||||
async def update_insights(self, file_id: str, insights: FileInsights) -> None:
|
||||
"""记录已过期或不存在时什么都不做"""
|
||||
...
|
||||
|
||||
async def live_ids(self) -> set[str]:
|
||||
"""未过期且未删除的 file_id，清理目录时用"""
|
||||
...
|
||||
@@ -83,6 +87,10 @@ class InMemoryDataFileStore:
|
||||
if (record := self._records.get(file_id)) is not None:
|
||||
self._records[file_id] = record.model_copy(update={"deleted": True})
|
||||
|
||||
async def update_insights(self, file_id: str, insights: FileInsights) -> None:
|
||||
if (record := self._records.get(file_id)) is not None:
|
||||
self._records[file_id] = record.model_copy(update={"insights": insights})
|
||||
|
||||
async def live_ids(self) -> set[str]:
|
||||
self._purge_expired()
|
||||
return {key for key, record in self._records.items() if not record.deleted}
|
||||
@@ -129,6 +137,11 @@ class MongoDataFileStore:
|
||||
async def mark_deleted(self, file_id: str) -> None:
|
||||
await self._files.update_one({"_id": file_id}, {"$set": {"deleted": True}})
|
||||
|
||||
async def update_insights(self, file_id: str, insights: FileInsights) -> None:
|
||||
await self._files.update_one(
|
||||
{"_id": file_id}, {"$set": {"insights": insights.model_dump()}}
|
||||
)
|
||||
|
||||
async def live_ids(self) -> set[str]:
|
||||
return set(
|
||||
await self._files.distinct(
|
||||
|
||||
Reference in New Issue
Block a user