diff --git a/app/agent/data_tools.py b/app/agent/data_tools.py index bf5646b..8ce5ec7 100644 --- a/app/agent/data_tools.py +++ b/app/agent/data_tools.py @@ -7,7 +7,7 @@ apply_actions 里的 $query 也通过 DataContext.query 执行,见 tools.py 上下文类型因 agent 而异,所以写 ToolRuntime[Any](不写类型参数时按 None 处理,序列化会告警) """ -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path from typing import Annotated, Any @@ -17,7 +17,8 @@ from pydantic import Field from app.config import settings from app.datafiles.query import MAX_ROWS, QueryError, QueryResult, run_query -from app.schemas.data_file import RequestDataFile, RequestDataTable +from app.datafiles.storage import table_path +from app.schemas.data_file import DataFileRecord, DataTable # 列超过这么多时,不指定 columns 只返回列名和类型,避免一次塞进太多概况 BRIEF_COLUMNS = 50 @@ -25,9 +26,12 @@ BRIEF_COLUMNS = 50 @dataclass(frozen=True) class DataContext: - files: tuple[RequestDataFile, ...] = () + """本次运行可用的数据文件;root 是数据文件的根目录,Parquet 路径由它和文件记录算出""" - def find(self, name: str) -> tuple[RequestDataFile, RequestDataTable] | None: + files: tuple[DataFileRecord, ...] = () + root: Path = field(default_factory=lambda: settings.data_file_dir) + + def find(self, name: str) -> tuple[DataFileRecord, DataTable] | None: for file in self.files: for table in file.tables: if table.table == name: @@ -36,7 +40,7 @@ class DataContext: def paths(self) -> dict[str, Path]: return { - table.table: Path(table.path) + table.table: table_path(self.root, file.file_id, table) for file in self.files for table in file.tables } @@ -86,7 +90,7 @@ def describe_data( if found is None: return _error(runtime, f"没有表 {table},可用的表:{_table_list(data)}") file, item = found - profiled: list[dict[str, Any]] = item.profile["columns"] + profiled = item.profile.columns result: dict[str, Any] = { "table": item.table, "label": item.label, @@ -94,20 +98,25 @@ def describe_data( } if columns is None and len(profiled) > BRIEF_COLUMNS: result["columns"] = [ - {key: column[key] for key in ("name", "type", "format") if key in column} + column.model_dump( + mode="json", include={"name", "type", "format"}, exclude_none=True + ) for column in profiled ] result["note"] = ( f"共 {len(profiled)} 列,只列出列名和类型;用 columns 指定列查看详细概况和样例" ) else: - index = {column["name"]: position for position, column in enumerate(profiled)} + index = {column.name: position for position, column in enumerate(profiled)} names = list(index) if columns is None else list(dict.fromkeys(columns)) positions = [index[name] for name in names if name in index] - result["columns"] = [profiled[position] for position in positions] + result["columns"] = [ + profiled[position].model_dump(mode="json", exclude_none=True) + for position in positions + ] # 样例行的值和 columns 一一对应 result["sample_rows"] = [ - [row[position] for position in positions] for row in item.profile["sample"] + [row[position] for position in positions] for row in item.profile.sample ] if missing := [name for name in names if name not in index]: result["missing_columns"] = missing diff --git a/app/api/files.py b/app/api/files.py index 4df4a1b..28b0517 100644 --- a/app/api/files.py +++ b/app/api/files.py @@ -3,13 +3,17 @@ from typing import Annotated, Literal, get_args from fastapi import APIRouter, Depends, File, Form, Query, Request, UploadFile +from fastapi.responses import JSONResponse +from starlette.types import ASGIApp, Receive, Scope, Send -from app.common.response import ApiResponse, success_response +from app.common.response import ApiResponse, error_content, success_response from app.core.auth import get_current_user +from app.datafiles.base import MAX_TABLE_ROWS from app.datafiles.service import ( MAX_FILES_PER_CONVERSATION, PARSE_CONCURRENCY, DataFileService, + file_too_large, ) from app.schemas.base import NonBlankString from app.schemas.data_file import ( @@ -31,18 +35,55 @@ from app.schemas.user import CurrentUser limits_router = APIRouter(prefix="/files", tags=["DataFile"]) router = APIRouter(prefix="/files", tags=["DataFile"]) -# 文件选择框的过滤条件;后端按文件内容识别格式,不看扩展名 -# 翻页取表数据:单页上限,和一张表最多的行数(100 万) +# 翻页取表数据:单页上限,offset 不超过一张表最多的行数 MAX_PAGE_ROWS = 500 -MAX_OFFSET = 1_000_000 +MAX_OFFSET = MAX_TABLE_ROWS +# 文件选择框的过滤条件;后端按文件内容识别格式,不看扩展名 ACCEPT = ".csv,.tsv,.txt,.xlsx,.xls,.xlsb,.json,.jsonl,.ndjson,.parquet" +# 上传的请求体比文件本身多出 multipart 的分隔行、字段名和会话 ID +UPLOAD_OVERHEAD_BYTES = 64 * 1024 + def _service(request: Request) -> DataFileService: return request.app.state.data_file_service +class UploadSizeLimit: + """上传接口在读取请求体之前按 Content-Length 拒绝过大的请求 + + FastAPI 会先把整个 multipart 请求体收完、写进临时文件,再调用接口函数; + 只在 DataFileService 里边写边查大小的话,超大的上传已经占满了带宽和磁盘 + """ + + def __init__(self, app: ASGIApp, path: str): + self.app = app + self.path = path + + async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None: + if ( + scope["type"] == "http" + and scope["method"] == "POST" + and scope["path"] == self.path + and (rejection := self._reject(scope)) is not None + ): + await rejection(scope, receive, send) + return + await self.app(scope, receive, send) + + @staticmethod + def _reject(scope: Scope) -> JSONResponse | None: + max_mb = scope["app"].state.data_file_service.max_mb + length = dict(scope["headers"]).get(b"content-length") + if length is None or not length.isdigit(): + return JSONResponse(error_content(411, "缺少 Content-Length"), 411) + if int(length) > max_mb * 1024 * 1024 + UPLOAD_OVERHEAD_BYTES: + error = file_too_large(max_mb) + return JSONResponse(error_content(error.code, error.msg), error.code) + return None + + def _detail(record: DataFileRecord) -> DataFileDetail: # 按接口模型重新校验,去掉 owner_id、sha256、列概况等只在后端用的字段 return DataFileDetail.model_validate(record.model_dump()) diff --git a/app/datafiles/aggregate.py b/app/datafiles/aggregate.py index aff0d07..513becb 100644 --- a/app/datafiles/aggregate.py +++ b/app/datafiles/aggregate.py @@ -3,9 +3,18 @@ 列名、别名一律通过 quote_identifier 引用,请求里的字符串不会拼进 SQL 的其他位置 """ +from pathlib import Path + +from app.config import settings from app.core.exceptions import BizException -from app.datafiles.query import quote_identifier -from app.schemas.data_file import AggregateRequest, DataColumn, GroupBy, Metric +from app.datafiles.query import quote_identifier, run_query +from app.schemas.data_file import ( + AggregateRequest, + AggregateResult, + DataColumn, + GroupBy, + Metric, +) # 只能用于数值列的聚合;其余聚合对列类型没有要求 NUMBER_ONLY = {"sum", "avg"} @@ -20,7 +29,28 @@ TEMPLATES = { TIME_TYPES = {"date", "datetime"} -def build_aggregate_sql( +def run_aggregate( + path: Path, request: AggregateRequest, columns: list[DataColumn] +) -> AggregateResult: + """对一张表的 Parquet 执行汇总;接口和上传后的分析共用 + + 请求不合法时抛 422;执行失败(包括文件已被删除)抛 QueryError + """ + sql, names = _build_sql(request, columns) + result = run_query( + sql, + {"t": path}, + timeout=settings.data_query_timeout, + max_rows=request.limit, + ) + return AggregateResult( + columns=names, + rows=[dict(zip(names, row)) for row in result.rows], + truncated=result.truncated, + ) + + +def _build_sql( request: AggregateRequest, columns: list[DataColumn] ) -> tuple[str, list[str]]: """返回 SQL(表名固定为 t)和结果列名;请求不合法时抛 422 diff --git a/app/datafiles/base.py b/app/datafiles/base.py index 4896947..4c355a6 100644 --- a/app/datafiles/base.py +++ b/app/datafiles/base.py @@ -1,4 +1,6 @@ -"""数据文件各环节共用的异常和上限""" +"""数据文件各环节共用的异常、上限和警告文案的写法""" + +from collections.abc import Sequence # 每张表的数据行数上限 MAX_TABLE_ROWS = 1_000_000 @@ -26,6 +28,11 @@ def where(sheet: str | None) -> str: return f"{sheet}:" if sheet else "" +def brief(items: Sequence[str], sep: str = "、", limit: int = 5) -> str: + """警告里列举的内容最多写 limit 个,多出的用“等”概括""" + return sep.join(items[:limit]) + ("等" if len(items) > limit else "") + + def rows_exceeded(sheet: str | None) -> DataFileError: return DataFileError( 422, f"{where(sheet)}超过 {MAX_TABLE_ROWS // 10_000} 万行,请拆分后上传" diff --git a/app/datafiles/columns.py b/app/datafiles/columns.py index 7077cb9..c2077ec 100644 --- a/app/datafiles/columns.py +++ b/app/datafiles/columns.py @@ -15,7 +15,7 @@ from typing import Any import pyarrow as pa -from app.datafiles.base import where +from app.datafiles.base import brief, where from app.schemas.data_file import DataColumn, ParseWarning NULL_TOKENS = {"", "-", "--", "—", "——", "N/A", "n/a", "NULL", "null", "#N/A"} @@ -63,7 +63,7 @@ def clean_names( """去掉空白和换行;空列名按位置命名,重复列名加后缀""" names, changes, seen = [], [], set() for position, value in enumerate(raw, start=1): - name = re.sub(r"\s*[\r\n]+\s*", " ", label_text(value)).strip() + name = re.sub(r"\s*[\r\n]+\s*", " ", as_text(value)).strip() if not name: name = f"列{position}" changes.append(f"第 {position} 列没有列名,命名为「{name}」") @@ -77,19 +77,9 @@ def clean_names( names.append(name) if not changes: return names, None - shown = ";".join(changes[:5]) + ("等" if len(changes) > 5 else "") - return names, ParseWarning(code="column_renamed", message=f"{where(sheet)}{shown}") - - -def label_text(value: Any) -> str: - """表头单元格转成列名;Excel 里年份等数字表头读出来是 2025.0""" - if value is None: - return "" - if isinstance(value, float) and value.is_integer(): - return str(int(value)) - if isinstance(value, (date, time)): - return value.isoformat() - return str(value) + return names, ParseWarning( + code="column_renamed", message=f"{where(sheet)}{brief(changes, ';')}" + ) def build_columns( @@ -257,29 +247,29 @@ def classify(value: Any) -> tuple[str, Any]: def parse_number(text: str) -> int | float | None: - """千分位和货币符号可以有;有前导零或 16 位以上的数字串是编号,不算数值""" - # 最常见的纯数字走快速路径,不用正则 - if text.isdigit(): - if len(text) >= MIN_ID_DIGITS or (len(text) > 1 and text[0] == "0"): - return None - return int(text) + """千分位和货币符号可以有;有前导零或 16 位以上的数字串是编号,不算数值 + + 判断数字一律用 isdecimal:isdigit 对 ²、① 也返回 True,但 int() 解析不了它们 + """ if text and text[0] in CURRENCY_SYMBOLS: text = text[1:].lstrip() - if not NUMBER.match(text): + # 最常见的纯数字不用正则 + plain = text.isdecimal() + if not plain and not NUMBER.match(text): return None digits = text.lstrip("+-") - if len(digits) > 1 and digits[0] == "0" and digits[1].isdigit(): + if len(digits) > 1 and digits[0] == "0" and digits[1].isdecimal(): return None - if digits.isdigit() and len(digits) >= MIN_ID_DIGITS: + if digits.isdecimal() and len(digits) >= MIN_ID_DIGITS: return None - text = text.replace(",", "") - if INTEGER.match(text): + if plain: return int(text) - return float(text) + text = text.replace(",", "") + return int(text) if INTEGER.match(text) else float(text) def parse_datetime(text: str) -> date | datetime | None: - if len(text) < 6 or not text[:4].isdigit(): + if len(text) < 6 or not text[:4].isdecimal(): return None # 标准 ISO 写法先用 fromisoformat,比 strptime 快一个数量级 if ":" in text: @@ -307,6 +297,9 @@ def parse_datetime(text: str) -> date | datetime | None: def as_text(value: Any) -> str: + """单元格转成文本:列名、保留为文本的值、警告里引用的内容都用它;Excel 里年份等数字读出来是 2025.0""" + if value is None: + return "" if isinstance(value, str): return value.strip() if isinstance(value, bool): @@ -318,7 +311,8 @@ def as_text(value: Any) -> str: if isinstance(value, (date, time)): return value.isoformat() if isinstance(value, (dict, list)): - return json.dumps(value, ensure_ascii=False) + # Parquet 的嵌套值里可能有日期、Decimal、bytes,JSON 表示不了的写成字符串 + return json.dumps(value, ensure_ascii=False, default=str) return str(value) diff --git a/app/datafiles/detect.py b/app/datafiles/detect.py index 163825d..27f4215 100644 --- a/app/datafiles/detect.py +++ b/app/datafiles/detect.py @@ -3,6 +3,7 @@ import codecs import csv import json +import re import zipfile from dataclasses import dataclass from pathlib import Path @@ -29,6 +30,10 @@ DELIMITERS = (",", "\t", ";", "|") # 判断分隔符只看开头这么多行 SNIFF_LINES = 50 +# 开头像 JSON 才按 JSON 处理:对象的 { 后面只能是 " 或 },数组的 [ 后面是一个值的开头; +# CSV 的列名也可能以 [ 开头,例如 [日期],[销售额] +JSON_START = re.compile(r'\{\s*["}]|\[\s*(?:[\[{"\]\d-]|(?:true|false|null)\s*[,\]])') + @dataclass(frozen=True) class Detected: @@ -56,7 +61,7 @@ def detect(path: Path) -> Detected: stripped = text.lstrip() if not stripped: raise DataFileError(422, "文件是空的") - if stripped[0] in "[{": + if JSON_START.match(stripped): return _json_format(text, encoding) delimiter = sniff_delimiter(text) return Detected("tsv" if delimiter == "\t" else "csv", encoding, text, delimiter) @@ -115,7 +120,7 @@ def _json_format(text: str, encoding: str) -> Detected: return Detected("jsonl", encoding, text) except ValueError: pass - # 以 [ 或 { 开头却不是合法 JSON,按 JSON 报错比当成 CSV 解析出一堆乱码更有用 + # 开头像 JSON 却不是合法 JSON,按 JSON 报错比当成 CSV 解析出一堆乱码更有用 line = getattr(json_error, "lineno", None) position = f"第 {line} 行" if line else "" raise DataFileError( diff --git a/app/datafiles/insights.py b/app/datafiles/insights.py index 1bd5105..715e90f 100644 --- a/app/datafiles/insights.py +++ b/app/datafiles/insights.py @@ -9,7 +9,6 @@ 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 @@ -18,9 +17,9 @@ 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.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, @@ -106,21 +105,16 @@ def build_input(record: DataFileRecord) -> dict: 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", []) + "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], } @@ -140,7 +134,7 @@ def check_draft( 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) + item, table, table_path(root, record.file_id, table) ): logger.info("丢弃不可用的图表推荐:{} {}", record.file_id, item.title) continue @@ -162,20 +156,11 @@ def _fits_chart(item: ChartSuggestion) -> bool: 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 +def _runs(item: ChartSuggestion, table: DataTable, path: Path) -> bool: + """真实跑一遍汇总:能编译、能执行(文件已被删除也算不能),并且有足够的行数画图""" 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): + result = run_aggregate(path, item.request, table.columns) + except (BizException, QueryError): return False return len(result.rows) >= MIN_RESULT_ROWS diff --git a/app/datafiles/parse.py b/app/datafiles/parse.py index 25f583f..3decf5a 100644 --- a/app/datafiles/parse.py +++ b/app/datafiles/parse.py @@ -3,7 +3,6 @@ 只读写本地文件,不碰数据库;上传流程在子进程里调用(worker.py),超时直接终止子进程 """ -from dataclasses import dataclass from pathlib import Path from typing import Any @@ -29,7 +28,7 @@ from app.datafiles.readers import ( read_parquet, ) from app.datafiles.tables import split_tables -from app.schemas.data_file import DataColumn, DataFormat, ParseWarning +from app.schemas.data_file import DataColumn, DataFormat, ParseWarning, TableProfile # 子进程把解析结果转成 JSON 交回上传流程,所以用 pydantic 模型 @@ -39,8 +38,9 @@ class ParsedTable(BaseModel): part: int | None rows: int columns: list[DataColumn] - path: Path - profile: dict[str, Any] + # out_dir 下的 Parquet 文件名 + parquet: str + profile: TableProfile preview: list[list[Any]] warnings: list[ParseWarning] @@ -53,40 +53,53 @@ class ParsedFile(BaseModel): warnings: list[ParseWarning] -@dataclass -class _Built: - sheet: str | None - part: int | None - typed: TypedColumns - warnings: list[ParseWarning] - - def parse_file(path: Path, out_dir: Path) -> ParsedFile: - """解析 path,Parquet 写到 out_dir/1.parquet、2.parquet……""" + """解析 path,每张表在 out_dir 下写一个 Parquet:1.parquet、2.parquet……""" detected = detect(path) + out_dir.mkdir(parents=True, exist_ok=True) warnings: list[ParseWarning] = [] - built: list[_Built] = [] + tables: list[ParsedTable] = [] + + def add_table( + sheet: str | None, + part: int | None, + typed: TypedColumns, + table_warnings: list[ParseWarning], + ) -> None: + target = out_dir / f"{len(tables) + 1}.parquet" + pq.write_table(typed.table, target) + profile, preview = profile_table(target, typed.columns, typed.notes) + tables.append( + ParsedTable( + sheet=sheet, + part=part, + rows=typed.table.num_rows, + columns=typed.columns, + parquet=target.name, + profile=profile, + preview=preview, + warnings=table_warnings, + ) + ) def add_grid(grid: Grid) -> None: - tables, dropped = split_tables(grid) + found, dropped = split_tables(grid) warnings.extend(dropped) - for table in tables: + for table in found: names, renamed = clean_names(table.header, table.sheet) typed, kept = build_columns(names, list(zip(*table.rows)), table.sheet) - built.append( - _Built( - table.sheet, - table.part, - typed, - [*table.warnings, *_present(renamed, kept)], - ) + add_table( + table.sheet, + table.part, + typed, + [*table.warnings, *_present(renamed, kept)], ) def add_records(records: Records) -> None: warnings.extend(records.warnings) names, renamed = clean_names(records.names, None) typed, kept = build_columns(names, records.columns, None) - built.append(_Built(None, None, typed, _present(renamed, kept))) + add_table(None, None, typed, _present(renamed, kept)) match detected.format: case "csv" | "tsv": @@ -107,28 +120,10 @@ def parse_file(path: Path, out_dir: Path) -> ParsedFile: case "parquet": typed, renamed = arrow_columns(read_parquet(path)) if typed.table.num_rows: - built.append(_Built(None, None, typed, _present(renamed))) + add_table(None, None, typed, _present(renamed)) - if not built: + if not tables: raise DataFileError(422, "文件里没有找到数据") - out_dir.mkdir(parents=True, exist_ok=True) - tables = [] - for index, item in enumerate(built, start=1): - target = out_dir / f"{index}.parquet" - pq.write_table(item.typed.table, target) - profile, preview = profile_table(target, item.typed.columns, item.typed.notes) - tables.append( - ParsedTable( - sheet=item.sheet, - part=item.part, - rows=item.typed.table.num_rows, - columns=item.typed.columns, - path=target, - profile=profile, - preview=preview, - warnings=item.warnings, - ) - ) return ParsedFile( format=detected.format, encoding=detected.encoding, diff --git a/app/datafiles/profile.py b/app/datafiles/profile.py index 3e59ca7..6ef2eef 100644 --- a/app/datafiles/profile.py +++ b/app/datafiles/profile.py @@ -1,4 +1,4 @@ -"""列概况和预览:上传时在 Parquet 上用 DuckDB 计算,存进文件记录,供 describe_data 和接口返回""" +"""列概况和预览:上传时在 Parquet 上用 DuckDB 计算,存进文件记录,供 describe_data、上传后的分析和接口返回""" import math from pathlib import Path @@ -6,8 +6,9 @@ from typing import Any import duckdb +from app.datafiles.columns import EXACT_FLOAT_INT from app.datafiles.query import connect, json_value, quote_identifier -from app.schemas.data_file import DataColumn +from app.schemas.data_file import ColumnProfile, DataColumn, TableProfile, TopValue PREVIEW_ROWS = 20 SAMPLE_ROWS = 5 @@ -23,11 +24,10 @@ PREVIEW_CELL_CHARS = 200 def profile_table( path: Path, columns: list[DataColumn], notes: dict[str, str] -) -> tuple[dict[str, Any], list[list[Any]]]: +) -> tuple[TableProfile, list[list[Any]]]: """返回列概况和接口用的预览(前 PREVIEW_ROWS 行)""" connection = connect({"t": path}) try: - rows = connection.execute("SELECT count(*) FROM t").fetchone()[0] connection.execute("SUMMARIZE t") fields = [column[0] for column in connection.description] summary = {row[0]: dict(zip(fields, row)) for row in connection.fetchall()} @@ -48,13 +48,12 @@ def profile_table( ).fetchall() finally: connection.close() - profile = { - "rows": rows, - "columns": profiled, - "sample": [ + profile = TableProfile( + columns=profiled, + sample=[ [json_value(value, PROFILE_CELL_CHARS) for value in row] for row in sample ], - } + ) preview = [ [json_value(value, PREVIEW_CELL_CHARS) for value in row] for row in head[:PREVIEW_ROWS] @@ -92,25 +91,19 @@ def _profile_column( head: list[tuple], position: int, note: str | None, -) -> dict[str, Any]: - item: dict[str, Any] = {"name": column.name, "type": column.type} - if column.format: - item["format"] = column.format - null_ratio = float(stats["null_percentage"]) / 100 - item["null_ratio"] = round(null_ratio, 4) - # approx_unique 是估算值,小表上可能比行数还大 - item["distinct"] = min( - stats["approx_unique"], round(stats["count"] * (1 - null_ratio)) - ) +) -> ColumnProfile: + item: dict[str, Any] = {} if column.type == "number": - for key in ("min", "max", "avg", "q25", "q50", "q75"): - item[key] = _number(stats[key]) + item = { + key: _number(stats[key]) + for key in ("min", "max", "avg", "q25", "q50", "q75") + } elif column.type in ("date", "datetime"): # SUMMARIZE 的时间用空格分隔,和预览里的 ISO 写法统一 - item["min"], item["max"] = ( - None if stats[key] is None else stats[key].replace(" ", "T", 1) + item = { + key: None if stats[key] is None else stats[key].replace(" ", "T", 1) for key in ("min", "max") - ) + } if ( column.type in ("text", "boolean") and stats["approx_unique"] <= TOP_VALUES_MAX_DISTINCT @@ -121,7 +114,7 @@ def _profile_column( f"GROUP BY 1 ORDER BY n DESC, 1 LIMIT {TOP_VALUES}" ).fetchall() item["top_values"] = [ - {"value": json_value(value, PROFILE_CELL_CHARS), "count": count} + TopValue(value=json_value(value, PROFILE_CELL_CHARS), count=count) for value, count in counts ] examples: list[Any] = [] @@ -131,10 +124,16 @@ def _profile_column( examples.append(value) if len(examples) == EXAMPLES: break - item["examples"] = examples - if note: - item["note"] = note - return item + null_ratio = float(stats["null_percentage"]) / 100 + return ColumnProfile( + **column.model_dump(), + null_ratio=round(null_ratio, 4), + # approx_unique 是估算值,小表上可能比行数还大 + distinct=min(stats["approx_unique"], round(stats["count"] * (1 - null_ratio))), + examples=examples, + note=note, + **item, + ) def _number(text: str | None) -> int | float | None: @@ -144,6 +143,6 @@ def _number(text: str | None) -> int | float | None: value = float(text) if not math.isfinite(value): return None - if value.is_integer() and abs(value) < 2**53: + if value.is_integer() and abs(value) < EXACT_FLOAT_INT: return int(value) return round(value, 6) diff --git a/app/datafiles/readers.py b/app/datafiles/readers.py index cb8bc22..f3d1835 100644 --- a/app/datafiles/readers.py +++ b/app/datafiles/readers.py @@ -19,7 +19,8 @@ from python_calamine import ( SheetVisibleEnum, ) -from app.datafiles.base import MAX_TABLE_ROWS, DataFileError, rows_exceeded +from app.datafiles.base import MAX_TABLE_ROWS, DataFileError, brief, rows_exceeded +from app.datafiles.columns import as_text from app.schemas.data_file import ParseWarning # 原始网格比数据多出的标题、表头、空行,读取时先粗略拦住明显超限的文件,精确的行数在拆表后检查 @@ -153,7 +154,7 @@ def read_jsonl(text: str) -> Records: result.warnings.append( ParseWarning( code="bad_lines_skipped", - message=f"跳过了 {len(bad)} 行无法解析的内容:{_line_list(bad)}", + message=f"跳过了 {len(bad)} 行无法解析的内容:{brief([f'第 {n} 行' for n in bad])}", ) ) return result @@ -230,11 +231,5 @@ def _flatten(record: dict, prefix: str, level: int, out: dict[str, Any]) -> None def _scalar(value: Any) -> Any: - if isinstance(value, (dict, list)): - return json.dumps(value, ensure_ascii=False) - return value - - -def _line_list(numbers: list[int]) -> str: - shown = "、".join(f"第 {number} 行" for number in numbers[:5]) - return shown + ("等" if len(numbers) > 5 else "") + # 嵌套的对象、数组存成 JSON 字符串,其余原样交给类型推断 + return as_text(value) if isinstance(value, (dict, list)) else value diff --git a/app/datafiles/service.py b/app/datafiles/service.py index eb83828..ed8fff6 100644 --- a/app/datafiles/service.py +++ b/app/datafiles/service.py @@ -14,12 +14,11 @@ 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.aggregate import run_aggregate from app.datafiles.insights import ( INSIGHTS_CONCURRENCY, INSIGHTS_TIMEOUT, @@ -27,9 +26,8 @@ from app.datafiles.insights import ( ) 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.datafiles.query import QueryError +from app.datafiles.storage import SOURCE_NAME, DataFileStore, table_path from app.schemas.data_file import ( AggregateRequest, AggregateResult, @@ -37,8 +35,6 @@ from app.schemas.data_file import ( DataTable, DataTableRows, FileInsights, - RequestDataFile, - RequestDataTable, ) from app.services.conversation import ConversationStore @@ -75,14 +71,18 @@ class DataFileService: self._conversations = conversations self._ttl = timedelta(seconds=ttl or settings.data_file_ttl) 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() + # 在这之前开始、记录里仍是 pending 的分析属于上一个进程,已随重启丢失; + # MongoDB 里的时间只精确到毫秒,这里也截到毫秒,刚启动就开始的分析读回后才不会早于它 + now = datetime.now(UTC) + self._started_at = now.replace(microsecond=now.microsecond // 1000 * 1000) async def upload( self, user_id: int, conversation_id: str, name: str | None, source: BinaryIO @@ -91,14 +91,16 @@ class DataFileService: # 先粗查一次,满了就不必上传和解析;写入前在锁里再查一次 _check_count(await self.store.find(conversation_id)) file_id = f"f_{uuid4().hex}" - directory = file_dir(self.root, file_id) + directory = self.root / file_id try: size, sha256 = await asyncio.to_thread(self._save, source, directory) async with self._parsing: parsed = await _parse(directory) async with self._saving: - existing = await self.store.find(conversation_id, include_deleted=True) - _check_count([record for record in existing if not record.deleted]) + _check_count(await self.store.find(conversation_id)) + first_number = await self.store.reserve_table_numbers( + conversation_id, len(parsed.tables) + ) record = self._record( file_id, user_id, @@ -107,18 +109,11 @@ class DataFileService: size, sha256, parsed, - _next_table_number(existing), + first_number, ) - 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: @@ -137,9 +132,7 @@ class DataFileService: await self._check_conversation(conversation_id, user_id) records = await self.store.find(conversation_id) return [ - await self._settle(record) - for record in records - if record.owner_id == user_id + 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: @@ -147,7 +140,7 @@ class DataFileService: # 他人的文件与不存在同样返回 404 if record is None or record.owner_id != user_id: raise BizException(404, NOT_FOUND) - return await self._settle(record) + return self._settle(record) async def regenerate_insights(self, user_id: int, file_id: str) -> DataFileRecord: """重新分析,用于上次失败之后;正在分析时 409""" @@ -157,7 +150,7 @@ class DataFileService: if file_id in self._analysing: raise BizException(409, "正在分析,请稍后") self._analysing.add(file_id) - pending = FileInsights(status="pending") + pending = FileInsights(status="pending", started_at=datetime.now(UTC)) try: await self.store.update_insights(file_id, pending) except BaseException: @@ -179,19 +172,12 @@ class DataFileService: ) -> 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, "数据表不存在") + table, path = self._table(record, table_name) 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 + rows = await asyncio.to_thread( + read_rows, path, offset, limit, sort_by, descending + ) return DataTableRows(total=table.rows, offset=offset, rows=rows) async def aggregate( @@ -199,81 +185,36 @@ class DataFileService: ) -> 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) + table, path = self._table(record, table_name) try: - result = await asyncio.to_thread( - run_query, - sql, - {"t": path}, - timeout=settings.data_query_timeout, - max_rows=request.limit, - ) + return await asyncio.to_thread(run_aggregate, path, request, table.columns) 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) # Windows 上正在查询的 Parquet 删不掉,留给定时清理 - await asyncio.to_thread( - shutil.rmtree, file_dir(self.root, file_id), ignore_errors=True - ) + await asyncio.to_thread(shutil.rmtree, self.root / file_id, ignore_errors=True) async def request_files( - self, user_id: int, conversation_id: str, attachments: list[JsonObject] - ) -> list[RequestDataFile]: - """本次请求可用的文件,调用方已校验会话归属 + self, user_id: int, conversation_id: str, attached: set[str] + ) -> list[DataFileRecord]: + """本次请求可用的文件:当前用户在本会话里未过期的文件,调用方已校验会话归属 - attachments 里 kind 为 data_file 的引用必须是本会话未过期的文件,否则 404;其他种类不处理 + attached 是本轮附件引用的文件,必须都在其中,否则 404 """ records = [ record for record in await self.store.find(conversation_id) if record.owner_id == user_id ] - attached = set() - for item in attachments: - if item.get("kind") != "data_file": - continue - if not isinstance(file_id := item.get("file_id"), str): - raise BizException(422, "data_file 附件缺少 file_id") - attached.add(file_id) if attached - {record.file_id for record in records}: raise BizException(404, NOT_FOUND) - return [ - RequestDataFile( - file_id=record.file_id, - name=record.name, - attached_this_turn=record.file_id in attached, - tables=[ - RequestDataTable.model_validate( - { - **table.model_dump(), - "path": str( - file_dir(self.root, record.file_id) / table.parquet - ), - } - ) - for table in record.tables - ], - warnings=record.warnings, - ) - for record in records - ] + return records async def aclose(self) -> None: - """应用退出时取消还在进行的分析,记录里留下的 pending 下次读取时会记为失败""" + """应用退出时取消还在进行的分析,记录里留下的 pending 下次启动后读取时按失败返回""" for task in list(self._tasks): task.cancel() await asyncio.gather(*self._tasks, return_exceptions=True) @@ -293,6 +234,7 @@ class DataFileService: await asyncio.sleep(SWEEP_INTERVAL) def _start_insights(self, record: DataFileRecord) -> None: + self._analysing.add(record.file_id) task = asyncio.create_task(self._analyse(record)) self._tasks.add(task) task.add_done_callback(self._tasks.discard) @@ -323,18 +265,32 @@ class DataFileService: finally: self._analysing.discard(record.file_id) - async def _settle(self, record: DataFileRecord) -> DataFileRecord: - """pending 但本进程没有在分析(重启丢了任务),记为失败,免得前端一直等""" + 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 + insights is not None + and insights.status == "pending" + and (insights.started_at is None or insights.started_at < self._started_at) ): - return record - failed = FileInsights(status="failed") - await self.store.update_insights(record.file_id, failed) - return record.model_copy(update={"insights": failed}) + return record.model_copy(update={"insights": FileInsights(status="failed")}) + return record + + def _table(self, record: DataFileRecord, name: str) -> tuple[DataTable, Path]: + """文件里的一张表和它的 Parquet + + 表不存在,或 Parquet 已被删除(删除时目录没删净、过期清理)时 404 + """ + table = next((item for item in record.tables if item.table == name), None) + if table is None: + raise BizException(404, "数据表不存在") + path = table_path(self.root, record.file_id, table) + if not path.is_file(): + raise BizException(404, NOT_FOUND) + return table, path async def _check_conversation(self, conversation_id: str, user_id: int) -> None: if not await self._conversations.exists(conversation_id, user_id): @@ -348,7 +304,7 @@ class DataFileService: while chunk := source.read(CHUNK_BYTES): size += len(chunk) if size > self.max_mb * 1024 * 1024: - raise BizException(413, f"文件不能超过 {self.max_mb} MB") + raise file_too_large(self.max_mb) digest.update(chunk) target.write(chunk) return size, digest.hexdigest() @@ -371,12 +327,7 @@ class DataFileService: DataTable( table=table_name, label=_label(name, table), - sheet=table.sheet, - rows=table.rows, - columns=table.columns, - preview=table.preview, - profile=table.profile, - parquet=table.path.name, + **table.model_dump(exclude={"part", "warnings"}), ) ) warnings.extend( @@ -397,6 +348,11 @@ class DataFileService: expires_at=now + self._ttl, tables=tables, warnings=warnings, + insights=( + None + if self._insights is None + else FileInsights(status="pending", started_at=now) + ), ) def _remove_orphans(self, live: set[str]) -> int: @@ -466,6 +422,10 @@ def _python() -> tuple[str, dict[str, str] | None]: return sys.executable, None +def file_too_large(max_mb: int) -> BizException: + return BizException(413, f"文件不能超过 {max_mb} MB") + + def _check_count(records: list[DataFileRecord]) -> None: if len(records) >= MAX_FILES_PER_CONVERSATION: raise BizException( @@ -473,12 +433,6 @@ def _check_count(records: list[DataFileRecord]) -> None: ) -def _next_table_number(records: list[DataFileRecord]) -> int: - """表名在会话内递增,已删除文件的表名也不再使用""" - used = (int(table.table[1:]) for record in records for table in record.tables) - return max(used, default=0) + 1 - - def _clean_name(name: str | None) -> str: # 有的浏览器会带上本地路径,只保留文件名 base = re.split(r"[\\/]", name or "")[-1].strip() diff --git a/app/datafiles/storage.py b/app/datafiles/storage.py index c5d83a3..568880e 100644 --- a/app/datafiles/storage.py +++ b/app/datafiles/storage.py @@ -2,6 +2,7 @@ 目录结构是 DATA_FILE_DIR//,里面有原文件 source 和每张表的 1.parquet、2.parquet…… DuckDB 只能读本地文件,所以文件不进 MongoDB;单实例部署,本地磁盘就够了 +会话内的表名 t1、t2……由每个会话一个的计数器分配,文件删除、过期后表名也不再使用 """ from datetime import UTC, datetime @@ -9,16 +10,16 @@ from pathlib import Path from typing import Any, Protocol from bson.codec_options import CodecOptions -from pymongo import ASCENDING +from pymongo import ASCENDING, ReturnDocument -from app.schemas.data_file import DataFileRecord, FileInsights +from app.schemas.data_file import DataFileRecord, DataTable, FileInsights # 原文件不带扩展名:calamine 见到 .xls 等扩展名会按扩展名打开,用户的扩展名和内容不符时会失败 SOURCE_NAME = "source" -def file_dir(root: Path, file_id: str) -> Path: - return root / file_id +def table_path(root: Path, file_id: str, table: DataTable) -> Path: + return root / file_id / table.parquet def _now() -> datetime: @@ -34,10 +35,12 @@ class DataFileStore(Protocol): """未过期且未删除的记录""" ... - async def find( - self, conversation_id: str, *, include_deleted: bool = False - ) -> list[DataFileRecord]: - """会话内未过期的记录,按上传时间排列""" + async def find(self, conversation_id: str) -> list[DataFileRecord]: + """会话内未过期且未删除的记录,按上传时间排列""" + ... + + async def reserve_table_numbers(self, conversation_id: str, count: int) -> int: + """为会话预留 count 个连续的表号,返回第一个;表号从 1 开始,不会重复分配""" ... async def mark_deleted(self, file_id: str) -> None: ... @@ -54,6 +57,8 @@ class DataFileStore(Protocol): class InMemoryDataFileStore: def __init__(self) -> None: self._records: dict[str, DataFileRecord] = {} + # 会话 → 已分配的最大表号 + self._table_numbers: dict[str, int] = {} def _purge_expired(self) -> None: now = _now() @@ -71,18 +76,20 @@ class InMemoryDataFileStore: record = self._records.get(file_id) return None if record is None or record.deleted else record - async def find( - self, conversation_id: str, *, include_deleted: bool = False - ) -> list[DataFileRecord]: + async def find(self, conversation_id: str) -> list[DataFileRecord]: self._purge_expired() records = [ record for record in self._records.values() - if record.conversation_id == conversation_id - and (include_deleted or not record.deleted) + if record.conversation_id == conversation_id and not record.deleted ] return sorted(records, key=lambda record: record.created_at) + async def reserve_table_numbers(self, conversation_id: str, count: int) -> int: + first = self._table_numbers.get(conversation_id, 0) + 1 + self._table_numbers[conversation_id] = first + count - 1 + return first + async def mark_deleted(self, file_id: str) -> None: if (record := self._records.get(file_id)) is not None: self._records[file_id] = record.model_copy(update={"deleted": True}) @@ -101,6 +108,8 @@ class MongoDataFileStore: # 读回的时间统一带 UTC 时区;集合名前缀用于测试隔离 options = CodecOptions(tz_aware=True, tzinfo=UTC) self._files = db.get_collection(f"{prefix}data_files", codec_options=options) + # 每个会话一条:_id 是会话 ID,tables 是已分配的最大表号;会话一直保留,计数器也不过期 + self._table_numbers = db.get_collection(f"{prefix}data_file_table_numbers") async def ensure_indexes(self) -> None: # 过期记录由 TTL 索引删除;TTL 每分钟才扫一次,查询时仍按 expires_at 过滤 @@ -118,15 +127,12 @@ class MongoDataFileStore: ) return None if document is None else _record(document) - async def find( - self, conversation_id: str, *, include_deleted: bool = False - ) -> list[DataFileRecord]: - query: dict[str, Any] = { + async def find(self, conversation_id: str) -> list[DataFileRecord]: + query = { "conversation_id": conversation_id, + "deleted": False, "expires_at": {"$gt": _now()}, } - if not include_deleted: - query["deleted"] = False documents = await ( self._files.find(query) .sort([("created_at", ASCENDING), ("_id", ASCENDING)]) @@ -134,6 +140,15 @@ class MongoDataFileStore: ) return [_record(document) for document in documents] + async def reserve_table_numbers(self, conversation_id: str, count: int) -> int: + document = await self._table_numbers.find_one_and_update( + {"_id": conversation_id}, + {"$inc": {"tables": count}}, + upsert=True, + return_document=ReturnDocument.AFTER, + ) + return document["tables"] - count + 1 + async def mark_deleted(self, file_id: str) -> None: await self._files.update_one({"_id": file_id}, {"$set": {"deleted": True}}) diff --git a/app/datafiles/tables.py b/app/datafiles/tables.py index 25caf5b..3ea1d4b 100644 --- a/app/datafiles/tables.py +++ b/app/datafiles/tables.py @@ -8,8 +8,8 @@ import re from dataclasses import dataclass, field from typing import Any -from app.datafiles.base import MAX_TABLE_ROWS, rows_exceeded, where -from app.datafiles.columns import classify, label_text +from app.datafiles.base import MAX_TABLE_ROWS, brief, rows_exceeded, where +from app.datafiles.columns import as_text, classify from app.datafiles.readers import Grid from app.schemas.data_file import ParseWarning @@ -209,9 +209,7 @@ def _finish(rows: list[list[Any]], piece: _Piece, sheet: str | None) -> RawTable totals = [index for index, row in data if _is_total(row)] if totals: data = [(index, row) for index, row in data if not _is_total(row)] - listed = "、".join(str(index + 1) for index in totals[:5]) + ( - "等" if len(totals) > 5 else "" - ) + listed = brief([str(index + 1) for index in totals]) warnings.append( ParseWarning( code="total_rows_removed", @@ -250,7 +248,7 @@ def _merge_header(top: list[Any], sub: list[Any]) -> list[str]: for upper, lower in zip(top, sub): if upper is not None: carry = upper - parts = [label_text(part) for part in (carry, lower) if part is not None] + parts = [as_text(part) for part in (carry, lower) if part is not None] merged.append("-".join(parts)) return merged @@ -280,7 +278,7 @@ def _quote(row: list[Any]) -> str: value = next((value for value in row if value is not None), None) if value is None: return "" - text = label_text(value) + text = as_text(value) if len(text) > QUOTE_CHARS: text = text[:QUOTE_CHARS] + "…" return f"「{text}」" diff --git a/app/main.py b/app/main.py index b5daa01..38e4040 100644 --- a/app/main.py +++ b/app/main.py @@ -8,6 +8,7 @@ from loguru import logger from app.api.agent import router from app.api.conversations import router as conversation_router +from app.api.files import UploadSizeLimit from app.api.files import limits_router as file_limits_router from app.api.files import router as file_router from app.api.materials import router as material_router @@ -98,6 +99,8 @@ def create_app() -> FastAPI: # 上传限制无需登录;要先于 /files/{file_id} 注册,否则会被当成文件 ID application.include_router(file_limits_router, prefix="/api") application.include_router(file_router, prefix="/api", dependencies=authenticated) + # 上传在读请求体之前先按 Content-Length 检查大小 + application.add_middleware(UploadSizeLimit, path=f"/api{file_router.prefix}") application.include_router(material_router, prefix="/api") return application diff --git a/app/observability/tracing.py b/app/observability/tracing.py index 113e34b..00118d4 100644 --- a/app/observability/tracing.py +++ b/app/observability/tracing.py @@ -55,14 +55,21 @@ def flush_tracing() -> None: logger.opt(exception=exc).warning("Langfuse flush 失败") -async def astream_screen(graph, payload, *, thread_id, session_id=None, user_id=None): - """同一轮流保持同一个追踪上下文,关闭流时结束追踪""" +async def astream_screen( + graph, payload, *, thread_id, session_id=None, user_id=None, context=None +): + """同一轮流保持同一个追踪上下文,关闭流时结束追踪;context 是图的运行上下文,不进检查点""" from contextlib import aclosing result = {} with _screen_trace(payload, thread_id, session_id, user_id) as (config, root): async with aclosing( - graph.astream(payload, config=config, stream_mode=["custom", "values"]) + graph.astream( + payload, + config=config, + context=context, + stream_mode=["custom", "values"], + ) ) as events: async for mode, data in events: if mode == "values": diff --git a/app/processor/graph.py b/app/processor/graph.py index 18372d9..4f5a303 100644 --- a/app/processor/graph.py +++ b/app/processor/graph.py @@ -11,7 +11,7 @@ from app.processor.nodes.intent_node import IntentNode from app.processor.nodes.qa_node import QANode from app.processor.nodes.screen_agent_node import ScreenAgentNode from app.processor.nodes.wait_for_user_node import WaitForUserNode -from app.processor.state import QueryState +from app.processor.state import QueryState, RunContext def build_checkpoint_serde() -> JsonPlusSerializer: @@ -47,7 +47,7 @@ def route_after_wait(state: QueryState) -> str: def build_graph(checkpointer=None): - workflow = StateGraph(QueryState) + workflow = StateGraph(QueryState, context_schema=RunContext) nodes = { "input_node": InputNode(), "intent_node": IntentNode(), diff --git a/app/processor/nodes/qa_node.py b/app/processor/nodes/qa_node.py index 56b127c..c52bd6b 100644 --- a/app/processor/nodes/qa_node.py +++ b/app/processor/nodes/qa_node.py @@ -10,11 +10,10 @@ from langchain_core.messages import AIMessageChunk, HumanMessage from langchain_core.output_parsers import StrOutputParser from langgraph.config import get_stream_writer -from app.agent.data_tools import DataContext from app.agent.qa_agent import QAContext, get_qa_agent from app.processor.base import BaseNode from app.processor.history import history_messages -from app.processor.state import QueryState +from app.processor.state import QueryState, run_data from app.processor.summary import data_file_summary # 模型在一次回答里分几次输出文字(例如查询前后各说一段)时,段落之间空一行 @@ -56,7 +55,7 @@ class QANode(BaseNode): HumanMessage(content=json.dumps(payload, ensure_ascii=False)), ] }, - context=QAContext(data=DataContext(tuple(context.data_files))), + context=QAContext(data=run_data()), stream_mode="messages", ): # 工具结果、工具调用和中间件插入的消息都不推送 diff --git a/app/processor/nodes/screen_agent_node.py b/app/processor/nodes/screen_agent_node.py index b5b0461..24b9bdc 100644 --- a/app/processor/nodes/screen_agent_node.py +++ b/app/processor/nodes/screen_agent_node.py @@ -7,12 +7,11 @@ import json from langchain_core.messages import HumanMessage -from app.agent.data_tools import DataContext from app.agent.screen_agent import AgentReply, get_screen_agent from app.agent.tools import ScreenAgentContext from app.processor.base import BaseNode from app.processor.history import history_messages -from app.processor.state import QueryState +from app.processor.state import QueryState, run_data from app.processor.summary import request_summary from app.resources.component import get_catalog_cache @@ -68,7 +67,7 @@ class ScreenAgentNode(BaseNode): context=ScreenAgentContext( screen=context.screen, catalog=catalog, - data=DataContext(tuple(context.data_files)), + data=run_data(), ), ) if result.get("structured_response") is None: diff --git a/app/processor/state.py b/app/processor/state.py index 30e35a9..38b6bca 100644 --- a/app/processor/state.py +++ b/app/processor/state.py @@ -1,5 +1,9 @@ +from dataclasses import dataclass, field from typing import Literal, NotRequired, TypedDict +from langgraph.runtime import get_runtime + +from app.agent.data_tools import DataContext from app.schemas.actions import Action from app.schemas.base import JsonObject from app.schemas.clarification import Clarification @@ -25,3 +29,21 @@ class QueryState(TypedDict): status: NotRequired[ Literal["running", "waiting_for_input", "completed", "cancelled", "failed"] ] # 后端处理、等待补充、处理完成、用户取消或失败,完成不代表前端已执行 + + +@dataclass(frozen=True) +class RunContext: + """图的运行上下文:本次运行要用、但不进检查点的依赖,每次启动和恢复时由调用方传入""" + + # 本会话的数据文件;状态里的 input_context.data_files 只有摘要,列概况和路径在这里 + data: DataContext = field(default_factory=DataContext) + + +def run_data() -> DataContext: + """当前运行可用的数据表;单独调用节点(不在图里执行)或调用方没传 context 时为空""" + try: + runtime = get_runtime(RunContext) + except RuntimeError: + return DataContext() + context = runtime.context if runtime is not None else None + return context.data if context is not None else DataContext() diff --git a/app/schemas/data_file.py b/app/schemas/data_file.py index 94884ca..aa839e6 100644 --- a/app/schemas/data_file.py +++ b/app/schemas/data_file.py @@ -44,6 +44,41 @@ class ParseWarning(BaseModel): message: str +class TopValue(BaseModel): + value: JsonValue + count: int + + +class ColumnProfile(DataColumn): + """一列的概况,上传时在全表上算出;不适用的统计为空,交给模型时省略""" + + null_ratio: float + # 去重值个数,是估算值 + distinct: int + # 数值列是数字,日期、日期时间列是 ISO 写法的字符串 + min: int | float | str | None = None + max: int | float | str | None = None + # 只有数值列有 + avg: int | float | None = None + q25: int | float | None = None + q50: int | float | None = None + q75: int | float | None = None + # 去重值不多的文本、布尔列才有,按出现次数从多到少 + top_values: list[TopValue] | None = None + # 开头几行里不同的取值 + examples: list[JsonValue] + # 保留为文本的原因等说明 + note: str | None = None + + +class TableProfile(BaseModel): + """表的列概况,给 describe_data 和上传后的分析用""" + + columns: list[ColumnProfile] + # 随机抽的几行,顺序与 columns 一致 + sample: list[list[JsonValue]] + + class DataTableInfo(BaseModel): # 会话内的表名 t1、t2……,模型写 SQL 时用 table: str @@ -60,8 +95,7 @@ class DataTableDetail(DataTableInfo): class DataTable(DataTableDetail): - # 列概况,给 describe_data 用 - profile: dict[str, JsonValue] + profile: TableProfile # 文件目录下的 Parquet 文件名 parquet: str @@ -162,6 +196,8 @@ class FileInsights(BaseModel): version: int = 1 # pending:正在分析;ready:已完成(推荐可能为空);failed:分析失败,文件本身可用 status: Literal["pending", "ready", "failed"] + # pending 时是开始分析的时间;服务启动前开始、至今仍是 pending 的分析已随重启丢失,读取时按失败返回 + started_at: datetime | None = None summary: str = "" suggestions: list[ChartSuggestion] = Field(default_factory=list) columns: list[ColumnInsight] = Field(default_factory=list) @@ -224,19 +260,14 @@ class DataFileLimits(BaseModel): aggregate_max_rows: int -class RequestDataTable(DataTableInfo): - """本次请求可用的表:多出工具要用的列概况和 Parquet 路径,不带预览""" - - profile: dict[str, JsonValue] - path: str - - class RequestDataFile(BaseModel): - """本次请求可用的文件:服务端按会话读取后放进请求,前端传的同名字段会被覆盖""" + """本次请求可用的文件:服务端按会话读取后放进请求,前端传的同名字段会被覆盖 + + 请求会进检查点,所以只放表名、行数和列;列概况和文件路径在图的运行上下文里,见 processor/state.py + """ file_id: str name: str # 本轮消息的 attachments 里带了这个文件,通常是刚上传的 attached_this_turn: bool - tables: list[RequestDataTable] - warnings: list[ParseWarning] + tables: list[DataTableInfo] diff --git a/app/schemas/request.py b/app/schemas/request.py index ede87c7..0eaf587 100644 --- a/app/schemas/request.py +++ b/app/schemas/request.py @@ -1,6 +1,6 @@ """前端设计请求:原始需求、当前大屏及可选补充上下文""" -from pydantic import Field +from pydantic import Field, field_validator from .base import ExtensibleModel, JsonObject, NonBlankString from .screen import Screen @@ -20,3 +20,22 @@ class DesignRequest(ExtensibleModel): # 本条消息附带的材料;{"kind": "data_file", "file_id": "..."} 引用本会话上传的数据文件, # 其他种类只原样进入摘要,不读取文件或远程资源 attachments: list[JsonObject] = Field(default_factory=list) + + @field_validator("attachments") + @classmethod + def _data_files_have_ids(cls, value: list[JsonObject]) -> list[JsonObject]: + for item in value: + if item.get("kind") == "data_file" and not isinstance( + item.get("file_id"), str + ): + raise ValueError("data_file 附件缺少 file_id") + return value + + @property + def data_file_ids(self) -> set[str]: + """本条消息附带的数据文件""" + return { + str(item["file_id"]) + for item in self.attachments + if item.get("kind") == "data_file" + } diff --git a/app/services/agent.py b/app/services/agent.py index e19f19a..852cb05 100644 --- a/app/services/agent.py +++ b/app/services/agent.py @@ -7,13 +7,16 @@ from uuid import uuid4 from langgraph.types import Command from loguru import logger +from app.agent.data_tools import DataContext from app.config import settings from app.core.exceptions import BizException from app.datafiles.service import DataFileService from app.observability.tracing import astream_screen from app.processor.history import trim_history +from app.processor.state import RunContext from app.schemas.agent import AgentResult, ResumeRequest from app.schemas.conversation import Turn +from app.schemas.data_file import DataFileRecord, RequestDataFile from app.schemas.request import DesignRequest from app.services.conversation import ConversationStore @@ -43,8 +46,9 @@ class AgentService: conversation_id, settings.conversation_history_turns ) # 附件引用的文件不属于本会话时在流开始前返回 404 - data_files = await self.data_files.request_files( - user_id, conversation_id, request.attachments + attached = request.data_file_ids + records = await self.data_files.request_files( + user_id, conversation_id, attached ) except BaseException: self.release(thread_id) @@ -64,7 +68,16 @@ class AgentService: settings.history_max_tokens, ) # 同样只信任服务端读取的结果;会话内的文件每一轮都可用,不附带也能查询 - payload["data_files"] = [item.model_dump(mode="json") for item in data_files] + # 请求会进检查点,只放摘要;工具要用的列概况和路径通过运行上下文传给图 + payload["data_files"] = [ + RequestDataFile( + file_id=record.file_id, + name=record.name, + attached_this_turn=record.file_id in attached, + tables=record.tables, + ).model_dump(mode="json") + for record in records + ] return ( thread_id, @@ -76,6 +89,7 @@ class AgentService: }, conversation_id, user_id, + self._run_context(records), ) async def prepare_resume(self, request: ResumeRequest, user_id: int): @@ -117,6 +131,10 @@ class AgentService: ) } ) + # 运行上下文不在检查点里,恢复时按会话重新读取;期间删除或过期的文件工具查不到 + records = await self.data_files.request_files( + user_id, conversation_id, set() + ) # 每次有效补充或取消生成新的输入 ID,任务和会话保持不变 return ( @@ -127,11 +145,15 @@ class AgentService: ), conversation_id, user_id, + self._run_context(records), ) except BaseException: self.release(thread_id) raise + def _run_context(self, records: list[DataFileRecord]) -> RunContext: + return RunContext(data=DataContext(tuple(records), self.data_files.root)) + def release(self, thread_id): self._active.discard(thread_id) self._conversation_of.pop(thread_id, None) @@ -199,7 +221,7 @@ class AgentService: clarification=state.get("clarification") if waiting else None, ) - async def stream(self, thread_id, payload, session_id, user_id=None): + async def stream(self, thread_id, payload, session_id, user_id=None, context=None): """消费图事件,完整状态只用于整理最终结果,不直接暴露给客户端""" state = None with logger.contextualize(thread_id=thread_id): @@ -210,6 +232,7 @@ class AgentService: thread_id=thread_id, session_id=session_id, user_id=user_id, + context=context, ) ) as events: async for mode, data in events: