feat:完成大部分数据文件上传解析功能

This commit is contained in:
2026-10-08 08:23:01 -07:00
parent 05c050d2ba
commit bb6a9405db
27 changed files with 2675 additions and 18 deletions
+126
View File
@@ -0,0 +1,126 @@
"""数据文件的本地工具:describe_data 查看列概况,query_data 执行只读 SQL;问答和 screen agent 共用
工具只能看到 DataContext 里的表,也就是当前用户在本会话里未过期的文件,查询引擎也只为这些表建视图
工具通过 runtime.context.data 读取 DataContext,用到它们的 agent 上下文都要有 data 字段;
上下文类型因 agent 而异,所以写 ToolRuntime[Any](不写类型参数时按 None 处理,序列化会告警)
"""
from dataclasses import dataclass
from pathlib import Path
from typing import Annotated, Any
from langchain.tools import ToolRuntime, tool
from langchain_core.messages import ToolMessage
from pydantic import Field
from app.config import settings
from app.datafiles.query import QueryError, run_query
from app.schemas.data_file import RequestDataFile, RequestDataTable
# 列超过这么多时,不指定 columns 只返回列名和类型,避免一次塞进太多概况
BRIEF_COLUMNS = 50
@dataclass(frozen=True)
class DataContext:
files: tuple[RequestDataFile, ...] = ()
def find(self, name: str) -> tuple[RequestDataFile, RequestDataTable] | None:
for file in self.files:
for table in file.tables:
if table.table == name:
return file, table
return None
def paths(self) -> dict[str, Path]:
return {
table.table: Path(table.path)
for file in self.files
for table in file.tables
}
def _error(runtime: ToolRuntime, message: str) -> ToolMessage:
# 返回给模型修正,不中断 agent
return ToolMessage(message, tool_call_id=runtime.tool_call_id, status="error")
def _table_list(data: DataContext) -> str:
return "、".join(data.paths()) or "无"
@tool
def describe_data(
table: str,
runtime: ToolRuntime[Any],
columns: Annotated[
list[str] | None,
Field(description="只看这些列;不传时看全部列"),
] = None,
) -> dict[str, Any] | ToolMessage:
"""查看一张数据表的列概况(类型、空值率、去重值个数、数值范围和四分位、常见值、样例)、随机样例行,以及解析时做的调整
分析或取数前先看要用的表;表名见请求 data_files 里的 table。统计值来自全表,样例只是其中几行
列很多时只返回列名和类型,再用 columns 指定要看的列
"""
data: DataContext = runtime.context.data
found = data.find(table)
if found is None:
return _error(runtime, f"没有表 {table},可用的表:{_table_list(data)}")
file, item = found
profiled: list[dict[str, Any]] = item.profile["columns"]
result: dict[str, Any] = {
"table": item.table,
"label": item.label,
"rows": item.rows,
}
if columns is None and len(profiled) > BRIEF_COLUMNS:
result["columns"] = [
{key: column[key] for key in ("name", "type", "format") if key in column}
for column in profiled
]
result["note"] = (
f"共 {len(profiled)} 列,只列出列名和类型;用 columns 指定列查看详细概况和样例"
)
else:
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]
# 样例行的值和 columns 一一对应
result["sample_rows"] = [
[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
# 针对这张表的警告,以及针对整个文件的警告(例如跳过的坏行)
result["warnings"] = [
warning.message
for warning in file.warnings
if warning.table in (item.table, None)
]
return result
@tool
def query_data(sql: str, runtime: ToolRuntime[Any]) -> dict[str, Any] | ToolMessage:
"""对上传的数据表执行一条只读 SQL(DuckDB 语法),表名是请求 data_files 里的 table(t1、t2……)
中文或带符号的列名用双引号括起来;宽表转长表用 UNPIVOT,PIVOT 要用 IN 写明取值
最多返回 200 行,truncated 为 true 时结果不完整,改用聚合或 LIMIT
回答和图表里的数字以查询结果为准,不要自己心算
"""
data: DataContext = runtime.context.data
paths = data.paths()
# 任务执行期间被删除或清理的文件跳过,不影响查询其他表
existing = {name: path for name, path in paths.items() if path.exists()}
try:
return run_query(sql, existing, timeout=settings.data_query_timeout).to_dict()
except QueryError as error:
message = str(error)
if gone := [name for name in paths if name not in existing]:
message += f";{'、'.join(gone)} 的文件已删除或过期,不能再查询"
return _error(runtime, message)
DATA_TOOLS = [describe_data, query_data]
+16 -10
View File
@@ -5,7 +5,7 @@
组件目录和物料在物料 MCP,各类任务的做法在 skill;系统提示只放每次都要用、别处放不下的原则
"""
from typing import Literal, Self
from typing import Any, Literal, Self
from langchain.agents import create_agent
from langchain.agents.middleware import ModelCallLimitMiddleware, wrap_model_call
@@ -14,6 +14,7 @@ from langchain_core.messages import AIMessage, HumanMessage
from loguru import logger
from pydantic import BaseModel, Field, model_validator
from app.agent.data_tools import DATA_TOOLS
from app.agent.skills import skills_prompt
from app.agent.tools import (
ScreenAgentContext,
@@ -98,16 +99,22 @@ async def require_tool_call(request, handler):
return response
_agent = None
# 是否挂数据工具 → agent
_agents: dict[bool, Any] = {}
async def get_screen_agent():
"""第一次使用时组装,之后复用;系统提示是常量,不随物料目录变化"""
global _agent
if _agent is None:
_agent = create_agent(
async def get_screen_agent(with_data: bool = False):
"""第一次使用时组装,之后复用;系统提示是常量,不随物料目录变化
会话里有数据文件时用挂了 describe_data、query_data 的实例,没有时不让模型看到用不了的工具
"""
if with_data not in _agents:
tools = [get_nodes, apply_actions, *await get_mcp_tools()]
if with_data:
tools.extend(DATA_TOOLS)
_agents[with_data] = create_agent(
model=get_llm(),
tools=[get_nodes, apply_actions, *await get_mcp_tools()],
tools=tools,
system_prompt=f"{SYSTEM_PROMPT}\n\n{skills_prompt()}",
# 用工具调用给出结论:试过原生结构化输出(ProviderStrategy),歧义时模型更常直接执行而不追问,评测多删了组件
response_format=ToolStrategy(AgentReply),
@@ -124,5 +131,4 @@ async def get_screen_agent():
checkpointer=False,
name="screen_agent",
)
return _agent
return _agents[with_data]
+3 -1
View File
@@ -3,7 +3,7 @@
它们读写的是本次请求的大屏和 agent 状态,所以放在本地;物料等能力由 MCP 提供
"""
from dataclasses import dataclass
from dataclasses import dataclass, field
from typing import Annotated, NotRequired
from langchain.agents import AgentState
@@ -13,6 +13,7 @@ from langgraph.types import Command
from pydantic import Field
from app.actions.batch import BatchError, run_batch
from app.agent.data_tools import DataContext
from app.schemas.actions import Action, AddAction, AgentAction
from app.schemas.component_resource import ComponentCatalog
from app.schemas.screen import Screen
@@ -22,6 +23,7 @@ from app.schemas.screen import Screen
class ScreenAgentContext:
screen: Screen # 本次请求的原始大屏,只读
catalog: ComponentCatalog # 当前物料目录,只用于补全和校验
data: DataContext = field(default_factory=DataContext) # 本会话可用的数据表
def _latest(_, new):
+95
View File
@@ -0,0 +1,95 @@
"""数据文件接口:上传限制、上传解析、列出、查看、删除,见 docs/screen-agent/data-files.md 第 5 节"""
from typing import Annotated, get_args
from fastapi import APIRouter, Depends, File, Form, Query, Request, UploadFile
from app.common.response import ApiResponse, success_response
from app.core.auth import get_current_user
from app.datafiles.service import MAX_FILES_PER_CONVERSATION, DataFileService
from app.schemas.base import NonBlankString
from app.schemas.data_file import (
DataFileDetail,
DataFileInfo,
DataFileLimits,
DataFileList,
DataFileRecord,
DataFormat,
)
from app.schemas.user import CurrentUser
# 上传限制不需要登录,单独一个路由;要先于 /files/{file_id} 注册
limits_router = APIRouter(prefix="/files", tags=["DataFile"])
router = APIRouter(prefix="/files", tags=["DataFile"])
# 文件选择框的过滤条件;后端按文件内容识别格式,不看扩展名
ACCEPT = ".csv,.tsv,.txt,.xlsx,.xls,.xlsb,.json,.jsonl,.ndjson,.parquet"
def _service(request: Request) -> DataFileService:
return request.app.state.data_file_service
def _detail(record: DataFileRecord) -> DataFileDetail:
# 按接口模型重新校验,去掉 owner_id、sha256、列概况等只在后端用的字段
return DataFileDetail.model_validate(record.model_dump())
@limits_router.get("/limits")
async def get_limits(request: Request) -> ApiResponse[DataFileLimits]:
return success_response(
DataFileLimits(
max_size_mb=_service(request).max_mb,
max_files_per_conversation=MAX_FILES_PER_CONVERSATION,
formats=list(get_args(DataFormat)),
accept=ACCEPT,
)
)
@router.post("")
async def upload_file(
request: Request,
conversation_id: Annotated[NonBlankString, Form()],
file: Annotated[UploadFile, File()],
user: Annotated[CurrentUser, Depends(get_current_user)],
) -> ApiResponse[DataFileDetail]:
record = await _service(request).upload(
user.id, conversation_id, file.filename, file.file
)
return success_response(_detail(record))
@router.get("")
async def list_files(
request: Request,
conversation_id: Annotated[NonBlankString, Query()],
user: Annotated[CurrentUser, Depends(get_current_user)],
) -> ApiResponse[DataFileList]:
records = await _service(request).list_files(user.id, conversation_id)
return success_response(
DataFileList(
files=[
DataFileInfo.model_validate(record.model_dump()) for record in records
]
)
)
@router.get("/{file_id}")
async def get_file(
file_id: str,
request: Request,
user: Annotated[CurrentUser, Depends(get_current_user)],
) -> ApiResponse[DataFileDetail]:
return success_response(_detail(await _service(request).get_file(user.id, file_id)))
@router.delete("/{file_id}")
async def delete_file(
file_id: str,
request: Request,
user: Annotated[CurrentUser, Depends(get_current_user)],
) -> ApiResponse[dict]:
await _service(request).delete_file(user.id, file_id)
return success_response()
+29 -1
View File
@@ -4,7 +4,9 @@ from pathlib import Path
from dotenv import load_dotenv
load_dotenv(Path(__file__).resolve().parent.parent / ".env")
PROJECT_ROOT = Path(__file__).resolve().parent.parent
load_dotenv(PROJECT_ROOT / ".env")
@dataclass(frozen=True)
@@ -87,6 +89,26 @@ class Settings:
default_factory=lambda: int(os.getenv("MONGODB_CHECKPOINT_TTL", "604800"))
)
# 上传的数据文件:原文件和解析出的 Parquet 放这个目录,相对路径相对项目根目录;容器部署时挂载出来
data_file_dir: Path = field(
default_factory=lambda: (
PROJECT_ROOT / (os.getenv("DATA_FILE_DIR") or "data/files")
)
)
# 单个数据文件的大小上限(MB)
data_file_max_mb: int = field(
default_factory=lambda: int(os.getenv("DATA_FILE_MAX_MB", "50"))
)
# 数据文件保留时长(秒),从上传时算起,默认 7 天
data_file_ttl: int = field(
default_factory=lambda: int(os.getenv("DATA_FILE_TTL", "604800"))
)
# 模型写的单条数据查询的超时(秒)
data_query_timeout: float = field(
default_factory=lambda: float(os.getenv("DATA_QUERY_TIMEOUT", "10"))
)
# 用户中心:请求携带的 token 交给它校验,换回 user_id
user_center_url: str = field(
default_factory=lambda: (
@@ -125,6 +147,12 @@ class Settings:
raise ValueError("MAX_AGENT_STEPS 必须是正整数")
if self.mcp_tool_timeout <= 0:
raise ValueError("MCP_TOOL_TIMEOUT 必须是正数")
if self.data_file_max_mb < 1:
raise ValueError("DATA_FILE_MAX_MB 必须是正整数")
if self.data_file_ttl < 1:
raise ValueError("DATA_FILE_TTL 必须是正整数")
if self.data_query_timeout <= 0:
raise ValueError("DATA_QUERY_TIMEOUT 必须是正数")
if not 0 <= self.intent_qa_min_probability <= 1:
raise ValueError("INTENT_QA_MIN_PROBABILITY 必须在 0 到 1 之间")
View File
+32
View File
@@ -0,0 +1,32 @@
"""数据文件各环节共用的异常和上限"""
# 每张表的数据行数上限
MAX_TABLE_ROWS = 1_000_000
UNSUPPORTED_FORMAT = "不支持的文件格式,支持 CSV、TSV、Excel、JSON、JSONL、Parquet"
class DataFileError(Exception):
"""文件无法接收或解析;status 是返回给调用方的 HTTP 状态码,message 可以直接展示
解析在子进程里执行,异常要能跨进程传递,所以参数原样放进 args
"""
def __init__(self, status: int, message: str):
super().__init__(status, message)
self.status = status
self.message = message
def __str__(self) -> str:
return self.message
def where(sheet: str | None) -> str:
"""警告和错误信息的位置前缀:Excel 写 sheet 名,其他格式不写"""
return f"{sheet}:" if sheet else ""
def rows_exceeded(sheet: str | None) -> DataFileError:
return DataFileError(
422, f"{where(sheet)}超过 {MAX_TABLE_ROWS // 10_000} 万行,请拆分后上传"
)
+349
View File
@@ -0,0 +1,349 @@
"""列名清理和类型推断,输出 Arrow 列
CSV 的单元格全是字符串;Excel 的数值、日期单元格直接采用,文本单元格再按规则推断
一列的非空值必须全部符合同一类型才转换,否则保留文本:宁可让模型在 SQL 里 TRY_CAST,也不悄悄丢掉值
"""
import json
import math
import re
from collections import Counter
from collections.abc import Sequence
from dataclasses import dataclass
from datetime import UTC, date, datetime, time
from typing import Any
import pyarrow as pa
from app.datafiles.base import where
from app.schemas.data_file import DataColumn, ParseWarning
NULL_TOKENS = {"", "-", "--", "—", "——", "N/A", "n/a", "NULL", "null", "#N/A"}
EXCEL_ERROR = re.compile(r"^#(?:NULL!|DIV/0!|VALUE!|REF!|NAME\?|NUM!|N/A)$")
CURRENCY_SYMBOLS = "¥¥$€"
NUMBER = re.compile(r"^[+-]?(?:\d+|\d{1,3}(?:,\d{3})+)(?:\.\d+)?(?:[eE][+-]?\d+)?$")
INTEGER = re.compile(r"^[+-]?\d+$")
# 16 位及以上的纯数字是证件号、卡号一类的编号,转成数值会丢失精度
MIN_ID_DIGITS = 16
# 按年份后的第一个字符选格式;%m、%d 也接受一位数
DATE_FORMATS = {
"-": ("%Y-%m-%d", "%Y-%m"),
"/": ("%Y/%m/%d", "%Y/%m"),
".": ("%Y.%m.%d",),
"年": ("%Y年%m月%d日", "%Y年%m月"),
}
DATETIME_FORMATS = (
"%Y-%m-%d %H:%M:%S",
"%Y-%m-%d %H:%M",
"%Y/%m/%d %H:%M:%S",
"%Y/%m/%d %H:%M",
)
INT64_RANGE = range(-(2**63), 2**63)
# 超过这个范围的整数存成 float 会丢精度
EXACT_FLOAT_INT = 2**53
NUMERIC_KINDS = ("int", "float", "percent")
DATE_KINDS = ("date", "datetime")
@dataclass
class TypedColumns:
columns: list[DataColumn]
table: pa.Table
# 列名 → 写进列概况的说明,例如保留为文本的原因
notes: dict[str, str]
def clean_names(
raw: Sequence[Any], sheet: str | None
) -> tuple[list[str], ParseWarning | None]:
"""去掉空白和换行;空列名按位置命名,重复列名加后缀"""
names, changes, seen = [], [], set()
for position, value in enumerate(raw, start=1):
name = re.sub(r"\s*[\r\n]+\s*", " ", label_text(value)).strip()
if not name:
name = f"列{position}"
changes.append(f"第 {position} 列没有列名,命名为「{name}」")
if name in seen:
suffix = 2
while f"{name}_{suffix}" in seen:
suffix += 1
changes.append(f"「{name}」重名,第 {position} 列改为「{name}_{suffix}」")
name = f"{name}_{suffix}"
seen.add(name)
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)
def build_columns(
names: list[str], values: list[Sequence[Any]], sheet: str | None
) -> tuple[TypedColumns, ParseWarning | None]:
"""按列推断类型;有列因少数无法解析的值保留为文本时,返回一条 kept_as_text 警告"""
columns, arrays, notes = [], [], {}
for name, column_values in zip(names, values):
column, array, note = infer_column(name, column_values)
columns.append(column)
arrays.append(array)
if note:
notes[name] = note
warning = None
if notes:
listed = "、".join(f"「{name}」" for name in notes)
warning = ParseWarning(
code="kept_as_text",
message=f"{where(sheet)}{listed}有无法解析的值,整列按文本处理",
)
return TypedColumns(
columns, pa.Table.from_arrays(arrays, names=names), notes
), warning
def arrow_columns(table: pa.Table) -> tuple[TypedColumns, ParseWarning | None]:
"""Parquet 自带类型,只映射到接口类型;列表、嵌套结构等存成 JSON 字符串"""
names, warning = clean_names(table.column_names, None)
columns, arrays = [], []
for name, array in zip(names, table.columns):
kind = _arrow_kind(array.type)
if kind is None:
values = array.to_pylist()
array = pa.array(
[None if v is None else as_text(v) for v in values], pa.string()
)
kind = "text"
elif kind == "text" and pa.types.is_dictionary(array.type):
array = array.cast(pa.string())
columns.append(DataColumn(name=name, type=kind))
arrays.append(array)
return TypedColumns(columns, pa.Table.from_arrays(arrays, names=names), {}), warning
def _arrow_kind(arrow_type: pa.DataType) -> str | None:
types = pa.types
if (
types.is_integer(arrow_type)
or types.is_floating(arrow_type)
or types.is_decimal(arrow_type)
):
return "number"
if types.is_boolean(arrow_type):
return "boolean"
if types.is_date(arrow_type):
return "date"
if types.is_timestamp(arrow_type):
return "datetime"
if (
types.is_string(arrow_type)
or types.is_large_string(arrow_type)
or types.is_string_view(arrow_type)
):
return "text"
if types.is_dictionary(arrow_type) and _arrow_kind(arrow_type.value_type) == "text":
return "text"
return None
def infer_column(
name: str, values: Sequence[Any]
) -> tuple[DataColumn, pa.Array, str | None]:
value_kinds, parsed = [], []
for value in values:
kind, result = classify(value)
value_kinds.append(kind)
parsed.append(result)
kinds = Counter(kind for kind in value_kinds if kind != "null")
present = set(kinds)
if present and present <= {"bool"}:
return DataColumn(name=name, type="boolean"), pa.array(parsed, pa.bool_()), None
if present and present <= {"int", "float"}:
if all(
value is None
or isinstance(value, int)
or (value.is_integer() and abs(value) < EXACT_FLOAT_INT)
for value in parsed
):
ints = [None if value is None else int(value) for value in parsed]
return (
DataColumn(name=name, type="number"),
pa.array(ints, pa.int64()),
None,
)
return (
DataColumn(name=name, type="number"),
pa.array(parsed, pa.float64()),
None,
)
if present == {"percent"}:
column = DataColumn(name=name, type="number", format="percent")
return column, pa.array(parsed, pa.float64()), None
if present and present <= {"date"}:
return DataColumn(name=name, type="date"), pa.array(parsed, pa.date32()), None
if present and present <= set(DATE_KINDS):
stamps = [
datetime.combine(value, time()) if type(value) is date else value
for value in parsed
]
return (
DataColumn(name=name, type="datetime"),
pa.array(stamps, pa.timestamp("us")),
None,
)
texts = [
None if result is None else as_text(value)
for value, result in zip(values, parsed)
]
note = _text_note(texts, value_kinds, kinds)
return DataColumn(name=name, type="text"), pa.array(texts, pa.string()), note
def classify(value: Any) -> tuple[str, Any]:
"""单个值的类型和转换后的值;类型是 null、bool、int、float、percent、date、datetime、text 之一"""
if value is None:
return "null", None
if isinstance(value, bool):
return "bool", value
if isinstance(value, int):
return ("int", value) if value in INT64_RANGE else ("text", str(value))
if isinstance(value, float):
return (
("null", None)
if math.isnan(value) or math.isinf(value)
else ("float", value)
)
if isinstance(value, datetime):
# Excel 只设了日期格式的单元格读出来是零点的 datetime
if value.time() == time() and value.tzinfo is None:
return "date", value.date()
return "datetime", _naive(value)
if isinstance(value, date):
return "date", value
if not isinstance(value, str):
return "text", as_text(value)
text = value.strip()
if text in NULL_TOKENS or EXCEL_ERROR.match(text):
return "null", None
if text.lower() in ("true", "false"):
return "bool", text.lower() == "true"
number = parse_number(text)
if number is not None:
return ("int" if isinstance(number, int) else "float"), number
if text.endswith("%"):
number = parse_number(text[:-1].rstrip())
if number is not None:
return "percent", number / 100
moment = parse_datetime(text)
if moment is not None:
return ("datetime" if isinstance(moment, datetime) else "date"), moment
return "text", text
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)
if text and text[0] in CURRENCY_SYMBOLS:
text = text[1:].lstrip()
if not NUMBER.match(text):
return None
digits = text.lstrip("+-")
if len(digits) > 1 and digits[0] == "0" and digits[1].isdigit():
return None
if digits.isdigit() and len(digits) >= MIN_ID_DIGITS:
return None
text = text.replace(",", "")
if INTEGER.match(text):
return int(text)
return float(text)
def parse_datetime(text: str) -> date | datetime | None:
if len(text) < 6 or not text[:4].isdigit():
return None
# 标准 ISO 写法先用 fromisoformat,比 strptime 快一个数量级
if ":" in text:
try:
return _naive(datetime.fromisoformat(text))
except ValueError:
pass
for pattern in DATETIME_FORMATS:
try:
return datetime.strptime(text, pattern)
except ValueError:
continue
return None
if len(text) == 10 and text[4] == "-":
try:
return date.fromisoformat(text)
except ValueError:
pass
for pattern in DATE_FORMATS.get(text[4], ()):
try:
return datetime.strptime(text, pattern).date()
except ValueError:
continue
return None
def as_text(value: Any) -> str:
if isinstance(value, str):
return value.strip()
if isinstance(value, bool):
return "true" if value else "false"
if isinstance(value, float) and value.is_integer():
return str(int(value))
if isinstance(value, datetime):
return value.isoformat(sep=" ")
if isinstance(value, (date, time)):
return value.isoformat()
if isinstance(value, (dict, list)):
return json.dumps(value, ensure_ascii=False)
return str(value)
def _naive(value: datetime) -> datetime:
# 带时区的时间统一换成 UTC 再去掉时区,同一列里的值才能比较
if value.tzinfo is None:
return value
return value.astimezone(UTC).replace(tzinfo=None)
def _text_note(
texts: list[str | None], value_kinds: list[str], kinds: Counter[str]
) -> str | None:
"""大部分值是数值或日期、少数无法解析时,说明原因,模型可以在 SQL 里用 TRY_CAST"""
failed = kinds["text"]
numeric = sum(kinds[kind] for kind in NUMERIC_KINDS)
dates = sum(kinds[kind] for kind in DATE_KINDS)
if not failed or max(numeric, dates) < failed:
return None
examples: list[str] = []
for text, kind in zip(texts, value_kinds):
if kind == "text" and text not in examples:
examples.append(text)
if len(examples) == 2:
break
shown = "、".join(f"「{example}」" for example in examples)
label = "数值" if numeric >= dates else "日期"
return f"看起来是{label},有 {failed} 个值无法解析,例如{shown}"
+143
View File
@@ -0,0 +1,143 @@
"""按文件内容识别格式和文本编码,不看扩展名"""
import codecs
import csv
import json
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from app.datafiles.base import UNSUPPORTED_FORMAT, DataFileError
from app.schemas.data_file import DataFormat
# xlsx/xlsb 解压后的总大小上限,防止解压炸弹
MAX_UNZIPPED_BYTES = 1024 * 1024 * 1024
# 带 BOM 的文本直接按 BOM 解码;UTF-32 的 BOM 以 UTF-16 的开头,解出来含 NUL,按二进制拒绝
BOMS = (
(codecs.BOM_UTF8, "utf-8", "utf-8"),
(codecs.BOM_UTF16_LE, "utf-16-le", "utf-16"),
(codecs.BOM_UTF16_BE, "utf-16-be", "utf-16"),
)
# 没有 BOM 时依次尝试;GB18030 兼容 GBK、GB2312,中文 Windows 上 Excel 另存的 CSV 多是这种
FALLBACK_ENCODINGS = ("utf-8", "gb18030")
DELIMITERS = (",", "\t", ";", "|")
# 判断分隔符只看开头这么多行
SNIFF_LINES = 50
@dataclass(frozen=True)
class Detected:
format: DataFormat
encoding: str | None = None
# 文本格式解码后的内容
text: str | None = None
# CSV/TSV 的分隔符
delimiter: str | None = None
# JSON 已解析的内容,避免读取时再解析一遍
value: Any = None
def detect(path: Path) -> Detected:
with path.open("rb") as file:
head = file.read(8)
if head.startswith(b"PK\x03\x04"):
return Detected(_zip_format(path))
if head.startswith(b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1"):
return Detected("xls")
if head.startswith(b"PAR1"):
return Detected("parquet")
text, encoding = decode_text(path.read_bytes())
stripped = text.lstrip()
if not stripped:
raise DataFileError(422, "文件是空的")
if stripped[0] in "[{":
return _json_format(text, encoding)
delimiter = sniff_delimiter(text)
return Detected("tsv" if delimiter == "\t" else "csv", encoding, text, delimiter)
def _zip_format(path: Path) -> DataFormat:
try:
with zipfile.ZipFile(path) as archive:
infos = archive.infolist()
except zipfile.BadZipFile:
raise DataFileError(415, UNSUPPORTED_FORMAT) from None
if sum(info.file_size for info in infos) > MAX_UNZIPPED_BYTES:
raise DataFileError(413, "Excel 文件解压后过大")
names = {info.filename for info in infos}
if "xl/workbook.xml" in names:
return "xlsx"
if "xl/workbook.bin" in names:
return "xlsb"
raise DataFileError(415, UNSUPPORTED_FORMAT)
def decode_text(data: bytes) -> tuple[str, str]:
"""返回解码后的文本和编码名"""
for bom, codec, name in BOMS:
if data.startswith(bom):
try:
text = data[len(bom) :].decode(codec)
except UnicodeDecodeError:
raise DataFileError(
422, "无法识别文件编码,请另存为 UTF-8 后重新上传"
) from None
return _reject_binary(text), name
for encoding in FALLBACK_ENCODINGS:
try:
return _reject_binary(data.decode(encoding)), encoding
except UnicodeDecodeError:
continue
raise DataFileError(422, "无法识别文件编码,请另存为 UTF-8 后重新上传")
def _reject_binary(text: str) -> str:
# 文本文件不会有 NUL;图片、PDF 等二进制文件偶尔能按 GB18030 解码成功,在这里拦下
if "\x00" in text:
raise DataFileError(415, UNSUPPORTED_FORMAT)
return text
def _json_format(text: str, encoding: str) -> Detected:
try:
return Detected("json", encoding, text, value=json.loads(text))
except ValueError as error:
json_error = error
first = next(line for line in text.splitlines() if line.strip())
try:
if isinstance(json.loads(first), dict):
return Detected("jsonl", encoding, text)
except ValueError:
pass
# 以 [ 或 { 开头却不是合法 JSON,按 JSON 报错比当成 CSV 解析出一堆乱码更有用
line = getattr(json_error, "lineno", None)
position = f"第 {line} 行" if line else ""
raise DataFileError(
422, f"JSON 格式错误:{position}{getattr(json_error, 'msg', json_error)}"
)
def sniff_delimiter(text: str) -> str:
"""选出让开头各行列数最一致、列数最多的分隔符;都只有一列时按逗号处理"""
lines = []
for line in text.splitlines():
if line.strip():
lines.append(line)
if len(lines) == SNIFF_LINES:
break
best, best_score = ",", (0.0, 1)
for delimiter in DELIMITERS:
widths = [len(row) for row in csv.reader(lines, delimiter=delimiter)]
width = max(set(widths), key=widths.count)
if width < 2:
continue
score = (widths.count(width) / len(widths), width)
if score > best_score:
best, best_score = delimiter, score
return best
+141
View File
@@ -0,0 +1,141 @@
"""解析入口:识别格式 → 读取 → 拆表 → 推断类型 → 每张表写一个 Parquet → 列概况
只读写本地文件,不碰数据库;上传流程在子进程里调用(worker.py),超时直接终止子进程
"""
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import pyarrow.parquet as pq
from pydantic import BaseModel
from app.datafiles.base import DataFileError
from app.datafiles.columns import (
TypedColumns,
arrow_columns,
build_columns,
clean_names,
)
from app.datafiles.detect import detect
from app.datafiles.profile import profile_table
from app.datafiles.readers import (
Grid,
Records,
read_delimited,
read_excel,
read_json,
read_jsonl,
read_parquet,
)
from app.datafiles.tables import split_tables
from app.schemas.data_file import DataColumn, DataFormat, ParseWarning
# 子进程把解析结果转成 JSON 交回上传流程,所以用 pydantic 模型
class ParsedTable(BaseModel):
sheet: str | None
# 同一 sheet 切出多张表时的序号,从 1 开始
part: int | None
rows: int
columns: list[DataColumn]
path: Path
profile: dict[str, Any]
preview: list[list[Any]]
warnings: list[ParseWarning]
class ParsedFile(BaseModel):
format: DataFormat
encoding: str | None
tables: list[ParsedTable]
# 不属于某张表的警告,例如跳过的 sheet
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……"""
detected = detect(path)
warnings: list[ParseWarning] = []
built: list[_Built] = []
def add_grid(grid: Grid) -> None:
tables, dropped = split_tables(grid)
warnings.extend(dropped)
for table in tables:
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)],
)
)
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)))
match detected.format:
case "csv" | "tsv":
add_grid(read_delimited(detected.text, detected.delimiter))
case "xlsx" | "xls" | "xlsb":
grids, sheet_warnings = read_excel(path)
warnings.extend(sheet_warnings)
for grid in grids:
add_grid(grid)
case "json":
result = read_json(detected.value)
if isinstance(result, Grid):
add_grid(result)
else:
add_records(result)
case "jsonl":
add_records(read_jsonl(detected.text))
case "parquet":
typed, renamed = arrow_columns(read_parquet(path))
if typed.table.num_rows:
built.append(_Built(None, None, typed, _present(renamed)))
if not built:
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,
tables=tables,
warnings=warnings,
)
def _present(*warnings: ParseWarning | None) -> list[ParseWarning]:
return [warning for warning in warnings if warning is not None]
+126
View File
@@ -0,0 +1,126 @@
"""列概况和预览:上传时在 Parquet 上用 DuckDB 计算,存进文件记录,供 describe_data 和接口返回"""
import math
from pathlib import Path
from typing import Any
import duckdb
from app.datafiles.query import connect, json_value, quote_identifier
from app.schemas.data_file import DataColumn
PREVIEW_ROWS = 20
SAMPLE_ROWS = 5
EXAMPLES = 3
TOP_VALUES = 5
# 去重值不超过这么多的文本、布尔列才给出现次数最多的值
TOP_VALUES_MAX_DISTINCT = 50
# 样例值从开头这么多行里取
EXAMPLE_SCAN_ROWS = 1000
PROFILE_CELL_CHARS = 50
PREVIEW_CELL_CHARS = 200
def profile_table(
path: Path, columns: list[DataColumn], notes: dict[str, str]
) -> tuple[dict[str, Any], 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()}
head = connection.execute(f"FROM t LIMIT {EXAMPLE_SCAN_ROWS}").fetchall()
profiled = [
_profile_column(
connection,
column,
summary[column.name],
head,
position,
notes.get(column.name),
)
for position, column in enumerate(columns)
]
sample = connection.execute(
f"FROM t USING SAMPLE reservoir({SAMPLE_ROWS} ROWS) REPEATABLE (42)"
).fetchall()
finally:
connection.close()
profile = {
"rows": rows,
"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]
]
return profile, preview
def _profile_column(
connection: duckdb.DuckDBPyConnection,
column: DataColumn,
stats: dict[str, Any],
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))
)
if column.type == "number":
for key in ("min", "max", "avg", "q25", "q50", "q75"):
item[key] = _number(stats[key])
elif column.type in ("date", "datetime"):
# SUMMARIZE 的时间用空格分隔,和预览里的 ISO 写法统一
item["min"], item["max"] = (
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
):
name = quote_identifier(column.name)
counts = connection.execute(
f"SELECT {name}, count(*) AS n FROM t WHERE {name} IS NOT NULL "
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}
for value, count in counts
]
examples: list[Any] = []
for row in head:
value = json_value(row[position], PROFILE_CELL_CHARS)
if value is not None and value not in examples:
examples.append(value)
if len(examples) == EXAMPLES:
break
item["examples"] = examples
if note:
item["note"] = note
return item
def _number(text: str | None) -> int | float | None:
"""SUMMARIZE 的统计值是字符串;整数不带小数点,其余保留 6 位小数;Parquet 里的 inf 返回空"""
if text is None:
return None
value = float(text)
if not math.isfinite(value):
return None
if value.is_integer() and abs(value) < 2**53:
return int(value)
return round(value, 6)
+168
View File
@@ -0,0 +1,168 @@
"""执行模型写的 SQL:连接只能读本次可见的 Parquet 文件,只允许单条 SELECT
每次查询新建一个内存连接:
- allowed_paths 只放这些 Parquet 文件;不用 allowed_directories,否则同目录下的其他文件能被 glob 读到,
COPY、ATTACH 也能往目录里写
- 关闭外部访问、扩展加载和 Python 变量替换扫描,最后锁定配置,SQL 里的 SET 一律失败
"""
import math
import threading
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import date, datetime, time, timedelta
from decimal import Decimal
from pathlib import Path
from typing import Any
import duckdb
MEMORY_LIMIT = "512MB"
THREADS = 2
MAX_ROWS = 200
MAX_CELL_CHARS = 200
class QueryError(ValueError):
"""查询失败的原因,原样返回给模型修正"""
@dataclass(frozen=True)
class QueryResult:
columns: list[str]
rows: list[list[Any]]
# 结果超过 max_rows 行,只返回了前 max_rows 行
truncated: bool
def to_dict(self) -> dict[str, Any]:
return {
"columns": self.columns,
"rows": self.rows,
"row_count": len(self.rows),
"truncated": self.truncated,
}
def run_query(
sql: str,
tables: Mapping[str, Path],
*,
timeout: float,
max_rows: int = MAX_ROWS,
) -> QueryResult:
"""tables 是表名到 Parquet 文件的映射,SQL 里只能用这些表名"""
try:
connection = connect(tables)
except duckdb.Error as error:
# 任务执行期间文件被删除或过期清理
raise QueryError(
f"数据表无法读取,文件可能已被删除:{_first_line(error)}"
) from None
try:
_check_statement(connection, sql, tables)
timer = threading.Timer(timeout, connection.interrupt)
timer.start()
try:
connection.execute(sql)
columns = [column[0] for column in connection.description]
rows = connection.fetchmany(max_rows + 1)
except duckdb.InterruptException:
raise QueryError(
f"查询超过 {timeout:g} 秒被中止,请缩小范围或先聚合"
) from None
except duckdb.PermissionException:
raise QueryError(f"只能查询上传的数据表:{_table_list(tables)}") from None
except duckdb.Error as error:
raise QueryError(
f"{_first_line(error)};可用的表:{_table_list(tables)}"
) from None
finally:
timer.cancel()
finally:
connection.close()
return QueryResult(
columns,
[[json_value(value) for value in row] for row in rows[:max_rows]],
len(rows) > max_rows,
)
def connect(tables: Mapping[str, Path]) -> duckdb.DuckDBPyConnection:
"""建好视图并锁定的连接;列概况也用它,和模型查询看到的数据一致"""
connection = duckdb.connect(
":memory:",
config={
"autoinstall_known_extensions": False,
"autoload_known_extensions": False,
"python_enable_replacements": False,
"memory_limit": MEMORY_LIMIT,
"threads": THREADS,
},
)
try:
paths = [str(Path(path).resolve()) for path in tables.values()]
for name, path in zip(tables, paths):
connection.execute(
f"CREATE VIEW {quote_identifier(name)} AS FROM read_parquet({_literal(path)})"
)
if paths:
connection.execute("SET allowed_paths = ?", [paths])
connection.execute("SET enable_external_access = false")
connection.execute("SET lock_configuration = true")
except BaseException:
connection.close()
raise
return connection
def quote_identifier(name: str) -> str:
return '"' + name.replace('"', '""') + '"'
def json_value(value: Any, max_chars: int = MAX_CELL_CHARS) -> Any:
"""查询结果转成 JSON 能表示的值;长文本截断"""
if value is None or isinstance(value, (bool, int)):
return value
if isinstance(value, float):
return None if math.isnan(value) or math.isinf(value) else value
if isinstance(value, Decimal):
return float(value)
if isinstance(value, str):
return value if len(value) <= max_chars else value[:max_chars] + "…"
if isinstance(value, (datetime, date, time)):
return value.isoformat()
if isinstance(value, (list, tuple)):
return [json_value(item, max_chars) for item in value]
if isinstance(value, dict):
return {str(key): json_value(item, max_chars) for key, item in value.items()}
if isinstance(value, timedelta):
return str(value)
return json_value(str(value), max_chars)
def _check_statement(
connection: duckdb.DuckDBPyConnection, sql: str, tables: Mapping[str, Path]
) -> None:
try:
statements = connection.extract_statements(sql)
except duckdb.Error as error:
raise QueryError(
f"{_first_line(error)};可用的表:{_table_list(tables)}"
) from None
# PIVOT 不写 IN 时会先建一个枚举类型,解析出来是两条语句
if len(statements) != 1 or statements[0].type != duckdb.StatementType.SELECT:
raise QueryError(
"只能执行一条 SELECT 查询(可以用 WITH、UNPIVOT;PIVOT 要用 IN 写明取值)"
)
def _literal(text: str) -> str:
return "'" + text.replace("'", "''") + "'"
def _first_line(error: Exception) -> str:
return str(error).strip().splitlines()[0]
def _table_list(tables: Mapping[str, Path]) -> str:
return "、".join(tables) if tables else "(无)"
+240
View File
@@ -0,0 +1,240 @@
"""把各种格式读成原始网格(CSV、Excel)或按列的记录(JSON、JSONL),Parquet 直接读成 Arrow 表
CSV 和 Excel 可能有标题行、多行表头、合计行,所以先读成不带表头的网格,再由 tables.py 识别表区域
"""
import csv
import io
import json
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any
import pyarrow as pa
import pyarrow.parquet as pq
from python_calamine import (
CalamineWorkbook,
PasswordError,
SheetTypeEnum,
SheetVisibleEnum,
)
from app.datafiles.base import MAX_TABLE_ROWS, DataFileError, rows_exceeded
from app.schemas.data_file import ParseWarning
# 原始网格比数据多出的标题、表头、空行,读取时先粗略拦住明显超限的文件,精确的行数在拆表后检查
ROW_SLACK = 1000
# JSON 嵌套对象展开的层数,更深的部分存成 JSON 字符串
JSON_FLATTEN_LEVELS = 2
# JSONL 坏行超过这个比例时整个文件判为失败
MAX_BAD_LINE_RATIO = 0.01
# 单元格最长的字符数;csv 模块默认 128KB,长文本单元格会报错
csv.field_size_limit(16 * 1024 * 1024)
@dataclass
class Grid:
"""不带表头的原始网格;rows[i] 是文件或 sheet 中的第 i + 1 行"""
sheet: str | None
rows: list[list[Any]]
@dataclass
class Records:
"""已有列结构的数据,按列存放"""
names: list[str]
columns: list[list[Any]]
warnings: list[ParseWarning] = field(default_factory=list)
def read_delimited(text: str, delimiter: str) -> Grid:
rows = []
for row in csv.reader(io.StringIO(text, newline=""), delimiter=delimiter):
rows.append(row)
if len(rows) > MAX_TABLE_ROWS + ROW_SLACK:
raise rows_exceeded(None)
return Grid(None, rows)
def read_excel(path: Path) -> tuple[list[Grid], list[ParseWarning]]:
try:
book = CalamineWorkbook.from_path(str(path))
except PasswordError:
raise DataFileError(422, "Excel 文件设置了密码,请去掉密码后重新上传") from None
except Exception:
raise DataFileError(422, "无法读取 Excel 文件,文件可能已损坏") from None
grids, warnings = [], []
try:
for meta in book.sheets_metadata:
# 图表页、宏表等没有单元格数据
if meta.typ != SheetTypeEnum.WorkSheet:
continue
if meta.visible != SheetVisibleEnum.Visible:
warnings.append(
ParseWarning(
code="hidden_sheet_skipped",
message=f"跳过了隐藏的 sheet「{meta.name}」",
)
)
continue
sheet = book.get_sheet_by_name(meta.name)
if sheet.total_height > MAX_TABLE_ROWS + ROW_SLACK:
raise rows_exceeded(meta.name)
# 保留左上角的空白区域,行号和 Excel 里看到的一致
rows = sheet.to_python(skip_empty_area=False)
if not any(cell != "" for row in rows for cell in row):
warnings.append(
ParseWarning(
code="empty_sheet_skipped",
message=f"sheet「{meta.name}」是空的,已跳过",
)
)
continue
grids.append(Grid(meta.name, rows))
finally:
book.close()
return grids, warnings
def read_json(value: Any) -> Grid | Records:
"""对象数组按记录读取;数组的数组当作网格,第一行通常是表头"""
warnings = []
if isinstance(value, dict):
records, path = _find_records(value)
if path:
warnings.append(
ParseWarning(
code="json_path_selected",
message=f"使用了 JSON 中「{path}」下的 {len(records)} 条记录",
)
)
else:
records = value
if not records:
raise DataFileError(422, "JSON 里没有找到数据")
if all(isinstance(item, list) for item in records):
return Grid(None, [[_scalar(cell) for cell in row] for row in records])
if all(isinstance(item, dict) for item in records):
result = _flatten_records(records)
result.warnings = warnings
return result
if not any(isinstance(item, (dict, list)) for item in records):
return Records(["值"], [list(records)], warnings)
raise DataFileError(422, "JSON 数组里的元素类型不一致,需要都是对象或都是数组")
def read_jsonl(text: str) -> Records:
records, bad, total = [], [], 0
for number, line in enumerate(text.splitlines(), start=1):
if not line.strip():
continue
total += 1
try:
record = json.loads(line)
except ValueError:
record = None
if isinstance(record, dict):
records.append(record)
else:
bad.append(number)
if not records:
raise DataFileError(422, "JSONL 里没有找到数据")
if len(bad) > total * MAX_BAD_LINE_RATIO:
raise DataFileError(
422, f"JSONL 有 {len(bad)} 行无法解析(例如第 {bad[0]} 行),超过 1%"
)
result = _flatten_records(records)
if bad:
result.warnings.append(
ParseWarning(
code="bad_lines_skipped",
message=f"跳过了 {len(bad)} 行无法解析的内容:{_line_list(bad)}",
)
)
return result
def read_parquet(path: Path) -> pa.Table:
try:
table = pq.read_table(path)
except Exception:
raise DataFileError(422, "无法读取 Parquet 文件,文件可能已损坏") from None
if table.num_rows > MAX_TABLE_ROWS:
raise rows_exceeded(None)
# struct 展开一层,列名写成 a.b;更深的嵌套由 columns.py 存成 JSON 字符串
return table.flatten()
def _find_records(value: dict) -> tuple[list, str | None]:
"""在对象的前两层里找最长的对象数组,例如 {"data": [...]};找不到时按列式或单条记录处理"""
candidates = []
for key, item in value.items():
if _is_record_list(item):
candidates.append((len(item), str(key), item))
elif isinstance(item, dict):
for sub_key, sub in item.items():
if _is_record_list(sub):
candidates.append((len(sub), f"{key}.{sub_key}", sub))
if candidates:
_, path, records = max(candidates, key=lambda candidate: candidate[0])
return records, path
lists = list(value.values())
# 列式:{"地区": [...], "销售额": [...]}
if (
lists
and all(isinstance(item, list) for item in lists)
and len({len(item) for item in lists}) == 1
):
return [dict(zip(value, row)) for row in zip(*lists)], None
return [value], None
def _is_record_list(value: Any) -> bool:
return (
isinstance(value, list)
and bool(value)
and all(isinstance(item, dict) for item in value)
)
def _flatten_records(records: list[dict]) -> Records:
# 列按首次出现的顺序排列,记录里没有的字段为空
index: dict[str, int] = {}
columns: list[list[Any]] = []
for row_number, record in enumerate(records):
flat: dict[str, Any] = {}
_flatten(record, "", 0, flat)
for name, value in flat.items():
if name not in index:
index[name] = len(columns)
columns.append([None] * row_number)
columns[index[name]].append(value)
for column in columns:
if len(column) == row_number:
column.append(None)
return Records(list(index), columns)
def _flatten(record: dict, prefix: str, level: int, out: dict[str, Any]) -> None:
for key, value in record.items():
name = f"{prefix}.{key}" if prefix else str(key)
if isinstance(value, dict) and value and level < JSON_FLATTEN_LEVELS:
_flatten(value, name, level + 1, out)
else:
out[name] = _scalar(value)
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 "")
+334
View File
@@ -0,0 +1,334 @@
"""上传流程:保存原文件 → 子进程解析(超时终止)→ 分配表名 → 保存记录;以及查看、删除和过期清理"""
import asyncio
import hashlib
import json
import os
import re
import shutil
import subprocess
import sys
import time
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import BinaryIO
from uuid import uuid4
from loguru import logger
from app.config import PROJECT_ROOT, settings
from app.core.exceptions import BizException
from app.datafiles.parse import ParsedFile, ParsedTable
from app.datafiles.storage import SOURCE_NAME, DataFileStore, file_dir
from app.schemas.base import JsonObject
from app.schemas.data_file import (
DataFileRecord,
DataTable,
RequestDataFile,
RequestDataTable,
)
from app.services.conversation import ConversationStore
MAX_FILES_PER_CONVERSATION = 10
# 解析超时(秒),超时终止子进程;同时最多解析的文件数,其余排队
PARSE_TIMEOUT = 60
PARSE_CONCURRENCY = 2
# 每小时清理一次目录;一小时内创建的目录不删,避免删掉正在上传的文件
SWEEP_INTERVAL = 3600
SWEEP_GRACE = 3600
# 只清理这个格式的目录,DATA_FILE_DIR 配错时也不会删到别的东西
FILE_ID = re.compile(r"^f_[0-9a-f]{32}$")
MAX_NAME_CHARS = 200
CHUNK_BYTES = 1024 * 1024
NOT_FOUND = "文件不存在或已过期"
PARSE_CRASHED = "解析文件时出错,请检查文件内容,或另存为 CSV 后重新上传"
class DataFileService:
def __init__(
self,
store: DataFileStore,
conversations: ConversationStore,
root: Path | None = None,
*,
max_mb: int | None = None,
ttl: int | None = None,
):
self.store = store
self.root = root or settings.data_file_dir
self.max_mb = max_mb or settings.data_file_max_mb
self._conversations = conversations
self._ttl = timedelta(seconds=ttl or settings.data_file_ttl)
self._parsing = asyncio.Semaphore(PARSE_CONCURRENCY)
# 文件数检查、分配表名和写入记录要一起完成,并发上传时才不会超限或重名;单实例部署,进程内的锁就够了
self._saving = asyncio.Lock()
async def upload(
self, user_id: int, conversation_id: str, name: str | None, source: BinaryIO
) -> DataFileRecord:
await self._check_conversation(conversation_id, user_id)
# 先粗查一次,满了就不必上传和解析;写入前在锁里再查一次
_check_count(await self.store.find(conversation_id))
file_id = f"f_{uuid4().hex}"
directory = file_dir(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])
record = self._record(
file_id,
user_id,
conversation_id,
_clean_name(name),
size,
sha256,
parsed,
_next_table_number(existing),
)
await self.store.insert(record)
except BaseException:
# 失败的上传不留文件,会话内的文件数不变
await asyncio.to_thread(shutil.rmtree, directory, ignore_errors=True)
raise
logger.info(
"数据文件已上传:{} 格式 {},{} 张表",
file_id,
record.format,
len(record.tables),
)
return record
async def list_files(
self, user_id: int, conversation_id: str
) -> 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]
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 record
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
)
async def request_files(
self, user_id: int, conversation_id: str, attachments: list[JsonObject]
) -> list[RequestDataFile]:
"""本次请求可用的文件,调用方已校验会话归属
attachments 里 kind 为 data_file 的引用必须是本会话未过期的文件,否则 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
]
async def sweep(self) -> int:
"""删除没有对应记录(过期、已删除、内存模式重启后丢失)的目录,返回删除的个数"""
live = await self.store.live_ids()
return await asyncio.to_thread(self._remove_orphans, live)
async def run_sweeper(self) -> None:
while True:
try:
if removed := await self.sweep():
logger.info("已清理 {} 个过期或删除的数据文件目录", removed)
except Exception as error:
logger.opt(exception=error).error("清理数据文件目录失败")
await asyncio.sleep(SWEEP_INTERVAL)
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, "会话不存在或已过期")
def _save(self, source: BinaryIO, directory: Path) -> tuple[int, str]:
"""边写边算大小和 sha256,超过上限立即停止"""
directory.mkdir(parents=True)
digest, size = hashlib.sha256(), 0
with (directory / SOURCE_NAME).open("wb") as target:
while chunk := source.read(CHUNK_BYTES):
size += len(chunk)
if size > self.max_mb * 1024 * 1024:
raise BizException(413, f"文件不能超过 {self.max_mb} MB")
digest.update(chunk)
target.write(chunk)
return size, digest.hexdigest()
def _record(
self,
file_id: str,
user_id: int,
conversation_id: str,
name: str,
size: int,
sha256: str,
parsed: ParsedFile,
first_number: int,
) -> DataFileRecord:
tables, warnings = [], list(parsed.warnings)
for number, table in enumerate(parsed.tables, start=first_number):
table_name = f"t{number}"
tables.append(
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,
)
)
warnings.extend(
warning.model_copy(update={"table": table_name})
for warning in table.warnings
)
now = datetime.now(UTC)
return DataFileRecord(
file_id=file_id,
owner_id=user_id,
conversation_id=conversation_id,
name=name,
format=parsed.format,
encoding=parsed.encoding,
size=size,
sha256=sha256,
created_at=now,
expires_at=now + self._ttl,
tables=tables,
warnings=warnings,
)
def _remove_orphans(self, live: set[str]) -> int:
if not self.root.is_dir():
return 0
cutoff = time.time() - SWEEP_GRACE
removed = 0
for entry in self.root.iterdir():
if (
FILE_ID.match(entry.name)
and entry.name not in live
and entry.is_dir()
and entry.stat().st_mtime < cutoff
):
shutil.rmtree(entry, ignore_errors=True)
removed += 1
return removed
async def _parse(directory: Path) -> ParsedFile:
python, env = _python()
command = [
python,
"-m",
"app.datafiles.worker",
str(directory / SOURCE_NAME),
str(directory),
]
try:
# 超时后 subprocess.run 会终止子进程
completed = await asyncio.to_thread(
subprocess.run,
command,
cwd=PROJECT_ROOT,
env=env,
stdin=subprocess.DEVNULL,
capture_output=True,
timeout=PARSE_TIMEOUT,
)
except subprocess.TimeoutExpired:
raise BizException(
422, f"解析超过 {PARSE_TIMEOUT} 秒,文件可能太大或太复杂,请拆分后上传"
) from None
if completed.returncode != 0:
logger.error(
"数据文件解析进程异常退出({}):{}",
completed.returncode,
completed.stderr.decode(errors="replace")[-2000:],
)
raise BizException(422, PARSE_CRASHED)
result = json.loads(completed.stdout)
if error := result.get("error"):
raise BizException(error["status"], error["message"])
return ParsedFile.model_validate(result["file"])
def _python() -> tuple[str, dict[str, str] | None]:
"""Windows 的 venv 里 sys.executable 是启动器,由它再拉起真正的解释器;超时时只杀得掉启动器,
刚启动就被杀时解释器会成为孤儿,一直占着输出管道。和 multiprocessing 一样直接启动基础解释器,
用 __PYVENV_LAUNCHER__ 告诉它所在的 venv
"""
if sys.platform == "win32" and sys._base_executable != sys.executable:
return sys._base_executable, {
**os.environ,
"__PYVENV_LAUNCHER__": sys.executable,
}
return sys.executable, None
def _check_count(records: list[DataFileRecord]) -> None:
if len(records) >= MAX_FILES_PER_CONVERSATION:
raise BizException(
409, f"每个会话最多 {MAX_FILES_PER_CONVERSATION} 个文件,请先删除不用的文件"
)
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()
return base[:MAX_NAME_CHARS] or "未命名文件"
def _label(name: str, table: ParsedTable) -> str:
label = f"{name} / {table.sheet}" if table.sheet else name
return f"{label}({table.part})" if table.part else label
+152
View File
@@ -0,0 +1,152 @@
"""数据文件的存储:文件放本地目录,元数据启用 MongoDB 时进集合 data_files,否则放进程内存
目录结构是 DATA_FILE_DIR/<file_id>/,里面有原文件 source 和每张表的 1.parquet、2.parquet……
DuckDB 只能读本地文件,所以文件不进 MongoDB;单实例部署,本地磁盘就够了
"""
from datetime import UTC, datetime
from pathlib import Path
from typing import Any, Protocol
from bson.codec_options import CodecOptions
from pymongo import ASCENDING
from app.schemas.data_file import DataFileRecord
# 原文件不带扩展名:calamine 见到 .xls 等扩展名会按扩展名打开,用户的扩展名和内容不符时会失败
SOURCE_NAME = "source"
def file_dir(root: Path, file_id: str) -> Path:
return root / file_id
def _now() -> datetime:
return datetime.now(UTC)
class DataFileStore(Protocol):
"""内存和 MongoDB 实现同一组方法;过期的记录一律当作不存在"""
async def insert(self, record: DataFileRecord) -> None: ...
async def get(self, file_id: str) -> DataFileRecord | None:
"""未过期且未删除的记录"""
...
async def find(
self, conversation_id: str, *, include_deleted: bool = False
) -> list[DataFileRecord]:
"""会话内未过期的记录,按上传时间排列"""
...
async def mark_deleted(self, file_id: str) -> None: ...
async def live_ids(self) -> set[str]:
"""未过期且未删除的 file_id,清理目录时用"""
...
class InMemoryDataFileStore:
def __init__(self) -> None:
self._records: dict[str, DataFileRecord] = {}
def _purge_expired(self) -> None:
now = _now()
for file_id in [
key for key, record in self._records.items() if record.expires_at <= now
]:
del self._records[file_id]
async def insert(self, record: DataFileRecord) -> None:
self._purge_expired()
self._records[record.file_id] = record
async def get(self, file_id: str) -> DataFileRecord | None:
self._purge_expired()
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]:
self._purge_expired()
records = [
record
for record in self._records.values()
if record.conversation_id == conversation_id
and (include_deleted or not record.deleted)
]
return sorted(records, key=lambda record: record.created_at)
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})
async def live_ids(self) -> set[str]:
self._purge_expired()
return {key for key, record in self._records.items() if not record.deleted}
class MongoDataFileStore:
def __init__(self, db: Any, prefix: str = ""):
# 读回的时间统一带 UTC 时区;集合名前缀用于测试隔离
options = CodecOptions(tz_aware=True, tzinfo=UTC)
self._files = db.get_collection(f"{prefix}data_files", codec_options=options)
async def ensure_indexes(self) -> None:
# 过期记录由 TTL 索引删除;TTL 每分钟才扫一次,查询时仍按 expires_at 过滤
await self._files.create_index("expires_at", expireAfterSeconds=0)
await self._files.create_index(
[("conversation_id", ASCENDING), ("created_at", ASCENDING)]
)
async def insert(self, record: DataFileRecord) -> None:
await self._files.insert_one({"_id": record.file_id, **record.model_dump()})
async def get(self, file_id: str) -> DataFileRecord | None:
document = await self._files.find_one(
{"_id": file_id, "deleted": False, "expires_at": {"$gt": _now()}}
)
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] = {
"conversation_id": conversation_id,
"expires_at": {"$gt": _now()},
}
if not include_deleted:
query["deleted"] = False
documents = await (
self._files.find(query)
.sort([("created_at", ASCENDING), ("_id", ASCENDING)])
.to_list()
)
return [_record(document) for document in documents]
async def mark_deleted(self, file_id: str) -> None:
await self._files.update_one({"_id": file_id}, {"$set": {"deleted": True}})
async def live_ids(self) -> set[str]:
return set(
await self._files.distinct(
"_id", {"deleted": False, "expires_at": {"$gt": _now()}}
)
)
def _record(document: dict[str, Any]) -> DataFileRecord:
return DataFileRecord.model_validate(
{key: value for key, value in document.items() if key != "_id"}
)
async def open_data_file_store(db: Any) -> DataFileStore:
"""db 为空(未启用 MongoDB)时用内存实现,和会话记录一致"""
if db is None:
return InMemoryDataFileStore()
store = MongoDataFileStore(db)
await store.ensure_indexes()
return store
+290
View File
@@ -0,0 +1,290 @@
"""在原始网格里识别表:按空行切块,找表头,合并两级表头,去掉标题、说明和合计行
只用规则;判断不了的地方照常产出表,同时记一条警告,调用方据此提示用户
"""
import math
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.readers import Grid
from app.schemas.data_file import ParseWarning
# 只在每块的前几行里找表头
HEADER_SEARCH_ROWS = 10
# 首个或第二个非空单元格是这些词(去掉空白和冒号后)的行是合计行
TOTAL_LABELS = {
"合计",
"总计",
"小计",
"共计",
"总和",
"平均",
"平均值",
"total",
"subtotal",
"grandtotal",
"sum",
"average",
}
# 例如"华东小计""一季度合计"
TOTAL_SUFFIXES = ("小计", "合计", "总计")
# 警告里引用的单元格文本最长显示这么多字
QUOTE_CHARS = 30
@dataclass
class RawTable:
sheet: str | None
# 同一 sheet 切出多张表时的序号,从 1 开始;只有一张表时为空
part: int | None
header: list[Any]
# 数据行,宽度与表头相同
rows: list[list[Any]]
warnings: list[ParseWarning] = field(default_factory=list)
@dataclass
class _Piece:
columns: list[int]
# 表头所在行在网格中的下标
header_row: int
header: list[Any]
# 数据行在网格中的下标
data: list[int]
warnings: list[ParseWarning]
def split_tables(grid: Grid) -> tuple[list[RawTable], list[ParseWarning]]:
"""返回识别出的表和不属于任何表的警告(被丢弃的块)"""
rows = _normalize(grid.rows)
sheet = grid.sheet
pieces: list[_Piece] = []
dropped: list[ParseWarning] = []
for block in _blocks(rows):
columns = [
c for c in range(len(rows[0])) if any(rows[r][c] is not None for r in block)
]
last = pieces[-1] if pieces else None
# 表中间夹着空行时,后面的块以数据行开头、列也落在上一张表内,接在上一张表后面
if (
last
and not _is_label_row([rows[block[0]][c] for c in columns])
and set(columns) <= set(last.columns)
):
last.data.extend(block)
continue
if len(block) < 2:
dropped.append(
ParseWarning(
code="block_dropped",
message=f"{where(sheet)}第 {block[0] + 1} 行的{_quote(rows[block[0]])}不是表格,已跳过",
)
)
continue
pieces.append(_read_block(rows, block, columns, sheet))
tables = []
for piece in pieces:
table = _finish(rows, piece, sheet)
if table is None:
dropped.append(
ParseWarning(
code="block_dropped",
message=f"{where(sheet)}第 {piece.header_row + 1} 行开始的表只有表头没有数据,已跳过",
)
)
else:
tables.append(table)
if len(tables) > 1:
for part, table in enumerate(tables, start=1):
table.part = part
return tables, dropped
def _normalize(grid: list[list[Any]]) -> list[list[Any]]:
"""去掉字符串两端空白,空字符串记为 None,补齐成同样宽度"""
width = max((len(row) for row in grid), default=0)
rows = []
for row in grid:
cells = [_cell(value) for value in row]
cells.extend([None] * (width - len(cells)))
rows.append(cells)
return rows
def _cell(value: Any) -> Any:
if isinstance(value, str):
value = value.strip()
return value or None
return value
def _blocks(rows: list[list[Any]]) -> list[list[int]]:
blocks, current = [], []
for index, row in enumerate(rows):
if any(value is not None for value in row):
current.append(index)
elif current:
blocks.append(current)
current = []
if current:
blocks.append(current)
return blocks
def _read_block(
rows: list[list[Any]], block: list[int], columns: list[int], sheet: str | None
) -> _Piece:
cells = [[rows[r][c] for c in columns] for r in block]
need = max(1, math.ceil(len(columns) / 2))
dense = [sum(value is not None for value in row) >= need for row in cells]
search = range(min(HEADER_SEARCH_ROWS, len(cells)))
warnings = []
# 表头:前几行里第一个够满、并且全是文字的行;标题行通常只有一个单元格,不够满
header_at = next((i for i in search if dense[i] and _is_label_row(cells[i])), None)
confident = header_at is not None
if header_at is None:
header_at = next((i for i in search if dense[i]), 0)
warnings.append(
ParseWarning(
code="header_uncertain",
message=f"{where(sheet)}没有找到明确的表头,用第 {block[header_at] + 1} 行作表头,可能识别错误",
)
)
if header_at > 0:
skipped = "".join(
_quote(cells[i]) for i in range(header_at) if _quote(cells[i])
)
warnings.append(
ParseWarning(
code="title_skipped",
message=f"{where(sheet)}{_row_range(block[0], block[header_at - 1])}是标题,已跳过:{skipped}",
)
)
header = cells[header_at]
data_at = header_at + 1
if (
confident
and header_at + 2 < len(cells)
and _is_subheader(cells[header_at], cells[header_at + 1], cells[header_at + 2])
):
header = _merge_header(cells[header_at], cells[header_at + 1])
data_at = header_at + 2
warnings.append(
ParseWarning(
code="header_merged",
message=f"{where(sheet)}{_row_range(block[header_at], block[header_at + 1])}合并为表头",
)
)
return _Piece(columns, block[header_at], header, block[data_at:], warnings)
def _finish(rows: list[list[Any]], piece: _Piece, sheet: str | None) -> RawTable | None:
data = [(index, [rows[index][c] for c in piece.columns]) for index in piece.data]
warnings = piece.warnings
# 表格下方只有一个文字单元格的行是说明,例如"注:数据来源于……"
notes = []
while data and len(piece.columns) >= 3 and _is_note(data[-1][1]):
notes.append(data.pop())
if notes:
notes.reverse()
shown = "".join(_quote(row) for _, row in notes)
warnings.append(
ParseWarning(
code="note_skipped",
message=f"{where(sheet)}{_row_range(notes[0][0], notes[-1][0])}是说明,已跳过:{shown}",
)
)
# 合计、小计行不是明细数据,留着会让求和重复计算
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 ""
)
warnings.append(
ParseWarning(
code="total_rows_removed",
message=f"{where(sheet)}已排除 {len(totals)} 行合计、小计(第 {listed} 行)",
)
)
if not data:
return None
if len(data) > MAX_TABLE_ROWS:
raise rows_exceeded(sheet)
return RawTable(sheet, None, piece.header, [row for _, row in data], warnings)
def _is_label(value: Any) -> bool:
return isinstance(value, str) and classify(value)[0] == "text"
def _is_label_row(row: list[Any]) -> bool:
values = [value for value in row if value is not None]
return bool(values) and all(_is_label(value) for value in values)
def _is_subheader(top: list[Any], sub: list[Any], first_data: list[Any]) -> bool:
"""第二行也全是文字、紧接着是带数值的数据行,并且两行里有合并单元格留下的空格"""
return (
_is_label_row(sub)
and not _is_label_row(first_data)
and (None in top or None in sub)
)
def _merge_header(top: list[Any], sub: list[Any]) -> list[str]:
"""合并单元格只有左上角有值,第一行横向向前填充后与第二行用 - 连接,例如 销售额-一月"""
merged, carry = [], None
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]
merged.append("-".join(parts))
return merged
def _is_note(row: list[Any]) -> bool:
values = [value for value in row if value is not None]
return len(values) == 1 and _is_label(values[0])
def _is_total(row: list[Any]) -> bool:
checked = 0
for value in row:
if value is None:
continue
if isinstance(value, str):
key = re.sub(r"[\s::]", "", value).lower()
if key in TOTAL_LABELS or key.endswith(TOTAL_SUFFIXES):
return True
checked += 1
if checked == 2:
return False
return False
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)
if len(text) > QUOTE_CHARS:
text = text[:QUOTE_CHARS] + "…"
return f"「{text}」"
def _row_range(first: int, last: int) -> str:
return f"第 {first + 1} 行" if first == last else f"第 {first + 1}–{last + 1} 行"
+27
View File
@@ -0,0 +1,27 @@
"""解析子进程:python -m app.datafiles.worker <原文件> <输出目录>
结果以一行 JSON 写到 stdout:成功是 {"file": ParsedFile},文件本身的问题是 {"error": {"status", "message"}}
其他异常让进程以非零状态退出,堆栈留在 stderr,由上传流程记日志
"""
import json
import sys
from pathlib import Path
from app.datafiles.base import DataFileError
from app.datafiles.parse import parse_file
def main(source: str, out_dir: str) -> None:
try:
result = {
"file": parse_file(Path(source), Path(out_dir)).model_dump(mode="json")
}
except DataFileError as error:
result = {"error": {"status": error.status, "message": error.message}}
# 只输出 ASCII,不受控制台编码影响
sys.stdout.write(json.dumps(result))
if __name__ == "__main__":
main(*sys.argv[1:])
+17 -1
View File
@@ -8,6 +8,8 @@ from loguru import logger
from app.api.agent import router
from app.api.conversations import router as conversation_router
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
from app.core.auth import get_current_user
from app.core.exceptions import register_exception_handlers
@@ -21,6 +23,8 @@ async def lifespan(application: FastAPI):
from app.clients.decision import close_decision_engine
from app.clients.mcp import start_mcp_servers, stop_mcp_servers
from app.clients.user_center import close_user_center
from app.datafiles.service import DataFileService
from app.datafiles.storage import open_data_file_store
from app.observability.tracing import flush_tracing
from app.processor.graph import build_graph
from app.services.agent import AgentService
@@ -31,7 +35,14 @@ async def lifespan(application: FastAPI):
persistence = await open_persistence()
application.state.conversation_store = persistence.store
mcp_processes = []
sweeper = None
try:
# 数据文件的元数据和会话记录用同一种存储;启动时先清理一遍过期目录,之后每小时一次
data_files = DataFileService(
await open_data_file_store(persistence.db), persistence.store
)
application.state.data_file_service = data_files
sweeper = asyncio.create_task(data_files.run_sweeper())
# 物料目录启动时载入内存,动作校验与前端接口共用,避免重复查询数据库
await warm_material_cache(persistence.db)
# 拉起 mcp_servers.json 中由应用管理的服务;单个服务失败只记录,用到它的请求执行时报错
@@ -42,10 +53,12 @@ async def lifespan(application: FastAPI):
except Exception as error:
logger.opt(exception=error).error("screen agent 预先组装失败")
application.state.agent_service = AgentService(
build_graph(persistence.checkpointer), persistence.store
build_graph(persistence.checkpointer), persistence.store, data_files
)
yield
finally:
if sweeper is not None:
sweeper.cancel()
# 共享客户端只在应用退出时关闭,追踪刷新放到线程避免阻塞事件循环
try:
await stop_mcp_servers(mcp_processes)
@@ -76,6 +89,9 @@ def create_app() -> FastAPI:
application.include_router(
conversation_router, prefix="/api", dependencies=authenticated
)
# 上传限制无需登录;要先于 /files/{file_id} 注册,否则会被当成文件 ID
application.include_router(file_limits_router, prefix="/api")
application.include_router(file_router, prefix="/api", dependencies=authenticated)
application.include_router(material_router, prefix="/api")
return application
+7 -2
View File
@@ -7,6 +7,7 @@ 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
@@ -53,7 +54,7 @@ class ScreenAgentNode(BaseNode):
return updated_state
context = state["input_context"]
agent = await get_screen_agent()
agent = await get_screen_agent(with_data=bool(context.data_files))
result = await agent.ainvoke(
{
# 会话历史作为此前的对话消息,本次请求是最后一条用户消息
@@ -64,7 +65,11 @@ class ScreenAgentNode(BaseNode):
),
]
},
context=ScreenAgentContext(screen=context.screen, catalog=catalog),
context=ScreenAgentContext(
screen=context.screen,
catalog=catalog,
data=DataContext(tuple(context.data_files)),
),
)
if result.get("structured_response") is None:
last = result["messages"][-1]
+29 -1
View File
@@ -4,11 +4,15 @@
"""
from app.schemas.base import JsonObject
from app.schemas.data_file import RequestDataFile
from app.schemas.input_context import InputContext
# 每张表在摘要里最多列出的列名,完整的列用 describe_data 查看
SUMMARY_COLUMNS = 30
def request_summary(context: InputContext) -> JsonObject:
"""用户输入、澄清记录、编辑上下文、附件和大屏摘要"""
"""用户输入、澄清记录、编辑上下文、附件、数据文件和大屏摘要"""
nodes: list[JsonObject] = []
for node in context.screen.nodes:
# 图表标题帮助判断用户指的是哪个组件;没有 option 或标题格式不符时按无标题处理
@@ -45,8 +49,32 @@ def request_summary(context: InputContext) -> JsonObject:
for source in context.screen.dataSource
],
},
# 数据文件的附件由 data_files 的 attached_this_turn 表示
"attachments": [
{key: item[key] for key in ("id", "kind", "name") if key in item}
for item in context.attachments
if item.get("kind") != "data_file"
],
"data_files": [_data_file(item) for item in context.data_files],
}
def _data_file(item: RequestDataFile) -> JsonObject:
"""每个文件只给表名、label、行数和列名;列概况由 describe_data 按需查看"""
tables: list[JsonObject] = []
for table in item.tables:
summary: JsonObject = {
"table": table.table,
"label": table.label,
"rows": table.rows,
"columns": [column.name for column in table.columns[:SUMMARY_COLUMNS]],
}
if len(table.columns) > SUMMARY_COLUMNS:
summary["column_count"] = len(table.columns)
tables.append(summary)
return {
"file_id": item.file_id,
"name": item.name,
"attached_this_turn": item.attached_this_turn,
"tables": tables,
}
+122
View File
@@ -0,0 +1,122 @@
"""数据文件的接口模型:文件记录、表、列和解析警告,见 docs/screen-agent/data-files.md 5.4、6.6、7
文件和表各有三层,后一层在前一层上加字段:列表接口返回的 Info、单个文件返回的 Detail、存储用的完整记录
接口返回时用对应的模型重新校验存储记录,多出的字段自然丢掉
"""
from datetime import datetime
from typing import Literal
from pydantic import BaseModel, JsonValue
DataFormat = Literal["csv", "tsv", "xlsx", "xls", "xlsb", "json", "jsonl", "parquet"]
ColumnType = Literal["text", "number", "date", "datetime", "boolean"]
WarningCode = Literal[
"title_skipped",
"note_skipped",
"block_dropped",
"header_merged",
"header_uncertain",
"total_rows_removed",
"column_renamed",
"kept_as_text",
"bad_lines_skipped",
"json_path_selected",
"empty_sheet_skipped",
"hidden_sheet_skipped",
]
class DataColumn(BaseModel):
name: str
type: ColumnType
# 原文是百分比时为 percent,值已存为小数
format: Literal["percent"] | None = None
class ParseWarning(BaseModel):
code: WarningCode
# 表名;为空表示针对整个文件。解析时还没有分配表名,由上传流程填写
table: str | None = None
# 可以直接展示的中文说明
message: str
class DataTableInfo(BaseModel):
# 会话内的表名 t1、t2……,模型写 SQL 时用
table: str
# 文件名 / sheet,同一 sheet 切出多张表时再加序号
label: str
sheet: str | None
rows: int
columns: list[DataColumn]
class DataTableDetail(DataTableInfo):
# 前 20 行,顺序与 columns 一致
preview: list[list[JsonValue]]
class DataTable(DataTableDetail):
# 列概况,给 describe_data 用
profile: dict[str, JsonValue]
# 文件目录下的 Parquet 文件名
parquet: str
class DataFileInfo(BaseModel):
file_id: str
conversation_id: str
name: str
format: DataFormat
# 文本文件识别出的编码,二进制格式为空
encoding: str | None
size: int
created_at: datetime
expires_at: datetime
tables: list[DataTableInfo]
warnings: list[ParseWarning]
class DataFileDetail(DataFileInfo):
tables: list[DataTableDetail]
class DataFileRecord(DataFileDetail):
owner_id: int
sha256: str
tables: list[DataTable]
# 删除只做标记,记录留到过期,会话内的表名不会被新文件重复使用
deleted: bool = False
class DataFileList(BaseModel):
files: list[DataFileInfo]
class DataFileLimits(BaseModel):
max_size_mb: int
max_files_per_conversation: int
formats: list[DataFormat]
# 文件选择框的 accept,只用于前端过滤;后端按文件内容识别格式
accept: str
class RequestDataTable(DataTableInfo):
"""本次请求可用的表:多出工具要用的列概况和 Parquet 路径,不带预览"""
profile: dict[str, JsonValue]
path: str
class RequestDataFile(BaseModel):
"""本次请求可用的文件:服务端按会话读取后放进请求,前端传的同名字段会被覆盖"""
file_id: str
name: str
# 本轮消息的 attachments 里带了这个文件,通常是刚上传的
attached_this_turn: bool
tables: list[RequestDataTable]
warnings: list[ParseWarning]
+3
View File
@@ -3,6 +3,7 @@
from pydantic import Field
from .base import JsonObject
from .data_file import RequestDataFile
from .request import DesignRequest
@@ -12,3 +13,5 @@ class InputContext(DesignRequest):
clarification_history: list[JsonObject] = Field(default_factory=list)
# 同一会话此前各轮的用户输入与助手回复,由服务端按 conversation_id 读取并按 token 预算裁剪,按时间从早到晚排列
history: list[JsonObject] = Field(default_factory=list)
# 本会话里未过期的数据文件,由服务端读取并校验 attachments 后填入
data_files: list[RequestDataFile] = Field(default_factory=list)
+2 -1
View File
@@ -17,5 +17,6 @@ class DesignRequest(ExtensibleModel):
prompt: NonBlankString
screen: Screen
context: RequestContext = Field(default_factory=RequestContext)
# 只预留 JSON 描述入口,暂不规定材料种类,也不读取文件或远程资源
# 本条消息附带的材料;{"kind": "data_file", "file_id": "..."} 引用本会话上传的数据文件,
# 其他种类只原样进入摘要,不读取文件或远程资源
attachments: list[JsonObject] = Field(default_factory=list)
+9 -1
View File
@@ -9,6 +9,7 @@ from loguru import logger
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.schemas.agent import AgentResult, ResumeRequest
@@ -18,9 +19,10 @@ from app.services.conversation import ConversationStore
class AgentService:
def __init__(self, graph, store: ConversationStore):
def __init__(self, graph, store: ConversationStore, data_files: DataFileService):
self.graph = graph
self.store = store
self.data_files = data_files
# 同一时刻,一个任务和一个会话都只允许一个执行,避免历史顺序错乱
self._active: set[str] = set()
self._conversation_of: dict[str, str] = {}
@@ -40,6 +42,10 @@ class AgentService:
turns = await self.store.load_turns(
conversation_id, settings.conversation_history_turns
)
# 附件引用的文件不属于本会话时在流开始前返回 404
data_files = await self.data_files.request_files(
user_id, conversation_id, request.attachments
)
except BaseException:
self.release(thread_id)
raise
@@ -57,6 +63,8 @@ class AgentService:
],
settings.history_max_tokens,
)
# 同样只信任服务端读取的结果;会话内的文件每一轮都可用,不附带也能查询
payload["data_files"] = [item.model_dump(mode="json") for item in data_files]
return (
thread_id,