Feat/0.2.0 #4
+7
-1
@@ -30,10 +30,16 @@ MONGODB_CHECKPOINT_TTL=604800
|
||||
|
||||
# 上传的数据文件：存放目录（相对路径相对项目根目录，容器部署时挂载出来）、单个文件大小上限（MB）、保留时长（秒，默认 7 天）
|
||||
DATA_FILE_DIR=data/files
|
||||
DATA_FILE_MAX_MB=50
|
||||
DATA_FILE_MAX_MB=500
|
||||
DATA_FILE_TTL=604800
|
||||
# 上传后用 LLM 分析文件（一句话说明、图表推荐、列的含义），false 关闭；没有配置 LLM_API_KEY 时同样不分析
|
||||
DATA_FILE_INSIGHTS=true
|
||||
# 大屏保存/发布的请求体上限（MB）、封面图目录（相对路径相对项目根目录）；大屏存取必须配置 MONGODB_URI，否则接口返回 503
|
||||
SCREEN_MAX_MB=10
|
||||
# 回收站里的大屏保留多少天，超过后自动彻底删除
|
||||
TRASH_RETENTION_DAYS=30
|
||||
THUMBNAIL_DIR=data/thumbnails
|
||||
|
||||
# 模型写的单条数据查询的超时（秒）
|
||||
DATA_QUERY_TIMEOUT=10
|
||||
# apply_actions 里一个 $query 写进组件的最多行数
|
||||
|
||||
@@ -0,0 +1,295 @@
|
||||
# Screen Agent · my-v0-server
|
||||
|
||||
面向可视化大屏编辑器的 AI 后端。用户通过自然语言创建大屏、调整组件和分析上传的数据，服务将需求转换为经过校验的结构化编辑动作，交给前端执行。
|
||||
|
||||
本仓库提供 API、Agent 工作流、组件物料和数据处理能力。接入时，前端负责提交当前画布快照、展示流式反馈，并应用返回的动作。
|
||||
|
||||
[快速开始](#快速开始) · [使用指南](#使用指南) · [接口概览](#接口概览) · [架构设计](#架构设计) · [开发与测试](#开发与测试)
|
||||
|
||||
## 功能特性
|
||||
|
||||
- **自然语言编辑大屏**：新建大屏，新增、删除或更换图表，修改组件属性、位置、尺寸与画布配置，支持对齐和分布等布局操作。
|
||||
- **数据驱动的问答与制图**：上传 CSV、TSV、Excel（XLSX / XLS / XLSB）、JSON、JSONL 或 Parquet，查询和汇总数据，用查询结果填充图表。
|
||||
- **动作校验与模拟执行**：按物料字段定义校验修改，并在画布副本上模拟每批动作；成功的批次累加在工作画布上，失败原因返回 Agent
|
||||
修正，通过的动作按顺序交付。
|
||||
- **多轮对话与澄清**：保存会话历史；需求需要补充时暂停任务，用户回答后继续，也可以取消等待中的任务。
|
||||
- **实时反馈**：通过 SSE 返回处理进度、问答文本增量和最终结果，区分业务失败与执行异常。
|
||||
- **组件能力扩展**：以 JSON 定义物料、以 Skills 描述任务做法，通过 MCP 提供组件目录与字段查询。
|
||||
|
||||
## 快速开始
|
||||
|
||||
### 环境要求
|
||||
|
||||
| 依赖 | 说明 |
|
||||
|----------------|----------------------------------------------------------------|
|
||||
| Python | 3.11 或以上，版本要求见 [pyproject.toml](pyproject.toml) |
|
||||
| uv | 安装依赖并运行项目，仓库提供 `uv.lock` |
|
||||
| MongoDB | 完整大屏编辑所需的物料目录，以及会话、检查点和文件元数据存储 |
|
||||
| 聊天模型服务 | 通过 `ChatOpenAI` 接入；需要支持工具调用、结构化输出与流式响应 |
|
||||
| TypeSafe / Jev | 判断请求是否进入问答流程，需要独立的 API Key |
|
||||
|
||||
用户中心和 Langfuse 可在本地开发时关闭。默认物料 MCP 服务由应用自动启动，使用本机 `8201` 端口；API 默认使用 `8000` 端口。
|
||||
|
||||
> 不配置 MongoDB 时，会话和检查点可使用进程内存，但物料目录不会回退到本地 JSON，新增、更换组件等能力无法正常使用。首次运行请完成物料初始化。
|
||||
|
||||
### 1. 安装依赖
|
||||
|
||||
获取仓库后，在项目根目录运行：
|
||||
|
||||
```sh
|
||||
uv sync --locked
|
||||
```
|
||||
|
||||
### 2. 配置环境
|
||||
|
||||
将 [.env.example](.env.example) 复制为 `.env`，已有配置无需覆盖。填写以下配置：
|
||||
|
||||
| 配置 | 说明 |
|
||||
|--------------------------------------------|----------------------------------------------|
|
||||
| `LLM_API_KEY`、`LLM_MODEL`、`LLM_BASE_URL` | 聊天模型密钥、模型名称和 API 地址 |
|
||||
| `TYPESAFE_API_KEY` | Jev 意图判断服务密钥 |
|
||||
| `MONGODB_URI`、`MONGODB_DB` | MongoDB 连接地址和数据库名 |
|
||||
| `AUTH_ENABLED` | 本地调试可设为 `false`，所有请求共用开发用户 |
|
||||
| `LANGFUSE_TRACING_ENABLED` | 未配置 Langfuse 时设为 `false` |
|
||||
|
||||
客户端默认发送 `reasoning_effort=low`；模型服务不支持该参数时，设置 `LLM_REASONING_EFFORT=`。
|
||||
|
||||
正式接入用户中心时设置 `AUTH_ENABLED=true` 和 `USER_CENTER_URL`，请求携带 `Authorization: Bearer <token>`。后端通过用户中心的
|
||||
`GET /api/users/me` 获取当前用户。
|
||||
|
||||
### 3. 初始化物料目录
|
||||
|
||||
首次使用需将 `app/resources/materials` 中的物料定义导入 MongoDB。工作区包含 `scripts/seed_materials.py` 时，可运行：
|
||||
|
||||
```sh
|
||||
uv run python -m scripts.seed_materials
|
||||
```
|
||||
|
||||
该脚本会整体替换目标数据库中的物料目录和共享字段分组，更新后需重启服务。当前 `scripts/` 未纳入版本控制，新克隆的工作区需另行取得初始化脚本。
|
||||
|
||||
### 4. 启动与验证
|
||||
|
||||
```sh
|
||||
uv run uvicorn app.main:app --host 127.0.0.1 --port 8000 --reload
|
||||
```
|
||||
|
||||
启动后访问：
|
||||
|
||||
- [Swagger UI](http://127.0.0.1:8000/docs)：查看请求模型和调试普通接口。
|
||||
- [OpenAPI](http://127.0.0.1:8000/openapi.json)：获取接口定义。
|
||||
- [物料目录](http://127.0.0.1:8000/api/materials)：确认物料已载入，正常返回 `code: 0`。
|
||||
|
||||
## 使用指南
|
||||
|
||||
### 创建会话并发起任务
|
||||
|
||||
调用 `POST /api/conversations` 创建会话，取得 `data.conversation_id`，再向 `POST /api/agent/run` 提交需求和当前画布：
|
||||
|
||||
```json
|
||||
{
|
||||
"prompt": "在画布顶部添加标题：销售数据概览",
|
||||
"screen": {
|
||||
"canvas": {
|
||||
"width": 1920,
|
||||
"height": 1080
|
||||
},
|
||||
"nodes": []
|
||||
},
|
||||
"context": {
|
||||
"conversation_id": "创建会话接口返回的 ID"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
开启鉴权后，创建会话和后续请求都需要携带 Bearer token。每次新任务都应提交 **当前画布的完整快照**；继续同一段对话时复用
|
||||
`conversation_id`，历史消息由服务端读取。
|
||||
|
||||
### 处理流式结果
|
||||
|
||||
`/api/agent/run` 和 `/api/agent/resume` 均返回 `text/event-stream`。前端使用支持 POST 和响应流读取的客户端，例如 `fetch`。
|
||||
|
||||
| 事件 | 客户端处理方式 |
|
||||
|----------------|------------------------------------------------------------------------|
|
||||
| `progress` | 展示当前处理阶段 |
|
||||
| `answer_delta` | 按顺序追加 `data.delta`；用于问答场景，最终以 `final.data.answer` 为准 |
|
||||
| `final` | 根据 `data.status` 展示结果、应用动作或发起澄清 |
|
||||
| `error` | 展示 `msg` 中的执行错误，结束本轮处理 |
|
||||
|
||||
每个事件的 `data:` 行统一为 `{code, data, msg}` JSON，`final.data.status` 的含义如下：
|
||||
|
||||
- `completed`：本轮完成；若有 `actions`，由前端按顺序应用。此状态不表示前端已经完成修改。
|
||||
- `waiting_for_input`：展示 `clarification`，保留 `thread_id` 等待用户回答。
|
||||
- `failed`：业务流程无法完成，`answer` 说明原因，不交付动作草稿。
|
||||
- `cancelled`：用户取消了等待补充的任务。
|
||||
|
||||
流开始前的错误通过 HTTP 状态码返回；流开始后的异常通过 `error` 事件返回。客户端不能仅凭 HTTP 200 判断成功，未收到 `final` 或
|
||||
`error` 就断开应视为结果不完整。
|
||||
|
||||
### 回答澄清问题
|
||||
|
||||
当结果为 `waiting_for_input` 时，向 `POST /api/agent/resume` 提交：
|
||||
|
||||
```json
|
||||
{
|
||||
"thread_id": "替换为等待补充的任务 UUID",
|
||||
"reply": {
|
||||
"answer": "使用月度销售额，放在画布左侧"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
也可使用 `reply.option_id` 提交返回的选项，或用 `reply: {"cancel": true}` 取消暂停中的任务。恢复接口沿用任务所属会话。
|
||||
|
||||
### 上传数据并生成图表
|
||||
|
||||
1. 创建会话，将文件与 `conversation_id` 一起以 `multipart/form-data` 提交到 `POST /api/files`。
|
||||
2. 接口返回 HTTP 202 和 `data.file_id`，轮询 `GET /api/files/{file_id}/progress`。
|
||||
3. 等待 `status=ready` 后，在同一会话发起 Agent 请求，例如“按月份汇总销售额，生成柱状图”。解析失败时状态为 `failed`，原因在
|
||||
`error.msg` 中。
|
||||
4. Agent 查询文件数据，将结果转换为图表配置，再返回经过校验的编辑动作。
|
||||
|
||||
会话中已解析且未过期的文件会自动提供给 Agent，无需每轮重新上传。
|
||||
|
||||
## 接口概览
|
||||
|
||||
所有业务接口统一使用 `/api` 前缀，普通 JSON 响应同样采用 `{code, data, msg}` 结构。除物料目录和上传限制外，下列接口默认需要登录。
|
||||
|
||||
| 方法与路径 | 用途 |
|
||||
|------------------------------------------------------|--------------------------------------|
|
||||
| `POST /api/conversations` | 创建会话 |
|
||||
| `GET /api/conversations` | 分页查询当前用户的会话 |
|
||||
| `GET /api/conversations/{conversation_id}` | 读取会话历史 |
|
||||
| `POST /api/agent/run` | 发起任务，返回 SSE |
|
||||
| `POST /api/agent/resume` | 恢复或取消暂停中的任务，返回 SSE |
|
||||
| `GET /api/materials` | 获取公开的组件物料目录 |
|
||||
| `GET /api/files/limits` | 获取当前上传限制与支持格式，无需认证 |
|
||||
| `POST /api/files` | 上传文件并后台解析 |
|
||||
| `GET /api/files?conversation_id=...` | 获取会话文件和进行中的上传 |
|
||||
| `GET /api/files/{file_id}` | 获取文件详情与表信息 |
|
||||
| `GET /api/files/{file_id}/progress` | 查询解析进度 |
|
||||
| `GET /api/files/{file_id}/tables/{table}/rows` | 分页读取并排序表数据 |
|
||||
| `POST /api/files/{file_id}/tables/{table}/aggregate` | 分组汇总表数据 |
|
||||
| `DELETE /api/files/{file_id}` | 删除文件 |
|
||||
|
||||
完整字段、约束和响应模型以运行中的 [Swagger UI](http://127.0.0.1:8000/docs)
|
||||
为准；动作协议定义见 [app/schemas/actions](app/schemas/actions)。
|
||||
|
||||
## 架构设计
|
||||
|
||||
```mermaid
|
||||
flowchart TD
|
||||
Client[大屏编辑器] --> API[FastAPI / AgentService]
|
||||
API --> Input[输入与会话上下文]
|
||||
Input --> Intent[Jev 意图判断]
|
||||
Intent -->|明确问答且达到阈值| QA[问答 Agent]
|
||||
Intent -->|其他请求| Screen[大屏 Agent]
|
||||
Screen --> Materials[MCP 物料与字段查询]
|
||||
Screen --> Batch[动作校验与整批模拟执行]
|
||||
Screen -->|需要补充| Wait[LangGraph 暂停]
|
||||
Wait -->|用户回答并恢复| Input
|
||||
QA --> Data[数据描述与 DuckDB 查询]
|
||||
Screen --> Data
|
||||
Data --> Files[会话文件 / Parquet]
|
||||
QA --> SSE[SSE 进度与结果]
|
||||
Screen --> SSE
|
||||
SSE --> Client
|
||||
```
|
||||
|
||||
外层 LangGraph 管理路由、任务状态与澄清恢复；内层 LangChain Agent 根据上下文调用工具。只有 Jev 选中问答且概率达到阈值时进入问答分支，其余请求交给大屏
|
||||
Agent。
|
||||
|
||||
大屏 Agent 通过 `get_nodes` 读取组件完整配置，通过 MCP 查询物料，通过 `apply_actions`
|
||||
提交修改。每批动作在工作画布（原始画布加上之前成功的修改）上模拟执行，读取和修改都基于它；字段和动作校验失败时这一批不生效，返回原因由
|
||||
Agent 修正。上传数据解析为 Parquet，查询结果可直接填入图表配置，避免模型逐行复写数据。
|
||||
|
||||
| 技术 | 职责 |
|
||||
|------------------------------------|------------------------------------------|
|
||||
| FastAPI / Pydantic | HTTP 接口、请求与响应模型、字段校验 |
|
||||
| LangGraph | 外层工作流、检查点、暂停与恢复 |
|
||||
| LangChain / langchain-openai | Agent 工具循环与聊天模型调用 |
|
||||
| TypeSafe / Jev | 意图判断 |
|
||||
| MCP | 组件物料和字段查询，接入额外工具服务 |
|
||||
| MongoDB | 物料目录、会话历史、图检查点、文件元数据 |
|
||||
| python-calamine / PyArrow / DuckDB | 表格读取、Parquet 存储、数据查询与汇总 |
|
||||
| Loguru / Langfuse | 应用日志与可选的调用追踪 |
|
||||
|
||||
## 项目结构
|
||||
|
||||
```text
|
||||
app/
|
||||
├── main.py # 应用入口与服务生命周期
|
||||
├── config.py # 环境变量与配置校验
|
||||
├── api/ # Agent、会话、物料与文件接口
|
||||
├── services/ # 任务调度、会话存储与持久化装配
|
||||
├── processor/ # LangGraph 工作流、节点、状态与历史裁剪
|
||||
├── agent/ # 大屏和问答 Agent、本地工具、Skills 加载
|
||||
├── actions/ # 动作执行、布局计算与整批模拟
|
||||
├── datafiles/ # 文件解析、存储、查询、概况与后台任务
|
||||
├── schemas/ # 请求、画布、动作、会话和数据文件模型
|
||||
├── resources/ # 物料定义、目录缓存与字段规则
|
||||
├── clients/ # LLM、Jev、MCP、用户中心客户端
|
||||
├── prompt/ # 领域约束、意图、问答与数据分析提示词
|
||||
├── core/ # 认证、异常处理和日志
|
||||
├── common/ # 统一响应与 SSE 编码
|
||||
└── observability/ # Langfuse 追踪
|
||||
mcp_servers/ # 应用管理的本地 MCP 服务
|
||||
skills/ # 大屏设计、编辑、图表替换与数据制图规则
|
||||
mcp_servers.json # MCP 服务连接配置
|
||||
.env.example # 环境变量示例
|
||||
pyproject.toml # Python 依赖与项目元数据
|
||||
uv.lock # 依赖锁文件
|
||||
```
|
||||
|
||||
## 配置与运行边界
|
||||
|
||||
主要可调参数如下，完整配置见 [.env.example](.env.example) 和 [app/config.py](app/config.py)。
|
||||
|
||||
| 配置 | 默认值 | 用途 |
|
||||
|-----------------------------------------------------|------------------|-----------------------------------------------------------------------------|
|
||||
| `MAX_AGENT_STEPS` / `MAX_QA_STEPS` | `12` / `8` | 大屏和问答 Agent 的单次模型调用上限；大屏 Agent 没有给出结论时另加 1 次收尾 |
|
||||
| `INTENT_QA_MIN_PROBABILITY` | `0.5` | 问答分支的判定阈值 |
|
||||
| `CONVERSATION_HISTORY_TURNS` / `HISTORY_MAX_TOKENS` | `20` / `4096` | 历史读取轮数与带入模型的 token 预算 |
|
||||
| `MONGODB_CHECKPOINT_TTL` | `604800` | 检查点保留秒数，`0` 表示永久保留 |
|
||||
| `DATA_FILE_DIR` | `data/files` | 原始文件和解析结果的磁盘目录 |
|
||||
| `DATA_FILE_MAX_MB` / `DATA_FILE_TTL` | `500` / `604800` | 单文件大小上限与保留秒数 |
|
||||
| `DATA_FILE_INSIGHTS` | `true` | 上传后生成文件说明、列含义和图表建议，可关闭 |
|
||||
| `DATA_QUERY_TIMEOUT` / `DATA_EMBED_MAX_ROWS` | `10` / `1000` | 单次 SQL 查询超时秒数与每处查询写入组件的行数上限 |
|
||||
| `AUTH_ENABLED` / `AUTH_CACHE_TTL_SECONDS` | `true` / `30` | 用户中心鉴权与 token 校验缓存秒数 |
|
||||
| `MCP_TOOL_TIMEOUT` | `5` | 单次 MCP 工具调用超时秒数 |
|
||||
|
||||
- **单 worker 运行**：任务与会话的并发保护、上传进度保存在进程内；即使配置 MongoDB，也不能据此直接扩展为多
|
||||
worker。同一会话同时只能执行一个任务。
|
||||
- **文件存储**：MongoDB 保存元数据，原始文件和 Parquet 位于 `DATA_FILE_DIR`。容器部署需保留该目录，并为反向代理配置匹配的上传大小限制。
|
||||
- **流式连接**：断开连接会取消当前执行，不会转为后台任务。当前没有 SSE 心跳、事件重放或断线续传；代理应关闭响应缓冲，并给模型调用留出足够的等待时间。
|
||||
- **前端接入**：保持 `/api` 前缀；当前应用未配置 CORS，中间开发代理或同源反向代理可将 `/api` 转发到后端。
|
||||
- **日志与追踪**：日志写入 `logs/app.log`。启用 Langfuse 时填写真实的公钥、私钥和服务地址；追踪故障不阻断主业务。
|
||||
|
||||
## 开发与测试
|
||||
|
||||
阅读主链路可从 [app/main.py](app/main.py)
|
||||
开始，依次查看 [Agent 接口](app/api/agent.py)、[AgentService](app/services/agent.py)、[外层工作流](app/processor/graph.py)、[大屏 Agent](app/agent/screen_agent.py)
|
||||
和 [动作模拟](app/actions/batch.py)。
|
||||
|
||||
扩展组件时修改 [物料定义](app/resources/materials)，导入数据库后重启服务；修改任务处理规则时查看 [skills](skills)；新增 MCP
|
||||
服务时修改 [mcp_servers.json](mcp_servers.json)，连接配置支持 `${环境变量名}` 展开。
|
||||
|
||||
当前 `.gitignore` 排除了 `tests/`、`docs/`、`scripts/` 等本地辅助目录，普通克隆不会包含这些内容。若工作区已有测试文件，在根目录使用标准库
|
||||
`unittest` 运行：
|
||||
|
||||
```sh
|
||||
uv run python -m unittest discover -s tests -v
|
||||
```
|
||||
|
||||
测试所需的 `openpyxl` 已列入开发依赖，`uv sync --locked` 默认安装。单元测试不能替代真实模型、Jev、MongoDB、用户中心和前端动作应用的联调验证。
|
||||
|
||||
## 常见问题
|
||||
|
||||
| 现象 | 检查项 |
|
||||
|---------------------------------------------|-------------------------------------------------------------------------------------------|
|
||||
| 会话或 Agent 接口返回 401 | 本地调试确认 `AUTH_ENABLED=false`；开启鉴权时携带有效 Bearer token |
|
||||
| 提示物料目录未加载，或 MCP 物料服务无法启动 | 检查 `MONGODB_URI`、物料初始化结果与 `8201` 端口，导入后重启服务 |
|
||||
| Agent 返回 422，提示缺少会话 | 先创建会话，并在 `context.conversation_id` 中传入返回的 ID |
|
||||
| 请求返回 409 | 同一会话或任务仍在运行，或尝试恢复未处于等待补充状态的任务 |
|
||||
| API 可访问，但 Agent 执行失败 | 查看流中的 `error`、`final.data.answer` 和 `logs/app.log`，核对模型、Jev 密钥与服务连通性 |
|
||||
| 上传返回 202，但文件暂时不可用 | 202 仅表示上传已接收，继续轮询解析进度，等待 `ready` |
|
||||
| 前端一直等不到流式反馈 | 检查是否持续读取响应流，以及代理是否缓冲 SSE 或提前超时 |
|
||||
+39
-2
@@ -7,7 +7,7 @@ from copy import deepcopy
|
||||
|
||||
from app.actions.layout import EXPANDERS
|
||||
from app.actions.registry import apply_action
|
||||
from app.actions.update import fit_existing_path
|
||||
from app.actions.update import fit_existing_path, same_value
|
||||
from app.resources.material_fields import check_update_field, find_material
|
||||
from app.schemas.actions import Action, AddAction, AgentAction, UpdateAction
|
||||
from app.schemas.capabilities import DATA_BINDING_ERROR, DATA_BINDING_SUPPORTED
|
||||
@@ -60,12 +60,47 @@ def complete_add(action: AddAction, catalog: ComponentCatalog) -> AddAction:
|
||||
return action.model_copy(update={"node": node, "sample_data": sample_data})
|
||||
|
||||
|
||||
def misplaced_props(actions: list[AgentAction], catalog: ComponentCatalog) -> list[str]:
|
||||
"""新增组件的 props 里物料没有的键：前端按物料的结构读取，多出来的键不会生效，多半是写错了层级
|
||||
|
||||
评测里模型把 series 写在 props 下而不是 props.option 下，物料默认的 series 留着，图表引用的列对不上；
|
||||
它看到的报错说 encode 用了别的列，和自己写的不一样，反复原样提交。返回给模型的说明，每项一个组件
|
||||
"""
|
||||
notes = []
|
||||
for action in actions:
|
||||
if not isinstance(action, AddAction):
|
||||
continue
|
||||
try:
|
||||
known = (
|
||||
find_material(catalog, action.node.type).template().get("props") or {}
|
||||
)
|
||||
except ValueError: # 物料不存在，补全时会报错
|
||||
continue
|
||||
written = (
|
||||
action.node.model_dump(mode="json", exclude_unset=True).get("props") or {}
|
||||
)
|
||||
for key in sorted(set(written) - set(known)):
|
||||
parent = next(
|
||||
(
|
||||
name
|
||||
for name, value in known.items()
|
||||
if isinstance(value, dict) and key in value
|
||||
),
|
||||
None,
|
||||
)
|
||||
where = f"（物料里在 props.{parent}.{key}）" if parent else ""
|
||||
notes.append(f"{action.node.id} 的 props.{key}{where}")
|
||||
return notes
|
||||
|
||||
|
||||
def run_batch(
|
||||
screen: Screen, catalog: ComponentCatalog, actions: list[AgentAction]
|
||||
) -> tuple[list[Action], Screen]:
|
||||
"""从原始大屏开始按顺序模拟执行，返回前端要执行的动作和执行后的大屏；第一个错误抛 BatchError
|
||||
"""从传入的大屏开始按顺序模拟执行，返回前端要执行的动作和执行后的大屏；第一个错误抛 BatchError
|
||||
|
||||
布局意图（arrange、align、distribute、set_layout）在这里按当时的大屏展开成 move 和 resize
|
||||
update 写入的值和字段现在的值完全一样时没有变化，跳过：不经过字段契约、也不交付；
|
||||
模型常把已经正确的值再写一遍（例如已经是 category 的轴类型），字段不在契约里时整批都会失败
|
||||
"""
|
||||
current = screen.model_copy(deep=True)
|
||||
completed: list[Action] = []
|
||||
@@ -86,6 +121,8 @@ def run_batch(
|
||||
None,
|
||||
)
|
||||
if node is not None:
|
||||
if same_value(node, action):
|
||||
continue
|
||||
check_update_field(catalog, node.type, action)
|
||||
# 字段定义允许、但组件中还没有的中间对象，折叠进值里再执行
|
||||
action = fit_existing_path(node, action)
|
||||
|
||||
+35
-2
@@ -1,5 +1,6 @@
|
||||
"""在大屏 JSON 副本上按路径更新组件属性，用于后端校验动作能否执行"""
|
||||
|
||||
import json
|
||||
from copy import deepcopy
|
||||
|
||||
from app.actions.base import BaseActionHandler
|
||||
@@ -7,6 +8,30 @@ from app.schemas.actions import UpdateAction
|
||||
from app.schemas.screen import Screen, ScreenNode
|
||||
|
||||
|
||||
def same_value(node: ScreenNode, action: UpdateAction) -> bool:
|
||||
"""字段已经存在、并且值和要写入的完全一样（按 JSON 比较，1 和 1.0、true 和 1 不算一样）"""
|
||||
if action.section == "name":
|
||||
current = node.name
|
||||
else:
|
||||
current = getattr(node, action.section)
|
||||
for segment in action.path:
|
||||
if (
|
||||
isinstance(current, dict)
|
||||
and isinstance(segment, str)
|
||||
and segment in current
|
||||
) or (
|
||||
isinstance(current, list)
|
||||
and type(segment) is int
|
||||
and segment < len(current)
|
||||
):
|
||||
current = current[segment]
|
||||
else:
|
||||
return False
|
||||
return json.dumps(current, sort_keys=True) == json.dumps(
|
||||
action.value, sort_keys=True
|
||||
)
|
||||
|
||||
|
||||
def fit_existing_path(node: ScreenNode, action: UpdateAction) -> UpdateAction:
|
||||
"""字段定义允许的路径在组件中可能还没有中间对象（例如未开启过的数值标签）
|
||||
|
||||
@@ -17,7 +42,11 @@ def fit_existing_path(node: ScreenNode, action: UpdateAction) -> UpdateAction:
|
||||
return action
|
||||
current = getattr(node, action.section) or {}
|
||||
for index, segment in enumerate(action.path[:-1]):
|
||||
if isinstance(current, dict) and isinstance(segment, str) and segment not in current:
|
||||
if (
|
||||
isinstance(current, dict)
|
||||
and isinstance(segment, str)
|
||||
and segment not in current
|
||||
):
|
||||
rest = action.path[index + 1 :]
|
||||
if not all(isinstance(item, str) for item in rest):
|
||||
return action
|
||||
@@ -29,7 +58,11 @@ def fit_existing_path(node: ScreenNode, action: UpdateAction) -> UpdateAction:
|
||||
)
|
||||
if isinstance(current, dict) and isinstance(segment, str):
|
||||
current = current[segment]
|
||||
elif isinstance(current, list) and type(segment) is int and segment < len(current):
|
||||
elif (
|
||||
isinstance(current, list)
|
||||
and type(segment) is int
|
||||
and segment < len(current)
|
||||
):
|
||||
current = current[segment]
|
||||
else:
|
||||
return action
|
||||
|
||||
@@ -10,6 +10,7 @@ from langchain.agents import create_agent
|
||||
from langchain.agents.middleware import ModelCallLimitMiddleware
|
||||
|
||||
from app.agent.data_tools import DATA_TOOLS, DataContext
|
||||
from app.agent.tool_args import normalize_tool_args
|
||||
from app.clients.llm import get_llm
|
||||
from app.config import settings
|
||||
from app.prompt.qa import QA_DATA_RULES, QA_SYSTEM_PROMPT
|
||||
@@ -37,6 +38,8 @@ def get_qa_agent(with_data: bool = False):
|
||||
),
|
||||
context_schema=QAContext,
|
||||
middleware=[
|
||||
# 有工具时才挂参数兼容：挂了 wrap_tool_call 中间件，没有工具也会多出一个工具节点
|
||||
*([normalize_tool_args] if with_data else []),
|
||||
# 超过上限时抛 ModelCallLimitExceededError，由节点按失败处理；
|
||||
# 写成 end 时中间件只插入一条英文说明，节点分不清是正常结束还是用尽了次数
|
||||
ModelCallLimitMiddleware(
|
||||
|
||||
@@ -16,6 +16,8 @@ from app.schemas.actions import AddAction, AgentAction, UpdateAction
|
||||
from app.schemas.screen import Screen
|
||||
|
||||
QUERY_KEY = "$query"
|
||||
# 图表物料的数据都在这里，series 的 encode 按它的列名引用；整个换成查询结果时数据就不再是示例数据
|
||||
SOURCE_PATH = "props.option.dataset.source"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -124,10 +126,30 @@ def _join(path: str, key: str | int) -> str:
|
||||
return f"{path}.{key}" if path else str(key)
|
||||
|
||||
|
||||
def encode_problems(screen: Screen, node_ids: set[str]) -> list[str]:
|
||||
"""用了 $query 的组件，series 的 encode 引用的列必须在 dataset.source 里
|
||||
def _binding(node) -> tuple[Any, Any]:
|
||||
option = node.props.get("option")
|
||||
if not isinstance(option, dict):
|
||||
return None, None
|
||||
return option.get("dataset"), option.get("series")
|
||||
|
||||
数据换成查询结果后列名变了，encode 还引用物料默认的列名时图表是空的，前端也不报错，所以在这里拦下
|
||||
|
||||
def rebound_nodes(before: Screen, after: Screen) -> set[str]:
|
||||
"""这一批改过数据或映射（option.dataset、option.series）的组件，新增的也算
|
||||
|
||||
只按前后状态比较，不看是哪种动作改的：$query、直接写数据、单独改 encode 都会让两者对不上；没动过的旧组件不核对
|
||||
"""
|
||||
old = {node.id: _binding(node) for node in before.nodes}
|
||||
return {
|
||||
node.id
|
||||
for node in after.nodes
|
||||
if node.id not in old or old[node.id] != _binding(node)
|
||||
}
|
||||
|
||||
|
||||
def encode_problems(screen: Screen, node_ids: set[str]) -> list[str]:
|
||||
"""指定组件的 series encode 引用的列必须在 dataset.source 里
|
||||
|
||||
数据换了列名、或 encode 改成了不存在的列时图表是空的，前端也不报错，所以在这里拦下
|
||||
"""
|
||||
problems = []
|
||||
for node in screen.nodes:
|
||||
@@ -161,11 +183,11 @@ def encode_problems(screen: Screen, node_ids: set[str]) -> list[str]:
|
||||
]
|
||||
if missing:
|
||||
wrong.append(
|
||||
f"series.{position}.encode.{key} 用了 {'、'.join(missing)}"
|
||||
f"props.option.series.{position}.encode.{key} 用了 {'、'.join(missing)}"
|
||||
)
|
||||
if wrong:
|
||||
problems.append(
|
||||
f"组件 {node.id}：{'；'.join(wrong)}，但 dataset.source 的列是 {'、'.join(columns)}；"
|
||||
"把 encode 改成查询结果的列名，或在 SQL 里用 AS 起对应的别名"
|
||||
f"组件 {node.id}：{'；'.join(wrong)}，但 props.option.dataset.source 的列是 {'、'.join(columns)}；"
|
||||
"把 encode 改成 dataset.source 里的列名；数据来自 $query 时也可以在 SQL 里用 AS 起对应的别名"
|
||||
)
|
||||
return problems
|
||||
|
||||
+96
-38
@@ -5,27 +5,31 @@
|
||||
组件目录和物料在物料 MCP，各类任务的做法在 skill；系统提示只放每次都要用、别处放不下的原则
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from typing import Any, Literal, Self
|
||||
|
||||
from langchain.agents import create_agent
|
||||
from langchain.agents.middleware import ModelCallLimitMiddleware, wrap_model_call
|
||||
from langchain.agents.middleware import ModelCallLimitMiddleware
|
||||
from langchain.agents.structured_output import ToolStrategy
|
||||
from langchain_core.messages import AIMessage, HumanMessage
|
||||
from langchain_core.messages import AIMessage, BaseMessage, SystemMessage
|
||||
from loguru import logger
|
||||
from pydantic import BaseModel, Field, model_validator
|
||||
from pydantic import BaseModel, Field, ValidationError, model_validator
|
||||
|
||||
from app.agent.data_tools import DATA_TOOLS
|
||||
from app.agent.skills import skills_prompt
|
||||
from app.agent.tool_args import normalize_tool_args
|
||||
from app.agent.tools import (
|
||||
ScreenAgentContext,
|
||||
ScreenAgentState,
|
||||
apply_actions,
|
||||
get_nodes,
|
||||
last_write_failed,
|
||||
)
|
||||
from app.clients.llm import get_llm
|
||||
from app.clients.mcp import get_mcp_tools
|
||||
from app.config import settings
|
||||
from app.prompt.domain import INPUT_AS_DATA
|
||||
from app.schemas.actions import Action, AddAction
|
||||
from app.schemas.clarification import Clarification
|
||||
|
||||
SYSTEM_PROMPT = f"""你是大屏编辑助手，按用户的要求新建大屏，新增、修改、删除组件，或调整画布
|
||||
@@ -45,13 +49,15 @@ SYSTEM_PROMPT = f"""你是大屏编辑助手，按用户的要求新建大屏，
|
||||
|
||||
class AgentReply(BaseModel):
|
||||
kind: Literal["done", "clarify", "unable"] = Field(
|
||||
description="done：最后一次 apply_actions 已成功，或请求不需要修改（例如只是提问）；"
|
||||
"clarify：需要用户回答才能继续；unable：无法完成"
|
||||
description="done：要求的修改都已写入草稿，且最后一次 apply_actions 成功，草稿交付给用户；"
|
||||
"或请求不需要修改（例如只是提问）；clarify：需要用户回答才能继续；unable：无法完成。"
|
||||
"clarify 和 unable 时草稿整个丢弃，用户看到的仍是原始大屏"
|
||||
)
|
||||
message: str = Field(
|
||||
default="",
|
||||
description="给用户看的一两句话。done：说明改了什么，用了默认值或示例数据时一并说明，没有修改时直接回答；"
|
||||
"unable：说明原因，有可行的替代做法时一并给出；clarify：不用填",
|
||||
description="给用户看的一两句话，不要提草稿、工具、参数这类内部说法。done：说明改了什么，用了默认值或示例数据时一并说明，没有修改时直接回答；"
|
||||
"unable：草稿整个丢弃，用户的大屏和请求前一样，不要说已经删除、新增或修改了什么；"
|
||||
"说明没能完成什么和原因，有可行的替代做法时一并给出；clarify：不用填",
|
||||
)
|
||||
clarification: Clarification | None = Field(
|
||||
default=None, description="kind=clarify 时填写，其他情况不填"
|
||||
@@ -64,39 +70,89 @@ class AgentReply(BaseModel):
|
||||
return self
|
||||
|
||||
|
||||
# 模型没有调用工具时最多纠正几次
|
||||
NO_TOOL_CALL_RETRIES = 2
|
||||
NO_TOOL_CALL_HINT = (
|
||||
"每一轮都要调用工具。要结束时调用 AgentReply，给用户的话写在 message 里"
|
||||
)
|
||||
FINALIZE_RULES = """<finalize>
|
||||
这次执行已经结束，不能再调用其他工具，也不能再修改大屏，只能调用 AgentReply 给出结论
|
||||
按上面的对话和执行情况判断：用户要求的修改没有全部完成时不能回复 done
|
||||
上面 apply_actions 返回的"执行成功""已删除"说的是这次的草稿，不是用户看到的大屏；回复 unable 时草稿整个丢弃，用户的大屏没有任何变化
|
||||
</finalize>"""
|
||||
|
||||
|
||||
@wrap_model_call
|
||||
async def require_tool_call(request, handler):
|
||||
"""ToolStrategy 要求每一轮都调用工具，结束时调用 AgentReply；模型偶尔只回文本（想直接回复用户）或空回复
|
||||
def reached_call_limit(messages: Sequence[BaseMessage]) -> bool:
|
||||
"""是不是因为调用次数用尽而结束：没结束的每次模型调用都带工具调用，数到上限就是被截断的"""
|
||||
calls = sum(
|
||||
isinstance(item, AIMessage) and bool(item.tool_calls) for item in messages
|
||||
)
|
||||
return calls >= settings.max_agent_steps
|
||||
|
||||
原样重试往往得到同样的结果，或者模型转而重复提交 apply_actions 直到调用次数用尽，
|
||||
所以把它的回复和纠正说明加在下一次调用里；只用于这一次调用，不写进 agent 的消息
|
||||
|
||||
def _describe(action: Action) -> str:
|
||||
if isinstance(action, AddAction):
|
||||
return f"add {action.node.id}"
|
||||
return f"{action.type} {getattr(action, 'target_id', '画布')}"
|
||||
|
||||
|
||||
def _run_state(result: dict, limit_hit: bool, texts: list[str]) -> str:
|
||||
"""收尾时告诉模型执行到了哪一步；对话里有每次工具结果，这里只给整体情况"""
|
||||
lines = [
|
||||
"调用次数已用完，执行被截断"
|
||||
if limit_hit
|
||||
else "上一轮没有调用 AgentReply 给出结论"
|
||||
]
|
||||
if actions := result.get("actions"):
|
||||
described = "；".join(_describe(item) for item in actions[:40])
|
||||
# 只写"已成功提交"时，模型回复 unable 也会说"已删除了某某"，用户以为修改已经生效
|
||||
lines.append(
|
||||
f"草稿里的动作 {len(actions)} 个，还没有交付：{described}；"
|
||||
"只有回复 done 才会交付，回复 unable 或 clarify 时草稿整个丢弃，用户的大屏保持原样"
|
||||
)
|
||||
else:
|
||||
lines.append("还没有成功提交过修改")
|
||||
if last_write_failed(result.get("messages", [])):
|
||||
lines.append("最后一次 apply_actions 失败，那一批没有写入草稿")
|
||||
if texts:
|
||||
lines.append(
|
||||
f"上一轮只输出了文本，没有调用 AgentReply：{' '.join(texts)[:1000]}"
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
async def finalize_reply(result: dict, *, limit_hit: bool) -> AgentReply | None:
|
||||
"""运行正常结束却没有结论时补一次结论：调用次数用尽，或模型只回了文本、空回复
|
||||
|
||||
只调用一次，只绑定 AgentReply 并强制调用，没有能修改大屏的工具，不会再进入修改循环；
|
||||
说明放在系统消息里，不追加用户消息，否则模型会把内部要求当成用户的话来回复
|
||||
强制调用也可能给出不合法的参数，校验不过就返回 None，由节点按失败处理，不重试
|
||||
"""
|
||||
response = await handler(request)
|
||||
for _ in range(NO_TOOL_CALL_RETRIES):
|
||||
message = response.result[0] if response.result else None
|
||||
if not isinstance(message, AIMessage) or message.tool_calls:
|
||||
break
|
||||
logger.warning(
|
||||
f"模型没有调用工具，纠正后重新调用: {str(message.content)[:200]}"
|
||||
)
|
||||
# 空回复不放回去，部分接口不接受内容为空的助手消息
|
||||
replied = [message] if message.content else []
|
||||
request = request.override(
|
||||
messages=[
|
||||
*request.messages,
|
||||
*replied,
|
||||
HumanMessage(content=NO_TOOL_CALL_HINT),
|
||||
]
|
||||
)
|
||||
response = await handler(request)
|
||||
return response
|
||||
messages = list(result["messages"])
|
||||
# 末尾没有工具调用的助手消息（模型的文本、空回复、调用上限的说明）不放进对话：
|
||||
# 以助手消息结尾时部分接口会把它当成续写的开头；模型的文本写进说明里
|
||||
trailing: list[AIMessage] = []
|
||||
while (
|
||||
messages and isinstance(messages[-1], AIMessage) and not messages[-1].tool_calls
|
||||
):
|
||||
trailing.insert(0, messages.pop())
|
||||
texts = (
|
||||
[] if limit_hit else [str(item.content) for item in trailing if item.content]
|
||||
)
|
||||
run_state = _run_state(result, limit_hit, texts)
|
||||
system = SystemMessage(
|
||||
f"{SYSTEM_PROMPT}\n\n{FINALIZE_RULES}\n\n<run_state>\n{run_state}\n</run_state>"
|
||||
)
|
||||
model = get_llm().bind_tools([AgentReply], tool_choice="AgentReply")
|
||||
response = await model.ainvoke(
|
||||
[system, *messages], config={"run_name": "finalize_reply", "tags": ["finalize"]}
|
||||
)
|
||||
call = next(
|
||||
(item for item in response.tool_calls if item["name"] == "AgentReply"), None
|
||||
)
|
||||
if call is None:
|
||||
logger.warning(f"收尾没有调用 AgentReply: {str(response.content)[:200]}")
|
||||
return None
|
||||
try:
|
||||
return AgentReply.model_validate(call["args"])
|
||||
except ValidationError as error:
|
||||
logger.warning(f"收尾给出的结论不合法: {error}")
|
||||
return None
|
||||
|
||||
|
||||
# 是否挂数据工具 → agent
|
||||
@@ -120,9 +176,11 @@ async def get_screen_agent(with_data: bool = False):
|
||||
response_format=ToolStrategy(AgentReply),
|
||||
state_schema=ScreenAgentState,
|
||||
context_schema=ScreenAgentContext,
|
||||
# 超过上限时插入说明并结束；和模型只回文本、空回复一样没有结论，由节点调用 finalize_reply 收尾
|
||||
# 模型没有调用工具时 agent 直接结束，不在循环里纠正：纠正说明会被当成用户的话，或让模型回到修改循环
|
||||
middleware=[
|
||||
require_tool_call,
|
||||
# 超过上限时插入说明并结束，没有结构化结果，由节点按失败处理
|
||||
# 列表、对象参数被写成 JSON 字符串时先解析回来，见 tool_args.py
|
||||
normalize_tool_args,
|
||||
ModelCallLimitMiddleware(
|
||||
run_limit=settings.max_agent_steps, exit_behavior="end"
|
||||
),
|
||||
|
||||
@@ -0,0 +1,264 @@
|
||||
"""工具参数兼容：模型把列表、对象参数写成 JSON 字符串时，按工具的参数 schema 解析回来
|
||||
|
||||
部分模型（例如 claude-haiku-5.5）在嵌套较深、内容较长时，常把 apply_actions 的 actions 整个写成 JSON 字符串，
|
||||
或把列表里的每一项写成字符串；参数校验报 "Input should be a valid list" 后它往往原样再交，评测里出现过
|
||||
连续 10 次直到调用次数用尽。这类问题和哪个工具无关，所以在工具调用入口统一处理，对本地工具和 MCP 工具都生效
|
||||
|
||||
- 合法的输入不动：schema 里任何一处接受字符串（包括联合类型的某个成员、数组某个成员的 items）时不解析；
|
||||
布尔 schema、解析不了的 $ref 这类看不准的情况按"什么都接受"处理，交给原有的参数校验
|
||||
- 只处理两层：参数本身，以及列表参数里的每一项
|
||||
- 写成字符串的长 JSON 多数只是结尾的括号多写或少写了一个：原文以括号结尾、模型输出没有因为长度被截断时，
|
||||
按还没闭合的括号重新补齐结尾；截断时丢掉的内容补不回来，补齐反而会得到一个意思不同的合法动作，所以不补
|
||||
- 仍然解析不出来时不执行，用中文说明错在哪里：参数校验的英文报错只说"应该是列表"，模型会照旧写成字符串
|
||||
- 解析或补齐后执行时，在工具结果末尾说明：补齐只能猜结尾，评测里模型把中间的 option 提前闭合、结尾又少一个括号，
|
||||
补齐后 series 落到了 props 下，执行结果对不上而它看不出原因，反复原样提交
|
||||
"""
|
||||
|
||||
import json
|
||||
from dataclasses import replace
|
||||
from typing import Any
|
||||
|
||||
from langchain.agents.middleware import wrap_tool_call
|
||||
from langchain_core.messages import AIMessage, ToolMessage
|
||||
from langchain_core.tools import BaseTool
|
||||
from langgraph.types import Command
|
||||
from loguru import logger
|
||||
|
||||
STRUCTURED = {"array", "object"}
|
||||
ANY = frozenset({"string", "number", "integer", "boolean", "null", "array", "object"})
|
||||
CLOSERS = {"[": "]", "{": "}"}
|
||||
# 各家接口表示"输出达到长度上限被截断"的写法
|
||||
TRUNCATED = {"length", "max_tokens"}
|
||||
|
||||
|
||||
def _resolve(schema: Any, defs: dict) -> Any:
|
||||
"""展开 $ref；指向的定义找不到时按 True（什么都接受）处理"""
|
||||
while isinstance(schema, dict) and "$ref" in schema:
|
||||
name = schema["$ref"].rsplit("/", 1)[-1]
|
||||
if name not in defs:
|
||||
return True
|
||||
schema = defs[name]
|
||||
return schema
|
||||
|
||||
|
||||
def _types(schema: Any, defs: dict) -> set[str]:
|
||||
"""schema 接受的 JSON 类型：anyOf、oneOf 取并集，allOf 取交集，没有类型约束时什么都接受"""
|
||||
schema = _resolve(schema, defs)
|
||||
if schema is False:
|
||||
return set()
|
||||
if not isinstance(schema, dict):
|
||||
return set(ANY)
|
||||
kind = schema.get("type")
|
||||
if kind:
|
||||
types = set(kind) if isinstance(kind, list) else {kind}
|
||||
elif "properties" in schema:
|
||||
types = {"object"}
|
||||
elif "items" in schema:
|
||||
types = {"array"}
|
||||
else:
|
||||
types = set(ANY)
|
||||
for member in schema.get("allOf", []):
|
||||
types &= _types(member, defs)
|
||||
if members := [*schema.get("anyOf", []), *schema.get("oneOf", [])]:
|
||||
types &= set().union(*(_types(member, defs) for member in members))
|
||||
return types
|
||||
|
||||
|
||||
def _items(schema: Any, defs: dict) -> list:
|
||||
"""各个数组成员的 items schema；有数组成员没有限定 items 时，其中有一个 True（什么都接受）"""
|
||||
schema = _resolve(schema, defs)
|
||||
if not isinstance(schema, dict):
|
||||
return [True] if schema is True else []
|
||||
kind = schema.get("type")
|
||||
own = set(kind) if isinstance(kind, list) else {kind}
|
||||
found = []
|
||||
if "prefixItems" in schema:
|
||||
# 每个位置各有 schema 的元组，不逐项解析
|
||||
found.append(True)
|
||||
elif "items" in schema:
|
||||
found.append(schema["items"])
|
||||
elif "array" in own:
|
||||
found.append(True)
|
||||
for key in ("anyOf", "oneOf", "allOf"):
|
||||
for member in schema.get(key, []):
|
||||
found.extend(_items(member, defs))
|
||||
return found
|
||||
|
||||
|
||||
def _reclose(text: str) -> str | None:
|
||||
"""去掉结尾连续的 ] 和 }，按还没闭合的括号重新补齐；结尾停在字符串中间时补不了"""
|
||||
body = text.rstrip().rstrip("]} \t\r\n")
|
||||
stack: list[str] = []
|
||||
in_string = escaped = False
|
||||
for char in body:
|
||||
if in_string:
|
||||
if escaped:
|
||||
escaped = False
|
||||
elif char == "\\":
|
||||
escaped = True
|
||||
elif char == '"':
|
||||
in_string = False
|
||||
elif char == '"':
|
||||
in_string = True
|
||||
elif char in CLOSERS:
|
||||
stack.append(char)
|
||||
elif char in "]}" and (not stack or CLOSERS[stack.pop()] != char):
|
||||
# 中间就有括号对不上，不是结尾的问题，补不了
|
||||
return None
|
||||
if in_string or not stack:
|
||||
return None
|
||||
return body + "".join(CLOSERS[item] for item in reversed(stack))
|
||||
|
||||
|
||||
def _matches(value: Any, types: set[str]) -> bool:
|
||||
return (isinstance(value, list) and "array" in types) or (
|
||||
isinstance(value, dict) and "object" in types
|
||||
)
|
||||
|
||||
|
||||
def _parse(value: Any, types: set[str], repair: bool) -> tuple[Any, str | None, bool]:
|
||||
"""schema 只接受列表或对象、却收到字符串时解析成 JSON；返回 (值, 解析不了时的说明, 是否补了结尾括号)"""
|
||||
if not isinstance(value, str) or "string" in types or not types & STRUCTURED:
|
||||
return value, None, False
|
||||
expected = "数组" if "array" in types else "对象"
|
||||
patched = False
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
except json.JSONDecodeError as error:
|
||||
# 原文不以括号结尾多半是被截断了，补齐结尾会丢掉后面的内容，不补
|
||||
ends_closed = value.rstrip().endswith(("]", "}"))
|
||||
repaired = _reclose(value) if repair and ends_closed else None
|
||||
try:
|
||||
parsed, patched = (json.loads(repaired) if repaired else None), True
|
||||
except json.JSONDecodeError:
|
||||
parsed = None
|
||||
if not _matches(parsed, types):
|
||||
problem = (
|
||||
f"应该直接写成 JSON {expected}，不要写成字符串；"
|
||||
f"这个字符串也不是合法的 JSON（第 {error.pos} 个字符附近：{error.msg}）"
|
||||
)
|
||||
return value, problem, False
|
||||
if not _matches(parsed, types):
|
||||
return value, f"应该直接写成 JSON {expected}，不要写成字符串", False
|
||||
return parsed, None, patched
|
||||
|
||||
|
||||
def normalize_args(
|
||||
tool: BaseTool, args: dict, *, repair: bool = True
|
||||
) -> tuple[dict, list[str], list[str]]:
|
||||
"""按工具的参数 schema 把写成字符串的列表、对象参数解析回来
|
||||
|
||||
返回新的参数（不修改传入的）、解析不了的参数说明、补了结尾括号的参数名；有说明时不该执行这次调用
|
||||
repair=False 时不补括号，用于模型输出被截断的情况
|
||||
"""
|
||||
schema = tool.tool_call_schema
|
||||
schema = schema if isinstance(schema, dict) else schema.model_json_schema()
|
||||
defs = schema.get("$defs", {})
|
||||
properties = schema.get("properties", {})
|
||||
fixed, problems, repaired = {}, [], []
|
||||
for name, value in args.items():
|
||||
prop = properties.get(name)
|
||||
if prop is None:
|
||||
fixed[name] = value
|
||||
continue
|
||||
value, problem, patched = _parse(value, _types(prop, defs), repair)
|
||||
if problem:
|
||||
problems.append(f"{name} {problem}")
|
||||
item_schemas = _items(prop, defs)
|
||||
if isinstance(value, list) and item_schemas:
|
||||
item_types = set().union(*(_types(item, defs) for item in item_schemas))
|
||||
parsed_items = []
|
||||
for index, item in enumerate(value):
|
||||
item, problem, item_patched = _parse(item, item_types, repair)
|
||||
if problem:
|
||||
problems.append(f"{name} 的第 {index + 1} 项{problem}")
|
||||
patched = patched or item_patched
|
||||
parsed_items.append(item)
|
||||
value = parsed_items
|
||||
if patched:
|
||||
repaired.append(name)
|
||||
fixed[name] = value
|
||||
return fixed, problems, repaired
|
||||
|
||||
|
||||
def _truncated(state: Any, call_id: str) -> bool:
|
||||
"""发起这次调用的模型输出是不是因为长度上限被截断"""
|
||||
messages = state.get("messages", []) if isinstance(state, dict) else []
|
||||
message = next(
|
||||
(
|
||||
item
|
||||
for item in reversed(messages)
|
||||
if isinstance(item, AIMessage)
|
||||
and any(call["id"] == call_id for call in item.tool_calls)
|
||||
),
|
||||
None,
|
||||
)
|
||||
metadata = message.response_metadata if message is not None else {}
|
||||
return bool(
|
||||
{metadata.get("finish_reason"), metadata.get("stop_reason")} & TRUNCATED
|
||||
)
|
||||
|
||||
|
||||
@wrap_tool_call
|
||||
async def normalize_tool_args(request, handler):
|
||||
"""在参数校验前统一修正参数；没有改动时原样执行，解析不了时返回中文说明，不执行"""
|
||||
tool = request.tool
|
||||
args = request.tool_call.get("args")
|
||||
if tool is None or not isinstance(args, dict):
|
||||
return await handler(request)
|
||||
truncated = _truncated(request.state, request.tool_call["id"])
|
||||
fixed, problems, repaired = normalize_args(tool, args, repair=not truncated)
|
||||
if problems:
|
||||
logger.info(
|
||||
f"工具 {tool.name} 的参数写成了不合法的 JSON 字符串"
|
||||
f"{'（模型输出被截断）' if truncated else ''}: {problems}"
|
||||
)
|
||||
return ToolMessage(
|
||||
"参数格式错误，这次调用没有执行：\n"
|
||||
+ "\n".join(f"- {item}" for item in problems)
|
||||
+ ("\n- 输出达到了长度上限被截断，分成几批提交" if truncated else ""),
|
||||
tool_call_id=request.tool_call["id"],
|
||||
name=tool.name,
|
||||
status="error",
|
||||
)
|
||||
if fixed == args:
|
||||
return await handler(request)
|
||||
changed = [name for name in args if fixed[name] != args[name]]
|
||||
strict = [name for name in changed if name not in repaired]
|
||||
notes = []
|
||||
if strict:
|
||||
logger.info(
|
||||
f"工具 {tool.name} 的参数 {'、'.join(strict)} 写成了 JSON 字符串，已解析"
|
||||
)
|
||||
notes.append(f"{'、'.join(strict)} 写成了 JSON 字符串，已解析后执行")
|
||||
if repaired:
|
||||
logger.info(
|
||||
f"工具 {tool.name} 的参数 {'、'.join(repaired)} 写成了 JSON 字符串，结尾括号对不上，补齐后解析"
|
||||
)
|
||||
# 补的是结尾，括号多写或少写在中间时，补齐后层级是错的，但它是合法的 JSON，执行时不一定报错
|
||||
notes.append(
|
||||
f"{'、'.join(repaired)} 写成了 JSON 字符串，结尾的括号数量不对，已按还没闭合的括号补齐后执行；"
|
||||
"如果中间的括号写错了，补齐后的层级也是错的，对照执行结果核对"
|
||||
)
|
||||
result = await handler(
|
||||
request.override(tool_call={**request.tool_call, "args": fixed})
|
||||
)
|
||||
# 不告诉模型时它会一直写成字符串，补齐猜错层级时它也看不出原因
|
||||
return _with_note(
|
||||
result, "参数说明：" + "；".join(notes) + "；以后直接写成 JSON，不要写成字符串"
|
||||
)
|
||||
|
||||
|
||||
def _with_note(result: Any, note: str) -> Any:
|
||||
"""在工具返回的消息末尾加一句说明；工具返回 Command 时改其中的工具消息"""
|
||||
if isinstance(result, ToolMessage):
|
||||
return result.model_copy(update={"content": f"{result.content}\n{note}"})
|
||||
if isinstance(result, Command) and isinstance(result.update, dict):
|
||||
messages = [
|
||||
item.model_copy(update={"content": f"{item.content}\n{note}"})
|
||||
if isinstance(item, ToolMessage)
|
||||
else item
|
||||
for item in result.update.get("messages", [])
|
||||
]
|
||||
return replace(result, update={**result.update, "messages": messages})
|
||||
return result
|
||||
+221
-69
@@ -1,6 +1,7 @@
|
||||
"""screen agent 的本地工具：读取本次请求的大屏、提交整批动作
|
||||
"""screen agent 的本地工具：读取当前大屏、在当前大屏上提交一批修改
|
||||
|
||||
它们读写的是本次请求的大屏和 agent 状态，所以放在本地；物料等能力由 MCP 提供
|
||||
读取、修改和校验都基于同一份草稿：原始大屏加上之前写入草稿的修改；回复 done 时草稿交付给用户
|
||||
"""
|
||||
|
||||
from collections.abc import Sequence
|
||||
@@ -9,44 +10,57 @@ from typing import Annotated, NotRequired
|
||||
|
||||
from langchain.agents import AgentState
|
||||
from langchain.tools import ToolRuntime, tool
|
||||
from langchain_core.messages import ToolMessage
|
||||
from langchain_core.messages import AIMessage, BaseMessage, ToolMessage
|
||||
from langgraph.types import Command
|
||||
from pydantic import Field
|
||||
|
||||
from app.actions.batch import BatchError, run_batch
|
||||
from app.actions.batch import BatchError, misplaced_props, run_batch
|
||||
from app.agent.data_tools import DataContext
|
||||
from app.agent.query_values import EmbeddedQuery, encode_problems, fill_queries
|
||||
from app.schemas.actions import Action, AddAction, AgentAction
|
||||
from app.agent.query_values import (
|
||||
SOURCE_PATH,
|
||||
EmbeddedQuery,
|
||||
encode_problems,
|
||||
fill_queries,
|
||||
rebound_nodes,
|
||||
)
|
||||
from app.schemas.actions import (
|
||||
Action,
|
||||
AddAction,
|
||||
AgentAction,
|
||||
DeleteAction,
|
||||
UpdateAction,
|
||||
)
|
||||
from app.schemas.component_resource import ComponentCatalog
|
||||
from app.schemas.screen import Screen
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ScreenAgentContext:
|
||||
screen: Screen # 本次请求的原始大屏，只读
|
||||
screen: Screen # 本次请求的原始大屏，只读，交付的动作从它开始执行
|
||||
catalog: ComponentCatalog # 当前物料目录，只用于补全和校验
|
||||
data: DataContext = field(default_factory=DataContext) # 本会话可用的数据表
|
||||
|
||||
|
||||
def _latest(_, new):
|
||||
# 同一步里有多次提交时取最后一个，避免并行工具调用写同一字段时报错
|
||||
return new
|
||||
|
||||
|
||||
# 两个字段只由 apply_actions 写，同一轮只执行第一次提交，所以一步里最多一次写入，不需要合并规则
|
||||
# 提交是否失败不另存字段，由 last_write_failed 从消息推出：参数解析、参数校验、查询、执行各处的失败都只体现在消息上
|
||||
class ScreenAgentState(AgentState):
|
||||
# 最近一次 apply_actions 成功时补全后的动作；失败时清空，以最后一次提交为准
|
||||
actions: NotRequired[Annotated[list[Action] | None, _latest]]
|
||||
# 最近一次 apply_actions 失败的原因，成功时清空
|
||||
apply_error: NotRequired[Annotated[str | None, _latest]]
|
||||
# 草稿：最近一次写入后的大屏；还没有写入时不存在，用原始大屏
|
||||
working_screen: NotRequired[Screen]
|
||||
# 交付记录：各次写入的动作（补全、展开后）按顺序拼接，在原始大屏上依次执行得到草稿
|
||||
actions: NotRequired[list[Action]]
|
||||
|
||||
|
||||
def current_screen(runtime: ToolRuntime[ScreenAgentContext]) -> Screen:
|
||||
return runtime.state.get("working_screen") or runtime.context.screen
|
||||
|
||||
|
||||
@tool
|
||||
def get_nodes(node_ids: list[str], runtime: ToolRuntime[ScreenAgentContext]) -> dict:
|
||||
"""读取组件的完整配置（layout、props、style 等）。修改属性或迁移数据前先读取
|
||||
|
||||
读取的是原始大屏，不包含本次已提交的修改
|
||||
读取的是草稿：原始大屏加上之前写入草稿的修改；和 apply_actions 在同一轮调用时读不到这次提交的结果，要核对就在下一轮读取
|
||||
"""
|
||||
nodes = {node.id: node for node in runtime.context.screen.nodes}
|
||||
nodes = {node.id: node for node in current_screen(runtime).nodes}
|
||||
return {
|
||||
"nodes": [
|
||||
nodes[item].model_dump(mode="json") for item in node_ids if item in nodes
|
||||
@@ -61,13 +75,38 @@ def _result_text(
|
||||
actions: list[Action],
|
||||
undeclared_sample: frozenset[str] = frozenset(),
|
||||
queries: Sequence[EmbeddedQuery] = (),
|
||||
unchanged: bool = False,
|
||||
repeated: bool = False,
|
||||
relabeled: Sequence[str] = (),
|
||||
) -> str:
|
||||
"""执行结果：除了动作数量，还给出执行后的画布和每个组件的布局
|
||||
|
||||
模型提交后才第一次看到结果，能自己发现位置超出画布、互相重叠、没排整齐等问题再修正，
|
||||
不用靠重复提交来试；本次动过布局或新增的组件前面标 *，已删除的单独列出
|
||||
模型提交后才第一次看到结果，能自己发现位置超出画布、互相重叠、没排整齐等问题再修正；
|
||||
这一批动过布局或新增的组件前面标 *，这一批删除的单独列出
|
||||
没有带来变化时也照样给出布局：只说"已经写入"会让模型以为达到了目标，可能当前大屏本来就是这样、并不符合要求
|
||||
统一说"写入草稿"、"回复 done 时交付"：说"已生效"时模型回复 unable 也会说已经改了什么，
|
||||
不说怎么交付时它会再交一批空的或重复的来"确认"
|
||||
"""
|
||||
lines = [f"执行成功，共 {len(actions)} 个动作已生效"]
|
||||
if unchanged and repeated:
|
||||
# 之前的提交已经改过：模型不确定是否写入才会重复提交，要写明已经在草稿里、不要重复，并说怎么结束，
|
||||
# 只说"没有变化"时它以为修改没成功，会换着字段反复提交直到调用次数用尽
|
||||
lines = [
|
||||
(
|
||||
"执行成功，但这一批没有改变草稿（0 个动作）：提交的内容和草稿里已有的一样，之前的提交已经改过。"
|
||||
"对照要求和下面的布局：已经达到就调用 AgentReply 回复 done，草稿会一起交付；"
|
||||
"没达到就提交不同的修改，不要原样重复提交"
|
||||
)
|
||||
]
|
||||
elif unchanged:
|
||||
# 这次运行还没有改过大屏：原始大屏本来就是这样，模型的方案并没有改动任何东西，只陈述事实；
|
||||
# 加上"完成就回复"之类的提示时，大屏并不符合要求它也会当成完成
|
||||
lines = [
|
||||
"执行成功，共 0 个动作写入草稿：提交的修改和原始大屏已有的一样，大屏原本就是这样"
|
||||
]
|
||||
else:
|
||||
lines = [
|
||||
f"执行成功，共 {len(actions)} 个动作已写入草稿；回复 done 时草稿交付给用户，不需要再提交确认"
|
||||
]
|
||||
sample = [
|
||||
action.node.id
|
||||
for action in actions
|
||||
@@ -78,10 +117,15 @@ def _result_text(
|
||||
if guessed := [item for item in sample if item in undeclared_sample]:
|
||||
lines.append(
|
||||
f"未写 sample_data，按物料默认标成了示例数据：{'、'.join(guessed)}；"
|
||||
"如果这些组件展示的数据都来自大屏上已有的组件或上传的数据文件，把 sample_data 设为 false 重新提交整批，否则不用处理"
|
||||
"如果这些组件展示的数据都来自大屏上已有的组件或上传的数据文件，"
|
||||
"用 $query 整体替换它们的 option.dataset.source，或再提交一批：先 delete，再用 sample_data=false 重新 add，否则不用处理"
|
||||
)
|
||||
if declared := [item for item in sample if item not in undeclared_sample]:
|
||||
lines.append(f"标为示例数据的新增组件：{'、'.join(declared)}")
|
||||
if relabeled:
|
||||
lines.append(
|
||||
f"数据整体换成了 $query 的查询结果，已改标为非示例数据：{'、'.join(relabeled)}"
|
||||
)
|
||||
# 模型提交时只写了 SQL，看不到结果；给出列名和行数，用来核对 encode 和数据量
|
||||
if queries:
|
||||
lines.append("$query 已替换为查询结果：")
|
||||
@@ -110,74 +154,182 @@ def _result_text(
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _reply(
|
||||
runtime: ToolRuntime, text: str, *, error: bool = False, **update
|
||||
) -> Command:
|
||||
"""返回工具消息；状态字段由调用方显式传入，没传的保持不变"""
|
||||
status = "error" if error else "success"
|
||||
message = ToolMessage(text, tool_call_id=runtime.tool_call_id, status=status)
|
||||
return Command(update={**update, "messages": [message]})
|
||||
|
||||
|
||||
def _first_write(message: AIMessage) -> str | None:
|
||||
"""一条模型消息里第一个 apply_actions 调用的 ID；同一轮只执行它"""
|
||||
return next(
|
||||
(call["id"] for call in message.tool_calls if call["name"] == "apply_actions"),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _is_first_submission(runtime: ToolRuntime) -> bool:
|
||||
"""按模型消息里 tool_calls 的顺序，这次调用是不是第一个 apply_actions
|
||||
|
||||
同一轮的多个调用都从同一份草稿出发，都执行会互相覆盖；只执行第一个，与并发谁先跑完无关
|
||||
直接调用工具、找不到发起调用的模型消息时，视为第一个
|
||||
"""
|
||||
message = next(
|
||||
(
|
||||
item
|
||||
for item in reversed(runtime.state.get("messages", []))
|
||||
if isinstance(item, AIMessage)
|
||||
and any(call["id"] == runtime.tool_call_id for call in item.tool_calls)
|
||||
),
|
||||
None,
|
||||
)
|
||||
return message is None or _first_write(message) == runtime.tool_call_id
|
||||
|
||||
|
||||
def last_write_failed(messages: Sequence[BaseMessage]) -> bool:
|
||||
"""最后一次被接受的提交是否失败：最后一条带 apply_actions 的模型消息里，第一个调用的结果是错误
|
||||
|
||||
参数解析、参数校验、查询、执行任何一处失败，结果都是 status="error" 的工具消息，都在这里判断，
|
||||
不靠各处各自记状态，以后多一种失败也不会漏；同一轮被拒绝的后续调用不是第一个，不算
|
||||
没有结果时也算失败：这一批没有执行，最后想做的修改没有写入草稿
|
||||
"""
|
||||
for position in range(len(messages) - 1, -1, -1):
|
||||
message = messages[position]
|
||||
if isinstance(message, AIMessage) and (call_id := _first_write(message)):
|
||||
reply = next(
|
||||
(
|
||||
item
|
||||
for item in messages[position + 1 :]
|
||||
if isinstance(item, ToolMessage) and item.tool_call_id == call_id
|
||||
),
|
||||
None,
|
||||
)
|
||||
return reply is None or reply.status == "error"
|
||||
return False
|
||||
|
||||
|
||||
@tool
|
||||
def apply_actions(
|
||||
actions: Annotated[
|
||||
list[AgentAction],
|
||||
# 空的一批会清掉之前成功的结果；参数校验失败不改状态，之前的结果保留
|
||||
Field(min_length=1, description="至少一个动作；不需要修改时不要调用，直接回复"),
|
||||
# 不限制至少一个：以前整批替换时空的一批会清掉之前的结果，增量执行下空的一批只是没有变化，
|
||||
# 按普通的无变化回复，比参数校验的英文报错更能让模型知道该结束了
|
||||
Field(description="这一批要执行的动作；不需要修改时不要调用，直接回复"),
|
||||
],
|
||||
runtime: ToolRuntime[ScreenAgentContext],
|
||||
) -> Command:
|
||||
"""提交本次要对大屏做的全部修改
|
||||
"""把一批修改写进这次的草稿；回复 done 时草稿交付给用户，不需要再提交确认
|
||||
|
||||
每次调用都是完整的一批，从原始大屏开始按顺序模拟执行，不在上一次调用的结果上累加
|
||||
成功时返回执行后的画布和各组件布局，用它核对结果，不对再重新提交整批
|
||||
失败时整批不生效，返回带动作序号的错误，修正后重新提交整批；以最后一次调用的结果为准
|
||||
第一次提交从原始大屏开始，之后每一批都在草稿上继续修改；回复 unable 或 clarify 时草稿整个丢弃
|
||||
成功时返回这一批执行后的画布和各组件布局，用它核对结果；结果不对再提交一批修正，已经写入的动作不用重复提交
|
||||
一批里任何一个动作失败，这一批都不写入草稿，之前写入的保留；按带序号的错误修正这一批后再提交
|
||||
同一轮只执行第一次调用，同一轮的 get_nodes 也读不到这次的结果；看到执行结果后再调用 AgentReply
|
||||
"""
|
||||
if not _is_first_submission(runtime):
|
||||
# 被拒绝的调用不写任何状态字段，不会覆盖第一次提交的结果
|
||||
return _reply(
|
||||
runtime,
|
||||
"同一轮只执行第一次 apply_actions，这一批没有执行；需要的话在下一轮根据执行结果再提交",
|
||||
error=True,
|
||||
)
|
||||
before = current_screen(runtime)
|
||||
misplaced = misplaced_props(actions, runtime.context.catalog)
|
||||
# 写错层级的属性不生效，失败和成功时都说明；它常是 encode 对不上的真正原因
|
||||
hint = (
|
||||
"\n新增组件写了物料没有的属性，前端不会使用："
|
||||
+ "；".join(misplaced)
|
||||
+ "；检查是不是写错了层级"
|
||||
if misplaced
|
||||
else ""
|
||||
)
|
||||
try:
|
||||
actions, queries = fill_queries(actions, runtime.context.data)
|
||||
completed, after = run_batch(
|
||||
runtime.context.screen, runtime.context.catalog, actions
|
||||
)
|
||||
problems = encode_problems(after, {item.node_id for item in queries})
|
||||
completed, after = run_batch(before, runtime.context.catalog, actions)
|
||||
# 这一批改过数据或映射的组件都核对，不论数据来自这一批的 $query、之前的提交还是直接写入
|
||||
problems = encode_problems(after, rebound_nodes(before, after))
|
||||
except BatchError as error:
|
||||
problems = [str(error)]
|
||||
if problems:
|
||||
message = "执行失败，整批未生效：\n" + "\n".join(
|
||||
# 失败只体现在工具消息上，由 last_write_failed 判断，不写状态字段
|
||||
message = "执行失败，这一批未写入草稿，之前写入的修改保留：\n" + "\n".join(
|
||||
f"- {item}" for item in problems
|
||||
)
|
||||
return Command(
|
||||
update={
|
||||
"actions": None,
|
||||
"apply_error": message,
|
||||
"messages": [
|
||||
ToolMessage(
|
||||
message, tool_call_id=runtime.tool_call_id, status="error"
|
||||
)
|
||||
],
|
||||
}
|
||||
)
|
||||
# 和上一次成功的提交完全一样：结果早已生效，再提交不会有变化，直接告诉模型结束，不再空转
|
||||
if completed == runtime.state.get("actions"):
|
||||
return Command(
|
||||
update={
|
||||
"messages": [
|
||||
ToolMessage(
|
||||
"这一批的结果和上一次提交完全相同，已经生效，不要重复提交；"
|
||||
"任务完成就调用 AgentReply 回复用户，没完成就改成不同的内容",
|
||||
tool_call_id=runtime.tool_call_id,
|
||||
)
|
||||
],
|
||||
}
|
||||
)
|
||||
return _reply(runtime, message + hint, error=True)
|
||||
undeclared = frozenset(
|
||||
action.node.id
|
||||
for action in actions
|
||||
if isinstance(action, AddAction)
|
||||
and "sample_data" not in action.model_fields_set
|
||||
)
|
||||
return Command(
|
||||
update={
|
||||
"actions": completed,
|
||||
"apply_error": None,
|
||||
"messages": [
|
||||
ToolMessage(
|
||||
_result_text(
|
||||
runtime.context.screen, after, completed, undeclared, queries
|
||||
),
|
||||
tool_call_id=runtime.tool_call_id,
|
||||
)
|
||||
],
|
||||
}
|
||||
log = runtime.state.get("actions", [])
|
||||
delivered, relabeled = _relabel_real_data(
|
||||
[*log, *completed], _replaced_sources(actions, queries)
|
||||
)
|
||||
# 按交付结果判断有没有变化，不比较参数：同样的 move 再执行一次会继续移动，同样的 update 再执行一次没有变化
|
||||
# 交付结果除了大屏还有只在动作上的示例数据标记，只比大屏会把"重新 add 改 sample_data"当成没变化而丢掉
|
||||
unchanged = after == before and _sample_marks(delivered) == _sample_marks(log)
|
||||
# 这次运行之前已有成功提交时，没有变化多半是在重复自己的修改；还没有时说明原始大屏本来就是这样
|
||||
text = _result_text(
|
||||
before,
|
||||
after,
|
||||
delivered[len(log) :],
|
||||
undeclared,
|
||||
queries,
|
||||
unchanged,
|
||||
repeated=bool(log),
|
||||
relabeled=relabeled,
|
||||
)
|
||||
if unchanged:
|
||||
return _reply(runtime, text + hint)
|
||||
return _reply(runtime, text + hint, working_screen=after, actions=delivered)
|
||||
|
||||
|
||||
def _replaced_sources(
|
||||
actions: Sequence[AgentAction], queries: Sequence[EmbeddedQuery]
|
||||
) -> set[str]:
|
||||
"""这一批用 update 把整个 dataset.source 换成 $query 结果的组件；只替换其中一部分时不算"""
|
||||
return {
|
||||
item.node_id
|
||||
for item in queries
|
||||
if item.path == SOURCE_PATH and isinstance(actions[item.index], UpdateAction)
|
||||
}
|
||||
|
||||
|
||||
def _relabel_real_data(
|
||||
actions: list[Action], node_ids: set[str]
|
||||
) -> tuple[list[Action], list[str]]:
|
||||
"""数据整体换成查询结果的组件，把它最近一次 add 的示例数据标记改为 false
|
||||
|
||||
标记只在 add 动作上、不影响大屏，改它不影响"交付记录依次执行得到草稿"；
|
||||
原始大屏上的组件没有 add，跨请求的标记要等大屏和前端契约支持
|
||||
"""
|
||||
result = list(actions)
|
||||
relabeled = []
|
||||
for node_id in sorted(node_ids):
|
||||
index = next(
|
||||
(
|
||||
position
|
||||
for position in range(len(result) - 1, -1, -1)
|
||||
if isinstance(result[position], AddAction)
|
||||
and result[position].node.id == node_id
|
||||
),
|
||||
None,
|
||||
)
|
||||
if index is not None and result[index].sample_data:
|
||||
result[index] = result[index].model_copy(update={"sample_data": False})
|
||||
relabeled.append(node_id)
|
||||
return result, relabeled
|
||||
|
||||
|
||||
def _sample_marks(actions: Sequence[Action]) -> dict[str, bool]:
|
||||
"""按交付记录，当前每个新增组件是否标为示例数据；这个标记只在 add 动作上，不在大屏里"""
|
||||
marks: dict[str, bool] = {}
|
||||
for action in actions:
|
||||
if isinstance(action, AddAction):
|
||||
marks[action.node.id] = action.sample_data
|
||||
elif isinstance(action, DeleteAction):
|
||||
marks.pop(action.target_id, None)
|
||||
return marks
|
||||
|
||||
@@ -8,22 +8,34 @@ from app.common.response import ApiResponse, success_response
|
||||
from app.core.auth import get_current_user
|
||||
from app.core.exceptions import BizException
|
||||
from app.schemas.conversation import (
|
||||
ConversationCreate,
|
||||
ConversationCreated,
|
||||
ConversationDetail,
|
||||
ConversationList,
|
||||
ConversationRename,
|
||||
)
|
||||
from app.schemas.user import CurrentUser
|
||||
|
||||
MAX_OFFSET = 10000
|
||||
|
||||
router = APIRouter(prefix="/conversations", tags=["Conversation"])
|
||||
|
||||
|
||||
@router.post("")
|
||||
async def create_conversation(
|
||||
request: Request, user: CurrentUser = Depends(get_current_user)
|
||||
request: Request,
|
||||
user: CurrentUser = Depends(get_current_user),
|
||||
body: ConversationCreate | None = None,
|
||||
) -> ApiResponse[ConversationCreated]:
|
||||
store = request.app.state.conversation_store
|
||||
screen_id = body.screen_id if body else None
|
||||
# 会话挂到大屏上，大屏必须是自己的；没配 MongoDB（没有大屏服务）时无从校验，直接放过
|
||||
screens = getattr(request.app.state, "screen_service", None)
|
||||
if screen_id is not None and screens is not None:
|
||||
if not await screens.exists(screen_id, user.id):
|
||||
raise BizException(404, "大屏不存在或已删除")
|
||||
return success_response(
|
||||
ConversationCreated(conversation_id=await store.create(user.id))
|
||||
ConversationCreated(conversation_id=await store.create(user.id, screen_id))
|
||||
)
|
||||
|
||||
|
||||
@@ -32,11 +44,15 @@ async def list_conversations(
|
||||
request: Request,
|
||||
user: Annotated[CurrentUser, Depends(get_current_user)],
|
||||
limit: Annotated[int, Query(ge=1, le=100)] = 20,
|
||||
offset: Annotated[int, Query(ge=0)] = 0,
|
||||
offset: Annotated[int, Query(ge=0, le=MAX_OFFSET)] = 0,
|
||||
screen_id: Annotated[str | None, Query(min_length=1)] = None,
|
||||
) -> ApiResponse[ConversationList]:
|
||||
store = request.app.state.conversation_store
|
||||
return success_response(
|
||||
ConversationList(conversations=await store.list_by_user(user.id, limit, offset))
|
||||
ConversationList(
|
||||
conversations=await store.list_by_user(user.id, limit, offset, screen_id),
|
||||
total=await store.count_by_user(user.id, screen_id),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@@ -54,3 +70,28 @@ async def get_conversation(
|
||||
return success_response(
|
||||
ConversationDetail(conversation_id=conversation_id, turns=turns)
|
||||
)
|
||||
|
||||
|
||||
@router.patch("/{conversation_id}")
|
||||
async def rename_conversation(
|
||||
conversation_id: str,
|
||||
body: ConversationRename,
|
||||
request: Request,
|
||||
user: CurrentUser = Depends(get_current_user),
|
||||
) -> ApiResponse[dict]:
|
||||
store = request.app.state.conversation_store
|
||||
if not await store.rename(conversation_id, user.id, body.title):
|
||||
raise BizException(404, "会话不存在或已过期")
|
||||
return success_response()
|
||||
|
||||
|
||||
@router.delete("/{conversation_id}")
|
||||
async def delete_conversation(
|
||||
conversation_id: str,
|
||||
request: Request,
|
||||
user: CurrentUser = Depends(get_current_user),
|
||||
) -> ApiResponse[dict]:
|
||||
store = request.app.state.conversation_store
|
||||
if not await store.delete(conversation_id, user.id):
|
||||
raise BizException(404, "会话不存在或已过期")
|
||||
return success_response()
|
||||
|
||||
@@ -0,0 +1,29 @@
|
||||
"""编辑态数据源取数：登录用户传 URL，服务端替浏览器请求，responsePath / transform 仍由前端做"""
|
||||
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
from fastapi import APIRouter, Depends
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from app.common.response import ApiResponse, success_response
|
||||
from app.core.auth import get_current_user
|
||||
from app.schemas.user import CurrentUser
|
||||
from app.services.datasource_proxy import fetch_json
|
||||
|
||||
router = APIRouter(prefix="/datasource", tags=["DataSource"])
|
||||
|
||||
|
||||
class DataSourceFetch(BaseModel):
|
||||
url: str = Field(min_length=1, max_length=4096)
|
||||
method: Literal["GET", "POST"] = "GET"
|
||||
params: dict[str, Any] | None = None
|
||||
headers: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@router.post("/fetch")
|
||||
async def fetch_datasource(
|
||||
body: DataSourceFetch,
|
||||
user: Annotated[CurrentUser, Depends(get_current_user)],
|
||||
) -> ApiResponse[dict[str, Any] | list[Any]]:
|
||||
data = await fetch_json(body.url, body.method, body.params, body.headers)
|
||||
return success_response(data)
|
||||
+6
-4
@@ -85,9 +85,13 @@ class UploadSizeLimit:
|
||||
return None
|
||||
|
||||
|
||||
def _public(record: DataFileRecord) -> dict:
|
||||
return {**record.model_dump(), "pinned": bool(record.pinned_by)}
|
||||
|
||||
|
||||
def _detail(record: DataFileRecord) -> DataFileDetail:
|
||||
# 按接口模型重新校验，去掉 owner_id、sha256、列概况等只在后端用的字段
|
||||
return DataFileDetail.model_validate(record.model_dump())
|
||||
return DataFileDetail.model_validate(_public(record))
|
||||
|
||||
|
||||
@limits_router.get("/limits")
|
||||
@@ -130,9 +134,7 @@ async def list_files(
|
||||
records = await service.list_files(user.id, conversation_id)
|
||||
return success_response(
|
||||
DataFileList(
|
||||
files=[
|
||||
DataFileInfo.model_validate(record.model_dump()) for record in records
|
||||
],
|
||||
files=[DataFileInfo.model_validate(_public(record)) for record in records],
|
||||
uploads=service.list_uploads(user.id, conversation_id),
|
||||
)
|
||||
)
|
||||
|
||||
@@ -0,0 +1,54 @@
|
||||
"""匿名接口，统一放在 /api/public 下，一眼能看出哪些不鉴权"""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter
|
||||
from fastapi.responses import FileResponse
|
||||
|
||||
from app.api.screens import Service
|
||||
from app.common.response import ApiResponse, success_response
|
||||
from app.core.exceptions import BizException
|
||||
from app.schemas.saved_screen import PublicScreen
|
||||
from app.services.screens import public_page
|
||||
|
||||
router = APIRouter(prefix="/public", tags=["Public"])
|
||||
|
||||
|
||||
@router.get("/screens/{share_id}")
|
||||
async def get_public_screen(
|
||||
share_id: str, service: Service
|
||||
) -> ApiResponse[PublicScreen]:
|
||||
"""线上版本：返回发布快照，不是草稿；已下线、已删除、不存在都是 404"""
|
||||
summary, snapshot = await service.get_live(share_id)
|
||||
return success_response(
|
||||
PublicScreen(
|
||||
name=summary.name,
|
||||
page=public_page(snapshot.page),
|
||||
version=snapshot.version,
|
||||
published_at=snapshot.published_at,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.get("/screens/{share_id}/datasources/{datasource_id}")
|
||||
async def fetch_public_datasource(
|
||||
share_id: str, datasource_id: str, service: Service
|
||||
) -> ApiResponse[dict[str, Any] | list[Any]]:
|
||||
"""线上取数：地址取自已发布快照里的数据源配置，前端只传 ID"""
|
||||
return success_response(
|
||||
await service.fetch_live_datasource(share_id, datasource_id)
|
||||
)
|
||||
|
||||
|
||||
@router.get("/thumbnails/{thumbnail_id}")
|
||||
async def get_thumbnail(thumbnail_id: str, service: Service) -> FileResponse:
|
||||
opened = service.open_thumbnail(thumbnail_id)
|
||||
if opened is None:
|
||||
raise BizException(404, "封面不存在")
|
||||
path, media_type = opened
|
||||
# 每次上传都换新 ID，内容永远不变，可以放心长期缓存
|
||||
return FileResponse(
|
||||
path,
|
||||
media_type=media_type,
|
||||
headers={"Cache-Control": "public, max-age=31536000, immutable"},
|
||||
)
|
||||
@@ -0,0 +1,249 @@
|
||||
"""大屏存取、发布和管理页接口
|
||||
|
||||
登录用户只能操作自己的大屏；别人的大屏与不存在同样返回 404
|
||||
"""
|
||||
|
||||
from datetime import timedelta
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from fastapi import APIRouter, Depends, File, Query, Request, UploadFile
|
||||
|
||||
from app.common.response import ApiResponse, success_response
|
||||
from app.core.auth import get_current_user
|
||||
from app.core.exceptions import BizException
|
||||
from app.schemas.saved_screen import (
|
||||
ListStatus,
|
||||
PublishBrief,
|
||||
PublishHistory,
|
||||
PublishHistoryItem,
|
||||
ScreenCopied,
|
||||
ScreenCopyRequest,
|
||||
ScreenCreate,
|
||||
ScreenCreated,
|
||||
ScreenDetail,
|
||||
ScreenList,
|
||||
ScreenListItem,
|
||||
ScreenPublished,
|
||||
ScreenPublishRequest,
|
||||
ScreenRename,
|
||||
ScreenRestored,
|
||||
ScreenSave,
|
||||
ScreenSaved,
|
||||
ScreenSummary,
|
||||
SortField,
|
||||
ThumbnailResult,
|
||||
TrashEmptied,
|
||||
)
|
||||
from app.schemas.user import CurrentUser
|
||||
from app.services import thumbnails
|
||||
from app.services.screens import ScreenService
|
||||
|
||||
router = APIRouter(prefix="/screens", tags=["Screen"])
|
||||
|
||||
# 深翻页的 skip 会让 MongoDB 空转，管理页一屏 20 条，这个上限足够
|
||||
MAX_OFFSET = 10000
|
||||
THUMBNAIL_URL = "/api/public/thumbnails/{}"
|
||||
|
||||
User = Annotated[CurrentUser, Depends(get_current_user)]
|
||||
|
||||
|
||||
def get_screen_service(request: Request) -> ScreenService:
|
||||
"""没配 MongoDB 时大屏数据只会存在进程内存里，重启就丢，所以直接 503"""
|
||||
service = getattr(request.app.state, "screen_service", None)
|
||||
if service is None:
|
||||
raise BizException(503, "未配置 MongoDB，大屏存取不可用")
|
||||
return service
|
||||
|
||||
|
||||
Service = Annotated[ScreenService, Depends(get_screen_service)]
|
||||
|
||||
|
||||
def thumbnail_url(thumbnail_id: str | None) -> str | None:
|
||||
return None if thumbnail_id is None else THUMBNAIL_URL.format(thumbnail_id)
|
||||
|
||||
|
||||
def _list_item(summary: ScreenSummary, trash_days: int) -> ScreenListItem:
|
||||
publish = summary.publish
|
||||
return ScreenListItem(
|
||||
screen_id=summary.screen_id,
|
||||
name=summary.name,
|
||||
thumbnail_url=thumbnail_url(summary.thumbnail),
|
||||
canvas=summary.canvas,
|
||||
status="draft" if publish is None else publish.status,
|
||||
has_unpublished_changes=publish is not None
|
||||
and summary.revision > publish.revision,
|
||||
share_id=None if publish is None else publish.share_id,
|
||||
updated_at=summary.updated_at,
|
||||
published_at=None if publish is None else publish.published_at,
|
||||
deleted_at=summary.deleted_at,
|
||||
purge_at=None
|
||||
if summary.deleted_at is None
|
||||
else summary.deleted_at + timedelta(days=trash_days),
|
||||
)
|
||||
|
||||
|
||||
@router.post("")
|
||||
async def create_screen(
|
||||
body: ScreenCreate, user: User, service: Service
|
||||
) -> ApiResponse[ScreenCreated]:
|
||||
record = await service.create(user.id, body.name, body.page)
|
||||
return success_response(
|
||||
ScreenCreated(screen_id=record.screen_id, revision=record.revision)
|
||||
)
|
||||
|
||||
|
||||
@router.get("")
|
||||
async def list_screens(
|
||||
user: User,
|
||||
service: Service,
|
||||
keyword: Annotated[str | None, Query(max_length=100)] = None,
|
||||
status: ListStatus | None = None,
|
||||
sort: SortField = "updated_at",
|
||||
order: Literal["asc", "desc"] = "desc",
|
||||
limit: Annotated[int, Query(ge=1, le=100)] = 20,
|
||||
offset: Annotated[int, Query(ge=0, le=MAX_OFFSET)] = 0,
|
||||
deleted: bool = False,
|
||||
) -> ApiResponse[ScreenList]:
|
||||
items, total = await service.list_screens(
|
||||
user.id,
|
||||
keyword=keyword.strip() or None if keyword else None,
|
||||
status=status,
|
||||
sort=sort,
|
||||
descending=order == "desc",
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
deleted=deleted,
|
||||
)
|
||||
return success_response(
|
||||
ScreenList(
|
||||
items=[_list_item(i, service.trash_days) for i in items], total=total
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.delete("/trash")
|
||||
async def empty_trash(user: User, service: Service) -> ApiResponse[TrashEmptied]:
|
||||
"""清空回收站：彻底删除，不可恢复。要先于 /{screen_id} 注册"""
|
||||
return success_response(TrashEmptied(count=await service.empty_trash(user.id)))
|
||||
|
||||
|
||||
@router.get("/{screen_id}")
|
||||
async def get_screen(
|
||||
screen_id: str, user: User, service: Service
|
||||
) -> ApiResponse[ScreenDetail]:
|
||||
record = await service.get(screen_id, user.id)
|
||||
publish = record.publish
|
||||
return success_response(
|
||||
ScreenDetail(
|
||||
screen_id=record.screen_id,
|
||||
name=record.name,
|
||||
page=record.page,
|
||||
revision=record.revision,
|
||||
updated_at=record.updated_at,
|
||||
publish=None
|
||||
if publish is None
|
||||
else PublishBrief.model_validate(publish.model_dump(exclude={"revision"})),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.put("/{screen_id}")
|
||||
async def save_screen(
|
||||
screen_id: str, body: ScreenSave, user: User, service: Service
|
||||
) -> ApiResponse[ScreenSaved]:
|
||||
saved = await service.save(screen_id, user.id, body.page, body.revision)
|
||||
return success_response(
|
||||
ScreenSaved(revision=saved.revision, updated_at=saved.updated_at)
|
||||
)
|
||||
|
||||
|
||||
@router.patch("/{screen_id}")
|
||||
async def rename_screen(
|
||||
screen_id: str, body: ScreenRename, user: User, service: Service
|
||||
) -> ApiResponse[dict]:
|
||||
await service.rename(screen_id, user.id, body.name)
|
||||
return success_response()
|
||||
|
||||
|
||||
@router.delete("/{screen_id}")
|
||||
async def delete_screen(
|
||||
screen_id: str, user: User, service: Service
|
||||
) -> ApiResponse[dict]:
|
||||
await service.delete(screen_id, user.id)
|
||||
return success_response()
|
||||
|
||||
|
||||
@router.post("/{screen_id}/restore")
|
||||
async def restore_deleted_screen(
|
||||
screen_id: str, user: User, service: Service
|
||||
) -> ApiResponse[dict]:
|
||||
await service.restore_deleted(screen_id, user.id)
|
||||
return success_response()
|
||||
|
||||
|
||||
@router.post("/{screen_id}/publish")
|
||||
async def publish_screen(
|
||||
screen_id: str, body: ScreenPublishRequest, user: User, service: Service
|
||||
) -> ApiResponse[ScreenPublished]:
|
||||
info = await service.publish(screen_id, user.id, body.revision)
|
||||
return success_response(
|
||||
ScreenPublished(
|
||||
version=info.version, share_id=info.share_id, published_at=info.published_at
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post("/{screen_id}/unpublish")
|
||||
async def unpublish_screen(
|
||||
screen_id: str, user: User, service: Service
|
||||
) -> ApiResponse[dict]:
|
||||
await service.unpublish(screen_id, user.id)
|
||||
return success_response()
|
||||
|
||||
|
||||
@router.post("/{screen_id}/copy")
|
||||
async def copy_screen(
|
||||
screen_id: str,
|
||||
user: User,
|
||||
service: Service,
|
||||
body: ScreenCopyRequest | None = None,
|
||||
) -> ApiResponse[ScreenCopied]:
|
||||
new_id = await service.copy_screen(screen_id, user.id, body.name if body else None)
|
||||
return success_response(ScreenCopied(screen_id=new_id))
|
||||
|
||||
|
||||
@router.put("/{screen_id}/thumbnail")
|
||||
async def upload_thumbnail(
|
||||
screen_id: str,
|
||||
file: Annotated[UploadFile, File()],
|
||||
user: User,
|
||||
service: Service,
|
||||
) -> ApiResponse[ThumbnailResult]:
|
||||
data = await file.read(thumbnails.MAX_BYTES + 1)
|
||||
thumbnail_id = await service.set_thumbnail(screen_id, user.id, data)
|
||||
return success_response(ThumbnailResult(thumbnail_url=thumbnail_url(thumbnail_id)))
|
||||
|
||||
|
||||
@router.get("/{screen_id}/publishes")
|
||||
async def list_publishes(
|
||||
screen_id: str,
|
||||
user: User,
|
||||
service: Service,
|
||||
limit: Annotated[int, Query(ge=1, le=100)] = 20,
|
||||
offset: Annotated[int, Query(ge=0, le=MAX_OFFSET)] = 0,
|
||||
) -> ApiResponse[PublishHistory]:
|
||||
entries, total = await service.list_publishes(screen_id, user.id, limit, offset)
|
||||
return success_response(
|
||||
PublishHistory(
|
||||
items=[PublishHistoryItem.model_validate(e.model_dump()) for e in entries],
|
||||
total=total,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@router.post("/{screen_id}/publishes/{version}/restore")
|
||||
async def restore_publish(
|
||||
screen_id: str, version: int, user: User, service: Service
|
||||
) -> ApiResponse[ScreenRestored]:
|
||||
revision = await service.restore_publish(screen_id, user.id, version)
|
||||
return success_response(ScreenRestored(revision=revision))
|
||||
+22
-2
@@ -55,7 +55,8 @@ class Settings:
|
||||
default_factory=lambda: float(os.getenv("INTENT_QA_MIN_PROBABILITY", "0.5"))
|
||||
)
|
||||
|
||||
# screen agent 单次运行的模型调用上限，包括读取 skill、查询物料、提交修正和最终回复
|
||||
# screen agent 单次运行的模型调用上限，包括读取 skill、查询物料、提交修正和最终回复；
|
||||
# 用尽或没有给出结论时节点再收尾 1 次，所以一次请求最多 MAX_AGENT_STEPS + 1 次
|
||||
max_agent_steps: int = field(
|
||||
default_factory=lambda: int(os.getenv("MAX_AGENT_STEPS", "12"))
|
||||
)
|
||||
@@ -101,7 +102,7 @@ class Settings:
|
||||
)
|
||||
# 单个数据文件的大小上限（MB）
|
||||
data_file_max_mb: int = field(
|
||||
default_factory=lambda: int(os.getenv("DATA_FILE_MAX_MB", "50"))
|
||||
default_factory=lambda: int(os.getenv("DATA_FILE_MAX_MB", "500"))
|
||||
)
|
||||
# 数据文件保留时长（秒），从上传时算起，默认 7 天
|
||||
data_file_ttl: int = field(
|
||||
@@ -116,6 +117,21 @@ class Settings:
|
||||
)
|
||||
)
|
||||
|
||||
# 大屏保存与发布：单次请求体上限（MB）；图表绑定数据文件时数据快照嵌在页面里，一屏很容易上 MB
|
||||
screen_max_mb: int = field(
|
||||
default_factory=lambda: int(os.getenv("SCREEN_MAX_MB", "10"))
|
||||
)
|
||||
# 回收站保留天数：软删的大屏超过这个时间由后台任务彻底删除（连同发布历史和封面）
|
||||
trash_retention_days: int = field(
|
||||
default_factory=lambda: int(os.getenv("TRASH_RETENTION_DAYS", "30"))
|
||||
)
|
||||
# 大屏封面图存放目录，相对路径相对项目根目录；和数据文件一样放本地磁盘
|
||||
thumbnail_dir: Path = field(
|
||||
default_factory=lambda: (
|
||||
PROJECT_ROOT / (os.getenv("THUMBNAIL_DIR") or "data/thumbnails")
|
||||
)
|
||||
)
|
||||
|
||||
# 模型写的单条数据查询的超时（秒）
|
||||
data_query_timeout: float = field(
|
||||
default_factory=lambda: float(os.getenv("DATA_QUERY_TIMEOUT", "10"))
|
||||
@@ -170,6 +186,10 @@ class Settings:
|
||||
raise ValueError("DATA_FILE_MAX_MB 必须是正整数")
|
||||
if self.data_file_ttl < 1:
|
||||
raise ValueError("DATA_FILE_TTL 必须是正整数")
|
||||
if self.trash_retention_days < 1:
|
||||
raise ValueError("TRASH_RETENTION_DAYS 必须是正整数")
|
||||
if self.screen_max_mb < 1:
|
||||
raise ValueError("SCREEN_MAX_MB 必须是正整数")
|
||||
if self.data_query_timeout <= 0:
|
||||
raise ValueError("DATA_QUERY_TIMEOUT 必须是正数")
|
||||
if self.data_embed_max_rows < 1:
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
"""限制请求体大小：声明的 Content-Length 超限直接 413，没声明（分块）或声明得不实的，边读边数
|
||||
|
||||
FastAPI 会先把整个请求体收完再调用接口函数，超大的请求体已经占满了带宽和内存
|
||||
和数据文件上传的 UploadSizeLimit 同一出发点；限制随路径而不同，所以用规则表，第一条命中的生效
|
||||
"""
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
|
||||
from fastapi.responses import JSONResponse
|
||||
from starlette.exceptions import HTTPException
|
||||
from starlette.types import ASGIApp, Message, Receive, Scope, Send
|
||||
|
||||
from app.common.response import error_content
|
||||
|
||||
|
||||
class BodyTooLarge(HTTPException):
|
||||
"""边读边数时超限；message 由全局 HTTP 异常处理放进信封，和声明超限的响应文案一致"""
|
||||
|
||||
def __init__(self, message: str):
|
||||
super().__init__(413)
|
||||
self.message = message
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BodyLimitRule:
|
||||
methods: frozenset[str]
|
||||
path: re.Pattern[str]
|
||||
max_bytes: int
|
||||
message: str
|
||||
|
||||
|
||||
class BodyLimit:
|
||||
def __init__(self, app: ASGIApp, rules: list[BodyLimitRule]):
|
||||
self.app = app
|
||||
self.rules = rules
|
||||
|
||||
async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
|
||||
rule = self._match(scope) if scope["type"] == "http" else None
|
||||
if rule is None:
|
||||
await self.app(scope, receive, send)
|
||||
return
|
||||
length = dict(scope["headers"]).get(b"content-length", b"")
|
||||
if length.isdigit() and int(length) > rule.max_bytes:
|
||||
response = JSONResponse(error_content(413, rule.message), 413)
|
||||
await response(scope, receive, send)
|
||||
return
|
||||
received = 0
|
||||
|
||||
async def counted() -> Message:
|
||||
nonlocal received
|
||||
message = await receive()
|
||||
if message["type"] == "http.request":
|
||||
received += len(message.get("body", b""))
|
||||
if received > rule.max_bytes:
|
||||
# 抛出的异常由应用里统一的 HTTP 异常处理转成 413 信封
|
||||
raise BodyTooLarge(rule.message)
|
||||
return message
|
||||
|
||||
await self.app(scope, counted, send)
|
||||
|
||||
def _match(self, scope: Scope) -> BodyLimitRule | None:
|
||||
return next(
|
||||
(
|
||||
rule
|
||||
for rule in self.rules
|
||||
if scope["method"] in rule.methods
|
||||
and rule.path.fullmatch(scope["path"])
|
||||
),
|
||||
None,
|
||||
)
|
||||
@@ -50,6 +50,8 @@ async def handle_http_exception(_request: Request, exc: HTTPException) -> JSONRe
|
||||
message = HTTPStatus(exc.status_code).phrase
|
||||
except ValueError:
|
||||
message = "请求失败"
|
||||
# 个别异常自带面向用户的文案，例如请求体超限
|
||||
message = getattr(exc, "message", None) or message
|
||||
return JSONResponse(
|
||||
status_code=exc.status_code,
|
||||
content=error_content(exc.status_code, message),
|
||||
|
||||
@@ -0,0 +1,516 @@
|
||||
"""单元格的类型规则：单个值用 classify 判断，一段 Arrow 数组用 classify_piece 批量判断
|
||||
|
||||
classify 是规则的唯一定义：一个值是什么类型、转换成什么值。大表不能逐个调用它，classify_piece 对常见写法
|
||||
（空值记号、布尔、ASCII 数字和百分数、编号、ISO 日期和常见的日期时间）用正则和 Arrow 转换批量得出结论，
|
||||
其余的值去重后逐个交给 classify。批量规则只认它有把握的写法，认不出的一律交给 classify，所以结论和逐个
|
||||
调用 classify 相同；测试里逐值对比两者
|
||||
"""
|
||||
|
||||
import json
|
||||
import math
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, date, datetime, time
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
|
||||
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 也接受一位数。没有对应格式的直接判为不是日期：
|
||||
# strptime 每次失败都要抛异常，一个值试完所有格式约 16 μs，还要抢一把模块级的锁，多线程时更慢
|
||||
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
|
||||
|
||||
# str.isspace() 为真的全部字符，也就是 str.strip() 去掉的字符；Arrow 自带的去空白函数按另一套定义，
|
||||
# 读取时显式传入这些字符，结果才和 Python 一致
|
||||
WHITESPACE = "\t\n\x0b\x0c\r\x1c\x1d\x1e\x1f \x85\xa0
"
|
||||
# 正则里匹配上面这些字符的写法（RE2 的 \s 只认 ASCII 空白）
|
||||
WHITESPACE_CLASS = "".join(f"\\x{{{ord(char):x}}}" for char in WHITESPACE)
|
||||
|
||||
# 类型代码，classify 返回的类型名按这个顺序编号
|
||||
KINDS = ("null", "bool", "int", "float", "percent", "date", "datetime", "text")
|
||||
NULL, BOOL, INT, FLOAT, PERCENT, DATE, DATETIME, TEXT = range(len(KINDS))
|
||||
KIND_CODES = {name: code for code, name in enumerate(KINDS)}
|
||||
|
||||
MICROS_PER_DAY = 86_400_000_000
|
||||
EPOCH = datetime(1970, 1, 1)
|
||||
|
||||
|
||||
def classify(value: Any) -> tuple[str, Any]:
|
||||
"""单个值的类型和转换后的值；类型是 KINDS 之一"""
|
||||
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:
|
||||
if isinstance(number, int):
|
||||
# 带千分位的超长整数和超出范围的原生整数一样保留文本
|
||||
return ("int", number) if number in INT64_RANGE else ("text", text)
|
||||
return "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 位以上的数字串是编号，不算数值
|
||||
|
||||
判断数字一律用 isdecimal：isdigit 对 ²、① 也返回 True，但 int() 解析不了它们
|
||||
"""
|
||||
if text and text[0] in CURRENCY_SYMBOLS:
|
||||
text = text[1:].lstrip()
|
||||
# 最常见的纯数字不用正则
|
||||
plain = text.isdecimal()
|
||||
if not plain and not NUMBER.match(text):
|
||||
return None
|
||||
digits = text.lstrip("+-")
|
||||
if len(digits) > 1 and digits[0] == "0" and digits[1].isdecimal():
|
||||
return None
|
||||
if digits.isdecimal() and len(digits) >= MIN_ID_DIGITS:
|
||||
return None
|
||||
if plain:
|
||||
return int(text)
|
||||
text = text.replace(",", "")
|
||||
return int(text) if INTEGER.match(text) else float(text)
|
||||
|
||||
|
||||
def parse_datetime(text: str) -> date | datetime | None:
|
||||
if len(text) < 6 or not text[:4].isdecimal():
|
||||
return None
|
||||
# 标准 ISO 写法先用 fromisoformat，比 strptime 快一个数量级
|
||||
if ":" in text:
|
||||
try:
|
||||
return _naive(datetime.fromisoformat(text))
|
||||
except ValueError:
|
||||
pass
|
||||
for pattern in DATETIME_FORMATS.get(text[4], ()):
|
||||
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:
|
||||
"""单元格转成文本：列名、保留为文本的值、警告里引用的内容都用它；Excel 里年份等数字读出来是 2025.0"""
|
||||
if value is None:
|
||||
return ""
|
||||
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)):
|
||||
# Parquet 的嵌套值里可能有日期、Decimal、bytes，JSON 表示不了的写成字符串
|
||||
return json.dumps(value, ensure_ascii=False, default=str)
|
||||
return str(value)
|
||||
|
||||
|
||||
def _naive(value: datetime) -> datetime:
|
||||
# 带时区的时间统一换成 UTC 再去掉时区，同一列里的值才能比较
|
||||
if value.tzinfo is None:
|
||||
return value
|
||||
return value.astimezone(UTC).replace(tzinfo=None)
|
||||
|
||||
|
||||
# ---- 批量判断 ----
|
||||
|
||||
# classify 只会把这些开头的字符串判成文本以外的类型（空值记号和布尔另外判断）：
|
||||
# 货币符号，或可选的正负号加数字（数值、百分数、日期都以数字开头）
|
||||
CANDIDATE = r"^(?:[¥￥$€]|[+-]?\p{Nd})"
|
||||
# 最多 15 位整数部分的 ASCII 数值，保证整数在 int64 内、百分数除以 100 与 Python 结果相同；
|
||||
# 有前导零的（编号）、货币符号后带空格的等写法不在这里，交给 classify
|
||||
_DIGITS = r"[¥￥$€]?[+-]?(?:0|[1-9][0-9]{0,14}|[1-9][0-9]{0,2}(?:,[0-9]{3}){1,4})(?:\.[0-9]+)?(?:[eE][+-]?[0-9]{1,3})?"
|
||||
FAST_NUMBER = f"^{_DIGITS}$"
|
||||
FAST_PERCENT = f"^{_DIGITS}%$"
|
||||
# 有前导零或 16 位以上的纯数字是编号：classify 判成文本
|
||||
FAST_ID = r"^0[0-9]+$|^[0-9]{16,}$"
|
||||
FAST_DATE = r"^[1-9][0-9]{3}-[0-9]{2}-[0-9]{2}$"
|
||||
# fromisoformat 能解析的写法，以及 DATETIME_FORMATS 里分隔符为 - 或 /、各部分一两位数字的写法
|
||||
FAST_DATETIMES = (
|
||||
r"^[1-9][0-9]{3}-[0-9]{2}-[0-9]{2}[T ][0-9]{2}:[0-9]{2}(?::[0-9]{2}(?:\.[0-9]{3}(?:[0-9]{3})?)?)?$",
|
||||
r"^[1-9][0-9]{3}-[0-9]{1,2}-[0-9]{1,2} [0-9]{1,2}:[0-9]{1,2}(?::[0-9]{1,2})?$",
|
||||
r"^[1-9][0-9]{3}/[0-9]{1,2}/[0-9]{1,2} [0-9]{1,2}:[0-9]{1,2}(?::[0-9]{1,2})?$",
|
||||
)
|
||||
|
||||
NULL_SET = pa.array(sorted(NULL_TOKENS))
|
||||
BOOL_SET = pa.array(["true", "false"])
|
||||
NULL_TEXT = pa.scalar(None, pa.string())
|
||||
|
||||
|
||||
@dataclass
|
||||
class Classified:
|
||||
"""一段单元格逐个分类的结果，和对每个值调用 classify 的结论相同
|
||||
|
||||
kinds 是每个值的类型代码；ints、reals、moments、flags 和这段等长，只在对应类型的位置上有意义，
|
||||
这段里没有这种类型时为空
|
||||
"""
|
||||
|
||||
piece: pa.Array
|
||||
kinds: np.ndarray
|
||||
# int
|
||||
ints: np.ndarray | None = None
|
||||
# float、percent（已除以 100）
|
||||
reals: np.ndarray | None = None
|
||||
# date、datetime：自 1970 年起的微秒数，日期按零点
|
||||
moments: np.ndarray | None = None
|
||||
# bool
|
||||
flags: np.ndarray | None = None
|
||||
|
||||
def values(self, name: str) -> np.ndarray:
|
||||
"""取某类转换后的值，没有时补零"""
|
||||
array = getattr(self, name)
|
||||
if array is None:
|
||||
array = np.zeros(len(self.kinds), _DTYPES[name])
|
||||
setattr(self, name, array)
|
||||
return array
|
||||
|
||||
def count(self) -> np.ndarray:
|
||||
return np.bincount(self.kinds, minlength=len(KINDS))
|
||||
|
||||
|
||||
_DTYPES = {"ints": np.int64, "reals": np.float64, "moments": np.int64, "flags": bool}
|
||||
|
||||
|
||||
def union_layout(piece: pa.UnionArray) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""dense union 每个位置属于第几个子数组、是子数组里的第几个
|
||||
|
||||
切片后的 union，pyarrow 的 type_codes、offsets 不计切片的偏移，所以直接读缓冲区
|
||||
"""
|
||||
start, stop = piece.offset, piece.offset + len(piece)
|
||||
_, codes, offsets = piece.buffers()[:3]
|
||||
return (
|
||||
np.frombuffer(codes, np.int8)[start:stop],
|
||||
np.frombuffer(offsets, np.int32)[start:stop],
|
||||
)
|
||||
|
||||
|
||||
def mask(values: pa.Array | pa.ChunkedArray) -> np.ndarray:
|
||||
"""Arrow 布尔数组转成 numpy，空值当作 False"""
|
||||
return pc.fill_null(values, False).to_numpy(zero_copy_only=False)
|
||||
|
||||
|
||||
def classify_piece(piece: pa.Array, order: list | None = None) -> Classified:
|
||||
"""批量判断一段单元格；段的类型是读取时产生的几种之一，见 grid.py
|
||||
|
||||
order 是字符串段依次尝试的批量规则，同一列的各段传同一个列表：这些规则认的写法互不重叠，先试哪个
|
||||
结论都一样，所以每段之后把认出最多的规则挪到前面，后面的段大多一次就认完
|
||||
"""
|
||||
kind = piece.type
|
||||
if pa.types.is_union(kind):
|
||||
return _classify_union(piece, list(FAST_RULES) if order is None else order)
|
||||
valid = mask(piece.is_valid())
|
||||
if pa.types.is_string(kind):
|
||||
return _classify_strings(
|
||||
piece, valid, list(FAST_RULES) if order is None else order
|
||||
)
|
||||
result = Classified(piece, np.where(valid, TEXT, NULL).astype(np.int8))
|
||||
if pa.types.is_null(kind):
|
||||
return result
|
||||
if pa.types.is_boolean(kind):
|
||||
result.kinds[valid] = BOOL
|
||||
result.flags = mask(piece)
|
||||
elif pa.types.is_integer(kind):
|
||||
result.kinds[valid] = INT
|
||||
result.ints = pc.fill_null(piece, 0).to_numpy()
|
||||
elif pa.types.is_floating(kind):
|
||||
reals = pc.fill_null(piece, 0.0).to_numpy()
|
||||
result.kinds[valid] = np.where(np.isfinite(reals[valid]), FLOAT, NULL)
|
||||
result.reals = reals
|
||||
elif pa.types.is_date32(kind):
|
||||
result.kinds[valid] = DATE
|
||||
result.moments = (
|
||||
pc.fill_null(piece, 0).cast(pa.int32()).to_numpy().astype(np.int64)
|
||||
* MICROS_PER_DAY
|
||||
)
|
||||
elif pa.types.is_timestamp(kind):
|
||||
moments = pc.fill_null(piece.cast(pa.int64()), 0).to_numpy()
|
||||
midnight = moments % MICROS_PER_DAY == 0
|
||||
result.kinds[valid] = np.where(midnight[valid], DATE, DATETIME)
|
||||
result.moments = moments
|
||||
else:
|
||||
raise TypeError(f"不支持的单元格类型：{kind}")
|
||||
return result
|
||||
|
||||
|
||||
def _classify_union(piece: pa.UnionArray, order: list) -> Classified:
|
||||
"""类型混杂的段（见 grid.py）：每个子数组各判断一次，再按位置拼回"""
|
||||
result = Classified(piece, np.full(len(piece), NULL, np.int8))
|
||||
codes, offsets = union_layout(piece)
|
||||
for code in range(piece.type.num_fields):
|
||||
at = np.flatnonzero(codes == code)
|
||||
if not len(at):
|
||||
continue
|
||||
child = classify_piece(piece.field(code), order)
|
||||
rows = offsets[at]
|
||||
result.kinds[at] = child.kinds[rows]
|
||||
for name in _DTYPES:
|
||||
values = getattr(child, name)
|
||||
if values is not None:
|
||||
result.values(name)[at] = values[rows]
|
||||
return result
|
||||
|
||||
|
||||
def _classify_strings(piece: pa.Array, valid: np.ndarray, order: list) -> Classified:
|
||||
"""字符串段：按 classify 的顺序（空值记号 → 布尔 → 数值 → 百分数 → 日期时间 → 文本）批量判断"""
|
||||
result = Classified(piece, np.where(valid, TEXT, NULL).astype(np.int8))
|
||||
# 空值记号、Excel 错误值、布尔都不以货币符号或数字开头，只在其余的值里找；
|
||||
# 这些值剩下的就是文本，类型已经记好。数值、日期列整段都是候选，省掉这几步
|
||||
pending = valid & mask(pc.match_substring_regex(piece, CANDIDATE))
|
||||
others = valid & ~pending
|
||||
if others.any():
|
||||
hit = others & mask(
|
||||
pc.or_(
|
||||
pc.is_in(piece, value_set=NULL_SET),
|
||||
pc.match_substring_regex(piece, EXCEL_ERROR.pattern),
|
||||
)
|
||||
)
|
||||
result.kinds[hit] = NULL
|
||||
others &= ~hit
|
||||
if others.any():
|
||||
lower = pc.utf8_lower(piece)
|
||||
hit = others & mask(pc.is_in(lower, value_set=BOOL_SET))
|
||||
if hit.any():
|
||||
result.kinds[hit] = BOOL
|
||||
result.values("flags")[hit] = mask(pc.equal(lower, "true"))[hit]
|
||||
claimed = {}
|
||||
for rule in order:
|
||||
left = int(pending.sum())
|
||||
if not left:
|
||||
break
|
||||
rule(piece, result, pending)
|
||||
claimed[rule] = left - int(pending.sum())
|
||||
order.sort(key=lambda rule: -claimed.get(rule, 0))
|
||||
if pending.any():
|
||||
_delegate(piece, result, pending)
|
||||
return result
|
||||
|
||||
|
||||
def _matched(piece: pa.Array, pattern: str, pending: np.ndarray) -> np.ndarray:
|
||||
"""还没有结论、并且符合 pattern 的值的位置"""
|
||||
return np.flatnonzero(pending & mask(pc.match_substring_regex(piece, pattern)))
|
||||
|
||||
|
||||
def _plain_number(strings: pa.Array) -> pa.Array:
|
||||
"""去掉货币符号、千分位和正号，剩下 Arrow 能直接转换的数字
|
||||
|
||||
只用于符合 FAST_NUMBER 的值：货币符号只在开头，正号在开头或指数里（1e+5 去掉后是 1e5，值不变）
|
||||
"""
|
||||
return pc.replace_substring_regex(strings, r"[¥￥$€,+]", "")
|
||||
|
||||
|
||||
def _fast_numbers(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
|
||||
rows = _matched(piece, FAST_NUMBER, pending)
|
||||
if not len(rows):
|
||||
return
|
||||
plain = _plain_number(piece.take(rows))
|
||||
integer = mask(pc.match_substring_regex(plain, r"^-?[0-9]+$"))
|
||||
if integer.any():
|
||||
at = rows[integer]
|
||||
result.kinds[at] = INT
|
||||
result.values("ints")[at] = pc.cast(
|
||||
plain.filter(integer), pa.int64()
|
||||
).to_numpy()
|
||||
if not integer.all():
|
||||
at = rows[~integer]
|
||||
result.kinds[at] = FLOAT
|
||||
result.values("reals")[at] = pc.cast(
|
||||
plain.filter(~integer), pa.float64()
|
||||
).to_numpy()
|
||||
pending[rows] = False
|
||||
|
||||
|
||||
def _fast_percents(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
|
||||
rows = _matched(piece, FAST_PERCENT, pending)
|
||||
if not len(rows):
|
||||
return
|
||||
plain = _plain_number(pc.utf8_slice_codeunits(piece.take(rows), 0, -1))
|
||||
result.kinds[rows] = PERCENT
|
||||
result.values("reals")[rows] = pc.cast(plain, pa.float64()).to_numpy() / 100
|
||||
pending[rows] = False
|
||||
|
||||
|
||||
def _fast_ids(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
|
||||
# 类型已经记为文本，只需标记为已判断
|
||||
pending[_matched(piece, FAST_ID, pending)] = False
|
||||
|
||||
|
||||
def _fast_dates(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
|
||||
rows = _matched(piece, FAST_DATE, pending)
|
||||
if len(rows):
|
||||
_put_moments(result, pending, rows, piece.take(rows), pa.date32(), DATE)
|
||||
|
||||
|
||||
def _fast_datetimes(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
|
||||
for pattern in FAST_DATETIMES:
|
||||
rows = _matched(piece, pattern, pending)
|
||||
if not len(rows):
|
||||
continue
|
||||
# 统一成 Arrow 认识的 ISO 写法：/ 换成 -，一位数补零（相邻的两位数各占一次，所以替换两遍）
|
||||
strings = pc.replace_substring(piece.take(rows), "/", "-")
|
||||
for _ in range(2):
|
||||
strings = pc.replace_substring_regex(
|
||||
strings, r"([-: ])([0-9])([-: ]|$)", r"\10\2\3"
|
||||
)
|
||||
_put_moments(result, pending, rows, strings, pa.timestamp("us"), DATETIME)
|
||||
|
||||
|
||||
def _put_moments(
|
||||
result: Classified,
|
||||
pending: np.ndarray,
|
||||
rows: np.ndarray,
|
||||
strings: pa.Array,
|
||||
arrow_type: pa.DataType,
|
||||
kind: int,
|
||||
) -> None:
|
||||
"""转换成日期或时间并填回；转换不了的（2 月 30 日、25 点）留给 classify"""
|
||||
ok = np.zeros(len(rows), bool)
|
||||
moments = np.zeros(len(rows), np.int64)
|
||||
_cast_parts(strings, arrow_type, 0, ok, moments)
|
||||
at = rows[ok]
|
||||
result.kinds[at] = kind
|
||||
result.values("moments")[at] = moments[ok]
|
||||
pending[at] = False
|
||||
|
||||
|
||||
# 二分到这么短还转换失败，就整段交给 classify
|
||||
MIN_CAST_PART = 64
|
||||
|
||||
|
||||
def _cast_parts(
|
||||
strings: pa.Array,
|
||||
arrow_type: pa.DataType,
|
||||
offset: int,
|
||||
ok: np.ndarray,
|
||||
moments: np.ndarray,
|
||||
) -> None:
|
||||
"""Arrow 的转换遇到一个坏值就整体失败；失败时对半分开重试，只把含坏值的小段留下"""
|
||||
try:
|
||||
values = pc.cast(strings, arrow_type)
|
||||
except pa.ArrowInvalid:
|
||||
if len(strings) <= MIN_CAST_PART:
|
||||
return
|
||||
half = len(strings) // 2
|
||||
_cast_parts(strings.slice(0, half), arrow_type, offset, ok, moments)
|
||||
_cast_parts(strings.slice(half), arrow_type, offset + half, ok, moments)
|
||||
return
|
||||
end = offset + len(strings)
|
||||
ok[offset:end] = True
|
||||
if pa.types.is_date32(arrow_type):
|
||||
moments[offset:end] = (
|
||||
values.cast(pa.int32()).to_numpy().astype(np.int64) * MICROS_PER_DAY
|
||||
)
|
||||
else:
|
||||
moments[offset:end] = values.cast(pa.int64()).to_numpy()
|
||||
|
||||
|
||||
# 批量规则；认的写法互不重叠，顺序不影响结论
|
||||
FAST_RULES = (_fast_numbers, _fast_percents, _fast_ids, _fast_dates, _fast_datetimes)
|
||||
|
||||
|
||||
def _delegate(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
|
||||
"""批量规则认不出的值去重后逐个交给 classify，再按位置填回"""
|
||||
rows = np.flatnonzero(pending)
|
||||
values = piece.take(rows)
|
||||
unique = pc.unique(values)
|
||||
codes = pc.index_in(values, value_set=unique).to_numpy()
|
||||
outcomes = [classify(value) for value in unique.to_pylist()]
|
||||
kinds = np.array([KIND_CODES[kind] for kind, _ in outcomes], np.int8)
|
||||
result.kinds[rows] = kinds[codes]
|
||||
for name, wanted, convert in _FIELDS:
|
||||
chosen = np.isin(kinds, wanted)
|
||||
if not chosen.any():
|
||||
continue
|
||||
converted = np.array(
|
||||
[
|
||||
convert(value) if pick else 0
|
||||
for (_, value), pick in zip(outcomes, chosen)
|
||||
],
|
||||
_DTYPES[name],
|
||||
)
|
||||
at = chosen[codes]
|
||||
result.values(name)[rows[at]] = converted[codes[at]]
|
||||
pending[rows] = False
|
||||
|
||||
|
||||
def _micros(value: date) -> int:
|
||||
if isinstance(value, datetime):
|
||||
return (value - EPOCH) // _MICROSECOND
|
||||
return (value - EPOCH.date()).days * MICROS_PER_DAY
|
||||
|
||||
|
||||
_MICROSECOND = datetime(1970, 1, 1, microsecond=1) - EPOCH
|
||||
# classify 的结果填进 Classified 的哪个字段：字段名、对应的类型代码、怎样转换
|
||||
_FIELDS = (
|
||||
("ints", (INT,), int),
|
||||
("reals", (FLOAT, PERCENT), float),
|
||||
("moments", (DATE, DATETIME), _micros),
|
||||
("flags", (BOOL,), bool),
|
||||
)
|
||||
+125
-211
@@ -1,52 +1,46 @@
|
||||
"""列名清理和类型推断，输出 Arrow 列
|
||||
|
||||
CSV 的单元格全是字符串；Excel 的数值、日期单元格直接采用，文本单元格再按规则推断
|
||||
一列的非空值必须全部符合同一类型才转换，否则保留文本：宁可让模型在 SQL 里 TRY_CAST，也不悄悄丢掉值
|
||||
单元格的类型规则在 cells.py；这里按列汇总：一列的非空值必须全部符合同一类型才转换，否则保留文本，
|
||||
宁可让模型在 SQL 里 TRY_CAST，也不悄悄丢掉值。各列互不相关，用线程并行推断（Arrow 的计算不占 GIL）
|
||||
"""
|
||||
|
||||
import json
|
||||
import math
|
||||
import os
|
||||
import re
|
||||
from collections import Counter
|
||||
from collections.abc import Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from datetime import UTC, date, datetime, time
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
|
||||
from app.datafiles.base import Reporter, brief, no_report, where
|
||||
from app.datafiles.cells import (
|
||||
BOOL,
|
||||
DATE,
|
||||
DATETIME,
|
||||
EXACT_FLOAT_INT,
|
||||
FAST_RULES,
|
||||
FLOAT,
|
||||
INT,
|
||||
KINDS,
|
||||
MICROS_PER_DAY,
|
||||
NULL,
|
||||
NULL_TEXT,
|
||||
PERCENT,
|
||||
TEXT,
|
||||
Classified,
|
||||
as_text,
|
||||
classify_piece,
|
||||
)
|
||||
from app.datafiles.grid import Column
|
||||
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")
|
||||
NUMERIC_KINDS = (INT, FLOAT, PERCENT)
|
||||
DATE_KINDS = (DATE, DATETIME)
|
||||
# 同时推断的列数
|
||||
INFER_THREADS = min(8, os.cpu_count() or 1)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -84,19 +78,28 @@ def clean_names(
|
||||
|
||||
def build_columns(
|
||||
names: list[str],
|
||||
values: list[Sequence[Any]],
|
||||
columns: list[Column],
|
||||
sheet: str | None,
|
||||
report: Reporter = no_report,
|
||||
) -> tuple[TypedColumns, ParseWarning | None]:
|
||||
"""按列推断类型；有列因少数无法解析的值保留为文本时，返回一条 kept_as_text 警告"""
|
||||
columns, arrays, notes = [], [], {}
|
||||
for number, (name, column_values) in enumerate(zip(names, values), start=1):
|
||||
column, array, note = infer_column(name, column_values)
|
||||
report(number / len(names))
|
||||
columns.append(column)
|
||||
arrays.append(array)
|
||||
if note:
|
||||
notes[name] = note
|
||||
"""按列推断类型；有列因少数无法解析的值保留为文本时，返回一条 kept_as_text 警告
|
||||
|
||||
columns 会被清空：每列推断完就释放它的原始单元格，大表不必同时留着原始和转换后的两份
|
||||
"""
|
||||
described, arrays, notes = [], [], {}
|
||||
with ThreadPoolExecutor(INFER_THREADS) as pool:
|
||||
futures = [
|
||||
pool.submit(infer_column, name, column)
|
||||
for name, column in zip(names, columns)
|
||||
]
|
||||
columns.clear()
|
||||
for number, (name, future) in enumerate(zip(names, futures), start=1):
|
||||
column, array, note = future.result()
|
||||
report(number / len(names))
|
||||
described.append(column)
|
||||
arrays.append(array)
|
||||
if note:
|
||||
notes[name] = note
|
||||
warning = None
|
||||
if notes:
|
||||
listed = "、".join(f"「{name}」" for name in notes)
|
||||
@@ -105,7 +108,7 @@ def build_columns(
|
||||
message=f"{where(sheet)}{listed}有无法解析的值，整列按文本处理",
|
||||
)
|
||||
return TypedColumns(
|
||||
columns, pa.Table.from_arrays(arrays, names=names), notes
|
||||
described, pa.Table.from_arrays(arrays, names=names), notes
|
||||
), warning
|
||||
|
||||
|
||||
@@ -154,194 +157,105 @@ def _arrow_kind(arrow_type: pa.DataType) -> str | 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)
|
||||
name: str, column: Column
|
||||
) -> tuple[DataColumn, pa.ChunkedArray, str | None]:
|
||||
order = list(FAST_RULES)
|
||||
parts = [classify_piece(piece, order) for piece in column.pieces]
|
||||
counts = sum((part.count() for part in parts), np.zeros(len(KINDS), np.int64))
|
||||
present = {kind for kind in range(len(KINDS)) if counts[kind] and kind != NULL}
|
||||
|
||||
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,
|
||||
def build(make, arrow_type: pa.DataType) -> pa.ChunkedArray:
|
||||
return pa.chunked_array([make(part) for part in parts], arrow_type)
|
||||
|
||||
if present and present <= {BOOL}:
|
||||
array = build(
|
||||
lambda part: _pick(part, "flags", (BOOL,), pa.bool_()), pa.bool_()
|
||||
)
|
||||
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
|
||||
return DataColumn(name=name, type="boolean"), array, None
|
||||
if present and present <= {INT, FLOAT}:
|
||||
if all(_integral(part) for part in parts):
|
||||
array = build(lambda part: _numbers(part, np.int64), pa.int64())
|
||||
else:
|
||||
array = build(lambda part: _numbers(part, np.float64), pa.float64())
|
||||
return DataColumn(name=name, type="number"), array, None
|
||||
if present == {PERCENT}:
|
||||
array = build(
|
||||
lambda part: _pick(part, "reals", (PERCENT,), pa.float64()), pa.float64()
|
||||
)
|
||||
return DataColumn(name=name, type="number", format="percent"), array, None
|
||||
if present and present <= {DATE}:
|
||||
array = build(_dates, pa.date32())
|
||||
return DataColumn(name=name, type="date"), array, 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,
|
||||
array = build(
|
||||
lambda part: _pick(part, "moments", DATE_KINDS, pa.timestamp("us")),
|
||||
pa.timestamp("us"),
|
||||
)
|
||||
return DataColumn(name=name, type="datetime"), array, 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
|
||||
array = build(_texts, pa.string())
|
||||
return DataColumn(name=name, type="text"), array, _text_note(counts, parts)
|
||||
|
||||
|
||||
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 _pick(
|
||||
part: Classified, field: str, kinds: tuple[int, ...], arrow_type: pa.DataType
|
||||
) -> pa.Array:
|
||||
"""这段里属于 kinds 的值，其余位置为空"""
|
||||
keep = np.isin(part.kinds, kinds)
|
||||
return pa.array(part.values(field), arrow_type, mask=~keep)
|
||||
|
||||
|
||||
def parse_number(text: str) -> int | float | None:
|
||||
"""千分位和货币符号可以有；有前导零或 16 位以上的数字串是编号，不算数值
|
||||
|
||||
判断数字一律用 isdecimal：isdigit 对 ²、① 也返回 True，但 int() 解析不了它们
|
||||
"""
|
||||
if text and text[0] in CURRENCY_SYMBOLS:
|
||||
text = text[1:].lstrip()
|
||||
# 最常见的纯数字不用正则
|
||||
plain = text.isdecimal()
|
||||
if not plain and not NUMBER.match(text):
|
||||
return None
|
||||
digits = text.lstrip("+-")
|
||||
if len(digits) > 1 and digits[0] == "0" and digits[1].isdecimal():
|
||||
return None
|
||||
if digits.isdecimal() and len(digits) >= MIN_ID_DIGITS:
|
||||
return None
|
||||
if plain:
|
||||
return int(text)
|
||||
text = text.replace(",", "")
|
||||
return int(text) if INTEGER.match(text) else float(text)
|
||||
def _integral(part: Classified) -> bool:
|
||||
"""小数都是整数值并且在 float 能精确表示的范围内（Excel 的整数读出来是 100.0）"""
|
||||
reals = part.values("reals")[part.kinds == FLOAT]
|
||||
return bool(np.all((np.floor(reals) == reals) & (np.abs(reals) < EXACT_FLOAT_INT)))
|
||||
|
||||
|
||||
def parse_datetime(text: str) -> date | datetime | None:
|
||||
if len(text) < 6 or not text[:4].isdecimal():
|
||||
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 _numbers(part: Classified, dtype) -> pa.Array:
|
||||
is_int, is_float = part.kinds == INT, part.kinds == FLOAT
|
||||
# 其他位置可能是 NaN（float 段里的空值），先清零再转换
|
||||
reals = np.where(is_float, part.values("reals"), 0.0)
|
||||
values = np.where(is_int, part.values("ints").astype(dtype), reals.astype(dtype))
|
||||
keep = is_int | is_float
|
||||
return pa.array(values, pa.from_numpy_dtype(dtype), mask=~keep)
|
||||
|
||||
|
||||
def as_text(value: Any) -> str:
|
||||
"""单元格转成文本：列名、保留为文本的值、警告里引用的内容都用它；Excel 里年份等数字读出来是 2025.0"""
|
||||
if value is None:
|
||||
return ""
|
||||
if isinstance(value, str):
|
||||
return value.strip()
|
||||
if isinstance(value, bool):
|
||||
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)):
|
||||
# Parquet 的嵌套值里可能有日期、Decimal、bytes，JSON 表示不了的写成字符串
|
||||
return json.dumps(value, ensure_ascii=False, default=str)
|
||||
return str(value)
|
||||
def _dates(part: Classified) -> pa.Array:
|
||||
days = (part.values("moments") // MICROS_PER_DAY).astype(np.int32)
|
||||
return pa.array(days, pa.date32(), mask=part.kinds != DATE)
|
||||
|
||||
|
||||
def _naive(value: datetime) -> datetime:
|
||||
# 带时区的时间统一换成 UTC 再去掉时区，同一列里的值才能比较
|
||||
if value.tzinfo is None:
|
||||
return value
|
||||
return value.astimezone(UTC).replace(tzinfo=None)
|
||||
def _texts(part: Classified) -> pa.Array:
|
||||
"""保留为文本：字符串是去掉空白后的原文，其他类型按 as_text 转换；空值记号仍为空"""
|
||||
keep = part.kinds != NULL
|
||||
if pa.types.is_string(part.piece.type):
|
||||
return pc.if_else(pa.array(keep), part.piece, NULL_TEXT)
|
||||
return pa.array(
|
||||
[
|
||||
as_text(value) if kept else None
|
||||
for value, kept in zip(part.piece.to_pylist(), keep)
|
||||
],
|
||||
pa.string(),
|
||||
)
|
||||
|
||||
|
||||
def _text_note(
|
||||
texts: list[str | None], value_kinds: list[str], kinds: Counter[str]
|
||||
) -> str | None:
|
||||
def _text_note(counts: np.ndarray, parts: list[Classified]) -> 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)
|
||||
failed = int(counts[TEXT])
|
||||
numeric = int(sum(counts[kind] for kind in NUMERIC_KINDS))
|
||||
dates = int(sum(counts[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)
|
||||
for part in parts:
|
||||
failures = _texts(part).filter(pa.array(part.kinds == TEXT))
|
||||
for text in pc.unique(failures).to_pylist():
|
||||
if text not in examples:
|
||||
examples.append(text)
|
||||
if len(examples) == 2:
|
||||
break
|
||||
if len(examples) == 2:
|
||||
break
|
||||
shown = "、".join(f"「{example}」" for example in examples)
|
||||
label = "数值" if numeric >= dates else "日期"
|
||||
return f"看起来是{label}，有 {failed} 个值无法解析，例如{shown}"
|
||||
|
||||
+122
-34
@@ -1,4 +1,7 @@
|
||||
"""按文件内容识别格式和文本编码，不看扩展名"""
|
||||
"""按文件内容识别格式和文本编码，不看扩展名
|
||||
|
||||
文本文件整个解码一遍确认编码，但只留开头一段用来判断格式和分隔符；读取时再按确认的编码边读边解码
|
||||
"""
|
||||
|
||||
import codecs
|
||||
import csv
|
||||
@@ -15,11 +18,16 @@ from app.schemas.data_file import DataFormat
|
||||
# xlsx/xlsb 解压后的总大小上限，防止解压炸弹
|
||||
MAX_UNZIPPED_BYTES = 1024 * 1024 * 1024
|
||||
|
||||
# 带 BOM 的文本直接按 BOM 解码；UTF-32 的 BOM 以 UTF-16 的开头，解出来含 NUL，按二进制拒绝
|
||||
# 整体是一个 JSON 值的文件只能整个读进内存解析，峰值内存约为文件大小的 6～7 倍；
|
||||
# 这个上限让它和最大的 CSV 一样峰值约 1 GB。JSONL 边读边解析，不受这个限制
|
||||
MAX_JSON_BYTES = 150 * 1024 * 1024
|
||||
|
||||
# 带 BOM 的文本直接按 BOM 解码：编码名、跳过 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"),
|
||||
(codecs.BOM_UTF8, "utf-8", "utf-8", "utf-8-sig"),
|
||||
(codecs.BOM_UTF16_LE, "utf-16", "utf-16-le", "utf-16"),
|
||||
(codecs.BOM_UTF16_BE, "utf-16", "utf-16-be", "utf-16"),
|
||||
)
|
||||
|
||||
# 没有 BOM 时依次尝试；GB18030 兼容 GBK、GB2312，中文 Windows 上 Excel 另存的 CSV 多是这种
|
||||
@@ -30,23 +38,42 @@ DELIMITERS = (",", "\t", ";", "|")
|
||||
# 判断分隔符只看开头这么多行
|
||||
SNIFF_LINES = 50
|
||||
|
||||
# 判断格式和分隔符用的开头文本长度；每次解码读入的字节数
|
||||
HEAD_CHARS = 1024 * 1024
|
||||
READ_BYTES = 1024 * 1024
|
||||
|
||||
# 开头像 JSON 才按 JSON 处理：对象的 { 后面只能是 " 或 }，数组的 [ 后面是一个值的开头；
|
||||
# CSV 的列名也可能以 [ 开头，例如 [日期],[销售额]
|
||||
JSON_START = re.compile(r'\{\s*["}]|\[\s*(?:[\[{"\]\d-]|(?:true|false|null)\s*[,\]])')
|
||||
|
||||
UNKNOWN_ENCODING = "无法识别文件编码，请另存为 UTF-8 后重新上传"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Detected:
|
||||
format: DataFormat
|
||||
# 文本格式的编码名（返回给调用方）和读取时传给 open 的编码
|
||||
encoding: str | None = None
|
||||
# 文本格式解码后的内容
|
||||
text: str | None = None
|
||||
codec: str | None = None
|
||||
# 文本格式开头的一段
|
||||
head: str | None = None
|
||||
# CSV/TSV 的分隔符
|
||||
delimiter: str | None = None
|
||||
# JSON 已解析的内容，避免读取时再解析一遍
|
||||
value: Any = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _Text:
|
||||
encoding: str
|
||||
codec: str
|
||||
head: str
|
||||
# 开头一段之后还有内容
|
||||
truncated: bool
|
||||
# 整个文件只有空白
|
||||
blank: bool
|
||||
|
||||
|
||||
def detect(path: Path) -> Detected:
|
||||
with path.open("rb") as file:
|
||||
head = file.read(8)
|
||||
@@ -57,14 +84,19 @@ def detect(path: Path) -> Detected:
|
||||
if head.startswith(b"PAR1"):
|
||||
return Detected("parquet")
|
||||
|
||||
text, encoding = decode_text(path.read_bytes())
|
||||
stripped = text.lstrip()
|
||||
if not stripped:
|
||||
text = scan_text(path)
|
||||
if text.blank:
|
||||
raise DataFileError(422, "文件是空的")
|
||||
if JSON_START.match(stripped):
|
||||
return _json_format(text, encoding)
|
||||
delimiter = sniff_delimiter(text)
|
||||
return Detected("tsv" if delimiter == "\t" else "csv", encoding, text, delimiter)
|
||||
if JSON_START.match(text.head.lstrip()):
|
||||
return _json_format(path, text)
|
||||
delimiter = sniff_delimiter(text.head)
|
||||
return Detected(
|
||||
"tsv" if delimiter == "\t" else "csv",
|
||||
text.encoding,
|
||||
text.codec,
|
||||
text.head,
|
||||
delimiter,
|
||||
)
|
||||
|
||||
|
||||
def _zip_format(path: Path) -> DataFormat:
|
||||
@@ -83,41 +115,76 @@ def _zip_format(path: Path) -> DataFormat:
|
||||
raise DataFileError(415, UNSUPPORTED_FORMAT)
|
||||
|
||||
|
||||
def decode_text(data: bytes) -> tuple[str, str]:
|
||||
"""返回解码后的文本和编码名"""
|
||||
for bom, codec, name in BOMS:
|
||||
if data.startswith(bom):
|
||||
def scan_text(path: Path) -> _Text:
|
||||
"""确认文本编码：BOM 优先，否则依次严格尝试 FALLBACK_ENCODINGS"""
|
||||
with path.open("rb") as file:
|
||||
start = file.read(4)
|
||||
for bom, name, decoding, codec in BOMS:
|
||||
if start.startswith(bom):
|
||||
try:
|
||||
text = data[len(bom) :].decode(codec)
|
||||
return _decode(path, name, decoding, codec, len(bom))
|
||||
except UnicodeDecodeError:
|
||||
raise DataFileError(
|
||||
422, "无法识别文件编码，请另存为 UTF-8 后重新上传"
|
||||
) from None
|
||||
return _reject_binary(text), name
|
||||
raise DataFileError(422, UNKNOWN_ENCODING) from None
|
||||
for encoding in FALLBACK_ENCODINGS:
|
||||
try:
|
||||
return _reject_binary(data.decode(encoding)), encoding
|
||||
return _decode(path, encoding, encoding, encoding, 0)
|
||||
except UnicodeDecodeError:
|
||||
continue
|
||||
raise DataFileError(422, "无法识别文件编码，请另存为 UTF-8 后重新上传")
|
||||
raise DataFileError(422, UNKNOWN_ENCODING)
|
||||
|
||||
|
||||
def _reject_binary(text: str) -> str:
|
||||
# 文本文件不会有 NUL；图片、PDF 等二进制文件偶尔能按 GB18030 解码成功，在这里拦下
|
||||
if "\x00" in text:
|
||||
raise DataFileError(415, UNSUPPORTED_FORMAT)
|
||||
return text
|
||||
def _decode(path: Path, name: str, decoding: str, codec: str, skip: int) -> _Text:
|
||||
"""严格解码整个文件，留下开头一段；含 NUL 时按二进制文件拒绝"""
|
||||
decoder = codecs.getincrementaldecoder(decoding)("strict")
|
||||
parts: list[str] = []
|
||||
kept, blank, truncated = 0, True, False
|
||||
with path.open("rb") as file:
|
||||
file.seek(skip)
|
||||
while True:
|
||||
chunk = file.read(READ_BYTES)
|
||||
text = decoder.decode(chunk, final=not chunk)
|
||||
# 文本文件不会有 NUL；图片、PDF 等二进制文件偶尔能按 GB18030 解码成功，在这里拦下
|
||||
if "\x00" in text:
|
||||
raise DataFileError(415, UNSUPPORTED_FORMAT)
|
||||
blank = blank and not text.strip()
|
||||
if kept < HEAD_CHARS:
|
||||
parts.append(text)
|
||||
kept += len(text)
|
||||
elif text:
|
||||
truncated = True
|
||||
if not chunk:
|
||||
break
|
||||
head = "".join(parts)
|
||||
if len(head) > HEAD_CHARS:
|
||||
head, truncated = head[:HEAD_CHARS], True
|
||||
cut = head.rfind("\n")
|
||||
if truncated and cut >= 0:
|
||||
# 截断处可能在一行中间，丢掉最后半行
|
||||
head = head[: cut + 1]
|
||||
return _Text(name, codec, head, truncated, blank)
|
||||
|
||||
|
||||
def _json_format(text: str, encoding: str) -> Detected:
|
||||
def _json_format(path: Path, text: _Text) -> Detected:
|
||||
if _first_line_is_record(text):
|
||||
return Detected("jsonl", text.encoding, text.codec, text.head)
|
||||
# 整体是一个 JSON 值，只能整个读进来解析
|
||||
if path.stat().st_size > MAX_JSON_BYTES:
|
||||
raise DataFileError(
|
||||
413,
|
||||
f"JSON 文件不能超过 {MAX_JSON_BYTES // 1024 // 1024} MB，"
|
||||
"请转成 JSONL（每行一个 JSON 对象）后上传",
|
||||
)
|
||||
content = path.read_bytes().decode(text.codec)
|
||||
try:
|
||||
return Detected("json", encoding, text, value=json.loads(text))
|
||||
return Detected(
|
||||
"json", text.encoding, text.codec, text.head, value=json.loads(content)
|
||||
)
|
||||
except ValueError as error:
|
||||
json_error = error
|
||||
first = next(line for line in text.splitlines() if line.strip())
|
||||
first = next(line for line in content.splitlines() if line.strip())
|
||||
try:
|
||||
if isinstance(json.loads(first), dict):
|
||||
return Detected("jsonl", encoding, text)
|
||||
return Detected("jsonl", text.encoding, text.codec, text.head)
|
||||
except ValueError:
|
||||
pass
|
||||
# 开头像 JSON 却不是合法 JSON，按 JSON 报错比当成 CSV 解析出一堆乱码更有用
|
||||
@@ -128,6 +195,27 @@ def _json_format(text: str, encoding: str) -> Detected:
|
||||
)
|
||||
|
||||
|
||||
def _first_line_is_record(text: _Text) -> bool:
|
||||
"""第一行是完整的 JSON 对象、后面还有内容：整体不可能是一个 JSON 值，是 JSONL
|
||||
|
||||
不用整个解析就能判断，大 JSONL 文件不必整个读进内存
|
||||
"""
|
||||
lines = text.head.splitlines(keepends=True)
|
||||
index = next((i for i, line in enumerate(lines) if line.strip()), None)
|
||||
if index is None:
|
||||
return False
|
||||
first, rest = lines[index], lines[index + 1 :]
|
||||
if first == first.rstrip("\r\n") and (rest or text.truncated):
|
||||
# 第一行在开头一段里没有结束
|
||||
return False
|
||||
if not (text.truncated or any(line.strip() for line in rest)):
|
||||
return False
|
||||
try:
|
||||
return isinstance(json.loads(first), dict)
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def sniff_delimiter(text: str) -> str:
|
||||
"""选出让开头各行列数最一致、列数最多的分隔符；都只有一列时按逗号处理"""
|
||||
lines = []
|
||||
|
||||
@@ -0,0 +1,306 @@
|
||||
"""按列存放的单元格，读取、拆表和类型推断共用
|
||||
|
||||
每列是若干段 Arrow 数组，每段一种类型：字符串（已去掉首尾空白，空串记为空值）、布尔、int64、float64、
|
||||
date32、timestamp[us]，或全是空值。读取时每 CHUNK_ROWS 行转一次，同一时间只有一个块是 Python 对象；
|
||||
块里同一列的类型混杂时（表头和数值在同一列、数值列里夹着"待定"），这一段存成 Arrow 的 dense union：
|
||||
每种类型一个子数组，原值和原类型都保留。Arrow 的 is_valid 等函数不认 union 的子数组里的空值，
|
||||
所以按段处理的地方都经过 each，对子数组分别计算再按位置拼回
|
||||
"""
|
||||
|
||||
from collections.abc import Iterable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from datetime import date, datetime, time
|
||||
from itertools import zip_longest
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
|
||||
from app.datafiles.cells import (
|
||||
EXACT_FLOAT_INT,
|
||||
INT64_RANGE,
|
||||
NULL_TEXT,
|
||||
WHITESPACE,
|
||||
as_text,
|
||||
mask,
|
||||
union_layout,
|
||||
)
|
||||
|
||||
CHUNK_ROWS = 65_536
|
||||
|
||||
NoneType = type(None)
|
||||
|
||||
|
||||
class Column:
|
||||
def __init__(self, pieces: Sequence[pa.Array]):
|
||||
self.pieces = [piece for piece in pieces if len(piece)]
|
||||
self.offsets = np.cumsum([0, *map(len, self.pieces)])
|
||||
|
||||
@classmethod
|
||||
def from_values(cls, values: Sequence[Any]) -> "Column":
|
||||
return cls([_convert(values)])
|
||||
|
||||
def __len__(self) -> int:
|
||||
return int(self.offsets[-1])
|
||||
|
||||
def present(self) -> np.ndarray:
|
||||
"""每个单元格是否有值"""
|
||||
return _join(
|
||||
[
|
||||
each(piece, lambda part: mask(part.is_valid()), bool)
|
||||
for piece in self.pieces
|
||||
],
|
||||
bool,
|
||||
)
|
||||
|
||||
def matches(self, test) -> np.ndarray:
|
||||
"""字符串单元格中 test（Arrow 字符串数组 → 布尔数组）为真的位置；其他类型的单元格都不算"""
|
||||
|
||||
def check(part: pa.Array) -> np.ndarray:
|
||||
if pa.types.is_string(part.type):
|
||||
return mask(test(part))
|
||||
return np.zeros(len(part), bool)
|
||||
|
||||
return _join([each(piece, check, bool) for piece in self.pieces], bool)
|
||||
|
||||
def slice(self, start: int, stop: int) -> "Column":
|
||||
pieces = []
|
||||
for piece, offset in zip(self.pieces, self.offsets):
|
||||
begin, end = max(start - offset, 0), min(stop - offset, len(piece))
|
||||
if begin < end:
|
||||
pieces.append(piece.slice(begin, end - begin))
|
||||
return Column(pieces)
|
||||
|
||||
def take(self, rows: np.ndarray) -> "Column":
|
||||
"""按行号（升序）取出若干行；行号连续时不复制数据"""
|
||||
if len(rows) and rows[-1] - rows[0] + 1 == len(rows):
|
||||
return self.slice(int(rows[0]), int(rows[-1]) + 1)
|
||||
bounds = np.searchsorted(rows, self.offsets)
|
||||
return Column(
|
||||
[
|
||||
piece.take(pa.array(rows[bounds[i] : bounds[i + 1]] - self.offsets[i]))
|
||||
for i, piece in enumerate(self.pieces)
|
||||
]
|
||||
)
|
||||
|
||||
def to_pylist(self) -> list[Any]:
|
||||
return [value for piece in self.pieces for value in piece.to_pylist()]
|
||||
|
||||
|
||||
@dataclass
|
||||
class Grid:
|
||||
"""不带表头的原始网格；第 i 行是文件或 sheet 中的第 i + 1 行，各列一样长"""
|
||||
|
||||
sheet: str | None
|
||||
columns: list[Column]
|
||||
height: int
|
||||
|
||||
@classmethod
|
||||
def from_rows(cls, sheet: str | None, rows: Iterable[Sequence[Any]]) -> "Grid":
|
||||
builder = GridBuilder()
|
||||
for row in rows:
|
||||
builder.add(row)
|
||||
return builder.finish(sheet)
|
||||
|
||||
def rows(
|
||||
self, start: int, stop: int, columns: Sequence[int] | None = None
|
||||
) -> list[list[Any]]:
|
||||
"""[start, stop) 行的 Python 值，只取 columns 这些列"""
|
||||
return self._rows(
|
||||
[column.slice(start, stop) for column in self._pick(columns)], stop - start
|
||||
)
|
||||
|
||||
def rows_at(
|
||||
self, rows: np.ndarray, columns: Sequence[int] | None = None
|
||||
) -> list[list[Any]]:
|
||||
return self._rows(
|
||||
[column.take(rows) for column in self._pick(columns)], len(rows)
|
||||
)
|
||||
|
||||
def _pick(self, columns: Sequence[int] | None) -> list[Column]:
|
||||
return self.columns if columns is None else [self.columns[c] for c in columns]
|
||||
|
||||
@staticmethod
|
||||
def _rows(columns: list[Column], count: int) -> list[list[Any]]:
|
||||
if not columns:
|
||||
return [[] for _ in range(count)]
|
||||
return [list(row) for row in zip(*(column.to_pylist() for column in columns))]
|
||||
|
||||
|
||||
class GridBuilder:
|
||||
"""逐行接收单元格，每 CHUNK_ROWS 行转成列；行长短不一时补空值
|
||||
|
||||
strings 为真表示单元格全是字符串（CSV），跳过类型检查；已经按列读好的字符串（Arrow 读 CSV）用 add_strings
|
||||
"""
|
||||
|
||||
def __init__(self, strings: bool = False):
|
||||
self._strings = strings
|
||||
self._columns: list[list[pa.Array]] = []
|
||||
self._pending: list[Sequence[Any]] = []
|
||||
self._done = 0
|
||||
|
||||
@property
|
||||
def height(self) -> int:
|
||||
return self._done + len(self._pending)
|
||||
|
||||
def add(self, row: Sequence[Any]) -> None:
|
||||
self._pending.append(row)
|
||||
if len(self._pending) >= CHUNK_ROWS:
|
||||
self._flush()
|
||||
|
||||
def add_rows(self, rows: Sequence[Sequence[Any]]) -> None:
|
||||
self._pending.extend(rows)
|
||||
if len(self._pending) >= CHUNK_ROWS:
|
||||
self._flush()
|
||||
|
||||
def add_strings(self, columns: Sequence[pa.Array]) -> None:
|
||||
"""一批按列存放的字符串单元格，各列等长"""
|
||||
self._flush()
|
||||
self._append([[_normalize(column)] for column in columns], len(columns[0]))
|
||||
|
||||
def finish(self, sheet: str | None) -> Grid:
|
||||
self._flush()
|
||||
return Grid(sheet, [Column(pieces) for pieces in self._columns], self._done)
|
||||
|
||||
def _flush(self) -> None:
|
||||
rows, count = self._pending, len(self._pending)
|
||||
if not count:
|
||||
return
|
||||
width = max(map(len, rows))
|
||||
self._append(
|
||||
[
|
||||
[
|
||||
_normalize(pa.array(values, pa.string()))
|
||||
if self._strings
|
||||
else _convert(values)
|
||||
]
|
||||
for values in (zip_longest(*rows) if width else ())
|
||||
],
|
||||
count,
|
||||
)
|
||||
self._pending = []
|
||||
|
||||
def _append(self, columns: list[list[pa.Array]], count: int) -> None:
|
||||
for position, pieces in enumerate(columns):
|
||||
if position == len(self._columns):
|
||||
# 这一列在前面的块里没有出现过
|
||||
self._columns.append([pa.nulls(self._done)])
|
||||
self._columns[position].extend(pieces)
|
||||
for pieces in self._columns[len(columns) :]:
|
||||
pieces.append(pa.nulls(count))
|
||||
self._done += count
|
||||
|
||||
|
||||
def each(piece: pa.Array, compute, dtype) -> np.ndarray:
|
||||
"""对一段计算 numpy 结果；union 段对每个子数组分别计算，再按位置拼回"""
|
||||
if not pa.types.is_union(piece.type):
|
||||
return compute(piece)
|
||||
codes, offsets = union_layout(piece)
|
||||
result = np.zeros(len(piece), dtype)
|
||||
for code in range(piece.type.num_fields):
|
||||
at = np.flatnonzero(codes == code)
|
||||
if len(at):
|
||||
result[at] = compute(piece.field(code))[offsets[at]]
|
||||
return result
|
||||
|
||||
|
||||
def _join(parts: list[np.ndarray], dtype) -> np.ndarray:
|
||||
return np.concatenate(parts) if parts else np.zeros(0, dtype)
|
||||
|
||||
|
||||
def _normalize(piece: pa.Array) -> pa.Array:
|
||||
"""字符串去掉首尾空白（和 str.strip 相同），空串记为空值"""
|
||||
if not pa.types.is_string(piece.type):
|
||||
return piece
|
||||
piece = pc.utf8_trim(piece, WHITESPACE)
|
||||
return pc.if_else(pc.equal(piece, ""), NULL_TEXT, piece)
|
||||
|
||||
|
||||
def _convert(values: Sequence[Any]) -> pa.Array:
|
||||
"""一段单元格转成 Arrow：类型一致时转成一个数组，混杂时转成 dense union"""
|
||||
types = set(map(type, values))
|
||||
types.discard(NoneType)
|
||||
if str in types and len(types) > 1:
|
||||
# 空白单元格（Excel 读出来是 ""）不算一种类型
|
||||
values = [
|
||||
None if type(v) is str and (not v or v.isspace()) else v for v in values
|
||||
]
|
||||
types = set(map(type, values))
|
||||
types.discard(NoneType)
|
||||
piece = _uniform(values, types)
|
||||
if piece is not None:
|
||||
return _normalize(piece)
|
||||
families: dict[type, int] = {}
|
||||
children: list[list[Any]] = []
|
||||
codes = np.empty(len(values), np.int8)
|
||||
offsets = np.empty(len(values), np.int32)
|
||||
for index, value in enumerate(values):
|
||||
family = _family(value)
|
||||
code = families.get(family)
|
||||
if code is None:
|
||||
code = families[family] = len(children)
|
||||
children.append([])
|
||||
codes[index] = code
|
||||
offsets[index] = len(children[code])
|
||||
children[code].append(value)
|
||||
return pa.UnionArray.from_dense(
|
||||
pa.array(codes),
|
||||
pa.array(offsets),
|
||||
[_normalize(_child(child)) for child in children],
|
||||
)
|
||||
|
||||
|
||||
def _uniform(values: Sequence[Any], types: set[type]) -> pa.Array | None:
|
||||
"""类型一致（或能无损地合成一种）时转成一个 Arrow 数组，否则返回空"""
|
||||
if not types:
|
||||
return pa.nulls(len(values))
|
||||
if types == {str}:
|
||||
return pa.array(values, pa.string())
|
||||
if types == {bool}:
|
||||
return pa.array(values, pa.bool_())
|
||||
if types <= {int, float}:
|
||||
if float not in types:
|
||||
try:
|
||||
return pa.array(values, pa.int64())
|
||||
except (OverflowError, pa.ArrowException):
|
||||
return None
|
||||
# 整数和小数混在一起时存成 float64；整数超出 float 能精确表示的范围就不合并
|
||||
if all(
|
||||
-EXACT_FLOAT_INT < v < EXACT_FLOAT_INT for v in values if type(v) is int
|
||||
):
|
||||
return pa.array(values, pa.float64())
|
||||
return None
|
||||
if types <= {date, datetime}:
|
||||
if any(type(v) is datetime and v.tzinfo is not None for v in values):
|
||||
return None
|
||||
if types == {date}:
|
||||
return pa.array(values, pa.date32())
|
||||
# 日期按零点存，类型推断时零点的时间仍算日期，和 classify 一致
|
||||
return pa.array(
|
||||
[datetime.combine(v, time()) if type(v) is date else v for v in values],
|
||||
pa.timestamp("us"),
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
def _family(value: Any) -> type:
|
||||
"""union 里归为同一个子数组的类型：日期和不带时区的时间合在一起，超出 int64 的整数等存成文本，其余按 Python 类型分开"""
|
||||
kind = type(value)
|
||||
if kind is date or (kind is datetime and value.tzinfo is None):
|
||||
return date
|
||||
if kind is int and value not in INT64_RANGE:
|
||||
return object
|
||||
if kind in (NoneType, str, bool, int, float):
|
||||
return kind
|
||||
return object
|
||||
|
||||
|
||||
def _child(values: list[Any]) -> pa.Array:
|
||||
types = set(map(type, values))
|
||||
types.discard(NoneType)
|
||||
piece = _uniform(values, types)
|
||||
if piece is not None:
|
||||
return piece
|
||||
# 超出 int64 的整数、带时区的时间、时长等：存成文本，classify 对它们的结论与文本相同
|
||||
return pa.array([None if v is None else as_text(v) for v in values], pa.string())
|
||||
+46
-22
@@ -3,6 +3,7 @@
|
||||
只读写本地文件，不碰数据库；上传流程在子进程里调用（worker.py），超时直接终止子进程
|
||||
"""
|
||||
|
||||
import gc
|
||||
from collections.abc import Callable
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
@@ -18,9 +19,9 @@ from app.datafiles.columns import (
|
||||
clean_names,
|
||||
)
|
||||
from app.datafiles.detect import detect
|
||||
from app.datafiles.grid import Grid
|
||||
from app.datafiles.profile import profile_table
|
||||
from app.datafiles.readers import (
|
||||
Grid,
|
||||
Records,
|
||||
read_delimited,
|
||||
read_excel,
|
||||
@@ -37,23 +38,28 @@ ParseStage = Literal["detecting", "reading", "building"]
|
||||
Progress = Callable[[ParseStage, float], None]
|
||||
|
||||
# 各格式识别、读取两个环节结束时的整体进度，剩下的是建表（拆表、推断类型、写 Parquet、算列概况）
|
||||
# 数字来自 scripts/measure_parse_progress.py 的实测：各环节真实耗时和进度大致成正比，预计时间才准
|
||||
# 数字来自 scripts/measure_parse_progress.py 在 100 万行上的实测：各环节真实耗时和进度大致成正比，预计时间才准
|
||||
STAGE_ENDS: dict[str, tuple[float, float]] = {
|
||||
"csv": (0.01, 0.08),
|
||||
"tsv": (0.01, 0.08),
|
||||
"xlsx": (0.01, 0.30),
|
||||
"xls": (0.01, 0.30),
|
||||
"xlsb": (0.01, 0.30),
|
||||
"json": (0.10, 0.35),
|
||||
"jsonl": (0.03, 0.47),
|
||||
"parquet": (0.02, 0.25),
|
||||
"csv": (0.06, 0.28),
|
||||
"tsv": (0.06, 0.28),
|
||||
"xlsx": (0.01, 0.84),
|
||||
"xls": (0.01, 0.84),
|
||||
"xlsb": (0.01, 0.84),
|
||||
"json": (0.27, 0.84),
|
||||
"jsonl": (0.03, 0.90),
|
||||
"parquet": (0.01, 0.75),
|
||||
}
|
||||
|
||||
# 建表各步骤完成时占其所在范围的比例：拆表（只有网格要）占网格的这么多，
|
||||
# 然后每张表推断完列类型、写完 Parquet，剩下的是算列概况
|
||||
SPLIT_AT = 0.35
|
||||
TYPED_AT = 0.8
|
||||
WRITTEN_AT = 0.85
|
||||
SPLIT_AT = 0.2
|
||||
TYPED_AT = 0.4
|
||||
WRITTEN_AT = 0.8
|
||||
|
||||
|
||||
# Parquet 每个行组的行数。DuckDB 按行组并行扫描，pyarrow 默认一个行组放 100 多万行，整张表只能单线程读；
|
||||
# 行组越小并行度越高，但字典编码按行组重复，文件变大、写得也慢，这个大小两头折中
|
||||
ROW_GROUP_ROWS = 128 * 1024
|
||||
|
||||
|
||||
def no_progress(stage: ParseStage, value: float) -> None:
|
||||
@@ -89,6 +95,18 @@ def parse_file(
|
||||
|
||||
progress 在识别、读取、建表过程中被调用，用来向调用方报告进度
|
||||
"""
|
||||
# 读取时成批创建短命的行列表，分代垃圾回收反复扫描它们却回收不到什么，关掉后读取快约三成；
|
||||
# 解析不产生循环引用，结束后照常回收
|
||||
enabled = gc.isenabled()
|
||||
gc.disable()
|
||||
try:
|
||||
return _parse_file(path, out_dir, progress)
|
||||
finally:
|
||||
if enabled:
|
||||
gc.enable()
|
||||
|
||||
|
||||
def _parse_file(path: Path, out_dir: Path, progress: Progress) -> ParsedFile:
|
||||
progress("detecting", 0.0)
|
||||
detected = detect(path)
|
||||
detect_end, read_end = STAGE_ENDS[detected.format]
|
||||
@@ -111,7 +129,7 @@ def parse_file(
|
||||
report: Reporter,
|
||||
) -> None:
|
||||
target = out_dir / f"{len(tables) + 1}.parquet"
|
||||
pq.write_table(typed.table, target)
|
||||
pq.write_table(typed.table, target, row_group_size=ROW_GROUP_ROWS)
|
||||
report(WRITTEN_AT)
|
||||
profile, preview = profile_table(target, typed.columns, typed.notes)
|
||||
tables.append(
|
||||
@@ -130,6 +148,8 @@ def parse_file(
|
||||
|
||||
def add_grid(grid: Grid, report: Reporter) -> None:
|
||||
found, dropped = split_tables(grid, scaled(report, 0.0, SPLIT_AT))
|
||||
# 拆出的表只引用用到的行；build_columns 推断完一列就释放一列，网格不能再留着
|
||||
del grid
|
||||
warnings.extend(dropped)
|
||||
tables = scaled(report, SPLIT_AT, 1.0)
|
||||
for number, table in enumerate(found):
|
||||
@@ -138,7 +158,7 @@ def parse_file(
|
||||
names, renamed = clean_names(table.header, table.sheet)
|
||||
typed, kept = build_columns(
|
||||
names,
|
||||
list(zip(*table.rows)),
|
||||
table.columns,
|
||||
table.sheet,
|
||||
scaled(share, 0.0, TYPED_AT),
|
||||
)
|
||||
@@ -163,18 +183,22 @@ def parse_file(
|
||||
|
||||
match detected.format:
|
||||
case "csv" | "tsv":
|
||||
grid = read_delimited(detected.text, detected.delimiter, reading)
|
||||
add_grid(grid, building)
|
||||
# 不留局部变量，网格才能在建表时用完即释放
|
||||
add_grid(
|
||||
read_delimited(path, detected.codec, detected.delimiter, reading),
|
||||
building,
|
||||
)
|
||||
case "xlsx" | "xls" | "xlsb":
|
||||
grids, sheet_warnings = read_excel(path, reading)
|
||||
warnings.extend(sheet_warnings)
|
||||
# 各 sheet 按行数分摊建表的进度
|
||||
total_rows = sum(len(grid.rows) for grid in grids)
|
||||
total_rows = sum(grid.height for grid in grids)
|
||||
done_rows = 0
|
||||
for grid in grids:
|
||||
while grids:
|
||||
start = done_rows / total_rows
|
||||
done_rows += len(grid.rows)
|
||||
add_grid(grid, scaled(building, start, done_rows / total_rows))
|
||||
done_rows += grids[0].height
|
||||
# 取出后不再引用，这个 sheet 的网格才能在建表时用完即释放
|
||||
add_grid(grids.pop(0), scaled(building, start, done_rows / total_rows))
|
||||
case "json":
|
||||
result = read_json(detected.value, reading)
|
||||
if isinstance(result, Grid):
|
||||
@@ -182,7 +206,7 @@ def parse_file(
|
||||
else:
|
||||
add_records(result, building)
|
||||
case "jsonl":
|
||||
add_records(read_jsonl(detected.text, reading), building)
|
||||
add_records(read_jsonl(path, detected.codec, reading), building)
|
||||
case "parquet":
|
||||
typed, renamed = arrow_columns(read_parquet(path))
|
||||
if typed.table.num_rows:
|
||||
|
||||
+55
-20
@@ -1,12 +1,13 @@
|
||||
"""列概况和预览：上传时在 Parquet 上用 DuckDB 计算，存进文件记录，供 describe_data、上传后的分析和接口返回"""
|
||||
|
||||
import math
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import duckdb
|
||||
|
||||
from app.datafiles.columns import EXACT_FLOAT_INT
|
||||
from app.datafiles.cells import EXACT_FLOAT_INT
|
||||
from app.datafiles.query import connect, json_value, quote_identifier
|
||||
from app.schemas.data_file import ColumnProfile, DataColumn, TableProfile, TopValue
|
||||
|
||||
@@ -20,28 +21,30 @@ TOP_VALUES_MAX_DISTINCT = 50
|
||||
EXAMPLE_SCAN_ROWS = 1000
|
||||
PROFILE_CELL_CHARS = 50
|
||||
PREVIEW_CELL_CHARS = 200
|
||||
# 算列概况的线程数：在解析子进程里、只读自己刚写的文件，不受模型查询的线程限制；
|
||||
# Parquet 按行组并行扫描，见 parse.py 的 ROW_GROUP_ROWS
|
||||
PROFILE_THREADS = min(8, os.cpu_count() or 1)
|
||||
|
||||
|
||||
def profile_table(
|
||||
path: Path, columns: list[DataColumn], notes: dict[str, str]
|
||||
) -> tuple[TableProfile, list[list[Any]]]:
|
||||
"""返回列概况和接口用的预览（前 PREVIEW_ROWS 行）"""
|
||||
connection = connect({"t": path})
|
||||
connection = connect({"t": path}, threads=PROFILE_THREADS)
|
||||
try:
|
||||
connection.execute("SUMMARIZE t")
|
||||
fields = [column[0] for column in connection.description]
|
||||
summary = {row[0]: dict(zip(fields, row)) for row in connection.fetchall()}
|
||||
total, statistics = _statistics(connection, columns)
|
||||
head = connection.execute(f"FROM t LIMIT {EXAMPLE_SCAN_ROWS}").fetchall()
|
||||
profiled = [
|
||||
_profile_column(
|
||||
connection,
|
||||
column,
|
||||
summary[column.name],
|
||||
stats,
|
||||
total,
|
||||
head,
|
||||
position,
|
||||
notes.get(column.name),
|
||||
)
|
||||
for position, column in enumerate(columns)
|
||||
for position, (column, stats) in enumerate(zip(columns, statistics))
|
||||
]
|
||||
sample = connection.execute(
|
||||
f"FROM t USING SAMPLE reservoir({SAMPLE_ROWS} ROWS) REPEATABLE (42)"
|
||||
@@ -84,24 +87,57 @@ def read_rows(
|
||||
return [[json_value(value, PREVIEW_CELL_CHARS) for value in row] for row in rows]
|
||||
|
||||
|
||||
def _statistics(
|
||||
connection: duckdb.DuckDBPyConnection, columns: list[DataColumn]
|
||||
) -> tuple[int, list[dict[str, Any]]]:
|
||||
"""总行数和每列的统计值，一次扫描算完：非空数、去重数（估算），数值列的最小、最大、平均和四分位数，
|
||||
日期列的最小、最大
|
||||
|
||||
四分位数用 quantile_disc 精确计算：近似算法在多线程下每次结果略有不同，同一个文件重新上传概况会变
|
||||
"""
|
||||
selects, layout = ["count(*)"], []
|
||||
for column in columns:
|
||||
name = quote_identifier(column.name)
|
||||
fields = {
|
||||
"count": f"count({name})",
|
||||
"approx_unique": f"approx_count_distinct({name})",
|
||||
}
|
||||
if column.type == "number":
|
||||
fields |= {
|
||||
"min": f"min({name})",
|
||||
"max": f"max({name})",
|
||||
"avg": f"avg({name})",
|
||||
"quartiles": f"quantile_disc({name}, [0.25, 0.5, 0.75])",
|
||||
}
|
||||
elif column.type in ("date", "datetime"):
|
||||
fields |= {"min": f"min({name})", "max": f"max({name})"}
|
||||
layout.append(list(fields))
|
||||
selects.extend(fields.values())
|
||||
row = connection.execute(f"SELECT {', '.join(selects)} FROM t").fetchone()
|
||||
values = iter(row[1:])
|
||||
return row[0], [dict(zip(fields, values)) for fields in layout]
|
||||
|
||||
|
||||
def _profile_column(
|
||||
connection: duckdb.DuckDBPyConnection,
|
||||
column: DataColumn,
|
||||
stats: dict[str, Any],
|
||||
total: int,
|
||||
head: list[tuple],
|
||||
position: int,
|
||||
note: str | None,
|
||||
) -> ColumnProfile:
|
||||
item: dict[str, Any] = {}
|
||||
if column.type == "number":
|
||||
item = {
|
||||
key: _number(stats[key])
|
||||
for key in ("min", "max", "avg", "q25", "q50", "q75")
|
||||
quartiles = stats["quartiles"] or [None] * 3
|
||||
item = {key: _number(stats[key]) for key in ("min", "max", "avg")}
|
||||
item |= {
|
||||
key: _number(value) for key, value in zip(("q25", "q50", "q75"), quartiles)
|
||||
}
|
||||
elif column.type in ("date", "datetime"):
|
||||
# SUMMARIZE 的时间用空格分隔，和预览里的 ISO 写法统一
|
||||
# 和预览里的时间写法一致
|
||||
item = {
|
||||
key: None if stats[key] is None else stats[key].replace(" ", "T", 1)
|
||||
key: None if stats[key] is None else stats[key].isoformat()
|
||||
for key in ("min", "max")
|
||||
}
|
||||
if (
|
||||
@@ -124,23 +160,22 @@ def _profile_column(
|
||||
examples.append(value)
|
||||
if len(examples) == EXAMPLES:
|
||||
break
|
||||
null_ratio = float(stats["null_percentage"]) / 100
|
||||
return ColumnProfile(
|
||||
**column.model_dump(),
|
||||
null_ratio=round(null_ratio, 4),
|
||||
# approx_unique 是估算值，小表上可能比行数还大
|
||||
distinct=min(stats["approx_unique"], round(stats["count"] * (1 - null_ratio))),
|
||||
null_ratio=round(1 - stats["count"] / total, 4),
|
||||
# approx_unique 是估算值，小表上可能比非空值的个数还大
|
||||
distinct=min(stats["approx_unique"], stats["count"]),
|
||||
examples=examples,
|
||||
note=note,
|
||||
**item,
|
||||
)
|
||||
|
||||
|
||||
def _number(text: str | None) -> int | float | None:
|
||||
"""SUMMARIZE 的统计值是字符串；整数不带小数点，其余保留 6 位小数；Parquet 里的 inf 返回空"""
|
||||
if text is None:
|
||||
def _number(raw: Any) -> int | float | None:
|
||||
"""整数不带小数点，其余保留 6 位小数；Parquet 里的 inf、NaN 返回空"""
|
||||
if raw is None:
|
||||
return None
|
||||
value = float(text)
|
||||
value = float(raw)
|
||||
if not math.isfinite(value):
|
||||
return None
|
||||
if value.is_integer() and abs(value) < EXACT_FLOAT_INT:
|
||||
|
||||
@@ -87,8 +87,13 @@ def run_query(
|
||||
)
|
||||
|
||||
|
||||
def connect(tables: Mapping[str, Path]) -> duckdb.DuckDBPyConnection:
|
||||
"""建好视图并锁定的连接；列概况也用它，和模型查询看到的数据一致"""
|
||||
def connect(
|
||||
tables: Mapping[str, Path], *, threads: int = THREADS
|
||||
) -> duckdb.DuckDBPyConnection:
|
||||
"""建好视图并锁定的连接；列概况也用它，和模型查询看到的数据一致
|
||||
|
||||
threads 只给解析子进程算列概况时放宽，模型写的查询一律用 THREADS
|
||||
"""
|
||||
connection = duckdb.connect(
|
||||
":memory:",
|
||||
config={
|
||||
@@ -96,7 +101,7 @@ def connect(tables: Mapping[str, Path]) -> duckdb.DuckDBPyConnection:
|
||||
"autoload_known_extensions": False,
|
||||
"python_enable_replacements": False,
|
||||
"memory_limit": MEMORY_LIMIT,
|
||||
"threads": THREADS,
|
||||
"threads": threads,
|
||||
},
|
||||
)
|
||||
try:
|
||||
|
||||
+182
-63
@@ -1,16 +1,20 @@
|
||||
"""把各种格式读成原始网格（CSV、Excel）或按列的记录（JSON、JSONL），Parquet 直接读成 Arrow 表
|
||||
|
||||
CSV 和 Excel 可能有标题行、多行表头、合计行，所以先读成不带表头的网格，再由 tables.py 识别表区域
|
||||
文本格式边读边转成列（grid.py），不把整个文件解码进内存
|
||||
"""
|
||||
|
||||
import codecs
|
||||
import csv
|
||||
import io
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from itertools import islice
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import pyarrow as pa
|
||||
import pyarrow.csv as pacsv
|
||||
import pyarrow.parquet as pq
|
||||
from python_calamine import (
|
||||
CalamineWorkbook,
|
||||
@@ -28,7 +32,8 @@ from app.datafiles.base import (
|
||||
rows_exceeded,
|
||||
scaled,
|
||||
)
|
||||
from app.datafiles.columns import as_text
|
||||
from app.datafiles.cells import as_text
|
||||
from app.datafiles.grid import CHUNK_ROWS, Column, Grid, GridBuilder
|
||||
from app.schemas.data_file import ParseWarning
|
||||
|
||||
# 原始网格比数据多出的标题、表头、空行，读取时先粗略拦住明显超限的文件，精确的行数在拆表后检查
|
||||
@@ -43,37 +48,120 @@ MAX_BAD_LINE_RATIO = 0.01
|
||||
# 逐行读取时每隔这么多行报告一次进度
|
||||
PROGRESS_ROWS = 20_000
|
||||
|
||||
# Arrow 读 CSV 每次处理的字节数；单行超过它时 Arrow 报错，改用标准库
|
||||
ARROW_BLOCK_BYTES = 4 * 1024 * 1024
|
||||
|
||||
# 单元格最长的字符数；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]]
|
||||
columns: list[Column]
|
||||
warnings: list[ParseWarning] = field(default_factory=list)
|
||||
|
||||
|
||||
def read_delimited(text: str, delimiter: str, report: Reporter = no_report) -> Grid:
|
||||
rows = []
|
||||
stream = io.StringIO(text, newline="")
|
||||
for row in csv.reader(stream, delimiter=delimiter):
|
||||
rows.append(row)
|
||||
if len(rows) > MAX_TABLE_ROWS + ROW_SLACK:
|
||||
raise rows_exceeded(None)
|
||||
if len(rows) % PROGRESS_ROWS == 0:
|
||||
report(stream.tell() / len(text))
|
||||
return Grid(None, rows)
|
||||
def read_delimited(
|
||||
path: Path, codec: str, delimiter: str, report: Reporter = no_report
|
||||
) -> Grid:
|
||||
"""codec 是 detect 确认过的编码，整个文件都能按它解码
|
||||
|
||||
先用 Arrow 读（C++ 实现，快一个数量级）；它要求每行列数相同，遇到标题行、说明行这类长短不齐的行，
|
||||
或引号不成对、单行过长，就报错。这时改用标准库 csv 从头读一遍。Arrow 读得出来时两者结果相同，测试里对比
|
||||
"""
|
||||
report = _increasing(report)
|
||||
try:
|
||||
return _read_delimited_arrow(path, codec, delimiter, report)
|
||||
except pa.ArrowInvalid:
|
||||
return _read_delimited_python(path, codec, delimiter, report)
|
||||
|
||||
|
||||
def _read_delimited_arrow(
|
||||
path: Path, codec: str, delimiter: str, report: Reporter
|
||||
) -> Grid:
|
||||
with path.open("rb") as raw, io.TextIOWrapper(raw, codec, newline="") as text:
|
||||
first = next(csv.reader(text, delimiter=delimiter), [])
|
||||
if not first:
|
||||
raise pa.ArrowInvalid("第一行是空行，列数要从后面的行推断")
|
||||
names = [f"c{i}" for i in range(len(first))]
|
||||
builder = GridBuilder(strings=True)
|
||||
size = max(path.stat().st_size, 1)
|
||||
with path.open("rb") as raw:
|
||||
if codec == "utf-8-sig" and raw.read(3) != codecs.BOM_UTF8:
|
||||
raw.seek(0)
|
||||
reader = pacsv.open_csv(
|
||||
raw,
|
||||
read_options=pacsv.ReadOptions(
|
||||
column_names=names,
|
||||
block_size=ARROW_BLOCK_BYTES,
|
||||
encoding="utf8" if codec.startswith("utf-8") else codec,
|
||||
),
|
||||
parse_options=pacsv.ParseOptions(
|
||||
delimiter=delimiter, newlines_in_values=True, ignore_empty_lines=False
|
||||
),
|
||||
convert_options=pacsv.ConvertOptions(
|
||||
column_types={name: pa.string() for name in names},
|
||||
strings_can_be_null=False,
|
||||
quoted_strings_can_be_null=False,
|
||||
),
|
||||
)
|
||||
batches: list[pa.RecordBatch] = []
|
||||
waiting = 0
|
||||
for batch in reader:
|
||||
batches.append(batch)
|
||||
waiting += len(batch)
|
||||
# 攒到一个块再转：每段太短时，类型推断的固定开销会变大
|
||||
if waiting >= CHUNK_ROWS:
|
||||
_add_batches(builder, batches)
|
||||
batches, waiting = [], 0
|
||||
if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
|
||||
raise rows_exceeded(None)
|
||||
report(raw.tell() / size)
|
||||
if batches:
|
||||
_add_batches(builder, batches)
|
||||
if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
|
||||
raise rows_exceeded(None)
|
||||
return builder.finish(None)
|
||||
|
||||
|
||||
def _add_batches(builder: GridBuilder, batches: list[pa.RecordBatch]) -> None:
|
||||
builder.add_strings(
|
||||
[
|
||||
pa.concat_arrays([batch.column(i) for batch in batches])
|
||||
for i in range(batches[0].num_columns)
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _read_delimited_python(
|
||||
path: Path, codec: str, delimiter: str, report: Reporter
|
||||
) -> Grid:
|
||||
builder = GridBuilder(strings=True)
|
||||
size = max(path.stat().st_size, 1)
|
||||
with path.open("rb") as raw, io.TextIOWrapper(raw, codec, newline="") as text:
|
||||
reader = csv.reader(text, delimiter=delimiter)
|
||||
while rows := list(islice(reader, PROGRESS_ROWS)):
|
||||
builder.add_rows(rows)
|
||||
if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
|
||||
raise rows_exceeded(None)
|
||||
report(raw.tell() / size)
|
||||
return builder.finish(None)
|
||||
|
||||
|
||||
def _increasing(report: Reporter) -> Reporter:
|
||||
"""Arrow 读到一半失败、换标准库从头读时，进度不往回走"""
|
||||
reached = 0.0
|
||||
|
||||
def inner(fraction: float) -> None:
|
||||
nonlocal reached
|
||||
if fraction > reached:
|
||||
reached = fraction
|
||||
report(fraction)
|
||||
|
||||
return inner
|
||||
|
||||
|
||||
def read_excel(
|
||||
@@ -104,9 +192,19 @@ def read_excel(
|
||||
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):
|
||||
# 解析 sheet 的 XML 占读取的大部分时间，逐行转成列占剩下的
|
||||
rows = scaled(
|
||||
report, (number + 0.8) / len(sheets), (number + 1) / len(sheets)
|
||||
)
|
||||
# iter_rows 保留顶部的空行，行号和 Excel 里看到的一致；左侧的空列不影响识别表区域；
|
||||
# 没有单元格的 sheet 上 iter_rows 会崩溃，不调用
|
||||
builder = GridBuilder()
|
||||
cells = sheet.iter_rows() if sheet.height else iter(())
|
||||
while batch := list(islice(cells, PROGRESS_ROWS)):
|
||||
builder.add_rows(batch)
|
||||
rows(builder.height / sheet.height)
|
||||
grid = builder.finish(meta.name)
|
||||
if not any(column.present().any() for column in grid.columns):
|
||||
warnings.append(
|
||||
ParseWarning(
|
||||
code="empty_sheet_skipped",
|
||||
@@ -114,7 +212,7 @@ def read_excel(
|
||||
)
|
||||
)
|
||||
continue
|
||||
grids.append(Grid(meta.name, rows))
|
||||
grids.append(grid)
|
||||
finally:
|
||||
book.close()
|
||||
return grids, warnings
|
||||
@@ -137,42 +235,53 @@ def read_json(value: Any, report: Reporter = no_report) -> Grid | Records:
|
||||
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 len(records) > MAX_TABLE_ROWS + ROW_SLACK:
|
||||
raise rows_exceeded(None)
|
||||
return Grid.from_rows(
|
||||
None, ([_scalar(cell) for cell in row] for row in records)
|
||||
)
|
||||
if all(isinstance(item, dict) for item in records):
|
||||
result = _flatten_records(records, report)
|
||||
flattener = _Flattener()
|
||||
for number, record in enumerate(records):
|
||||
if number % PROGRESS_ROWS == 0:
|
||||
report(number / len(records))
|
||||
flattener.add(record)
|
||||
result = flattener.finish()
|
||||
result.warnings = warnings
|
||||
return result
|
||||
if not any(isinstance(item, (dict, list)) for item in records):
|
||||
return Records(["值"], [list(records)], warnings)
|
||||
if len(records) > MAX_TABLE_ROWS:
|
||||
raise rows_exceeded(None)
|
||||
return Records(["值"], [Column.from_values(records)], warnings)
|
||||
raise DataFileError(422, "JSON 数组里的元素类型不一致，需要都是对象或都是数组")
|
||||
|
||||
|
||||
def read_jsonl(text: str, report: Reporter = no_report) -> Records:
|
||||
records, bad, total = [], [], 0
|
||||
lines = text.splitlines()
|
||||
# 解析每行 JSON 和展开成列各占一半
|
||||
parsing = scaled(report, 0.0, 0.5)
|
||||
for number, line in enumerate(lines, start=1):
|
||||
if number % PROGRESS_ROWS == 0:
|
||||
parsing(number / len(lines))
|
||||
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:
|
||||
def read_jsonl(path: Path, codec: str, report: Reporter = no_report) -> Records:
|
||||
"""逐行读取；codec 是 detect 确认过的编码"""
|
||||
flattener, bad, total = _Flattener(), [], 0
|
||||
size = max(path.stat().st_size, 1)
|
||||
with path.open("rb") as raw, io.TextIOWrapper(raw, codec) as text:
|
||||
for number, line in enumerate(text, start=1):
|
||||
if number % PROGRESS_ROWS == 0:
|
||||
report(raw.tell() / size)
|
||||
if not line.strip():
|
||||
continue
|
||||
total += 1
|
||||
try:
|
||||
record = json.loads(line)
|
||||
except ValueError:
|
||||
record = None
|
||||
if isinstance(record, dict):
|
||||
flattener.add(record)
|
||||
else:
|
||||
bad.append(number)
|
||||
if not flattener.count:
|
||||
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, scaled(report, 0.5, 1.0))
|
||||
result = flattener.finish()
|
||||
if bad:
|
||||
result.warnings.append(
|
||||
ParseWarning(
|
||||
@@ -226,24 +335,34 @@ def _is_record_list(value: Any) -> bool:
|
||||
)
|
||||
|
||||
|
||||
def _flatten_records(records: list[dict], report: Reporter = no_report) -> Records:
|
||||
# 列按首次出现的顺序排列，记录里没有的字段为空
|
||||
index: dict[str, int] = {}
|
||||
columns: list[list[Any]] = []
|
||||
for row_number, record in enumerate(records):
|
||||
if row_number % PROGRESS_ROWS == 0:
|
||||
report(row_number / len(records))
|
||||
class _Flattener:
|
||||
"""逐条展开记录，按列存放：列按首次出现的顺序排列，记录里没有的字段为空"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._index: dict[str, int] = {}
|
||||
self._builder = GridBuilder()
|
||||
|
||||
@property
|
||||
def count(self) -> int:
|
||||
return self._builder.height
|
||||
|
||||
def add(self, record: dict) -> None:
|
||||
if self._builder.height >= MAX_TABLE_ROWS:
|
||||
raise rows_exceeded(None)
|
||||
flat: dict[str, Any] = {}
|
||||
_flatten(record, "", 0, flat)
|
||||
row: list[Any] = [None] * len(self._index)
|
||||
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)
|
||||
position = self._index.setdefault(name, len(self._index))
|
||||
if position < len(row):
|
||||
row[position] = value
|
||||
else:
|
||||
row.append(value)
|
||||
self._builder.add(row)
|
||||
|
||||
def finish(self) -> Records:
|
||||
grid = self._builder.finish(None)
|
||||
return Records(list(self._index), grid.columns)
|
||||
|
||||
|
||||
def _flatten(record: dict, prefix: str, level: int, out: dict[str, Any]) -> None:
|
||||
|
||||
@@ -195,7 +195,8 @@ class DataFileService:
|
||||
|
||||
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)
|
||||
# 同时撤销大屏对它的固定：文件没了，记录不该再停在"永不过期"
|
||||
await self.store.mark_deleted(file_id, datetime.now(UTC) + self._ttl)
|
||||
# Windows 上正在查询的 Parquet 删不掉，留给定时清理
|
||||
await asyncio.to_thread(shutil.rmtree, self.root / file_id, ignore_errors=True)
|
||||
|
||||
@@ -215,6 +216,21 @@ class DataFileService:
|
||||
raise BizException(404, NOT_FOUND)
|
||||
return records
|
||||
|
||||
async def sync_screen_files(
|
||||
self, user_id: int, screen_id: str, file_ids: set[str]
|
||||
) -> None:
|
||||
"""大屏保存后调用：它引用的文件固定住，不再引用的放开（重新按保留时长计时）"""
|
||||
current = await self.store.pinned_ids(screen_id)
|
||||
if added := file_ids - current:
|
||||
await self.store.pin(added, user_id, screen_id)
|
||||
if removed := current - file_ids:
|
||||
await self.store.unpin(screen_id, removed, datetime.now(UTC) + self._ttl)
|
||||
|
||||
async def release_screen_files(self, screen_id: str) -> None:
|
||||
"""大屏进回收站或被彻底删除：它固定的文件全部放开"""
|
||||
if current := await self.store.pinned_ids(screen_id):
|
||||
await self.store.unpin(screen_id, current, datetime.now(UTC) + self._ttl)
|
||||
|
||||
async def aclose(self) -> None:
|
||||
"""应用退出时取消还在进行的分析，记录里留下的 pending 下次启动后读取时按失败返回"""
|
||||
for task in list(self._tasks):
|
||||
|
||||
+121
-5
@@ -18,6 +18,10 @@ from app.schemas.data_file import DataFileRecord, DataTable, FileInsights
|
||||
SOURCE_NAME = "source"
|
||||
|
||||
|
||||
# 被大屏固定的文件的占位过期时间：MongoDB 的 TTL 索引只认 expires_at，推到遥远的将来就不会被删
|
||||
PINNED_EXPIRES = datetime(9999, 1, 1, tzinfo=UTC)
|
||||
|
||||
|
||||
def table_path(root: Path, file_id: str, table: DataTable) -> Path:
|
||||
return root / file_id / table.parquet
|
||||
|
||||
@@ -43,7 +47,11 @@ class DataFileStore(Protocol):
|
||||
"""为会话预留 count 个连续的表号，返回第一个；表号从 1 开始，不会重复分配"""
|
||||
...
|
||||
|
||||
async def mark_deleted(self, file_id: str) -> None: ...
|
||||
async def mark_deleted(
|
||||
self, file_id: str, expires_at: datetime | None = None
|
||||
) -> None:
|
||||
"""给了 expires_at 就同时撤销所有大屏对它的固定，并把过期时间改成它（被固定的文件不能停在 9999 年）"""
|
||||
...
|
||||
|
||||
async def update_insights(self, file_id: str, insights: FileInsights) -> None:
|
||||
"""记录已过期或不存在时什么都不做"""
|
||||
@@ -53,6 +61,20 @@ class DataFileStore(Protocol):
|
||||
"""未过期且未删除的 file_id，清理目录时用"""
|
||||
...
|
||||
|
||||
async def pin(self, file_ids: set[str], owner_id: int, screen_id: str) -> None:
|
||||
"""让 screen_id 固定这些文件（只对 owner 名下、未过期、未删除的生效），被固定的文件不再过期"""
|
||||
...
|
||||
|
||||
async def unpin(
|
||||
self, screen_id: str, file_ids: set[str], expires_at: datetime
|
||||
) -> None:
|
||||
"""撤销 screen_id 对这些文件的固定；不再被任何大屏固定的文件，过期时间改为 expires_at"""
|
||||
...
|
||||
|
||||
async def pinned_ids(self, screen_id: str) -> set[str]:
|
||||
"""被 screen_id 固定的文件，包括已删除的（它们也要撤销固定）"""
|
||||
...
|
||||
|
||||
|
||||
class InMemoryDataFileStore:
|
||||
def __init__(self) -> None:
|
||||
@@ -90,9 +112,14 @@ class InMemoryDataFileStore:
|
||||
self._table_numbers[conversation_id] = first + count - 1
|
||||
return first
|
||||
|
||||
async def mark_deleted(self, file_id: str) -> None:
|
||||
async def mark_deleted(
|
||||
self, file_id: str, expires_at: datetime | None = None
|
||||
) -> None:
|
||||
if (record := self._records.get(file_id)) is not None:
|
||||
self._records[file_id] = record.model_copy(update={"deleted": True})
|
||||
changes: dict = {"deleted": True}
|
||||
if expires_at is not None:
|
||||
changes.update(pinned_by=[], expires_at=expires_at)
|
||||
self._records[file_id] = record.model_copy(update=changes)
|
||||
|
||||
async def update_insights(self, file_id: str, insights: FileInsights) -> None:
|
||||
if (record := self._records.get(file_id)) is not None:
|
||||
@@ -102,6 +129,37 @@ class InMemoryDataFileStore:
|
||||
self._purge_expired()
|
||||
return {key for key, record in self._records.items() if not record.deleted}
|
||||
|
||||
async def pin(self, file_ids, owner_id, screen_id):
|
||||
self._purge_expired()
|
||||
for file_id in file_ids:
|
||||
record = self._records.get(file_id)
|
||||
if record is None or record.deleted or record.owner_id != owner_id:
|
||||
continue
|
||||
pinned_by = record.pinned_by
|
||||
if screen_id not in pinned_by:
|
||||
pinned_by = [*pinned_by, screen_id]
|
||||
self._records[file_id] = record.model_copy(
|
||||
update={"pinned_by": pinned_by, "expires_at": PINNED_EXPIRES}
|
||||
)
|
||||
|
||||
async def unpin(self, screen_id, file_ids, expires_at):
|
||||
for file_id in file_ids:
|
||||
record = self._records.get(file_id)
|
||||
if record is None or screen_id not in record.pinned_by:
|
||||
continue
|
||||
pinned_by = [item for item in record.pinned_by if item != screen_id]
|
||||
changes = {"pinned_by": pinned_by}
|
||||
if not pinned_by:
|
||||
changes["expires_at"] = expires_at
|
||||
self._records[file_id] = record.model_copy(update=changes)
|
||||
|
||||
async def pinned_ids(self, screen_id):
|
||||
return {
|
||||
key
|
||||
for key, record in self._records.items()
|
||||
if screen_id in record.pinned_by
|
||||
}
|
||||
|
||||
|
||||
class MongoDataFileStore:
|
||||
def __init__(self, db: Any, prefix: str = ""):
|
||||
@@ -117,6 +175,8 @@ class MongoDataFileStore:
|
||||
await self._files.create_index(
|
||||
[("conversation_id", ASCENDING), ("created_at", ASCENDING)]
|
||||
)
|
||||
# 按大屏找它固定的文件
|
||||
await self._files.create_index("pinned_by")
|
||||
|
||||
async def insert(self, record: DataFileRecord) -> None:
|
||||
await self._files.insert_one({"_id": record.file_id, **record.model_dump()})
|
||||
@@ -149,8 +209,13 @@ class MongoDataFileStore:
|
||||
)
|
||||
return document["tables"] - count + 1
|
||||
|
||||
async def mark_deleted(self, file_id: str) -> None:
|
||||
await self._files.update_one({"_id": file_id}, {"$set": {"deleted": True}})
|
||||
async def mark_deleted(
|
||||
self, file_id: str, expires_at: datetime | None = None
|
||||
) -> None:
|
||||
changes: dict = {"deleted": True}
|
||||
if expires_at is not None:
|
||||
changes.update(pinned_by=[], expires_at=expires_at)
|
||||
await self._files.update_one({"_id": file_id}, {"$set": changes})
|
||||
|
||||
async def update_insights(self, file_id: str, insights: FileInsights) -> None:
|
||||
await self._files.update_one(
|
||||
@@ -164,6 +229,57 @@ class MongoDataFileStore:
|
||||
)
|
||||
)
|
||||
|
||||
async def pin(self, file_ids, owner_id, screen_id):
|
||||
if not file_ids:
|
||||
return
|
||||
await self._files.update_many(
|
||||
{
|
||||
"_id": {"$in": list(file_ids)},
|
||||
"owner_id": owner_id,
|
||||
"deleted": False,
|
||||
"expires_at": {"$gt": _now()},
|
||||
},
|
||||
{
|
||||
"$addToSet": {"pinned_by": screen_id},
|
||||
"$set": {"expires_at": PINNED_EXPIRES},
|
||||
},
|
||||
)
|
||||
|
||||
async def unpin(self, screen_id, file_ids, expires_at):
|
||||
if not file_ids:
|
||||
return
|
||||
# 一条聚合管道更新：去掉这个大屏，且只有数组因此变空时才重新计时。
|
||||
# 单文档原子，不会停在"数组空了但仍是 9999 年"；只改确实固定了它的文档
|
||||
await self._files.update_many(
|
||||
{"_id": {"$in": list(file_ids)}, "pinned_by": screen_id},
|
||||
[
|
||||
{
|
||||
"$set": {
|
||||
"pinned_by": {
|
||||
"$filter": {
|
||||
"input": "$pinned_by",
|
||||
"cond": {"$ne": ["$$this", screen_id]},
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
{
|
||||
"$set": {
|
||||
"expires_at": {
|
||||
"$cond": [
|
||||
{"$eq": [{"$size": "$pinned_by"}, 0]},
|
||||
expires_at,
|
||||
"$expires_at",
|
||||
]
|
||||
}
|
||||
}
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
async def pinned_ids(self, screen_id):
|
||||
return set(await self._files.distinct("_id", {"pinned_by": screen_id}))
|
||||
|
||||
|
||||
def _record(document: dict[str, Any]) -> DataFileRecord:
|
||||
return DataFileRecord.model_validate(
|
||||
|
||||
+109
-82
@@ -1,13 +1,17 @@
|
||||
"""在原始网格里识别表：按空行切块，找表头，合并两级表头，去掉标题、说明和合计行
|
||||
|
||||
只用规则；判断不了的地方照常产出表，同时记一条警告，调用方据此提示用户
|
||||
要看每一行的规则（空行切块、合计行）对整列批量计算；表头、标题、说明只看每块开头或末尾的几行，取成 Python 值判断
|
||||
"""
|
||||
|
||||
import math
|
||||
import re
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pyarrow as pa
|
||||
import pyarrow.compute as pc
|
||||
|
||||
from app.datafiles.base import (
|
||||
MAX_TABLE_ROWS,
|
||||
Reporter,
|
||||
@@ -16,8 +20,8 @@ from app.datafiles.base import (
|
||||
rows_exceeded,
|
||||
where,
|
||||
)
|
||||
from app.datafiles.columns import as_text, classify
|
||||
from app.datafiles.readers import Grid
|
||||
from app.datafiles.cells import WHITESPACE_CLASS, as_text, classify, mask
|
||||
from app.datafiles.grid import Column, Grid
|
||||
from app.schemas.data_file import ParseWarning
|
||||
|
||||
# 只在每块的前几行里找表头
|
||||
@@ -40,6 +44,11 @@ TOTAL_LABELS = {
|
||||
}
|
||||
# 例如"华东小计""一季度合计"
|
||||
TOTAL_SUFFIXES = ("小计", "合计", "总计")
|
||||
_TOTAL_SET = pa.array(sorted(TOTAL_LABELS))
|
||||
_TOTAL_STRIP = f"[{WHITESPACE_CLASS}:：]"
|
||||
_TOTAL_SUFFIX = f"(?:{'|'.join(TOTAL_SUFFIXES)})$"
|
||||
# 每个合计字样（含后缀）至少有其中一个字符；英文按转小写之前的大小写都列出
|
||||
_TOTAL_HINT = "[计和均tTmMeE]"
|
||||
|
||||
# 警告里引用的单元格文本最长显示这么多字
|
||||
QUOTE_CHARS = 30
|
||||
@@ -51,8 +60,8 @@ class RawTable:
|
||||
# 同一 sheet 切出多张表时的序号，从 1 开始；只有一张表时为空
|
||||
part: int | None
|
||||
header: list[Any]
|
||||
# 数据行，宽度与表头相同
|
||||
rows: list[list[Any]]
|
||||
# 数据列，和表头一一对应
|
||||
columns: list[Column]
|
||||
warnings: list[ParseWarning] = field(default_factory=list)
|
||||
|
||||
|
||||
@@ -62,8 +71,8 @@ class _Piece:
|
||||
# 表头所在行在网格中的下标
|
||||
header_row: int
|
||||
header: list[Any]
|
||||
# 数据行在网格中的下标
|
||||
data: list[int]
|
||||
# 数据行在网格中的范围 [start, stop)
|
||||
data: list[tuple[int, int]]
|
||||
warnings: list[ParseWarning]
|
||||
|
||||
|
||||
@@ -74,40 +83,41 @@ def split_tables(
|
||||
|
||||
几步都要扫一遍所有行，每步结束报告一次进度，比例来自大表上的实测耗时
|
||||
"""
|
||||
rows = _normalize(grid.rows)
|
||||
report(0.25)
|
||||
sheet = grid.sheet
|
||||
present = [column.present() for column in grid.columns]
|
||||
filled = np.logical_or.reduce(present) if present else np.zeros(grid.height, bool)
|
||||
blocks = _blocks(filled)
|
||||
used = _used_columns(present, blocks)
|
||||
firsts = grid.rows_at(np.array([start for start, _ in blocks], np.int64))
|
||||
report(0.3)
|
||||
pieces: list[_Piece] = []
|
||||
dropped: list[ParseWarning] = []
|
||||
blocks = _blocks(rows)
|
||||
report(0.3)
|
||||
for number, block in enumerate(blocks, start=1):
|
||||
columns = [
|
||||
c for c in range(len(rows[0])) if any(rows[r][c] is not None for r in block)
|
||||
]
|
||||
for number, ((start, stop), columns, first) in enumerate(
|
||||
zip(blocks, used, firsts), start=1
|
||||
):
|
||||
last = pieces[-1] if pieces else None
|
||||
# 表中间夹着空行时，后面的块以数据行开头、列也落在上一张表内，接在上一张表后面
|
||||
if (
|
||||
last
|
||||
and not _is_label_row([rows[block[0]][c] for c in columns])
|
||||
and not _is_label_row([first[c] for c in columns])
|
||||
and set(columns) <= set(last.columns)
|
||||
):
|
||||
last.data.extend(block)
|
||||
last.data.append((start, stop))
|
||||
continue
|
||||
if len(block) < 2:
|
||||
if stop - start < 2:
|
||||
dropped.append(
|
||||
ParseWarning(
|
||||
code="block_dropped",
|
||||
message=f"{where(sheet)}第 {block[0] + 1} 行的{_quote(rows[block[0]])}不是表格，已跳过",
|
||||
message=f"{where(sheet)}第 {start + 1} 行的{_quote(first)}不是表格，已跳过",
|
||||
)
|
||||
)
|
||||
continue
|
||||
pieces.append(_read_block(rows, block, columns, sheet))
|
||||
pieces.append(_read_block(grid, start, stop, columns, sheet))
|
||||
report(0.3 + 0.25 * number / len(blocks))
|
||||
|
||||
tables = []
|
||||
for number, piece in enumerate(pieces, start=1):
|
||||
table = _finish(rows, piece, sheet)
|
||||
table = _finish(grid, piece, sheet)
|
||||
report(0.55 + 0.45 * number / len(pieces))
|
||||
if table is None:
|
||||
dropped.append(
|
||||
@@ -125,41 +135,28 @@ def split_tables(
|
||||
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 _blocks(filled: np.ndarray) -> list[tuple[int, int]]:
|
||||
"""连续的非空行组成的块 [start, stop)"""
|
||||
edges = np.flatnonzero(np.diff(np.concatenate(([0], filled, [0])).astype(np.int8)))
|
||||
return list(zip(edges[0::2].tolist(), edges[1::2].tolist()))
|
||||
|
||||
|
||||
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 _used_columns(
|
||||
present: list[np.ndarray], blocks: list[tuple[int, int]]
|
||||
) -> list[list[int]]:
|
||||
"""每块里有值的列；块之间只隔着空行，按块的起点分段求或即可"""
|
||||
if not blocks:
|
||||
return []
|
||||
starts = np.array([start for start, _ in blocks])
|
||||
used = np.array([np.logical_or.reduceat(column, starts) for column in present])
|
||||
return [np.flatnonzero(used[:, index]).tolist() for index in range(len(blocks))]
|
||||
|
||||
|
||||
def _read_block(
|
||||
rows: list[list[Any]], block: list[int], columns: list[int], sheet: str | None
|
||||
grid: Grid, start: int, stop: int, columns: list[int], sheet: str | None
|
||||
) -> _Piece:
|
||||
cells = [[rows[r][c] for c in columns] for r in block]
|
||||
# 表头最多在第 HEADER_SEARCH_ROWS 行，两级表头还要看它下面两行
|
||||
cells = grid.rows(start, min(stop, start + HEADER_SEARCH_ROWS + 2), columns)
|
||||
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)))
|
||||
@@ -173,7 +170,7 @@ def _read_block(
|
||||
warnings.append(
|
||||
ParseWarning(
|
||||
code="header_uncertain",
|
||||
message=f"{where(sheet)}没有找到明确的表头，用第 {block[header_at] + 1} 行作表头，可能识别错误",
|
||||
message=f"{where(sheet)}没有找到明确的表头，用第 {start + header_at + 1} 行作表头，可能识别错误",
|
||||
)
|
||||
)
|
||||
if header_at > 0:
|
||||
@@ -183,7 +180,7 @@ def _read_block(
|
||||
warnings.append(
|
||||
ParseWarning(
|
||||
code="title_skipped",
|
||||
message=f"{where(sheet)}{_row_range(block[0], block[header_at - 1])}是标题，已跳过：{skipped}",
|
||||
message=f"{where(sheet)}{_row_range(start, start + header_at - 1)}是标题，已跳过：{skipped}",
|
||||
)
|
||||
)
|
||||
|
||||
@@ -191,7 +188,7 @@ def _read_block(
|
||||
data_at = header_at + 1
|
||||
if (
|
||||
confident
|
||||
and header_at + 2 < len(cells)
|
||||
and header_at + 2 < stop - start
|
||||
and _is_subheader(cells[header_at], cells[header_at + 1], cells[header_at + 2])
|
||||
):
|
||||
header = _merge_header(cells[header_at], cells[header_at + 1])
|
||||
@@ -199,20 +196,28 @@ def _read_block(
|
||||
warnings.append(
|
||||
ParseWarning(
|
||||
code="header_merged",
|
||||
message=f"{where(sheet)}{_row_range(block[header_at], block[header_at + 1])}合并为表头",
|
||||
message=f"{where(sheet)}{_row_range(start + header_at, start + header_at + 1)}合并为表头",
|
||||
)
|
||||
)
|
||||
return _Piece(columns, block[header_at], header, block[data_at:], warnings)
|
||||
return _Piece(
|
||||
columns, start + header_at, header, [(start + data_at, stop)], 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]
|
||||
def _finish(grid: Grid, piece: _Piece, sheet: str | None) -> RawTable | None:
|
||||
ranges = [np.arange(start, stop) for start, stop in piece.data if start < stop]
|
||||
rows = np.concatenate(ranges) if ranges else np.zeros(0, np.int64)
|
||||
warnings = piece.warnings
|
||||
|
||||
# 表格下方只有一个文字单元格的行是说明，例如"注：数据来源于……"
|
||||
notes = []
|
||||
while data and len(piece.columns) >= 3 and _is_note(data[-1][1]):
|
||||
notes.append(data.pop())
|
||||
while len(rows) and len(piece.columns) >= 3:
|
||||
last = int(rows[-1])
|
||||
row = grid.rows(last, last + 1, piece.columns)[0]
|
||||
if not _is_note(row):
|
||||
break
|
||||
notes.append((last, row))
|
||||
rows = rows[:-1]
|
||||
if notes:
|
||||
notes.reverse()
|
||||
shown = "".join(_quote(row) for _, row in notes)
|
||||
@@ -224,22 +229,59 @@ def _finish(rows: list[list[Any]], piece: _Piece, sheet: str | None) -> RawTable
|
||||
)
|
||||
|
||||
# 合计、小计行不是明细数据，留着会让求和重复计算
|
||||
totals = [index for index, row in data if _is_total(row)]
|
||||
if totals:
|
||||
data = [(index, row) for index, row in data if not _is_total(row)]
|
||||
listed = brief([str(index + 1) for index in totals])
|
||||
totals = _totals(grid, piece.columns, rows)
|
||||
if totals.any():
|
||||
listed = brief([str(index + 1) for index in rows[totals].tolist()])
|
||||
warnings.append(
|
||||
ParseWarning(
|
||||
code="total_rows_removed",
|
||||
message=f"{where(sheet)}已排除 {len(totals)} 行合计、小计（第 {listed} 行）",
|
||||
message=f"{where(sheet)}已排除 {int(totals.sum())} 行合计、小计（第 {listed} 行）",
|
||||
)
|
||||
)
|
||||
rows = rows[~totals]
|
||||
|
||||
if not data:
|
||||
if not len(rows):
|
||||
return None
|
||||
if len(data) > MAX_TABLE_ROWS:
|
||||
if len(rows) > MAX_TABLE_ROWS:
|
||||
raise rows_exceeded(sheet)
|
||||
return RawTable(sheet, None, piece.header, [row for _, row in data], warnings)
|
||||
columns = [grid.columns[c].take(rows) for c in piece.columns]
|
||||
return RawTable(sheet, None, piece.header, columns, warnings)
|
||||
|
||||
|
||||
def _totals(grid: Grid, columns: list[int], rows: np.ndarray) -> np.ndarray:
|
||||
"""rows 里哪些是合计行：每行只看前两个非空单元格，其中有合计字样的文字就是"""
|
||||
if not len(rows):
|
||||
return np.zeros(0, bool)
|
||||
low, high = int(rows[0]), int(rows[-1]) + 1
|
||||
seen = np.zeros(high - low, np.int64)
|
||||
total = np.zeros(high - low, bool)
|
||||
for c in columns:
|
||||
column = grid.columns[c].slice(low, high)
|
||||
present = column.present()
|
||||
open_ = present & (seen < 2)
|
||||
if open_.any():
|
||||
total |= open_ & column.matches(_is_total_label)
|
||||
seen += present
|
||||
if (seen >= 2).all():
|
||||
break
|
||||
return total[rows - low]
|
||||
|
||||
|
||||
def _is_total_label(strings: pa.Array) -> pa.Array:
|
||||
"""去掉空白和冒号、转小写后是合计字样；合计字样都含 _TOTAL_HINT 里的字符，先用它筛掉绝大多数值"""
|
||||
found = np.zeros(len(strings), bool)
|
||||
rows = np.flatnonzero(mask(pc.match_substring_regex(strings, _TOTAL_HINT)))
|
||||
if len(rows):
|
||||
key = pc.utf8_lower(
|
||||
pc.replace_substring_regex(strings.take(rows), _TOTAL_STRIP, "")
|
||||
)
|
||||
found[rows] = mask(
|
||||
pc.or_(
|
||||
pc.is_in(key, value_set=_TOTAL_SET),
|
||||
pc.match_substring_regex(key, _TOTAL_SUFFIX),
|
||||
)
|
||||
)
|
||||
return pa.array(found)
|
||||
|
||||
|
||||
def _is_label(value: Any) -> bool:
|
||||
@@ -276,21 +318,6 @@ def _is_note(row: list[Any]) -> bool:
|
||||
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)
|
||||
|
||||
+58
@@ -1,6 +1,7 @@
|
||||
"""FastAPI 应用入口"""
|
||||
|
||||
import asyncio
|
||||
import re
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from fastapi import Depends, FastAPI
|
||||
@@ -8,10 +9,15 @@ from loguru import logger
|
||||
|
||||
from app.api.agent import router
|
||||
from app.api.conversations import router as conversation_router
|
||||
from app.api.datasource import router as datasource_router
|
||||
from app.api.files import UploadSizeLimit
|
||||
from app.api.files import limits_router as file_limits_router
|
||||
from app.api.files import router as file_router
|
||||
from app.api.materials import router as material_router
|
||||
from app.api.public import router as public_router
|
||||
from app.api.screens import router as screen_router
|
||||
from app.config import settings
|
||||
from app.core.body_limit import BodyLimit, BodyLimitRule
|
||||
from app.core.auth import get_current_user
|
||||
from app.core.exceptions import register_exception_handlers
|
||||
from app.core.logging import setup_logging
|
||||
@@ -32,12 +38,16 @@ async def lifespan(application: FastAPI):
|
||||
from app.services.agent import AgentService
|
||||
from app.services.material_mongo import warm_material_cache
|
||||
from app.services.persistence import open_persistence
|
||||
from app.services.screen_mongo import MongoScreenStore
|
||||
from app.services.screens import ScreenService
|
||||
from app.services.thumbnails import ThumbnailStorage
|
||||
|
||||
# 按配置选择 MongoDB 或进程内存，会话记录与图检查点使用同一种存储
|
||||
persistence = await open_persistence()
|
||||
application.state.conversation_store = persistence.store
|
||||
mcp_processes = []
|
||||
sweeper = None
|
||||
trash_sweeper = None
|
||||
data_files = None
|
||||
try:
|
||||
# 数据文件的元数据和会话记录用同一种存储；启动时先清理一遍过期目录，之后每小时一次
|
||||
@@ -48,6 +58,21 @@ async def lifespan(application: FastAPI):
|
||||
)
|
||||
application.state.data_file_service = data_files
|
||||
sweeper = asyncio.create_task(data_files.run_sweeper())
|
||||
# 大屏存取必须落 MongoDB：没配时不创建服务，screens 相关接口返回 503，而不是存在内存里重启就丢
|
||||
screen_service = None
|
||||
if persistence.db is not None:
|
||||
screen_store = MongoScreenStore(persistence.db)
|
||||
await screen_store.ensure_indexes()
|
||||
screen_service = ScreenService(
|
||||
screen_store,
|
||||
ThumbnailStorage(settings.thumbnail_dir),
|
||||
data_files,
|
||||
settings.trash_retention_days,
|
||||
)
|
||||
trash_sweeper = asyncio.create_task(screen_service.run_trash_sweeper())
|
||||
else:
|
||||
logger.warning("未配置 MONGODB_URI，大屏存取接口不可用")
|
||||
application.state.screen_service = screen_service
|
||||
# 物料目录启动时载入内存，动作校验与前端接口共用，避免重复查询数据库
|
||||
await warm_material_cache(persistence.db)
|
||||
# 拉起 mcp_servers.json 中由应用管理的服务；单个服务失败只记录，用到它的请求执行时报错
|
||||
@@ -64,6 +89,8 @@ async def lifespan(application: FastAPI):
|
||||
finally:
|
||||
if sweeper is not None:
|
||||
sweeper.cancel()
|
||||
if trash_sweeper is not None:
|
||||
trash_sweeper.cancel()
|
||||
if data_files is not None:
|
||||
await data_files.aclose()
|
||||
# 共享客户端只在应用退出时关闭，追踪刷新放到线程避免阻塞事件循环
|
||||
@@ -102,6 +129,37 @@ def create_app() -> FastAPI:
|
||||
# 上传在读请求体之前先按 Content-Length 检查大小
|
||||
application.add_middleware(UploadSizeLimit, path=f"/api{file_router.prefix}")
|
||||
application.include_router(material_router, prefix="/api")
|
||||
application.include_router(screen_router, prefix="/api", dependencies=authenticated)
|
||||
application.include_router(
|
||||
datasource_router, prefix="/api", dependencies=authenticated
|
||||
)
|
||||
# 匿名接口统一放 /api/public 下，不挂鉴权
|
||||
application.include_router(public_router, prefix="/api")
|
||||
# 保存 / 发布的请求体比普通接口大得多，单独放宽；封面图 1 MB
|
||||
application.add_middleware(
|
||||
BodyLimit,
|
||||
rules=[
|
||||
BodyLimitRule(
|
||||
frozenset({"POST", "PUT"}),
|
||||
re.compile(r"/api/screens(/[^/]+)?"),
|
||||
settings.screen_max_mb * 1024 * 1024,
|
||||
f"大屏数据不能超过 {settings.screen_max_mb} MB",
|
||||
),
|
||||
BodyLimitRule(
|
||||
frozenset({"PUT"}),
|
||||
re.compile(r"/api/screens/[^/]+/thumbnail"),
|
||||
1024 * 1024 + 64 * 1024,
|
||||
"封面图不能超过 1 MB",
|
||||
),
|
||||
# 其余带请求体的接口（发布、复制、改名、数据源取数、会话）都只有几个字段
|
||||
BodyLimitRule(
|
||||
frozenset({"POST", "PUT", "PATCH"}),
|
||||
re.compile(r"/api/(screens|datasource|conversations)(/.*)?"),
|
||||
64 * 1024,
|
||||
"请求体过大",
|
||||
),
|
||||
],
|
||||
)
|
||||
return application
|
||||
|
||||
|
||||
|
||||
@@ -1,39 +1,49 @@
|
||||
"""调用 screen agent，按它的回复和最后一次提交的结果写回状态
|
||||
"""调用 screen agent，按它的回复和成功提交的动作写回状态
|
||||
|
||||
kind 由模型决定；代码只在"回复 done 但最后一次提交失败"时不采纳，因为此时没有能交付的动作
|
||||
kind 由模型决定；代码只在"回复 done 但最后一次提交失败"时不采纳：最后想做的修改没有写入草稿，
|
||||
模型却说已完成。这是保守做法，last_write_failed 只看工具结果，任务是否完成仍以模型的结论和评测为准
|
||||
agent 正常结束却没有结论时（调用次数用尽、只回了文本或空回复）补一次只读收尾；取消和异常照常抛出，不收尾
|
||||
"""
|
||||
|
||||
import json
|
||||
|
||||
from langchain_core.messages import HumanMessage
|
||||
|
||||
from app.agent.screen_agent import AgentReply, get_screen_agent
|
||||
from app.agent.tools import ScreenAgentContext
|
||||
from app.agent.screen_agent import (
|
||||
AgentReply,
|
||||
finalize_reply,
|
||||
get_screen_agent,
|
||||
reached_call_limit,
|
||||
)
|
||||
from app.agent.tools import ScreenAgentContext, last_write_failed
|
||||
from app.processor.base import BaseNode
|
||||
from app.processor.history import history_messages
|
||||
from app.processor.state import QueryState, run_data
|
||||
from app.processor.summary import request_summary
|
||||
from app.resources.component import get_catalog_cache
|
||||
|
||||
# 评测按这句话识别"草稿因最后一次提交失败没有交付"，单独核对草稿本身是否满足要求
|
||||
DRAFT_REJECTED = "修改没有通过校验，这次没有生效，可以换个说法再试"
|
||||
|
||||
def reply_update(result: dict) -> QueryState:
|
||||
|
||||
def reply_update(result: dict, *, limit_hit: bool = False) -> QueryState:
|
||||
"""把 agent 的运行结果转换成 QueryState 的更新"""
|
||||
reply: AgentReply | None = result.get("structured_response")
|
||||
if reply is None:
|
||||
# 调用次数用尽等情况，agent 没有给出结论
|
||||
return {
|
||||
"status": "failed",
|
||||
"answer": "处理步骤超出上限，这次没有完成，可以换个说法再试",
|
||||
}
|
||||
# 收尾也没有给出合法的结论；按真实原因说明，不再一律说超出上限
|
||||
answer = (
|
||||
"处理步骤超出上限，这次没有完成，修改没有生效，可以换个说法再试"
|
||||
if limit_hit
|
||||
else "这次没有得出结论，修改没有生效，可以换个说法再试"
|
||||
)
|
||||
return {"status": "failed", "answer": answer}
|
||||
if reply.kind == "clarify":
|
||||
return {"status": "waiting_for_input", "clarification": reply.clarification}
|
||||
if reply.kind == "unable":
|
||||
return {"status": "failed", "answer": reply.message}
|
||||
if result.get("apply_error"):
|
||||
return {
|
||||
"status": "failed",
|
||||
"answer": "修改没有通过校验，这次没有生效，可以换个说法再试",
|
||||
}
|
||||
if last_write_failed(result.get("messages", [])):
|
||||
return {"status": "failed", "answer": DRAFT_REJECTED}
|
||||
# 各次成功提交的动作按顺序拼接；clarify、unable 不交付，下次请求从原始大屏重新执行
|
||||
return {
|
||||
"status": "completed",
|
||||
"actions": result.get("actions") or [],
|
||||
@@ -70,12 +80,16 @@ class ScreenAgentNode(BaseNode):
|
||||
data=run_data(),
|
||||
),
|
||||
)
|
||||
limit_hit = False
|
||||
if result.get("structured_response") is None:
|
||||
limit_hit = reached_call_limit(result["messages"])
|
||||
last = result["messages"][-1]
|
||||
self.logger.warning(
|
||||
f"agent 没有给出结论: {last.type} {str(last.content)[:500]} "
|
||||
f"tool_calls={getattr(last, 'tool_calls', None)} "
|
||||
f"agent 没有给出结论（{'调用次数用尽' if limit_hit else '没有调用 AgentReply'}），收尾一次: "
|
||||
f"{last.type} {str(last.content)[:500]} "
|
||||
f"metadata={getattr(last, 'response_metadata', None)}"
|
||||
)
|
||||
updated_state.update(reply_update(result))
|
||||
reply = await finalize_reply(result, limit_hit=limit_hit)
|
||||
result = {**result, "structured_response": reply}
|
||||
updated_state.update(reply_update(result, limit_hit=limit_hit))
|
||||
return updated_state
|
||||
|
||||
@@ -20,7 +20,9 @@ class QueryState(TypedDict):
|
||||
input_context: NotRequired[InputContext] # 校验后的上下文，包含最新大屏和澄清历史
|
||||
intent_decision: NotRequired[DecisionResponse | None] # Jev 原始意图判断及分数
|
||||
intent: NotRequired[Route | None] # 门控后的去向
|
||||
actions: NotRequired[list[Action] | None] # screen agent 最后一次成功提交的动作
|
||||
actions: NotRequired[
|
||||
list[Action] | None
|
||||
] # screen agent 各次成功提交的动作，按执行顺序拼接
|
||||
clarification: NotRequired[Clarification | None] # 当前澄清问题及选项
|
||||
clarification_reply: NotRequired[
|
||||
ClarificationReply | None
|
||||
|
||||
@@ -15,7 +15,7 @@ class AddAction(ActionModel):
|
||||
node: ScreenNode = Field(
|
||||
description="新组件。id、type、layout 必填，其余只写需要覆盖物料默认配置的字段；"
|
||||
"后端按 type 找到物料，用默认配置补全：对象逐层覆盖，数组整体替换，null 明确覆盖。"
|
||||
"id 不能与现有组件重复，同一批里先 delete 的 ID 可以再用；type 用 list_components 返回的 type。"
|
||||
"id 不能与当前大屏上的组件重复（包括之前提交里新增的），先 delete 的 ID 可以再用；type 用 list_components 返回的 type。"
|
||||
'任何字段值都可以写成 {"$query": "SELECT ..."}，后端替换成上传数据的查询结果（对象数组，键是结果的列名）'
|
||||
)
|
||||
sample_data: bool = Field(
|
||||
|
||||
@@ -1,9 +1,15 @@
|
||||
"""会话及对话轮次，用于多轮上下文和后续持久化"""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Literal
|
||||
from typing import Annotated, Literal
|
||||
|
||||
from pydantic import BaseModel
|
||||
from pydantic import BaseModel, StringConstraints
|
||||
|
||||
from .base import NonBlankString
|
||||
|
||||
Title = Annotated[
|
||||
str, StringConstraints(strip_whitespace=True, min_length=1, max_length=100)
|
||||
]
|
||||
|
||||
|
||||
class Turn(BaseModel):
|
||||
@@ -18,6 +24,15 @@ class Turn(BaseModel):
|
||||
created_at: datetime
|
||||
|
||||
|
||||
class ConversationCreate(BaseModel):
|
||||
# 会话归属的大屏；不传则不归属任何大屏（兼容旧调用）
|
||||
screen_id: NonBlankString | None = None
|
||||
|
||||
|
||||
class ConversationRename(BaseModel):
|
||||
title: Title
|
||||
|
||||
|
||||
class ConversationCreated(BaseModel):
|
||||
conversation_id: str
|
||||
|
||||
@@ -31,9 +46,11 @@ class ConversationSummary(BaseModel):
|
||||
"""历史会话列表项，只包含已产生对话的会话"""
|
||||
|
||||
conversation_id: str
|
||||
title: str # 取自首轮用户输入，过长时截断
|
||||
title: str # 用户改过名用改的，否则取自首轮用户输入，过长时截断
|
||||
screen_id: str | None = None
|
||||
updated_at: datetime # 最近一轮对话的创建时间
|
||||
|
||||
|
||||
class ConversationList(BaseModel):
|
||||
conversations: list[ConversationSummary]
|
||||
total: int # 符合条件的会话总数（不受 limit / offset 影响），管理页翻页用
|
||||
|
||||
@@ -212,7 +212,9 @@ class DataFileInfo(BaseModel):
|
||||
encoding: str | None
|
||||
size: int
|
||||
created_at: datetime
|
||||
# 被已保存的大屏引用时固定为 True，此时 expires_at 是个遥远的占位时间，不会过期
|
||||
expires_at: datetime
|
||||
pinned: bool = False
|
||||
tables: list[DataTableInfo]
|
||||
warnings: list[ParseWarning]
|
||||
# 未启用智能分析，或这是功能上线前上传的文件时为空
|
||||
@@ -229,6 +231,9 @@ class DataFileRecord(DataFileDetail):
|
||||
tables: list[DataTable]
|
||||
# 删除只做标记，记录留到过期，会话内的表名不会被新文件重复使用
|
||||
deleted: bool = False
|
||||
# 引用这个文件的大屏 ID。非空时文件不过期（expires_at 被推到 PINNED_EXPIRES），
|
||||
# 最后一个引用消失后重新按保留时长计时
|
||||
pinned_by: list[str] = Field(default_factory=list)
|
||||
|
||||
|
||||
UploadStatus = Literal["queued", "parsing", "ready", "failed"]
|
||||
|
||||
@@ -0,0 +1,175 @@
|
||||
"""保存的大屏：存储记录与接口模型。页面本身（PageSchema）按原始 JSON 存取，校验见 app/schemas/screen.py"""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Annotated, Any, Literal
|
||||
|
||||
from pydantic import BaseModel, Field, StringConstraints
|
||||
|
||||
NAME_MAX_LENGTH = 100
|
||||
DEFAULT_NAME = "未命名大屏"
|
||||
|
||||
ScreenName = Annotated[
|
||||
str,
|
||||
StringConstraints(strip_whitespace=True, min_length=1, max_length=NAME_MAX_LENGTH),
|
||||
]
|
||||
PublishStatus = Literal["published", "offline"]
|
||||
ListStatus = Literal["draft", "published", "offline"]
|
||||
SortField = Literal["updated_at", "created_at", "name"]
|
||||
|
||||
|
||||
class PublishInfo(BaseModel):
|
||||
"""screens.publish：线上版本的指针，快照本身在 screen_publishes"""
|
||||
|
||||
status: PublishStatus
|
||||
version: int # 当前线上版本号
|
||||
revision: int # 线上版本是从哪个 revision 发出去的
|
||||
share_id: str
|
||||
published_at: datetime
|
||||
|
||||
|
||||
class ScreenSummary(BaseModel):
|
||||
"""不含页面的大屏记录，列表、元数据更新和线上读取用"""
|
||||
|
||||
screen_id: str
|
||||
owner_id: int
|
||||
name: str
|
||||
thumbnail: str | None = None
|
||||
canvas: dict[str, Any] | None = None
|
||||
revision: int
|
||||
created_at: datetime
|
||||
updated_at: datetime
|
||||
deleted_at: datetime | None = None
|
||||
publish: PublishInfo | None = None
|
||||
|
||||
|
||||
class ScreenRecord(ScreenSummary):
|
||||
page: dict[str, Any]
|
||||
|
||||
|
||||
class PublishEntry(BaseModel):
|
||||
screen_id: str
|
||||
version: int
|
||||
revision: int
|
||||
published_at: datetime
|
||||
published_by: int
|
||||
|
||||
|
||||
class PublishSnapshot(PublishEntry):
|
||||
"""发布快照，只增不改"""
|
||||
|
||||
page: dict[str, Any]
|
||||
|
||||
|
||||
# ---- 请求 ----
|
||||
|
||||
|
||||
class ScreenCreate(BaseModel):
|
||||
name: ScreenName = DEFAULT_NAME
|
||||
page: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class ScreenSave(BaseModel):
|
||||
page: dict[str, Any]
|
||||
revision: int = Field(ge=0)
|
||||
|
||||
|
||||
class ScreenRename(BaseModel):
|
||||
name: ScreenName
|
||||
|
||||
|
||||
class ScreenPublishRequest(BaseModel):
|
||||
revision: int = Field(ge=0)
|
||||
|
||||
|
||||
class ScreenCopyRequest(BaseModel):
|
||||
name: ScreenName | None = None
|
||||
|
||||
|
||||
# ---- 响应 ----
|
||||
|
||||
|
||||
class ScreenCreated(BaseModel):
|
||||
screen_id: str
|
||||
revision: int
|
||||
|
||||
|
||||
class PublishBrief(BaseModel):
|
||||
status: PublishStatus
|
||||
version: int
|
||||
share_id: str
|
||||
published_at: datetime
|
||||
|
||||
|
||||
class ScreenDetail(BaseModel):
|
||||
screen_id: str
|
||||
name: str
|
||||
page: dict[str, Any]
|
||||
revision: int
|
||||
updated_at: datetime
|
||||
publish: PublishBrief | None
|
||||
|
||||
|
||||
class ScreenSaved(BaseModel):
|
||||
revision: int
|
||||
updated_at: datetime
|
||||
|
||||
|
||||
class ScreenPublished(BaseModel):
|
||||
version: int
|
||||
share_id: str
|
||||
published_at: datetime
|
||||
|
||||
|
||||
class ScreenListItem(BaseModel):
|
||||
screen_id: str
|
||||
name: str
|
||||
thumbnail_url: str | None
|
||||
canvas: dict[str, Any] | None
|
||||
status: ListStatus
|
||||
has_unpublished_changes: bool
|
||||
share_id: str | None
|
||||
updated_at: datetime
|
||||
published_at: datetime | None
|
||||
deleted_at: datetime | None
|
||||
# 在回收站里的大屏，到这个时间会被自动彻底删除
|
||||
purge_at: datetime | None
|
||||
|
||||
|
||||
class ScreenList(BaseModel):
|
||||
items: list[ScreenListItem]
|
||||
total: int
|
||||
|
||||
|
||||
class TrashEmptied(BaseModel):
|
||||
count: int
|
||||
|
||||
|
||||
class ScreenCopied(BaseModel):
|
||||
screen_id: str
|
||||
|
||||
|
||||
class ThumbnailResult(BaseModel):
|
||||
thumbnail_url: str
|
||||
|
||||
|
||||
class PublishHistoryItem(BaseModel):
|
||||
version: int
|
||||
revision: int
|
||||
published_at: datetime
|
||||
published_by: int
|
||||
|
||||
|
||||
class PublishHistory(BaseModel):
|
||||
items: list[PublishHistoryItem]
|
||||
total: int
|
||||
|
||||
|
||||
class ScreenRestored(BaseModel):
|
||||
revision: int
|
||||
|
||||
|
||||
class PublicScreen(BaseModel):
|
||||
name: str
|
||||
page: dict[str, Any]
|
||||
version: int
|
||||
published_at: datetime
|
||||
@@ -1,7 +1,7 @@
|
||||
"""异步执行与恢复大屏任务，当前使用单进程内存检查点"""
|
||||
|
||||
from contextlib import aclosing
|
||||
from datetime import datetime, timezone
|
||||
from datetime import UTC, datetime
|
||||
from uuid import uuid4
|
||||
|
||||
from langgraph.types import Command
|
||||
@@ -201,7 +201,7 @@ class AgentService:
|
||||
prompt=state["request"]["prompt"],
|
||||
outcome=result.status,
|
||||
summary=self._summarize(result),
|
||||
created_at=datetime.now(timezone.utc),
|
||||
created_at=datetime.now(UTC),
|
||||
)
|
||||
)
|
||||
except Exception:
|
||||
@@ -210,7 +210,7 @@ class AgentService:
|
||||
@staticmethod
|
||||
def _result(state, thread_id):
|
||||
waiting = bool(state.get("__interrupt__"))
|
||||
# 只交付完成状态下的动作；actions 是 agent 最后一次成功提交、已通过模拟执行的整批动作
|
||||
# 只交付完成状态下的动作；actions 是 agent 各次成功提交的动作按顺序拼接，在原始大屏上依次执行得到最终大屏
|
||||
deliver = state["status"] == "completed"
|
||||
|
||||
return AgentResult(
|
||||
|
||||
@@ -21,7 +21,7 @@ def make_title(prompt: str) -> str:
|
||||
class ConversationStore(Protocol):
|
||||
"""内存和 MongoDB 实现同一组方法，业务代码不需要关心具体存储"""
|
||||
|
||||
async def create(self, user_id: int) -> str: ...
|
||||
async def create(self, user_id: int, screen_id: str | None = None) -> str: ...
|
||||
|
||||
async def exists(self, conversation_id: str, user_id: int) -> bool:
|
||||
"""会话已签发、未过期且属于该用户；不属于时与不存在同样返回 False"""
|
||||
@@ -34,9 +34,21 @@ class ConversationStore(Protocol):
|
||||
async def save_turn(self, turn: Turn) -> None: ...
|
||||
|
||||
async def list_by_user(
|
||||
self, user_id: int, limit: int, offset: int
|
||||
self, user_id: int, limit: int, offset: int, screen_id: str | None = None
|
||||
) -> list[ConversationSummary]:
|
||||
"""该用户已产生对话的会话，按最近更新时间倒序；没有对话的空会话不返回"""
|
||||
"""该用户已产生对话的会话，按最近更新时间倒序；没有对话的空会话不返回；给了 screen_id 只返回该大屏的"""
|
||||
...
|
||||
|
||||
async def count_by_user(self, user_id: int, screen_id: str | None = None) -> int:
|
||||
"""与 list_by_user 同一口径的总数"""
|
||||
...
|
||||
|
||||
async def rename(self, conversation_id: str, user_id: int, title: str) -> bool:
|
||||
"""会话不存在或不属于该用户返回 False"""
|
||||
...
|
||||
|
||||
async def delete(self, conversation_id: str, user_id: int) -> bool:
|
||||
"""连同对话轮次一起删除；会话不存在或不属于该用户返回 False"""
|
||||
...
|
||||
|
||||
|
||||
@@ -51,6 +63,8 @@ class InMemoryConversationStore:
|
||||
# 会话 ID -> 创建时间；轮次按 thread_id 去重并保持首次写入顺序
|
||||
self._created: dict[str, datetime] = {}
|
||||
self._owner: dict[str, int] = {}
|
||||
self._screen: dict[str, str | None] = {}
|
||||
self._titles: dict[str, str] = {}
|
||||
self._turns: dict[str, dict[str, Turn]] = {}
|
||||
|
||||
def _purge_expired(self) -> None:
|
||||
@@ -61,15 +75,18 @@ class InMemoryConversationStore:
|
||||
for key, created in self._created.items()
|
||||
if created < deadline and not self._turns.get(key)
|
||||
]:
|
||||
del self._created[conversation_id]
|
||||
self._owner.pop(conversation_id, None)
|
||||
self._turns.pop(conversation_id, None)
|
||||
self._forget(conversation_id)
|
||||
|
||||
async def create(self, user_id: int) -> str:
|
||||
def _forget(self, conversation_id: str) -> None:
|
||||
for table in (self._created, self._owner, self._screen, self._titles, self._turns):
|
||||
table.pop(conversation_id, None)
|
||||
|
||||
async def create(self, user_id: int, screen_id: str | None = None) -> str:
|
||||
self._purge_expired()
|
||||
conversation_id = str(uuid4())
|
||||
self._created[conversation_id] = _now()
|
||||
self._owner[conversation_id] = user_id
|
||||
self._screen[conversation_id] = screen_id
|
||||
return conversation_id
|
||||
|
||||
async def exists(self, conversation_id: str, user_id: int) -> bool:
|
||||
@@ -94,8 +111,8 @@ class InMemoryConversationStore:
|
||||
turn = turn.model_copy(update={"created_at": existing.created_at})
|
||||
self._turns[turn.conversation_id][turn.thread_id] = turn
|
||||
|
||||
async def list_by_user(
|
||||
self, user_id: int, limit: int, offset: int
|
||||
def _summaries(
|
||||
self, user_id: int, screen_id: str | None
|
||||
) -> list[ConversationSummary]:
|
||||
summaries = []
|
||||
for conversation_id, owner in self._owner.items():
|
||||
@@ -105,14 +122,37 @@ class InMemoryConversationStore:
|
||||
)
|
||||
if owner != user_id or not turns:
|
||||
continue
|
||||
if screen_id is not None and self._screen.get(conversation_id) != screen_id:
|
||||
continue
|
||||
summaries.append(
|
||||
ConversationSummary(
|
||||
conversation_id=conversation_id,
|
||||
title=make_title(turns[0].prompt),
|
||||
title=self._titles.get(conversation_id) or make_title(turns[0].prompt),
|
||||
screen_id=self._screen.get(conversation_id),
|
||||
updated_at=turns[-1].created_at,
|
||||
)
|
||||
)
|
||||
summaries.sort(
|
||||
key=lambda item: (item.updated_at, item.conversation_id), reverse=True
|
||||
)
|
||||
return summaries[offset : offset + limit]
|
||||
return summaries
|
||||
|
||||
async def list_by_user(
|
||||
self, user_id: int, limit: int, offset: int, screen_id: str | None = None
|
||||
) -> list[ConversationSummary]:
|
||||
return self._summaries(user_id, screen_id)[offset : offset + limit]
|
||||
|
||||
async def count_by_user(self, user_id: int, screen_id: str | None = None) -> int:
|
||||
return len(self._summaries(user_id, screen_id))
|
||||
|
||||
async def rename(self, conversation_id: str, user_id: int, title: str) -> bool:
|
||||
if not await self.exists(conversation_id, user_id):
|
||||
return False
|
||||
self._titles[conversation_id] = title
|
||||
return True
|
||||
|
||||
async def delete(self, conversation_id: str, user_id: int) -> bool:
|
||||
if not await self.exists(conversation_id, user_id):
|
||||
return False
|
||||
self._forget(conversation_id)
|
||||
return True
|
||||
|
||||
@@ -45,6 +45,8 @@ class MongoConversationStore:
|
||||
await self._conversations.create_index("created_at")
|
||||
# 按用户查询会话
|
||||
await self._conversations.create_index("user_id")
|
||||
# 按大屏过滤会话
|
||||
await self._conversations.create_index([("user_id", ASCENDING), ("screen_id", ASCENDING)])
|
||||
|
||||
async def _purge_expired(self) -> None:
|
||||
# 只清理没有任何轮次且超过保留时长的空会话，有对话的会话不受影响
|
||||
@@ -66,12 +68,17 @@ class MongoConversationStore:
|
||||
if empty:
|
||||
await self._conversations.delete_many({"_id": {"$in": empty}})
|
||||
|
||||
async def create(self, user_id: int) -> str:
|
||||
async def create(self, user_id: int, screen_id: str | None = None) -> str:
|
||||
await self._purge_expired()
|
||||
conversation_id = str(uuid4())
|
||||
# user_id 来自用户中心，会话只对创建者可见
|
||||
await self._conversations.insert_one(
|
||||
{"_id": conversation_id, "created_at": _now(), "user_id": user_id}
|
||||
{
|
||||
"_id": conversation_id,
|
||||
"created_at": _now(),
|
||||
"user_id": user_id,
|
||||
"screen_id": screen_id,
|
||||
}
|
||||
)
|
||||
return conversation_id
|
||||
|
||||
@@ -121,12 +128,13 @@ class MongoConversationStore:
|
||||
upsert=True,
|
||||
)
|
||||
|
||||
async def list_by_user(
|
||||
self, user_id: int, limit: int, offset: int
|
||||
) -> list[ConversationSummary]:
|
||||
def _with_stats(self, user_id: int, screen_id: str | None) -> list[dict]:
|
||||
match: dict = {"user_id": user_id}
|
||||
if screen_id is not None:
|
||||
match["screen_id"] = screen_id
|
||||
# $unwind 会丢弃没有任何轮次的空会话
|
||||
pipeline = [
|
||||
{"$match": {"user_id": user_id}},
|
||||
return [
|
||||
{"$match": match},
|
||||
{
|
||||
"$lookup": {
|
||||
"from": self._turns.name,
|
||||
@@ -146,6 +154,13 @@ class MongoConversationStore:
|
||||
}
|
||||
},
|
||||
{"$unwind": "$stats"},
|
||||
]
|
||||
|
||||
async def list_by_user(
|
||||
self, user_id: int, limit: int, offset: int, screen_id: str | None = None
|
||||
) -> list[ConversationSummary]:
|
||||
pipeline = [
|
||||
*self._with_stats(user_id, screen_id),
|
||||
{"$sort": {"stats.updated_at": -1, "_id": -1}},
|
||||
{"$skip": offset},
|
||||
{"$limit": limit},
|
||||
@@ -154,8 +169,30 @@ class MongoConversationStore:
|
||||
return [
|
||||
ConversationSummary(
|
||||
conversation_id=item["_id"],
|
||||
title=make_title(item["stats"]["first_prompt"]),
|
||||
title=item.get("title") or make_title(item["stats"]["first_prompt"]),
|
||||
screen_id=item.get("screen_id"),
|
||||
updated_at=item["stats"]["updated_at"],
|
||||
)
|
||||
for item in await cursor.to_list()
|
||||
]
|
||||
|
||||
async def count_by_user(self, user_id: int, screen_id: str | None = None) -> int:
|
||||
pipeline = [*self._with_stats(user_id, screen_id), {"$count": "total"}]
|
||||
cursor = await self._conversations.aggregate(pipeline)
|
||||
counted = await cursor.to_list()
|
||||
return counted[0]["total"] if counted else 0
|
||||
|
||||
async def rename(self, conversation_id: str, user_id: int, title: str) -> bool:
|
||||
if not await self.exists(conversation_id, user_id):
|
||||
return False
|
||||
await self._conversations.update_one(
|
||||
{"_id": conversation_id, "user_id": user_id}, {"$set": {"title": title}}
|
||||
)
|
||||
return True
|
||||
|
||||
async def delete(self, conversation_id: str, user_id: int) -> bool:
|
||||
if not await self.exists(conversation_id, user_id):
|
||||
return False
|
||||
await self._turns.delete_many({"conversation_id": conversation_id})
|
||||
await self._conversations.delete_one({"_id": conversation_id, "user_id": user_id})
|
||||
return True
|
||||
|
||||
@@ -0,0 +1,300 @@
|
||||
"""数据源代理：服务端替浏览器请求第三方接口
|
||||
|
||||
两个入口共用 fetch_json：编辑态由登录用户传 URL；线上态只传数据源 ID，配置取自已发布快照
|
||||
防护：只允许 http/https；DNS 解析后每个地址都必须是公网地址，并把连接钉在已检查过的地址上
|
||||
（不再让 httpx 二次解析，防 DNS 重绑定）；不读环境代理；重定向逐跳重新检查，跨域时丢掉调用方的密钥头；
|
||||
总超时 5 秒，响应体上限 2 MB，按解压后的字节计，解压时带输出上限（防压缩炸弹）
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import ipaddress
|
||||
import json
|
||||
import re
|
||||
import socket
|
||||
import time
|
||||
import zlib
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, Literal
|
||||
from urllib.parse import urljoin
|
||||
|
||||
import httpx
|
||||
|
||||
from app.core.exceptions import BizException
|
||||
|
||||
TIMEOUT = 5.0
|
||||
MAX_BYTES = 2 * 1024 * 1024
|
||||
MAX_REDIRECTS = 3
|
||||
REDIRECT_STATUS = {301, 302, 303, 307, 308}
|
||||
# 这些头由 httpx 或连接本身决定，不让调用方覆盖
|
||||
FORBIDDEN_HEADERS = {
|
||||
"host",
|
||||
"content-length",
|
||||
"connection",
|
||||
"transfer-encoding",
|
||||
"upgrade",
|
||||
"te",
|
||||
"keep-alive",
|
||||
"proxy-authorization",
|
||||
"proxy-connection",
|
||||
"accept-encoding",
|
||||
}
|
||||
# 只接受自己能带上限解压的编码；httpx 默认还会协商 br / zstd，那些解码器没有输出上限
|
||||
ACCEPT_ENCODING = "gzip, deflate"
|
||||
# is_global 会放过的几段：NAT64（可以嵌入任意 IPv4）、IPv4 兼容地址、已废弃的站点本地地址
|
||||
EXTRA_BLOCKED = [
|
||||
ipaddress.ip_network(net)
|
||||
for net in ("64:ff9b::/96", "64:ff9b:1::/48", "::/96", "fec0::/10")
|
||||
]
|
||||
# 十进制、十六进制、省略写法的 IPv4（2130706433、0x7f.1、127.1）：不是标准写法，
|
||||
# 各平台的解析器解释不一，一律当作内网地址拒绝，不交给 DNS
|
||||
NUMERIC_LABEL = re.compile(r"0x[0-9a-f]+|[0-9]+", re.IGNORECASE)
|
||||
|
||||
Method = Literal["GET", "POST"]
|
||||
|
||||
|
||||
def _blocked() -> BizException:
|
||||
return BizException(422, "不允许访问内网或本机地址")
|
||||
|
||||
|
||||
def _is_public(address: ipaddress.IPv4Address | ipaddress.IPv6Address) -> bool:
|
||||
# IPv4 映射的 IPv6 地址按内嵌的 IPv4 判断
|
||||
if isinstance(address, ipaddress.IPv6Address) and address.ipv4_mapped is not None:
|
||||
address = address.ipv4_mapped
|
||||
return (
|
||||
address.is_global
|
||||
and not address.is_multicast
|
||||
and not any(address in net for net in EXTRA_BLOCKED)
|
||||
)
|
||||
|
||||
|
||||
def _checked(
|
||||
address: ipaddress.IPv4Address | ipaddress.IPv6Address,
|
||||
) -> ipaddress.IPv4Address | ipaddress.IPv6Address:
|
||||
if not _is_public(address):
|
||||
raise _blocked()
|
||||
return address
|
||||
|
||||
|
||||
async def _resolve(host: str, port: int) -> str:
|
||||
"""解析主机名，所有地址都是公网才放行，返回第一个地址"""
|
||||
host = host.rstrip(".")
|
||||
try:
|
||||
literal = ipaddress.ip_address(host)
|
||||
except ValueError:
|
||||
pass
|
||||
else:
|
||||
return str(_checked(literal))
|
||||
if host.lower() == "localhost" or host.lower().endswith(".localhost"):
|
||||
raise _blocked()
|
||||
if all(NUMERIC_LABEL.fullmatch(label) for label in host.split(".")):
|
||||
raise _blocked()
|
||||
try:
|
||||
infos = await asyncio.get_running_loop().getaddrinfo(
|
||||
host, port, type=socket.SOCK_STREAM
|
||||
)
|
||||
except socket.gaierror:
|
||||
raise BizException(502, "数据源域名无法解析") from None
|
||||
addresses = [ipaddress.ip_address(info[4][0].split("%")[0]) for info in infos]
|
||||
if not addresses:
|
||||
raise BizException(502, "数据源域名无法解析")
|
||||
for address in addresses:
|
||||
_checked(address)
|
||||
return str(addresses[0])
|
||||
|
||||
|
||||
def _clean_headers(headers: dict[str, Any] | None) -> dict[str, str]:
|
||||
cleaned = {"Accept": "application/json", "Accept-Encoding": ACCEPT_ENCODING}
|
||||
for name, value in (headers or {}).items():
|
||||
if name.lower() in FORBIDDEN_HEADERS or not isinstance(
|
||||
value, (str, int, float)
|
||||
):
|
||||
continue
|
||||
cleaned[name] = str(value)
|
||||
return cleaned
|
||||
|
||||
|
||||
def _query_value(value: Any) -> Any:
|
||||
return value if isinstance(value, (str, int, float, bool)) else json.dumps(value)
|
||||
|
||||
|
||||
def _parse_url(raw: str) -> httpx.URL:
|
||||
try:
|
||||
url = httpx.URL(raw)
|
||||
except httpx.InvalidURL:
|
||||
raise BizException(422, "数据源地址不合法") from None
|
||||
if url.scheme not in ("http", "https") or not url.host or url.userinfo:
|
||||
raise BizException(422, "数据源地址只支持 http / https，且不能带账号密码")
|
||||
return url
|
||||
|
||||
|
||||
def _too_large() -> BizException:
|
||||
return BizException(502, "数据源响应超过 2 MB")
|
||||
|
||||
|
||||
async def _read_body(response: httpx.Response) -> bytes:
|
||||
"""读响应体：原始字节和解压后的字节都不超过 MAX_BYTES
|
||||
|
||||
不用 aiter_bytes：它对每个 64 KB 的块整块解压，没有输出上限，一个块就能膨胀到上 GB
|
||||
"""
|
||||
encoding = response.headers.get("content-encoding", "identity").strip().lower()
|
||||
if encoding in ("", "identity"):
|
||||
decoder = None
|
||||
elif encoding in ("gzip", "x-gzip"):
|
||||
decoder = zlib.decompressobj(16 + zlib.MAX_WBITS)
|
||||
elif encoding == "deflate":
|
||||
decoder = zlib.decompressobj()
|
||||
else:
|
||||
raise BizException(502, "数据源响应的压缩格式不支持")
|
||||
declared = response.headers.get("content-length", "")
|
||||
if declared.isdigit() and int(declared) > MAX_BYTES:
|
||||
raise _too_large()
|
||||
chunks: list[bytes] = []
|
||||
raw_size = size = 0
|
||||
try:
|
||||
async for chunk in response.aiter_raw():
|
||||
raw_size += len(chunk)
|
||||
if raw_size > MAX_BYTES:
|
||||
raise _too_large()
|
||||
if decoder is None:
|
||||
chunks.append(chunk)
|
||||
size += len(chunk)
|
||||
continue
|
||||
pending = chunk
|
||||
while pending:
|
||||
# max_length 为 0 表示不限，这里恒大于 0
|
||||
part = decoder.decompress(pending, MAX_BYTES - size + 1)
|
||||
size += len(part)
|
||||
if size > MAX_BYTES:
|
||||
raise _too_large()
|
||||
chunks.append(part)
|
||||
pending = decoder.unconsumed_tail
|
||||
if not part:
|
||||
break
|
||||
# 压缩流没有正常结束（被截断）也算失败，别把半截数据当成功
|
||||
if decoder is not None and not decoder.eof:
|
||||
raise BizException(502, "数据源响应解压失败")
|
||||
except zlib.error:
|
||||
raise BizException(502, "数据源响应解压失败") from None
|
||||
return b"".join(chunks)
|
||||
|
||||
|
||||
async def _request(
|
||||
client: httpx.AsyncClient,
|
||||
method: str,
|
||||
url: httpx.URL,
|
||||
headers: dict[str, str],
|
||||
body: Any,
|
||||
) -> tuple[int, httpx.Headers, bytes]:
|
||||
port = url.port or (443 if url.scheme == "https" else 80)
|
||||
pinned = url.copy_with(host=await _resolve(url.host, port))
|
||||
request = client.build_request(
|
||||
method,
|
||||
pinned,
|
||||
headers={**headers, "Host": url.netloc.decode()},
|
||||
json=body if method == "POST" else None,
|
||||
# 连接用 IP，TLS 的 SNI 和证书校验仍然用原主机名
|
||||
extensions={"sni_hostname": url.host},
|
||||
)
|
||||
response = await client.send(request, stream=True)
|
||||
try:
|
||||
return response.status_code, response.headers, await _read_body(response)
|
||||
finally:
|
||||
await response.aclose()
|
||||
|
||||
|
||||
def _origin(url: httpx.URL) -> tuple[str, str, int | None]:
|
||||
return url.scheme, url.host, url.port or (443 if url.scheme == "https" else 80)
|
||||
|
||||
|
||||
async def fetch_json(
|
||||
url: str,
|
||||
method: Method = "GET",
|
||||
params: dict[str, Any] | None = None,
|
||||
headers: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any] | list[Any]:
|
||||
"""GET 的 params 作为查询参数，POST 的 params 作为 JSON 请求体，和前端 fetchData 一致"""
|
||||
target = _parse_url(url)
|
||||
request_headers = _clean_headers(headers)
|
||||
body = None
|
||||
if method == "GET":
|
||||
if params:
|
||||
target = target.copy_merge_params(
|
||||
{key: _query_value(value) for key, value in params.items()}
|
||||
)
|
||||
else:
|
||||
body = params or {}
|
||||
try:
|
||||
async with asyncio.timeout(TIMEOUT):
|
||||
# trust_env=False：环境里的代理会绕过上面的地址检查
|
||||
# max_keepalive_connections=0：连接按 IP 复用，不关的话同一 IP 上的下一跳会沿用上一跳的 TLS 校验结果
|
||||
async with httpx.AsyncClient(
|
||||
trust_env=False,
|
||||
timeout=TIMEOUT,
|
||||
limits=httpx.Limits(max_keepalive_connections=0),
|
||||
) as client:
|
||||
for _ in range(MAX_REDIRECTS + 1):
|
||||
status, response_headers, content = await _request(
|
||||
client, method, target, request_headers, body
|
||||
)
|
||||
location = response_headers.get("location")
|
||||
if status not in REDIRECT_STATUS or not location:
|
||||
break
|
||||
previous = target
|
||||
target = _parse_url(urljoin(str(target), location))
|
||||
if _origin(target) != _origin(previous):
|
||||
# 密钥头只发给调用方指定的那个源
|
||||
request_headers = _clean_headers(None)
|
||||
if status not in (307, 308):
|
||||
method, body = "GET", None
|
||||
else:
|
||||
raise BizException(502, "数据源重定向次数过多")
|
||||
except (TimeoutError, httpx.TimeoutException):
|
||||
raise BizException(504, "数据源请求超时") from None
|
||||
except httpx.HTTPError:
|
||||
raise BizException(502, "数据源请求失败") from None
|
||||
if not 200 <= status < 300:
|
||||
raise BizException(502, f"数据源返回 HTTP {status}")
|
||||
try:
|
||||
data = json.loads(content)
|
||||
except (ValueError, RecursionError):
|
||||
# 2 MB 以内的深层嵌套会打满递归，同样当作非法 JSON
|
||||
raise BizException(502, "数据源返回的不是 JSON") from None
|
||||
if not isinstance(data, (dict, list)):
|
||||
raise BizException(502, "数据源返回的不是 JSON 对象或数组")
|
||||
return data
|
||||
|
||||
|
||||
class FetchCache:
|
||||
"""线上大屏多人同时开、按 interval 轮询：同一个 key 几秒内只真正请求一次，并发的请求共用同一次
|
||||
|
||||
失败不缓存；单实例部署，缓存放进程内存。单条最大 2 MB，条数上限乘出来就是内存上限，所以条数要小
|
||||
"""
|
||||
|
||||
def __init__(self, ttl: float = 3.0, max_entries: int = 64):
|
||||
self._ttl = ttl
|
||||
self._max_entries = max_entries
|
||||
self._entries: dict[Any, tuple[float, Any]] = {}
|
||||
self._inflight: dict[Any, asyncio.Future] = {}
|
||||
|
||||
async def get(self, key: Any, fetch: Callable[[], Awaitable[Any]]) -> Any:
|
||||
now = time.monotonic()
|
||||
hit = self._entries.get(key)
|
||||
if hit is not None and hit[0] > now:
|
||||
return hit[1]
|
||||
future = self._inflight.get(key)
|
||||
if future is None:
|
||||
future = asyncio.ensure_future(fetch())
|
||||
self._inflight[key] = future
|
||||
future.add_done_callback(lambda done: self._finish(key, done))
|
||||
# shield：某个请求者断开不能取消共用的那次请求
|
||||
return await asyncio.shield(future)
|
||||
|
||||
def _finish(self, key: Any, future: asyncio.Future) -> None:
|
||||
self._inflight.pop(key, None)
|
||||
if future.cancelled() or future.exception() is not None:
|
||||
return
|
||||
now = time.monotonic()
|
||||
# 每次写入先丢掉过期项；清完还满就不缓存这一条
|
||||
self._entries = {k: v for k, v in self._entries.items() if v[0] > now}
|
||||
if len(self._entries) < self._max_entries:
|
||||
self._entries[key] = (now + self._ttl, future.result())
|
||||
@@ -0,0 +1,238 @@
|
||||
"""大屏存储的 MongoDB 实现，方法语义与内存实现一致
|
||||
|
||||
集合：screens（草稿，一屏一条）、screen_publishes（发布快照，只增不改）
|
||||
"""
|
||||
|
||||
import re
|
||||
from contextlib import contextmanager
|
||||
from datetime import UTC
|
||||
from typing import Any
|
||||
|
||||
from bson.codec_options import CodecOptions
|
||||
from bson.errors import InvalidDocument
|
||||
from pydantic import BaseModel
|
||||
from pymongo import ASCENDING, DESCENDING, ReturnDocument
|
||||
from pymongo.errors import DocumentTooLarge, DuplicateKeyError
|
||||
|
||||
from app.core.exceptions import BizException
|
||||
from app.schemas.saved_screen import (
|
||||
PublishEntry,
|
||||
PublishSnapshot,
|
||||
ScreenRecord,
|
||||
ScreenSummary,
|
||||
)
|
||||
from app.services.screen_store import canvas_size
|
||||
|
||||
# 列表和元数据更新不需要页面，只取画布；Mongo 的投影不能同时"排除 page"和"包含 page.canvas"，所以列字段
|
||||
SUMMARY_PROJECTION = {
|
||||
"owner_id": 1,
|
||||
"name": 1,
|
||||
"thumbnail": 1,
|
||||
"page.canvas": 1,
|
||||
"revision": 1,
|
||||
"created_at": 1,
|
||||
"updated_at": 1,
|
||||
"deleted_at": 1,
|
||||
"publish": 1,
|
||||
}
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _storable():
|
||||
"""页面里的超大整数、含 NUL 的字段名、编码后超过 16 MB 的文档，驱动会抛各种非业务异常
|
||||
|
||||
10 MB 的 JSON 编码成 BSON 后会涨（数组的下标要当字符串键存），所以体积上限靠这里兜底
|
||||
"""
|
||||
try:
|
||||
yield
|
||||
except (OverflowError, InvalidDocument, DocumentTooLarge):
|
||||
raise BizException(
|
||||
422, "大屏数据无法保存：数字超出范围、字段名不合法或体积过大"
|
||||
) from None
|
||||
|
||||
|
||||
# update 自己拼的条件，调用方的 where 不许覆盖
|
||||
RESERVED_FILTER_KEYS = {"_id", "owner_id", "deleted_at", "revision"}
|
||||
|
||||
|
||||
def _owned(screen_id: str, owner_id: int, deleted: bool) -> dict[str, Any]:
|
||||
return {
|
||||
"_id": screen_id,
|
||||
"owner_id": owner_id,
|
||||
"deleted_at": {"$ne": None} if deleted else None,
|
||||
}
|
||||
|
||||
|
||||
def _encode(value: Any) -> Any:
|
||||
return value.model_dump() if isinstance(value, BaseModel) else value
|
||||
|
||||
|
||||
def _summary(document: dict[str, Any]) -> ScreenSummary:
|
||||
canvas = (document.get("page") or {}).get("canvas")
|
||||
data = {k: v for k, v in document.items() if k not in ("_id", "page")}
|
||||
return ScreenSummary.model_validate(
|
||||
{**data, "screen_id": document["_id"], "canvas": canvas_size(canvas)}
|
||||
)
|
||||
|
||||
|
||||
class MongoScreenStore:
|
||||
def __init__(self, db: Any, prefix: str = ""):
|
||||
# 读回的时间统一带 UTC 时区；集合名前缀用于测试隔离
|
||||
options = CodecOptions(tz_aware=True, tzinfo=UTC)
|
||||
self._screens = db.get_collection(f"{prefix}screens", codec_options=options)
|
||||
self._publishes = db.get_collection(
|
||||
f"{prefix}screen_publishes", codec_options=options
|
||||
)
|
||||
|
||||
async def ensure_indexes(self) -> None:
|
||||
await self._screens.create_index(
|
||||
[
|
||||
("owner_id", ASCENDING),
|
||||
("deleted_at", ASCENDING),
|
||||
("updated_at", DESCENDING),
|
||||
]
|
||||
)
|
||||
# 分享链接唯一；没发布过的大屏没有 publish.share_id，不进索引
|
||||
await self._screens.create_index(
|
||||
"publish.share_id",
|
||||
unique=True,
|
||||
partialFilterExpression={"publish.share_id": {"$type": "string"}},
|
||||
)
|
||||
await self._publishes.create_index(
|
||||
[("screen_id", ASCENDING), ("version", DESCENDING)], unique=True
|
||||
)
|
||||
# 回收站过期清理按 deleted_at 跨用户扫描；没删除的不进索引
|
||||
await self._screens.create_index(
|
||||
"deleted_at", partialFilterExpression={"deleted_at": {"$type": "date"}}
|
||||
)
|
||||
|
||||
async def insert(self, record: ScreenRecord) -> None:
|
||||
document = record.model_dump(exclude={"screen_id", "canvas"})
|
||||
with _storable():
|
||||
await self._screens.insert_one({"_id": record.screen_id, **document})
|
||||
|
||||
async def get(self, screen_id, owner_id, *, deleted=False):
|
||||
document = await self._screens.find_one(_owned(screen_id, owner_id, deleted))
|
||||
if document is None:
|
||||
return None
|
||||
return ScreenRecord.model_validate(
|
||||
{**_summary(document).model_dump(), "page": document["page"]}
|
||||
)
|
||||
|
||||
async def get_summary(self, screen_id, owner_id, *, deleted=False):
|
||||
document = await self._screens.find_one(
|
||||
_owned(screen_id, owner_id, deleted), SUMMARY_PROJECTION
|
||||
)
|
||||
return None if document is None else _summary(document)
|
||||
|
||||
async def update(
|
||||
self, screen_id, owner_id, changes, *, revision=None, where=None, deleted=False
|
||||
):
|
||||
query = _owned(screen_id, owner_id, deleted)
|
||||
if revision is not None:
|
||||
query["revision"] = revision
|
||||
if clash := RESERVED_FILTER_KEYS & (where or {}).keys():
|
||||
raise ValueError(f"where 不能覆盖所有权和版本条件：{sorted(clash)}")
|
||||
query.update(where or {})
|
||||
with _storable():
|
||||
document = await self._screens.find_one_and_update(
|
||||
query,
|
||||
{"$set": {key: _encode(value) for key, value in changes.items()}},
|
||||
projection=SUMMARY_PROJECTION,
|
||||
return_document=ReturnDocument.AFTER,
|
||||
)
|
||||
return None if document is None else _summary(document)
|
||||
|
||||
async def search(
|
||||
self, owner_id, *, keyword, status, sort, descending, limit, offset, deleted
|
||||
):
|
||||
query: dict[str, Any] = {
|
||||
"owner_id": owner_id,
|
||||
"deleted_at": {"$ne": None} if deleted else None,
|
||||
}
|
||||
if keyword:
|
||||
query["name"] = {"$regex": re.escape(keyword), "$options": "i"}
|
||||
if status == "draft":
|
||||
query["publish"] = None
|
||||
elif status is not None:
|
||||
query["publish.status"] = status
|
||||
direction = DESCENDING if descending else ASCENDING
|
||||
documents = await (
|
||||
self._screens.find(query, SUMMARY_PROJECTION)
|
||||
.sort([(sort, direction), ("_id", direction)])
|
||||
.skip(offset)
|
||||
.limit(limit)
|
||||
.to_list()
|
||||
)
|
||||
return [_summary(d) for d in documents], await self._screens.count_documents(
|
||||
query
|
||||
)
|
||||
|
||||
async def get_by_share_id(self, share_id):
|
||||
document = await self._screens.find_one(
|
||||
{"publish.share_id": share_id, "deleted_at": None}, SUMMARY_PROJECTION
|
||||
)
|
||||
return None if document is None else _summary(document)
|
||||
|
||||
async def list_expired_trash(self, deleted_before, limit):
|
||||
documents = await (
|
||||
self._screens.find(
|
||||
# 带 $type 才蕴含部分索引的过滤条件，查询才用得上索引
|
||||
{"deleted_at": {"$type": "date", "$lt": deleted_before}},
|
||||
SUMMARY_PROJECTION,
|
||||
)
|
||||
.sort("deleted_at", ASCENDING)
|
||||
.limit(limit)
|
||||
.to_list()
|
||||
)
|
||||
return [_summary(d) for d in documents]
|
||||
|
||||
async def purge(self, screen_id, owner_id):
|
||||
# 带"仍在回收站"的条件：清理期间被还原的大屏不能删
|
||||
result = await self._screens.delete_one(_owned(screen_id, owner_id, True))
|
||||
return result.deleted_count == 1
|
||||
|
||||
async def delete_publishes(self, screen_id):
|
||||
await self._publishes.delete_many({"screen_id": screen_id})
|
||||
|
||||
async def last_publish_version(self, screen_id):
|
||||
document = await self._publishes.find_one(
|
||||
{"screen_id": screen_id}, {"version": 1}, sort=[("version", DESCENDING)]
|
||||
)
|
||||
return 0 if document is None else document["version"]
|
||||
|
||||
async def insert_publish(self, snapshot):
|
||||
try:
|
||||
with _storable():
|
||||
await self._publishes.insert_one(snapshot.model_dump())
|
||||
except DuplicateKeyError:
|
||||
return False
|
||||
return True
|
||||
|
||||
async def get_publish(self, screen_id, version):
|
||||
document = await self._publishes.find_one(
|
||||
{"screen_id": screen_id, "version": version}
|
||||
)
|
||||
if document is None:
|
||||
return None
|
||||
return PublishSnapshot.model_validate(
|
||||
{k: v for k, v in document.items() if k != "_id"}
|
||||
)
|
||||
|
||||
async def list_publishes(self, screen_id, limit, offset):
|
||||
query = {"screen_id": screen_id}
|
||||
documents = await (
|
||||
self._publishes.find(query, {"page": 0})
|
||||
.sort("version", DESCENDING)
|
||||
.skip(offset)
|
||||
.limit(limit)
|
||||
.to_list()
|
||||
)
|
||||
entries = [
|
||||
PublishEntry.model_validate({k: v for k, v in d.items() if k != "_id"})
|
||||
for d in documents
|
||||
]
|
||||
return entries, await self._publishes.count_documents(query)
|
||||
|
||||
async def delete_publish(self, screen_id, version):
|
||||
await self._publishes.delete_one({"screen_id": screen_id, "version": version})
|
||||
@@ -0,0 +1,241 @@
|
||||
"""大屏存储：协议和进程内存实现；MongoDB 实现见 screen_mongo
|
||||
|
||||
内存实现只给测试用——大屏数据不能丢，没配 MongoDB 时 screens 接口直接返回 503
|
||||
所有查询都带 owner_id：别人的大屏与不存在同样返回 None
|
||||
"""
|
||||
|
||||
from datetime import datetime
|
||||
from typing import Any, Protocol
|
||||
|
||||
from app.schemas.saved_screen import (
|
||||
ListStatus,
|
||||
PublishEntry,
|
||||
PublishSnapshot,
|
||||
ScreenRecord,
|
||||
ScreenSummary,
|
||||
SortField,
|
||||
)
|
||||
|
||||
|
||||
def canvas_size(canvas: Any) -> dict[str, Any] | None:
|
||||
"""列表卡片只需要画布宽高"""
|
||||
if not isinstance(canvas, dict):
|
||||
return None
|
||||
return {key: canvas[key] for key in ("width", "height") if key in canvas}
|
||||
|
||||
|
||||
def summarize(record: ScreenRecord) -> ScreenSummary:
|
||||
"""去掉页面，只保留画布宽高"""
|
||||
data = record.model_dump(exclude={"page"})
|
||||
data["canvas"] = canvas_size(record.page.get("canvas"))
|
||||
return ScreenSummary.model_validate(data)
|
||||
|
||||
|
||||
class ScreenStore(Protocol):
|
||||
async def insert(self, record: ScreenRecord) -> None: ...
|
||||
|
||||
async def get(
|
||||
self, screen_id: str, owner_id: int, *, deleted: bool = False
|
||||
) -> ScreenRecord | None:
|
||||
"""带页面的记录；deleted 为 True 时只找回收站里的"""
|
||||
...
|
||||
|
||||
async def get_summary(
|
||||
self, screen_id: str, owner_id: int, *, deleted: bool = False
|
||||
) -> ScreenSummary | None: ...
|
||||
|
||||
async def update(
|
||||
self,
|
||||
screen_id: str,
|
||||
owner_id: int,
|
||||
changes: dict[str, Any],
|
||||
*,
|
||||
revision: int | None = None,
|
||||
where: dict[str, Any] | None = None,
|
||||
deleted: bool = False,
|
||||
) -> ScreenSummary | None:
|
||||
"""一步完成"匹配并更新"；不匹配返回 None
|
||||
|
||||
给了 revision 就要求库里的 revision 与之相等；where 是额外的等值条件，键可以带点，
|
||||
如 {"publish.version": 3}、{"publish": None}——更新是基于先前读到的值做的，值变了就不能盖上去
|
||||
"""
|
||||
...
|
||||
|
||||
async def search(
|
||||
self,
|
||||
owner_id: int,
|
||||
*,
|
||||
keyword: str | None,
|
||||
status: ListStatus | None,
|
||||
sort: SortField,
|
||||
descending: bool,
|
||||
limit: int,
|
||||
offset: int,
|
||||
deleted: bool,
|
||||
) -> tuple[list[ScreenSummary], int]: ...
|
||||
|
||||
async def get_by_share_id(self, share_id: str) -> ScreenSummary | None:
|
||||
"""未删除的大屏；是否已下线由调用方看 publish.status"""
|
||||
...
|
||||
|
||||
async def list_expired_trash(
|
||||
self, deleted_before: datetime, limit: int
|
||||
) -> list[ScreenSummary]:
|
||||
"""所有人的回收站里 deleted_at 早于 deleted_before 的大屏"""
|
||||
...
|
||||
|
||||
async def purge(self, screen_id: str, owner_id: int) -> bool:
|
||||
"""彻底删除回收站里的大屏；已经被还原（不在回收站）或不存在时返回 False"""
|
||||
...
|
||||
|
||||
async def delete_publishes(self, screen_id: str) -> None: ...
|
||||
|
||||
async def last_publish_version(self, screen_id: str) -> int: ...
|
||||
|
||||
async def insert_publish(self, snapshot: PublishSnapshot) -> bool:
|
||||
"""同一大屏同一版本号已存在时返回 False"""
|
||||
...
|
||||
|
||||
async def get_publish(
|
||||
self, screen_id: str, version: int
|
||||
) -> PublishSnapshot | None: ...
|
||||
|
||||
async def list_publishes(
|
||||
self, screen_id: str, limit: int, offset: int
|
||||
) -> tuple[list[PublishEntry], int]:
|
||||
"""版本号倒序"""
|
||||
...
|
||||
|
||||
async def delete_publish(self, screen_id: str, version: int) -> None: ...
|
||||
|
||||
|
||||
class InMemoryScreenStore:
|
||||
def __init__(self) -> None:
|
||||
self._screens: dict[str, ScreenRecord] = {}
|
||||
self._publishes: dict[tuple[str, int], PublishSnapshot] = {}
|
||||
|
||||
def _find(
|
||||
self, screen_id: str, owner_id: int, deleted: bool
|
||||
) -> ScreenRecord | None:
|
||||
record = self._screens.get(screen_id)
|
||||
if record is None or record.owner_id != owner_id:
|
||||
return None
|
||||
return record if (record.deleted_at is not None) == deleted else None
|
||||
|
||||
async def insert(self, record: ScreenRecord) -> None:
|
||||
self._screens[record.screen_id] = record.model_copy(deep=True)
|
||||
|
||||
async def get(self, screen_id, owner_id, *, deleted=False):
|
||||
record = self._find(screen_id, owner_id, deleted)
|
||||
return None if record is None else record.model_copy(deep=True)
|
||||
|
||||
async def get_summary(self, screen_id, owner_id, *, deleted=False):
|
||||
record = self._find(screen_id, owner_id, deleted)
|
||||
return None if record is None else summarize(record)
|
||||
|
||||
async def update(
|
||||
self, screen_id, owner_id, changes, *, revision=None, where=None, deleted=False
|
||||
):
|
||||
record = self._find(screen_id, owner_id, deleted)
|
||||
if record is None or (revision is not None and record.revision != revision):
|
||||
return None
|
||||
if where is not None and not _matches(record, where):
|
||||
return None
|
||||
updated = record.model_copy(update=changes, deep=True)
|
||||
self._screens[screen_id] = updated
|
||||
return summarize(updated)
|
||||
|
||||
async def search(
|
||||
self, owner_id, *, keyword, status, sort, descending, limit, offset, deleted
|
||||
):
|
||||
records = [
|
||||
record
|
||||
for record in self._screens.values()
|
||||
if record.owner_id == owner_id
|
||||
and (record.deleted_at is not None) == deleted
|
||||
and (keyword is None or keyword.lower() in record.name.lower())
|
||||
and _status_matches(record, status)
|
||||
]
|
||||
# 排序字段相同时按 ID 兜底，翻页才稳定
|
||||
records.sort(key=lambda r: (getattr(r, sort), r.screen_id), reverse=descending)
|
||||
return [summarize(r) for r in records[offset : offset + limit]], len(records)
|
||||
|
||||
async def get_by_share_id(self, share_id):
|
||||
for record in self._screens.values():
|
||||
if (
|
||||
record.deleted_at is None
|
||||
and record.publish is not None
|
||||
and record.publish.share_id == share_id
|
||||
):
|
||||
return summarize(record)
|
||||
return None
|
||||
|
||||
async def list_expired_trash(self, deleted_before, limit):
|
||||
expired = sorted(
|
||||
(
|
||||
r
|
||||
for r in self._screens.values()
|
||||
if r.deleted_at is not None and r.deleted_at < deleted_before
|
||||
),
|
||||
key=lambda r: r.deleted_at,
|
||||
)
|
||||
return [summarize(r) for r in expired[:limit]]
|
||||
|
||||
async def purge(self, screen_id, owner_id):
|
||||
if self._find(screen_id, owner_id, True) is None:
|
||||
return False
|
||||
del self._screens[screen_id]
|
||||
return True
|
||||
|
||||
async def delete_publishes(self, screen_id):
|
||||
for key in [key for key in self._publishes if key[0] == screen_id]:
|
||||
del self._publishes[key]
|
||||
|
||||
async def last_publish_version(self, screen_id):
|
||||
return max(
|
||||
(version for sid, version in self._publishes if sid == screen_id), default=0
|
||||
)
|
||||
|
||||
async def insert_publish(self, snapshot):
|
||||
key = (snapshot.screen_id, snapshot.version)
|
||||
if key in self._publishes:
|
||||
return False
|
||||
self._publishes[key] = snapshot.model_copy(deep=True)
|
||||
return True
|
||||
|
||||
async def get_publish(self, screen_id, version):
|
||||
snapshot = self._publishes.get((screen_id, version))
|
||||
return None if snapshot is None else snapshot.model_copy(deep=True)
|
||||
|
||||
async def list_publishes(self, screen_id, limit, offset):
|
||||
snapshots = sorted(
|
||||
(s for (sid, _), s in self._publishes.items() if sid == screen_id),
|
||||
key=lambda s: s.version,
|
||||
reverse=True,
|
||||
)
|
||||
entries = [
|
||||
PublishEntry.model_validate(s.model_dump(exclude={"page"}))
|
||||
for s in snapshots[offset : offset + limit]
|
||||
]
|
||||
return entries, len(snapshots)
|
||||
|
||||
async def delete_publish(self, screen_id, version):
|
||||
self._publishes.pop((screen_id, version), None)
|
||||
|
||||
|
||||
def _matches(record: ScreenRecord, where: dict[str, Any]) -> bool:
|
||||
for path, expected in where.items():
|
||||
value: Any = record
|
||||
for part in path.split("."):
|
||||
value = None if value is None else getattr(value, part, None)
|
||||
if value != expected:
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _status_matches(record: ScreenSummary, status: ListStatus | None) -> bool:
|
||||
if status is None:
|
||||
return True
|
||||
if status == "draft":
|
||||
return record.publish is None
|
||||
return record.publish is not None and record.publish.status == status
|
||||
@@ -0,0 +1,562 @@
|
||||
"""大屏的业务逻辑：存取、乐观锁、发布、历史、复制、封面、线上读取和数据源取数
|
||||
|
||||
存储层只提供"匹配并更新"这类原语，404 / 409 的判断都在这里
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import copy
|
||||
import secrets
|
||||
import weakref
|
||||
from collections.abc import Callable
|
||||
from datetime import UTC, datetime, timedelta
|
||||
from typing import Any, Protocol
|
||||
from uuid import uuid4
|
||||
|
||||
from loguru import logger
|
||||
from pydantic import ValidationError
|
||||
|
||||
from app.core.exceptions import BizException
|
||||
from app.schemas.saved_screen import (
|
||||
NAME_MAX_LENGTH,
|
||||
ListStatus,
|
||||
PublishEntry,
|
||||
PublishInfo,
|
||||
PublishSnapshot,
|
||||
ScreenRecord,
|
||||
ScreenSummary,
|
||||
SortField,
|
||||
)
|
||||
from app.schemas.screen import Screen
|
||||
from app.services.datasource_proxy import FetchCache, fetch_json
|
||||
from app.services.screen_store import ScreenStore
|
||||
from app.services.thumbnails import ThumbnailStorage
|
||||
|
||||
# 事件代码和数据源转换函数是字符串形式的 JS，线上会在浏览器里执行；写入权限只给 owner，再限制长度
|
||||
SCRIPT_MAX_CHARS = 64 * 1024
|
||||
COPY_SUFFIX = " 副本"
|
||||
EMPTY_PAGE = {
|
||||
"canvas": {"width": 1920, "height": 1080, "backgroundColor": "#ffffff10"},
|
||||
"nodes": [],
|
||||
"dataSource": [],
|
||||
}
|
||||
|
||||
# 回收站过期清理：一批最多清多少个，两次扫描的间隔（秒）
|
||||
PURGE_BATCH = 100
|
||||
TRASH_SWEEP_INTERVAL = 3600
|
||||
|
||||
STALE = "大屏已在别处修改"
|
||||
# 下线 / 删除 / 换封面遇到并发修改时的重读次数
|
||||
UPDATE_ATTEMPTS = 3
|
||||
|
||||
|
||||
def _now() -> datetime:
|
||||
return datetime.now(UTC)
|
||||
|
||||
|
||||
class FilePinner(Protocol):
|
||||
"""数据文件服务里和大屏有关的两个方法"""
|
||||
|
||||
async def sync_screen_files(
|
||||
self, user_id: int, screen_id: str, file_ids: set[str]
|
||||
) -> None: ...
|
||||
|
||||
async def release_screen_files(self, screen_id: str) -> None: ...
|
||||
|
||||
|
||||
def referenced_file_ids(page: dict[str, Any]) -> set[str]:
|
||||
"""页面里图表绑定的数据文件：nodes[].fileQuery.file_id"""
|
||||
ids = set()
|
||||
nodes = page.get("nodes")
|
||||
for node in nodes if isinstance(nodes, list) else []:
|
||||
query = node.get("fileQuery") if isinstance(node, dict) else None
|
||||
file_id = query.get("file_id") if isinstance(query, dict) else None
|
||||
if isinstance(file_id, str) and file_id:
|
||||
ids.add(file_id)
|
||||
return ids
|
||||
|
||||
|
||||
def _publish_is(publish: PublishInfo | None) -> dict[str, Any]:
|
||||
""" "线上指针还是我读到的那个"的更新条件"""
|
||||
if publish is None:
|
||||
return {"publish": None}
|
||||
return {"publish.version": publish.version, "publish.status": publish.status}
|
||||
|
||||
|
||||
def _not_found() -> BizException:
|
||||
# 别人的大屏与不存在同样返回 404
|
||||
return BizException(404, "大屏不存在或已删除")
|
||||
|
||||
|
||||
def validate_page(page: dict[str, Any]) -> dict[str, Any]:
|
||||
"""按 Screen 校验结构，但返回的仍是原始 JSON：model_dump 会给没传的字段补出 null，前端读回来就变了"""
|
||||
try:
|
||||
Screen.model_validate(page)
|
||||
except ValidationError as error:
|
||||
first = error.errors()[0]
|
||||
where = ".".join(str(part) for part in first["loc"])
|
||||
raise BizException(422, f"大屏数据不合法：{where}：{first['msg']}") from None
|
||||
_check_script_length(page)
|
||||
return page
|
||||
|
||||
|
||||
def _check_script_length(page: dict[str, Any]) -> None:
|
||||
scripts = [
|
||||
event.get("code")
|
||||
for node in page["nodes"]
|
||||
for event in (node.get("events") or [])
|
||||
if isinstance(event, dict)
|
||||
] + [
|
||||
source.get("transform")
|
||||
for source in page.get("dataSource", [])
|
||||
if isinstance(source, dict)
|
||||
]
|
||||
if any(isinstance(text, str) and len(text) > SCRIPT_MAX_CHARS for text in scripts):
|
||||
raise BizException(
|
||||
422, f"事件代码和数据转换函数不能超过 {SCRIPT_MAX_CHARS // 1024} KB"
|
||||
)
|
||||
|
||||
|
||||
def public_page(page: dict[str, Any]) -> dict[str, Any]:
|
||||
"""线上页面里去掉数据源的 headers：它们常放 API key，线上取数由服务端按快照里的完整配置发起，浏览器用不着"""
|
||||
sources = page.get("dataSource")
|
||||
if not isinstance(sources, list):
|
||||
return {**page}
|
||||
return {
|
||||
**page,
|
||||
"dataSource": [
|
||||
{key: value for key, value in item.items() if key != "headers"}
|
||||
if isinstance(item, dict)
|
||||
else item
|
||||
for item in sources
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
class ScreenService:
|
||||
def __init__(
|
||||
self,
|
||||
store: ScreenStore,
|
||||
thumbnails: ThumbnailStorage,
|
||||
files: FilePinner | None = None,
|
||||
trash_days: int = 30,
|
||||
):
|
||||
self._store = store
|
||||
self._thumbnails = thumbnails
|
||||
# 已保存的大屏引用的数据文件不能过期，否则"刷新数据"会 404；没有文件服务时不管
|
||||
self._files = files
|
||||
if trash_days < 1:
|
||||
raise ValueError("trash_days 必须是正整数")
|
||||
self.trash_days = trash_days
|
||||
# 按大屏的锁：保存 / 回滚 / 删除 / 还原都是"改库再同步固定的文件"两步，
|
||||
# 两个请求交错会让固定的文件跟着先开始、后结束的那个请求的旧页面走。单实例部署，进程内锁就够
|
||||
self._locks: weakref.WeakValueDictionary[str, asyncio.Lock] = (
|
||||
weakref.WeakValueDictionary()
|
||||
)
|
||||
self._datasource_cache = FetchCache()
|
||||
|
||||
def _lock(self, screen_id: str) -> asyncio.Lock:
|
||||
lock = self._locks.get(screen_id)
|
||||
if lock is None:
|
||||
lock = self._locks[screen_id] = asyncio.Lock()
|
||||
return lock
|
||||
|
||||
async def _pin_files(self, owner_id: int, screen_id: str, page: dict[str, Any]):
|
||||
"""让数据文件跟着大屏的引用走。主操作已经成功，这一步失败只记日志，不让保存报错"""
|
||||
if self._files is None:
|
||||
return
|
||||
try:
|
||||
await self._files.sync_screen_files(
|
||||
owner_id, screen_id, referenced_file_ids(page)
|
||||
)
|
||||
except Exception as error:
|
||||
logger.opt(exception=error).error(
|
||||
"同步大屏 {} 引用的数据文件失败", screen_id
|
||||
)
|
||||
|
||||
async def _release_files(self, screen_id: str) -> None:
|
||||
if self._files is None:
|
||||
return
|
||||
try:
|
||||
await self._files.release_screen_files(screen_id)
|
||||
except Exception as error:
|
||||
logger.opt(exception=error).error(
|
||||
"放开大屏 {} 固定的数据文件失败", screen_id
|
||||
)
|
||||
|
||||
# ---- 草稿 ----
|
||||
|
||||
async def create(
|
||||
self, owner_id: int, name: str, page: dict[str, Any] | None
|
||||
) -> ScreenRecord:
|
||||
now = _now()
|
||||
record = ScreenRecord(
|
||||
screen_id=f"s_{uuid4().hex}",
|
||||
owner_id=owner_id,
|
||||
name=name,
|
||||
page=validate_page(copy.deepcopy(EMPTY_PAGE) if page is None else page),
|
||||
revision=1,
|
||||
created_at=now,
|
||||
updated_at=now,
|
||||
)
|
||||
await self._store.insert(record)
|
||||
await self._pin_files(owner_id, record.screen_id, record.page)
|
||||
return record
|
||||
|
||||
async def get(self, screen_id: str, owner_id: int) -> ScreenRecord:
|
||||
record = await self._store.get(screen_id, owner_id)
|
||||
if record is None:
|
||||
raise _not_found()
|
||||
return record
|
||||
|
||||
async def save(
|
||||
self, screen_id: str, owner_id: int, page: dict[str, Any], revision: int
|
||||
) -> ScreenSummary:
|
||||
validate_page(page)
|
||||
async with self._lock(screen_id):
|
||||
saved = await self._update_with_revision(
|
||||
screen_id,
|
||||
owner_id,
|
||||
{"page": page, "revision": revision + 1, "updated_at": _now()},
|
||||
revision,
|
||||
)
|
||||
await self._pin_files(owner_id, screen_id, page)
|
||||
return saved
|
||||
|
||||
async def _update_with_revision(
|
||||
self, screen_id: str, owner_id: int, changes: dict[str, Any], revision: int
|
||||
) -> ScreenSummary:
|
||||
updated = await self._store.update(
|
||||
screen_id, owner_id, changes, revision=revision
|
||||
)
|
||||
if updated is not None:
|
||||
return updated
|
||||
# 没更新成功：大屏不存在是 404，存在但 revision 不一致是 409
|
||||
if await self._store.get_summary(screen_id, owner_id) is None:
|
||||
raise _not_found()
|
||||
raise BizException(409, STALE)
|
||||
|
||||
async def rename(self, screen_id: str, owner_id: int, name: str) -> None:
|
||||
# 改名不算改内容，revision 不动
|
||||
changes = {"name": name, "updated_at": _now()}
|
||||
if await self._store.update(screen_id, owner_id, changes) is None:
|
||||
raise _not_found()
|
||||
|
||||
async def _update_from_summary(
|
||||
self,
|
||||
screen_id: str,
|
||||
owner_id: int,
|
||||
build: Callable[[ScreenSummary], tuple[dict[str, Any], dict[str, Any]] | None],
|
||||
) -> None:
|
||||
"""读出当前状态，据此算出修改，再带着"状态没变"的条件写回；并发改了就重读再来
|
||||
|
||||
下线、删除要整段改写 publish，不带条件就会把并发发布写进去的新指针盖回旧值
|
||||
build 返回 (changes, where)；返回 None 表示已经是目标状态，什么都不用做
|
||||
"""
|
||||
for _ in range(UPDATE_ATTEMPTS):
|
||||
summary = await self._store.get_summary(screen_id, owner_id)
|
||||
if summary is None:
|
||||
raise _not_found()
|
||||
built = build(summary)
|
||||
if built is None:
|
||||
return
|
||||
changes, where = built
|
||||
if await self._store.update(screen_id, owner_id, changes, where=where):
|
||||
return
|
||||
raise BizException(409, "操作冲突，请重试")
|
||||
|
||||
async def delete(self, screen_id: str, owner_id: int) -> None:
|
||||
def build(summary: ScreenSummary):
|
||||
changes: dict[str, Any] = {"deleted_at": _now()}
|
||||
if summary.publish is None:
|
||||
return changes, {"publish": None}
|
||||
# 线上版本同时下线，分享链接立即 404
|
||||
changes["publish"] = summary.publish.model_copy(
|
||||
update={"status": "offline"}
|
||||
)
|
||||
return changes, _publish_is(summary.publish)
|
||||
|
||||
async with self._lock(screen_id):
|
||||
await self._update_from_summary(screen_id, owner_id, build)
|
||||
# 进了回收站的大屏不再固定文件，还原时再固定回来
|
||||
await self._release_files(screen_id)
|
||||
|
||||
async def restore_deleted(self, screen_id: str, owner_id: int) -> None:
|
||||
async with self._lock(screen_id):
|
||||
record = await self._store.get(screen_id, owner_id, deleted=True)
|
||||
if record is None:
|
||||
raise _not_found()
|
||||
changes = {"deleted_at": None}
|
||||
if (
|
||||
await self._store.update(screen_id, owner_id, changes, deleted=True)
|
||||
is None
|
||||
):
|
||||
raise _not_found()
|
||||
await self._pin_files(owner_id, screen_id, record.page)
|
||||
|
||||
# ---- 回收站 ----
|
||||
|
||||
async def _purge(self, summary: ScreenSummary) -> bool:
|
||||
"""彻底删除回收站里的一个大屏
|
||||
|
||||
先删大屏本身（带"仍在回收站"的条件，清理期间被还原的不会误删），再清理附属数据
|
||||
"""
|
||||
if not await self._store.purge(summary.screen_id, summary.owner_id):
|
||||
return False
|
||||
await self._store.delete_publishes(summary.screen_id)
|
||||
self._thumbnails.delete(summary.thumbnail)
|
||||
await self._release_files(summary.screen_id)
|
||||
return True
|
||||
|
||||
async def empty_trash(self, owner_id: int) -> int:
|
||||
"""清空自己的回收站，返回彻底删除的个数"""
|
||||
count = 0
|
||||
previous: set[str] = set()
|
||||
while True:
|
||||
batch, _ = await self._store.search(
|
||||
owner_id,
|
||||
keyword=None,
|
||||
status=None,
|
||||
sort="updated_at",
|
||||
descending=True,
|
||||
limit=PURGE_BATCH,
|
||||
offset=0,
|
||||
deleted=True,
|
||||
)
|
||||
ids = {summary.screen_id for summary in batch}
|
||||
purged = [await self._purge(summary) for summary in batch]
|
||||
count += sum(purged)
|
||||
# 空批才算清完；整批都没删成（被并发还原了）就再查一次，
|
||||
# 但同一批连续两次都删不掉说明有东西卡住，别空转
|
||||
if not batch or (not any(purged) and ids == previous):
|
||||
return count
|
||||
previous = ids
|
||||
|
||||
async def purge_expired_trash(self) -> int:
|
||||
"""彻底删除在回收站里超过保留天数的大屏（所有人的）"""
|
||||
deadline = _now() - timedelta(days=self.trash_days)
|
||||
count = 0
|
||||
previous: set[str] = set()
|
||||
while True:
|
||||
batch = await self._store.list_expired_trash(deadline, PURGE_BATCH)
|
||||
ids = {summary.screen_id for summary in batch}
|
||||
purged = [await self._purge(summary) for summary in batch]
|
||||
count += sum(purged)
|
||||
if not batch or (not any(purged) and ids == previous):
|
||||
return count
|
||||
previous = ids
|
||||
|
||||
async def run_trash_sweeper(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
if removed := await self.purge_expired_trash():
|
||||
logger.info(
|
||||
"已彻底删除 {} 个超过 {} 天的回收站大屏",
|
||||
removed,
|
||||
self.trash_days,
|
||||
)
|
||||
except Exception as error:
|
||||
logger.opt(exception=error).error("清理回收站失败")
|
||||
await asyncio.sleep(TRASH_SWEEP_INTERVAL)
|
||||
|
||||
# ---- 管理页 ----
|
||||
|
||||
async def list_screens(
|
||||
self,
|
||||
owner_id: int,
|
||||
*,
|
||||
keyword: str | None,
|
||||
status: ListStatus | None,
|
||||
sort: SortField,
|
||||
descending: bool,
|
||||
limit: int,
|
||||
offset: int,
|
||||
deleted: bool,
|
||||
) -> tuple[list[ScreenSummary], int]:
|
||||
return await self._store.search(
|
||||
owner_id,
|
||||
keyword=keyword,
|
||||
status=status,
|
||||
sort=sort,
|
||||
descending=descending,
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
deleted=deleted,
|
||||
)
|
||||
|
||||
async def copy_screen(self, screen_id: str, owner_id: int, name: str | None) -> str:
|
||||
"""只复制草稿和封面，不复制发布状态和历史"""
|
||||
source = await self.get(screen_id, owner_id)
|
||||
base = source.name[: NAME_MAX_LENGTH - len(COPY_SUFFIX)]
|
||||
copied = await self.create(
|
||||
owner_id,
|
||||
name or f"{base}{COPY_SUFFIX}",
|
||||
copy.deepcopy(source.page),
|
||||
)
|
||||
if source.thumbnail is not None:
|
||||
thumbnail = self._thumbnails.copy(source.thumbnail)
|
||||
if thumbnail is not None:
|
||||
try:
|
||||
attached = await self._store.update(
|
||||
copied.screen_id, owner_id, {"thumbnail": thumbnail}
|
||||
)
|
||||
except BaseException:
|
||||
self._thumbnails.delete(thumbnail)
|
||||
raise
|
||||
if not attached:
|
||||
self._thumbnails.delete(thumbnail)
|
||||
return copied.screen_id
|
||||
|
||||
async def set_thumbnail(self, screen_id: str, owner_id: int, data: bytes) -> str:
|
||||
thumbnail = self._thumbnails.save(data)
|
||||
replaced: list[str | None] = []
|
||||
|
||||
def build(summary: ScreenSummary):
|
||||
# 封面只是展示用的，不算改内容：不动 revision 和 updated_at
|
||||
replaced[:] = [summary.thumbnail]
|
||||
return {"thumbnail": thumbnail}, {"thumbnail": summary.thumbnail}
|
||||
|
||||
try:
|
||||
await self._update_from_summary(screen_id, owner_id, build)
|
||||
except BaseException:
|
||||
self._thumbnails.delete(thumbnail)
|
||||
raise
|
||||
# 删的是这次真正换掉的那张；并发上传时各自换掉各自读到的，不会留下孤儿文件
|
||||
self._thumbnails.delete(replaced[0])
|
||||
return thumbnail
|
||||
|
||||
def open_thumbnail(self, thumbnail_id: str):
|
||||
return self._thumbnails.open(thumbnail_id)
|
||||
|
||||
# ---- 发布 ----
|
||||
|
||||
async def publish(
|
||||
self, screen_id: str, owner_id: int, revision: int
|
||||
) -> PublishInfo:
|
||||
record = await self.get(screen_id, owner_id)
|
||||
if record.revision != revision:
|
||||
raise BizException(409, "大屏有未保存或已在别处修改的内容，请先保存")
|
||||
previous = record.publish
|
||||
snapshot = PublishSnapshot(
|
||||
screen_id=screen_id,
|
||||
# 版本号取已有快照的最大值加一：写快照后没来得及更新指针留下的孤儿快照也不会撞号
|
||||
version=await self._store.last_publish_version(screen_id) + 1,
|
||||
revision=record.revision,
|
||||
page=record.page,
|
||||
published_at=_now(),
|
||||
published_by=owner_id,
|
||||
)
|
||||
if not await self._store.insert_publish(snapshot):
|
||||
raise BizException(409, "发布冲突，请重试")
|
||||
info = PublishInfo(
|
||||
status="published",
|
||||
version=snapshot.version,
|
||||
revision=snapshot.revision,
|
||||
# 第一次发布生成 share_id，之后沿用；下线再发布也沿用
|
||||
share_id=previous.share_id
|
||||
if previous
|
||||
else f"p_{secrets.token_urlsafe(16)}",
|
||||
published_at=snapshot.published_at,
|
||||
)
|
||||
# 条件：草稿没被再保存过（revision），线上指针还是读到的那个——并发发布只有一个能成，
|
||||
# 否则第一次发布时两边各生成一个 share_id，先返回给前端的那个会失效
|
||||
try:
|
||||
updated = await self._store.update(
|
||||
screen_id,
|
||||
owner_id,
|
||||
{"publish": info},
|
||||
revision=revision,
|
||||
where=_publish_is(previous),
|
||||
)
|
||||
except BaseException:
|
||||
await self._store.delete_publish(screen_id, snapshot.version)
|
||||
raise
|
||||
if updated is None:
|
||||
await self._store.delete_publish(screen_id, snapshot.version)
|
||||
raise BizException(409, STALE)
|
||||
return info
|
||||
|
||||
async def unpublish(self, screen_id: str, owner_id: int) -> None:
|
||||
def build(summary: ScreenSummary):
|
||||
if summary.publish is None or summary.publish.status == "offline":
|
||||
return None
|
||||
offline = summary.publish.model_copy(update={"status": "offline"})
|
||||
return {"publish": offline}, _publish_is(summary.publish)
|
||||
|
||||
await self._update_from_summary(screen_id, owner_id, build)
|
||||
|
||||
async def exists(self, screen_id: str, owner_id: int) -> bool:
|
||||
return await self._store.get_summary(screen_id, owner_id) is not None
|
||||
|
||||
async def list_publishes(
|
||||
self, screen_id: str, owner_id: int, limit: int, offset: int
|
||||
) -> tuple[list[PublishEntry], int]:
|
||||
if await self._store.get_summary(screen_id, owner_id) is None:
|
||||
raise _not_found()
|
||||
return await self._store.list_publishes(screen_id, limit, offset)
|
||||
|
||||
async def restore_publish(self, screen_id: str, owner_id: int, version: int) -> int:
|
||||
"""把历史版本的页面覆盖回草稿，不自动发布"""
|
||||
async with self._lock(screen_id):
|
||||
record = await self.get(screen_id, owner_id)
|
||||
snapshot = await self._store.get_publish(screen_id, version)
|
||||
if snapshot is None:
|
||||
raise BizException(404, "发布版本不存在")
|
||||
updated = await self._update_with_revision(
|
||||
screen_id,
|
||||
owner_id,
|
||||
{
|
||||
"page": snapshot.page,
|
||||
"revision": record.revision + 1,
|
||||
"updated_at": _now(),
|
||||
},
|
||||
record.revision,
|
||||
)
|
||||
await self._pin_files(owner_id, screen_id, snapshot.page)
|
||||
return updated.revision
|
||||
|
||||
# ---- 线上（匿名） ----
|
||||
|
||||
async def _live(self, share_id: str) -> tuple[ScreenSummary, PublishSnapshot]:
|
||||
"""线上版本：已下线 / 已删除 / 不存在都是 404；返回的是发布快照，不是草稿"""
|
||||
summary = await self._store.get_by_share_id(share_id)
|
||||
if (
|
||||
summary is None
|
||||
or summary.publish is None
|
||||
or summary.publish.status != "published"
|
||||
):
|
||||
raise _not_found()
|
||||
snapshot = await self._store.get_publish(
|
||||
summary.screen_id, summary.publish.version
|
||||
)
|
||||
if snapshot is None:
|
||||
raise _not_found()
|
||||
return summary, snapshot
|
||||
|
||||
async def get_live(self, share_id: str) -> tuple[ScreenSummary, PublishSnapshot]:
|
||||
return await self._live(share_id)
|
||||
|
||||
async def fetch_live_datasource(self, share_id: str, datasource_id: str):
|
||||
"""前端只传数据源 ID，请求地址取自已发布快照；不能做成"前端传 URL 后端转发"，否则是匿名开放代理"""
|
||||
_, snapshot = await self._live(share_id)
|
||||
source = next(
|
||||
(
|
||||
item
|
||||
for item in snapshot.page.get("dataSource", [])
|
||||
if isinstance(item, dict) and item.get("id") == datasource_id
|
||||
),
|
||||
None,
|
||||
)
|
||||
if source is None or source.get("type") != "api" or not source.get("url"):
|
||||
raise BizException(404, "数据源不存在")
|
||||
method = "POST" if source.get("method") == "POST" else "GET"
|
||||
params = source.get("params")
|
||||
headers = source.get("headers")
|
||||
return await self._datasource_cache.get(
|
||||
(share_id, snapshot.version, datasource_id),
|
||||
lambda: fetch_json(
|
||||
str(source["url"]),
|
||||
method,
|
||||
params if isinstance(params, dict) else None,
|
||||
headers if isinstance(headers, dict) else None,
|
||||
),
|
||||
)
|
||||
@@ -0,0 +1,63 @@
|
||||
"""大屏封面图：放本地磁盘，文件名就是封面 ID；每次上传换新 ID，旧图删掉
|
||||
|
||||
图片类 URL 要能被 <img src> 匿名访问，靠不可猜的 ID 兜底
|
||||
"""
|
||||
|
||||
import re
|
||||
import secrets
|
||||
from pathlib import Path
|
||||
|
||||
from app.core.exceptions import BizException
|
||||
|
||||
MAX_BYTES = 1024 * 1024
|
||||
# 按文件头判断类型，不信任客户端声明的 Content-Type
|
||||
MAGIC = {b"\x89PNG\r\n\x1a\n": "image/png", b"\xff\xd8\xff": "image/jpeg"}
|
||||
# ID 进了文件路径，读写前必须过这个格式检查，防路径穿越
|
||||
ID_PATTERN = re.compile(r"th_[A-Za-z0-9_-]{16,64}")
|
||||
|
||||
|
||||
def sniff(head: bytes) -> str | None:
|
||||
for magic, media_type in MAGIC.items():
|
||||
if head.startswith(magic):
|
||||
return media_type
|
||||
return None
|
||||
|
||||
|
||||
class ThumbnailStorage:
|
||||
def __init__(self, root: Path):
|
||||
self._root = root
|
||||
|
||||
def _path(self, thumbnail_id: str) -> Path | None:
|
||||
if not ID_PATTERN.fullmatch(thumbnail_id):
|
||||
return None
|
||||
return self._root / thumbnail_id
|
||||
|
||||
def save(self, data: bytes) -> str:
|
||||
if len(data) > MAX_BYTES:
|
||||
raise BizException(413, "封面图不能超过 1 MB")
|
||||
if sniff(data[:8]) is None:
|
||||
raise BizException(422, "封面图只支持 PNG 或 JPEG")
|
||||
thumbnail_id = f"th_{secrets.token_urlsafe(16)}"
|
||||
self._root.mkdir(parents=True, exist_ok=True)
|
||||
self._path(thumbnail_id).write_bytes(data)
|
||||
return thumbnail_id
|
||||
|
||||
def open(self, thumbnail_id: str) -> tuple[Path, str] | None:
|
||||
"""文件路径和图片类型；ID 不合法、文件不存在或内容不是图片时返回 None"""
|
||||
path = self._path(thumbnail_id)
|
||||
if path is None or not path.is_file():
|
||||
return None
|
||||
with path.open("rb") as file:
|
||||
media_type = sniff(file.read(8))
|
||||
return None if media_type is None else (path, media_type)
|
||||
|
||||
def copy(self, thumbnail_id: str) -> str | None:
|
||||
path = self._path(thumbnail_id)
|
||||
if path is None or not path.is_file():
|
||||
return None
|
||||
return self.save(path.read_bytes())
|
||||
|
||||
def delete(self, thumbnail_id: str | None) -> None:
|
||||
path = None if thumbnail_id is None else self._path(thumbnail_id)
|
||||
if path is not None:
|
||||
path.unlink(missing_ok=True)
|
||||
Reference in New Issue
Block a user