feat:完成大部分数据文件上传解析功能
This commit is contained in:
@@ -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
@@ -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
@@ -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):
|
||||
|
||||
@@ -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
@@ -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 之间")
|
||||
|
||||
|
||||
@@ -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} 万行,请拆分后上传"
|
||||
)
|
||||
@@ -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}"
|
||||
@@ -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
|
||||
@@ -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]
|
||||
@@ -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)
|
||||
@@ -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 "(无)"
|
||||
@@ -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 "")
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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} 行"
|
||||
@@ -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
@@ -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,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]
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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,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)
|
||||
|
||||
@@ -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,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,
|
||||
|
||||
Reference in New Issue
Block a user