150 Commits
dev ... ft_wyt

Author SHA1 Message Date
wuyongtao
0b59c1ebee feat: 训练任务支持从最新 checkpoint 继续训练
- 后端: 新增 resume/checkpoints 端点,compute 客户端与适配器支持续训
- 计算节点: llama_factory 适配器断点续训支持及测试
- 前端: fine-tune API 封装与列表页续训/断点展示
2026-08-21 14:48:55 +08:00
wuyongtao
0233755859 feat: 评测任务完成进度收尾与指标维度汇总展示
- platform_store: 评测任务完成时置进度 100% 并落地 completed_time
- eval_runner: 新增指标维度汇总,供雷达图等维度展示
- 前端: 评测详情/列表展示优化、模型管理增强
2026-08-21 10:49:50 +08:00
caoxiaozhu
03254f8196 feat(data_process): 模型缓存统一仓库根 .cache 并修复 PDF 页眉页脚清理
- 新增 app/core/cache_paths.py:HF_HOME / tiktoken 缓存统一指向 <repo>/.cache,
  本地与 Docker 路径一致,离线部署打包 .cache 即可
- Dockerfile.backend 的 tiktoken 词表改用官方 SHA 文件名,避免运行时回退重建
- 修复 layout_hybrid 路径不运行 detect_pdf_document_noise 的缺陷:
  needs_pdf_noise 不再与 needs_layout_raw 互斥,PDF 智能预处理在版面切分下也生效
- 新增 layout_noise.py:识别跨页重复的页眉表格标签组并按行剔除,
  解决 docling layout 模型把中文企业 PDF 页眉识别成普通 Table 导致清不掉的问题
- 回收 HybridChunker 丢弃的末尾孤立标题,找回章节标题内容
2026-08-21 10:16:23 +08:00
wuyongtao
3b9361c237 merge: 合并远程 ft_wyt 分支,解决权限与日志模块冲突
- 冲突解决原则:本地权限治理(require_admin/current_user/资源ACL)与远程 op_log 日志装饰器双向保留
- platform.py: 9 处冲突,@op_log 与管理员校验叠加,避免远程丢失 require_admin 的安全回归
- logging.py: 合并 get_client_ip 与 user_id_var,X-Trace-Id 优先 + ContextVar 卫生处理
- op_log.py: 采纳远程将变量计算上移到函数顶部的结构
- compute_poller.py: 中文日志 + 失败去重限流/断连重置逻辑
- data_process.py: 保留租户归属字段 + biz_logger 成功日志
2026-08-21 09:58:43 +08:00
wuyongtao
6f0e82f351 feat: 平台治理与权限体系完善,存储进度/GPU预留/审批中心与日志整合
- 平台治理: 租户用户权限层次、资源ACL、审批中心与审批模板、访问申请
- 存储: MinIO 存储进度迁移、对象存储安全加固与测试
- 计算: GPU 资源预留、compute 轮询与同步增强
- 权限: permission v2 迁移、权限安全验收测试
- 日志: 后端运行日志中文说明、操作日志整合
- 数据处理/评测: 数据转换与模型评测优化

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-21 09:49:48 +08:00
wangjiming
29d8a506cd 更新日志的中文说明 2026-08-21 09:42:03 +08:00
wangjiming
8455d99d49 merge: 合并远程 ft_wyt 分支,解决冲突 2026-08-19 17:39:18 +08:00
wangjiming
78680a5c5c 新增后端日志,且更改了平台名字 2026-08-19 17:26:49 +08:00
wuyongtao
080ef6ab00 fix: 评测任务软删除健壮性增强,删除不存在任务返回 404
- delete_eval_task 用 RETURNING 校验记录存在性与重复删除,否则抛 KeyError
- eval_tasks/eval_task 列表与详情过滤已软删除记录(deleted_at IS NULL)
- model_eval_delete 捕获 KeyError 返回 404,避免接口崩溃

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-19 16:43:36 +08:00
wuyongtao
78e3baa9ba feat: 平台治理与对象存储增强,审批中心与运行日志整合
- 新增 storage/policy.py 落盘策略:按大小/类型决定文件存 MinIO 或内联数据库
- 数据处理源文件与生成结果写入 MinIO 并登记 storage_objects,支持失败回滚
- 算力节点训练产物按版本归档到 MinIO,登记 model_artifacts
- 数据转换任务输入输出对象化,支持从 MinIO 读写
- 新增审批中心(申请/我的/策略)、组织与权限、运行日志整合页面
- schema 与 docker 配置、前端路由侧边栏、治理文档同步更新

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-19 16:10:24 +08:00
caoxiaozhu
81c2f85c3a feat(data_process): 问答对数据评测体系与质量分雷达图
- 三层评测:规则层沿用原五维规则分,语义层用本地 BGE 向量算问答/来源
  相关性,评审层复用生成模型按 rubric 打分(忠实度/正确性/清晰度等,
  区分 standard/reasoning/dpo 输出类型),任一层失败自动降级
- 组合分 = 规则 35% + 语义 20% + 评审 45%,缺层自动重归一
- 新增 results/evaluate-batch 批量评测接口,镜像批量重生成的并发、
  乐观锁与部分成功语义;生成阶段不再展示质量分
- 详情页与结果编辑页新增"数据评测"按钮和批量进度;质量分列悬停弹出
  雷达图浮窗(评审 5 维 + 语义 2 维、三层分项、评审理由)
- 手动编辑/恢复后重算规则与语义层并丢弃过期评审分,雷达图不再展示
  失效数据
2026-08-19 14:22:27 +08:00
caoxiaozhu
f47611020a fix(data_process): 修复 Word 切分行号断档与预览标题缺失
- 正文抽取下钻 SDT 内容控件,目录等内容不再整段丢失
- 切片投影匹配剥离序列化插入的列表自动编号,新增行锚点兜底与游标防回退
- fixed/semantic 切分路径定位失败时保留切片,不再静默丢弃
- Word 预览按段落大纲级别识别标题,未套标题样式的小节正常渲染
- 本地嵌入模型抽为共享单例,供语义分块与质量评分共用
2026-08-19 14:22:27 +08:00
wuyongtao
93449b7f0e fix(frontend): 修复操作日志与数据转换类型错误,修正普通用户角色取值
- OperationLogView: showDetail 调用补充 OperationLog 类型断言
- DataConvertView: beforeUpload 参数类型改为 UploadRawFile 并直接读取 file.type
- UserCreateView: 普通用户角色值 user 修正为 operator

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-19 11:33:45 +08:00
wuyongtao
9379f93633 feat: 推理与评测支持多卡 GPU 选择,训练任务实时 GPU 监控
- 新增 GPU 选择归一化 helper,统一前端各形态的选择(gpu_indices/gpus/gpu_id)
- 评测与推理支持同一节点内多卡选择,校验所选 GPU 空闲后再派发
- 新增 /fine-tune/{id}/gpu-status 接口,训练日志页展示实时 GPU 指标
- 算力节点推理加载支持 CUDA_VISIBLE_DEVICES 多卡可见,nvidia-smi 进程级监控
- GPU 占用跟踪细化为按卡记录,覆盖评测任务、推理模型与对比任务

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-19 11:12:17 +08:00
wangjiming
c3e96ae61b Merge branch 'ft_wyt' of http://www.caoxiaozhu.com:13001/YG-Soft/YG_FT into ft_wyt 2026-08-19 10:33:06 +08:00
wangjiming
b254fa0985 修复普通用户对模型训练的权限操作 2026-08-19 10:32:58 +08:00
caoxiaozhu
f3625f994c fix(frontend): 表格页画布允许收缩,修复分页条被裁掉
内容自然高度超出视口时,画布 flex 不收缩导致分页条被裁掉且无滚动
条可及;为 .data-table-page 的画布补充 flex: 1 1 auto 与 min-height: 0。
2026-08-18 15:49:00 +08:00
caoxiaozhu
78fb894307 feat(data_process): 优化问题生成提示语,提升问题自然度
- 系统提示语注入问题风格规则:像真实用户自然提问、避免"请描述/
  请说明/根据文档"等模板化开头、多条问题交替句式,并附正反示例;
  服务端注入对存量任务即时生效。
- 未配置提示语时的后端兜底从一句话充实为与前端同信息量的完整默认。
- 前端标准/思维链/DPO 三套内置默认提示语新增"问题表述自然"要求,
  旧版归档为 _4 常量并注册进 isBuiltInGenerationPrompt 迁移映射,
  自定义提示语不受影响。
- 新增兜底提示语专项测试与风格规则断言,生成测试 32/32 通过。
2026-08-18 15:49:00 +08:00
caoxiaozhu
91b4ae2287 feat(data_process): 显式关闭 docling OCR 并拒绝无文本层 PDF
- docling 转换器通过 PdfPipelineOptions 显式设置 do_ocr=False,
  混合 PDF 的图片页不再产出 OCR 文字;不提供重新开启 OCR 的参数。
- 无文本层 PDF 的错误文案改为"扫描版或图片型 PDF 不支持",
  上传阶段整批拒绝,保留混合 PDF 的可处理判定。
- 新增 test_layout_converter_disables_ocr 守护开关状态,
  同步设计文档与 disable-ocr 实施计划/设计说明。
2026-08-18 15:49:00 +08:00
caoxiaozhu
2f48934e66 fix(data_process): 修复 tiktoken 加载器的变量作用域错误
函数内部的 import tiktoken.core 会把 tiktoken 变成局部变量,导致
tiktoken.get_encoding 在联网下载前就抛 UnboundLocalError,被 except
吞掉后联网与本地缓存两条路径全部失效。移除该内部 import(模块顶部
已引入 tiktoken),切分测试全部通过,编码器成功下载并缓存到
~/.cache/tiktoken,此后离线环境也可正常加载。
2026-08-18 15:49:00 +08:00
wangjiming
fed5694796 Merge branch 'ft_wyt' of http://www.caoxiaozhu.com:13001/YG-Soft/YG_FT into ft_wyt 2026-08-18 14:49:34 +08:00
wangjiming
04c3c1412c 修改普通用户的数据类型转换在数据集管理看不见的问题 2026-08-18 14:49:12 +08:00
wuyongtao
2adf78a6ed feat: 数据集下载鉴权、MinIO 发布与数据库兼容增强
- 数据集下载接口鉴权:单文件直接返回、多文件打包 ZIP,前端统一走请求客户端携带 Token
- 预检/同步支持 MinIO 对象补建与资源副本记录,兼容接入 MinIO 前的历史数据集
- create_dataset 增加名称重复校验,软删记录释放名称并保留墓碑
- 数据处理发布结果支持写入 MinIO storage_objects 并关联 dataset_file
- 数据存储根目录支持 YG_FT_DATA_ROOT 环境变量
- SQL 迁移兼容旧版 approval_steps / retention_policies / data_process_results
- 训练日志图表按曲线独立过滤缺失采样点

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-18 14:45:16 +08:00
wangjiming
4e27b98a84 修改用户设置的新增用户的权限点击操作 2026-08-17 16:04:04 +08:00
wangjiming
6c1bf61ff7 更新平台治理 2026-08-17 09:15:36 +08:00
wuyongtao
de9c8e4ffe docs: 精简 README 与 docker 部署说明,归档 20260812 文档快照
- 重写根 README 为架构、端口、启动速览,替代原详细开发指南
- 重写 docker/README.md 聚焦跨服务器部署边界与数据安全
- 新增平台测试用例文档,归档 docs/20260812 快照副本

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-12 15:38:34 +08:00
wuyongtao
5ecca9f0bc feat: 权限与日志治理完善,MinIO 独立部署与 tiktoken 离线打包适配
- 后端:强化平台/审批/资源/系统接口权限校验与操作日志,更新权限设计文档与测试用例
- 存储:新增 MinIO 独立部署适配(端口 19000/19001),外部端点与 host-gateway 互通
- 离线:打包 tiktoken cl100k_base 词表进镜像,避免无网环境联网下载
- 其他:算力节点接口微调,前端微调创建页小修,忽略 MinIO 运行时数据

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-12 15:21:42 +08:00
caoxiaozhu
2f64086177 docs(data_process): 在 __init__ 补充模块职责速查与分层说明 2026-08-12 14:46:56 +08:00
caoxiaozhu
bdaf72d58b chore: 删除误存的 document_chunking.txt(macOS 压缩的 zip 残留,代码未引用) 2026-08-12 14:39:11 +08:00
caoxiaozhu
ca012893f7 refactor: 完整重构 data_process 模块并修复拆分遗留缺陷
将 algorithms.py / store.py 拆分为 algorithms/ 与 store/ 子包,并修复
机械拆分造成的导入与辅助函数缺失:

- algorithms/: 补全各子模块依赖与 17 个私有辅助函数、8 个常量;重写
  __init__.py 移除坏的 importlib 兜底,分层导入并以局部 import 断开
  text_utils<->parsers、quality<->structured_processing 循环依赖。
- store/: 补回 DataProcessStoreError / hashlib / _serialize_value /
  estimate_token_count 等缺失导入,包入口导出测试与调用方依赖的私有
  辅助函数。
- 删除旧单文件 algorithms.py / store.py 及重构残留(_algorithms_old、
  backups、refactor 脚本、REFACTORING 文档)。

algorithms 与 store 测试套件 91 项全部通过。
2026-08-12 14:29:39 +08:00
wuyongtao
ae39c45a73 feat: 前端离线部署适配与在线模型测试延迟展示
- 移除 index.html 中 Google Fonts 外链,适配离线部署环境
- 侧边栏用户头像改为显示姓名首字母,去掉 picsum 占位图外链
- 在线模型测试补充 latency_ms 延迟字段展示
- api_key 为空时传空字符串避免校验报错

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-11 16:34:06 +08:00
wuyongtao
b5d2cd7935 feat: 新增 MinIO 对象存储与算力节点缓存预下载
- 后端新增 storage 模块(minio_store),支持 MinIO 预签名 URL 上传与对象管理
- config 新增 MinIO 及存储等待相关配置项
- 算力节点新增 /compute/cache/prepare 缓存预下载接口(带校验和原子落盘)
- 算力节点健康接口增加存储可用性探针
- SQL 迁移补充资源存储相关表结构
- Docker 新增 minio 服务与后端 minio 配置
- 补充 minio-compute-cache-plan 设计文档

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-11 16:25:18 +08:00
wangjiming
a12f80492d Merge branch 'ft_wyt' of http://www.caoxiaozhu.com:13001/YG-Soft/YG_FT into ft_wyt 2026-08-11 15:45:23 +08:00
wangjiming
71405def14 新增配置模型测试按钮 2026-08-11 15:41:51 +08:00
caoxiaozhu
5f6e7523cf feat: 新增外部数据源拉取与 DPO 输出格式支持
- 支持从 PostgreSQL 数据库拉取结构化数据作为训练来源
- 新增 DPO (Direct Preference Optimization) 输出类型
- 支持 chosen/rejected 字段的编辑、校验和发布
- 完善数据预处理切分逻辑和元数据管理
- 移除 OCR 扫描 PDF 功能,保持基础文本解析能力

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-11 14:17:45 +08:00
wangjiming
f809825a7d 完善部分平台治理功能,及修改看板缺陷 2026-08-10 11:41:16 +08:00
wuyongtao
75cc105ebc chore: 忽略离线部署包,提交安全加固、数据库初始化与文档
- .gitignore: 忽略 docker/offline 离线部署包(镜像/运行时等大文件)
- 安全加固: 新增 compute/api/security.py 及各端安全测试,补充 docs/security-hardening.md
- 数据库: 新增完整初始化 SQL 与 docs/database-config.md
- 数据转换与评测: 修复类型检查、增强校验并补充测试
- Docker 配置与环境变量更新

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-07 09:24:35 +08:00
wuyongtao
e397bcc2ca fix: 修复数据转换页面类型检查错误
- statusTag 返回类型改用 TagProps['type'],修复 el-tag type 属性 TS 报错
- handleDelete 参数加 DataConvertTask 类型断言,修复 DefaultRow 类型不匹配

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-05 16:47:53 +08:00
wangjiming
c64fa1cd61 新增数据类型转换 2026-08-05 16:23:00 +08:00
wangjiming
b5c6557341 Merge branch 'ft_wyt' of http://www.caoxiaozhu.com:13001/YG-Soft/YG_FT into ft_wyt 2026-08-05 08:34:55 +08:00
wuyongtao
ec7d8c0a3d fix: 算力节点任务/GPU 统计纳入评测与推理占用,补充测试依赖
- compute_nodes() 的 current_running_jobs 纳入运行中的评测任务(eval_tasks)
  与已加载的推理模型(compare_tasks 持久化状态 + 内存标记兜底),
  多节点时 GPU 被评测/推理占用不再显示 0/1
- gpus() 直接按 eval_tasks / compare_tasks 派生 GPU busy/reserved 状态,
  修复 gpu_id=0 时 or -1 导致匹配失败;删除评测后 GPU 不再残留 busy
- 评测不再复用 mark_inference_loaded 内存标记,GPU 占用由任务数据驱动
- requirements.txt 补充 pytest / ruff(此前仅声明在 pyproject dev 可选依赖)

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-04 18:46:53 +08:00
wuyongtao
0292bf5138 fix: 模型评测异步加载等待与多节点路由修复
- eval_runner 等待异步模型加载完成(InferenceSession.wait_until_loaded),
  修复 "model load failed: unknown"
- 评测算力节点选择:优先页面选择的节点 / 模型所在节点(_select_eval_node),
  多节点时不再派发到不可达节点导致连接超时
- 前端评测 GPU 选择改为节点感知(节点:GPU 复合值),透传 compute_node_id,
  并检查 startEval 结果展示真实错误
- 大模型评价(judge)使用模型记录的真实 API 模型名(api_model),
  避免用平台内部名调用 LLM API 导致 HTTP 400
- 新增后端节点选择与 compute wait_until_loaded 单元测试

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-04 18:21:16 +08:00
wangjiming
7f93ed6d09 feat: 更新 request.ts 2026-08-04 17:18:50 +08:00
wangjiming
5ec950cc9f feat: 更新 request.ts 2026-08-04 17:02:57 +08:00
wuyongtao
0271942ba5 feat: 模型推理异步加载与对话链路修复,同步基线
模型推理全异步化改造:
- 计算节点 InferenceSession 改为后台线程异步加载模型,load 立即返回,
  加载期间事件循环保持响应(/inference/status 与 /health 不阻塞)
- 后端模型加载改为异步派发 + 轮询对账器(reconcile_inference_loads),
  任务状态由 starting 自动推进到 ready/error,解决多节点启动超时
  (timeout of 120000ms exceeded)
- 推理删除/卸载改为任务感知 + 短超时,删除先删记录再 best-effort 卸载,
  不再被不可达节点阻塞;同节点新模型替换旧任务标记失效
- 流式对话透传 task_id/node_id 路由到真正加载模型的算力节点,
  useStreamChat 解析 SSE 错误帧以干净文案展示
- 对话历史按任务 id 本地持久化,退出重进可恢复;移除页脚提示文本
- 新增后端推理异步加载与计算节点异步状态机单元测试

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-04 16:59:34 +08:00
wuyongtao
250e060271 feat: 看板服务状态优化,移除前端访问统计埋点,修正列表页文案与类型
- 看板:补齐各服务模块图标,服务状态表格改为自适应行高并支持滚动
- 移除前端请求访问统计埋点及对应 audit-visit API 模块
- 平台接口:清理 service_checks 循环中未使用的路径变量
- 系统模块:导出 SystemUser 类型
- 项目/租户列表:统一“项目名称/编码 ID/租户 ID”文案,补充表格行类型断言修复

Co-Authored-By: Claude <noreply@anthropic.com>
2026-08-04 11:34:04 +08:00
wuyongtao
7b36bc774e Merge branch 'ft_wyt' of http://www.caoxiaozhu.com:13001/YG-Soft/YG_FT into ft_wyt 2026-08-03 17:34:23 +08:00
wuyongtao
4e5c43fad5 feat: 评测任务元数据增强与前端展示优化
后端:
- 新增评测指标标签生成(_build_eval_metric_label/_basic_metric_labels)
- eval 任务 payload 富化(_enrich_eval_payload): 解析模型名/数据集名/维度配置, 生成 metric_label
- 任务创建时解析 trained_models/模型名, metric 回退为 metric_label
- 运行时状态刷新增加 2s 节流, 避免频繁触发

前端:
- EvalTask 类型增加 dataset_id/metric_label
- 评测列表与详情展示 metric_label 指标标签, 模型名/指标超长 tooltip 省略
- 创建页改用 getComputeGpus 过滤空闲 GPU, 移除 getSystemInfo 依赖
2026-08-03 17:34:21 +08:00
wangjiming
62a1d03eac 更新前端看板1 2026-08-03 17:24:45 +08:00
wangjiming
94230cad16 更新前端看板 2026-08-03 16:33:08 +08:00
wangjiming
0c601934a0 更新前端看板 2026-08-03 16:20:21 +08:00
wuyongtao
5cc306eb0a fix: 推理/评测结果同步、数据集统计与算力节点管理增强
后端:
- 抽取 fetch_eval_result_content 复用函数,model_eval_detail 直接应用评测任务结果
- health 接口移除数据库依赖,返回静态指标
- 数据集: count_dataset_records JSON 感知计数; 文件统计改为从 dataset_files 聚合重算; 在线编辑记录 size/record_count/version_no; 上传同步批处理
- 算力节点: 调度支持 requested GPU 子集校验与容量计算; 新增 delete_compute_node(含活动任务保护)及 DELETE 接口; 连接池 connect_timeout
- 评测任务落库 basic_metrics/score/completed_time, failed/stopped 记录 error

评测引擎:
- _load_dataset 支持 JSON/JSONL 文件
- 新增 exact match 与文本相似度指标, 余弦相似度去掉 2 样本限制

前端:
- 算力节点列表「维护」改为「删除」(带确认弹窗), compute.ts 新增 deleteComputeNode
- 数据集上传超时调整为 120s; FineTuneTask 增加 compute_node_id; GpuInfo 状态增加 reserved
2026-08-03 15:49:21 +08:00
wuyongtao
cc08b164d0 fix: 前端表格行类型断言修复并补充 psycopg-pool 依赖
- 审批/项目/租户/用户设置等视图新增 asXxx 类型断言辅助函数
- search-fields 由字符串改为数组传参
- backend 依赖新增 psycopg-pool,Dockerfile 依赖校验同步更新
2026-08-03 11:02:41 +08:00
wuyongtao
24c77a990a Merge branch 'ft_wyt' of http://www.caoxiaozhu.com:13001/YG-Soft/YG_FT into ft_wyt
# Conflicts:
#	backend/app/api/v1/endpoints/platform.py
#	compute/requirements.txt
2026-08-03 09:42:49 +08:00
wangjiming
15c4223f2c update 2026-08-03 09:34:08 +08:00
caoxiaozhu
b975de02da fix: 完善数据预处理与 JSON 上传链路 2026-07-30 16:54:00 +08:00
wuyongtao
46d343fb63 chore: compute requirements.txt 新增评测依赖 sacrebleu / rouge-score / scikit-learn
Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-29 14:27:48 +08:00
wuyongtao
0c39f2f5b9 feat: 模型评测端到端闭环 — EvalRunner引擎 + 算力节点Job执行 + 结果回写
算力节点 (compute):
- 新建 eval_runner.py: 评测执行引擎,作为subprocess运行
  - 加载模型 + JSONL数据集 + 逐样本推理
  - BLEU/ROUGE/Cosine基础指标计算
  - LLM Judge评分(OpenAI兼容API调用)
  - 结果写入eval_results.json
- adapter.py: build_command新增engine=eval分支
- main.py: 新增/json模块导入,新增/compute/files/read端点,eval job校验

后端:
- platform.py: 重写startEval提交eval job到算力节点
  - 支持models表和trained_models表查找
  - 已合并模型不传adapter路径
- platform_store.py: 新增update_eval_task/running_eval_tasks/apply_eval_job_result
- sync.py: poller新增eval job同步,异步读取eval_results.json回写结果

前端:
- EvalCreateView/DimensionCreateView: eval模型过滤扩展(API类型+api_url)
- EvalCreateView: GPU过滤在线节点空闲GPU
- EvalTaskSetupStep: GPU value从数组index改为gpu.id
- BasicMetricSetupStep: ROUGE方法名修正(rouge_1→rouge1)
- EvalView: 新增5秒轮询刷新

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-28 19:34:41 +08:00
wuyongtao
c7c9ed925b feat: 模型推理端到端闭环 — 真实流式推理 + 释放/删除 + GPU 状态同步
后端 (platform.py + platform_store.py):
- 新增 _build_messages_payload() 转换前端格式为 OpenAI messages
- 新增 _stream_chat_proxy() SSE 流式代理到算力节点
- 新增 _unload_from_compute_node() 真正释放算力节点 GPU 显存
- 重写 model_compare_load: 从假 PID/端口改为真正调用算力节点加载模型
- 修复 model_compare_unload: 调用 _unload_from_compute_node 释放 GPU
- 修复 model_compare_delete: 先释放 GPU 再删除记录
- 修复 model_compare_stream_chat: 从 mock 改为 StreamingResponse 代理
- 修复 model_chat_local/stream: 消息格式转换 + 路径修正
- PlatformStore 新增 _inference_nodes 追踪,gpus() 同步推理占用状态
- preload/unload 端点标记/清除推理节点占用

算力节点 (compute):
- inference.py: 适配新版 LLaMA-Factory API (get_infer_args 4 返回值、ChatModel args dict、stream_chat 新签名)
- inference.py: unload() 增加 gc.collect + torch.cuda.empty_cache + synchronize 彻底释放显存
- main.py: inference/load 移除 HTTPException(500),错误以 200 正常返回

前端:
- InferenceChatView: 真实模式下走 SSE 流式推理,mock 模式保留兼容
- InferenceCreateView: 调用 preloadLocalModel + createCompare 真实创建推理任务,失败回退 mock
- InferenceListView: 「停止」改为「释放」,删除前先释放算力节点,改进错误提示
- compare.ts: 新增 streamChatReal() fetch SSE,preload 超时提升至 5 分钟
- useStreamChat.ts: send() 支持 useMock 参数,真实模式调用 streamChatReal
- GPU 选择过滤: 仅显示在线算力节点上的空闲 GPU

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-28 17:29:16 +08:00
wuyongtao
f917a025e1 feat: 基于 LLaMA-Factory 实现模型推理引擎
利用容器内已有的 LLaMA-Factory ChatModel (huggingface 后端) 实现真实
模型推理,无需额外安装 vLLM。

Compute 端新增:
- compute/engines/llama_factory/inference.py
  InferenceSession: 模型加载/卸载/对话/流式输出
  支持 base model 和 LoRA adapter,线程安全
- compute/api/main.py 新增 5 个推理端点:
  POST /inference/load      - 加载模型
  POST /inference/unload    - 卸载释放 GPU 显存
  GET  /inference/status    - 查询会话状态
  POST /inference/chat      - 非流式对话
  POST /inference/chat/stream - SSE 流式对话

后端新增:
- platform.py 推理代理端点(local/chat, local/chat/stream,
  local/preload, local/unload, local/status, trained/preload)
- ComputeNodeClient._request 通用请求方法
- _select_first_online_node 自动选择在线节点

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-28 13:49:10 +08:00
wuyongtao
a9ab130d43 feat: P0 训练闭环核心功能实现
P0-1 模型路径治理:
- 新增 003_model_path_governance.sql 迁移,models 表增加 can_train 字段
- create_model/update_model 自动计算 can_train(非API+有路径=可训练)
- _compute_job_payload_from_task_node 拒绝 API 模型和无可训练路径模型
- 平台诊断规则增加 API 模型/路径缺失检测

P0-2 数据集格式校验:
- 新增 dataset_format.py,支持 Alpaca/ShareGPT/DPO/CPT 格式校验
- 训练预检时自动根据 train_type 匹配格式并校验内容字段
- llama_dataset_info 增加 DPO/CPT 格式列映射

P0-3 训练完成产物入库:
- _ensure_trained_model 使用 compute 节点返回的真实 artifacts
- 注册 per-file artifact 记录(含 size_bytes/checksum_sha256)
- trained_models 表增加 artifact_dir 字段

P0-4 失败日志拉取:
- poll_compute_jobs_once 检测到 failed/stopped 时强制拉取最后 200 行日志
- apply_compute_job 持久化失败日志片段到任务 payload

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-28 13:10:53 +08:00
wuyongtao
525fc55cef fix: 修复 GPU 选择索引错误及 CUDA 不可用问题,优化 GPU 硬件单选
- 修复 FineTuneCreateView GPU 选择使用 v-for idx 替代真实 gpu.id 导致
  多节点环境下 GPU 索引错误(如 gpu-node-02 仅 GPU 0 但请求 GPU 1)
- GPU 选择改为单选模式,已离线节点自动过滤不展示
- GPU 卡片增加节点编号展示
- 修复 CUDA_VISIBLE_DEVICES=all 无效值导致 torch.cuda.is_available() False
  改为 CUDA_VISIBLE_DEVICES=0
- .gitignore 新增 .claude/ CLAUDE.md 排除规则
- GpuInfo 类型增加 node_id/node_code/node_name 字段

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-28 12:43:20 +08:00
wuyongtao
3fd9cf9100 fix: 移除 frontend/.gitignore 中对 dist 的例外规则
之前 frontend/.gitignore 中 !dist/ 和 !dist/** 规则导致构建产物
仍然被 Git 追踪,现已删除,与根 .gitignore 保持一致。

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-28 12:10:54 +08:00
wuyongtao
a72e2a2520 chore: 将 frontend/dist 从版本控制中移除,补充 ElMessageBox 导入
- 移除 .gitignore 中对 frontend/dist/ 的例外规则,构建产物不再纳入版本控制
- DataProcessDetailView.vue 补充 ElMessageBox 导入修复 TS2552 错误

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-28 12:07:57 +08:00
caoxiaozhu
01d2e6c76a feat(data-process): 完善后台生成与失败重试 2026-07-28 10:56:05 +08:00
caoxiaozhu
8a6a6574bb feat(data-process): 完善 Word 与 Excel 原文件预览 2026-07-27 16:24:53 +08:00
caoxiaozhu
9025437a37 feat(data-process): 增加 Office 原文件预览接口 2026-07-27 16:12:50 +08:00
caoxiaozhu
4623e3fa1c fix(data-process): 完善默认生成提示语 2026-07-27 14:58:02 +08:00
caoxiaozhu
f97245b814 feat(data-process): 增加可配置思维链生成 2026-07-27 14:41:38 +08:00
caoxiaozhu
97cdb5cc68 fix(data-process): 收准结构化预处理能力文案 2026-07-27 14:36:16 +08:00
caoxiaozhu
d8a11e4949 fix(data-process): 对齐输出类型控件宽度 2026-07-27 13:58:15 +08:00
caoxiaozhu
05c1a5c1e1 fix(data-process): 优化输出类型选择样式 2026-07-27 13:49:21 +08:00
caoxiaozhu
f3fa5f1a68 fix(data-process): 调整输出类型选项位置 2026-07-27 13:43:30 +08:00
caoxiaozhu
de2e8952b5 feat(data-process): 支持思维链输出类型 2026-07-27 13:08:44 +08:00
caoxiaozhu
ecafb7eb13 fix(dataset): 展示训练任务名称 2026-07-27 12:44:40 +08:00
caoxiaozhu
b82897ca3a fix(dataset): 限制长文本提示层尺寸 2026-07-27 12:35:09 +08:00
caoxiaozhu
e6bbb0bb49 feat(dataset): 补全样本原文与删除操作 2026-07-27 12:26:19 +08:00
caoxiaozhu
bce586697b fix(dataset): 统一大小与版本元数据 2026-07-27 12:26:09 +08:00
caoxiaozhu
e486d36a80 feat(dataset): 支持数据任务批量删除 2026-07-27 11:36:30 +08:00
caoxiaozhu
d6e325fe9e fix(data-process): 按发布状态控制结果编辑 2026-07-27 11:23:17 +08:00
caoxiaozhu
b08a771a61 fix(data-process): 保留可验证的任务耗时 2026-07-27 11:16:48 +08:00
caoxiaozhu
8caaaa5bbc fix(data-process): 恢复提前中断的重新生成任务 2026-07-27 11:07:18 +08:00
caoxiaozhu
3f5fedb9ed build(frontend): 更新重新生成流程构建 2026-07-27 10:43:53 +08:00
caoxiaozhu
53844a3a09 fix(frontend): 保持重新生成退出前详情 2026-07-27 10:43:48 +08:00
caoxiaozhu
53014bb381 fix(data-process): 延迟重新生成破坏性变更 2026-07-27 10:43:42 +08:00
caoxiaozhu
03bd0b6d03 build(frontend): 更新数据任务详情构建 2026-07-27 10:03:58 +08:00
caoxiaozhu
b4927a8952 fix(frontend): 对齐数据任务详情状态 2026-07-27 10:03:52 +08:00
caoxiaozhu
b14b2ecf22 fix(data-process): 修正任务详情数据契约 2026-07-27 10:03:46 +08:00
caoxiaozhu
42c0e4f5c2 build(frontend): 更新数据任务列表构建 2026-07-27 09:50:33 +08:00
caoxiaozhu
88a82ed771 feat(frontend): 完善数据任务数量与状态 2026-07-27 09:50:26 +08:00
caoxiaozhu
895983ac20 feat(data-process): 返回任务文档数量 2026-07-27 09:50:19 +08:00
caoxiaozhu
680fa905f8 build(frontend): 更新数据集保留提示构建 2026-07-27 09:39:29 +08:00
caoxiaozhu
ce0f908d20 fix(frontend): 标明重新生成前的数据集状态 2026-07-27 09:39:23 +08:00
caoxiaozhu
e4ea1f168c fix(data-process): 保留重新生成前的已发布数据集 2026-07-27 09:39:19 +08:00
caoxiaozhu
915f994c45 build(frontend): 更新五十条生成上限构建 2026-07-27 09:12:03 +08:00
caoxiaozhu
762f866175 feat(data-process): 放宽问答生成数量上限 2026-07-27 09:11:58 +08:00
caoxiaozhu
9428c6b785 feat(data-process): 支持单项生成五十条数据 2026-07-27 09:11:51 +08:00
caoxiaozhu
25d75f40c7 build(frontend): 更新重新生成流程构建 2026-07-25 22:41:16 +08:00
caoxiaozhu
06e7455630 feat(data-process): 接入重新生成配置流程 2026-07-25 22:41:06 +08:00
caoxiaozhu
396d3f6f47 feat(data-process): 支持任务重新生成 2026-07-25 22:40:55 +08:00
caoxiaozhu
64d7414b04 build(frontend): 更新跳转图标对齐构建 2026-07-25 22:00:23 +08:00
caoxiaozhu
9a5282f39c fix(data-process): 对齐数据集跳转图标 2026-07-25 22:00:23 +08:00
caoxiaozhu
f21a4c954f build(frontend): 更新数据集名称对齐构建 2026-07-25 21:59:02 +08:00
caoxiaozhu
2e2cbb4976 fix(data-process): 对齐数据集名称起点 2026-07-25 21:59:02 +08:00
caoxiaozhu
2e8278636b build(frontend): 更新数据集对齐修复构建 2026-07-25 21:55:24 +08:00
caoxiaozhu
bbc0df29bf fix(data-process): 对齐输出数据集列表 2026-07-25 21:55:16 +08:00
caoxiaozhu
b20e7aa595 build(frontend): 更新结果提示修复构建 2026-07-25 18:25:58 +08:00
caoxiaozhu
17615aa17d fix(data-process): 限制结果提示浮层尺寸 2026-07-25 18:25:34 +08:00
caoxiaozhu
07e2999323 fix(data-process): 修复三数据集发布参数错误 2026-07-25 18:19:37 +08:00
caoxiaozhu
749c84a62b build(frontend): 更新会话续期修复构建 2026-07-25 18:16:02 +08:00
caoxiaozhu
ea0013b99c fix(auth): 长任务活跃期间续期会话 2026-07-25 18:15:49 +08:00
caoxiaozhu
4f8aff5fc4 build(frontend): 更新文档切分页面构建 2026-07-25 18:00:55 +08:00
caoxiaozhu
9193f10e3e feat(frontend): 更新三种切分模式与PDF定位 2026-07-25 18:00:44 +08:00
caoxiaozhu
ea08478a37 feat(data-process): 接入三种文档切分引擎 2026-07-25 18:00:21 +08:00
caoxiaozhu
4782981169 build(frontend): 更新三数据集发布页面构建 2026-07-25 17:04:27 +08:00
caoxiaozhu
939a7f8e8f fix(frontend): 展示三个独立发布数据集 2026-07-25 17:04:21 +08:00
caoxiaozhu
e9a121cfeb fix(data-process): 发布三个独立切分数据集 2026-07-25 17:04:14 +08:00
caoxiaozhu
d4b9a76aa5 build(frontend): 更新数据切分页面构建 2026-07-24 21:11:06 +08:00
caoxiaozhu
9114f3d4c7 fix(dataset): 展示并切换三路数据切分 2026-07-24 21:11:01 +08:00
caoxiaozhu
9cb77c251a fix(data-process): 发布精确三路数据切分 2026-07-24 20:43:47 +08:00
caoxiaozhu
e6a5a36bc0 build(frontend): 更新数据处理详情构建 2026-07-24 16:31:38 +08:00
caoxiaozhu
b2c570f607 fix(data-process): 补齐原文参照与详情统计 2026-07-24 16:30:29 +08:00
caoxiaozhu
994ec6644a fix(data-process): 修正详情统计字段契约 2026-07-24 16:28:47 +08:00
caoxiaozhu
d6d3d27b2d chore: 恢复 .gitignore 对 .env 文件的忽略规则
上次合并 rebase 时远端引入的 !.env 规则会强制追踪环境文件,
存在将含敏感信息的 .env 误提交的风险。恢复为 .env 忽略规则,
同时保留 !.env.example 以便示例文件继续入库。
2026-07-24 16:13:04 +08:00
caoxiaozhu
eb6ff93150 build(frontend): 更新智能预处理生产构建 2026-07-24 15:10:48 +08:00
caoxiaozhu
215b4074e0 build(frontend): 同步PDF预览依赖锁文件 2026-07-24 15:06:37 +08:00
caoxiaozhu
4544483fc5 feat(data-process): 启用智能文档清理预览 2026-07-24 15:06:01 +08:00
caoxiaozhu
3266a6fc09 feat(data-process): 清理PDF文档级噪声 2026-07-24 15:05:39 +08:00
caoxiaozhu
a9b06140d0 docs: 补充评估工作台设计文档
从根目录评测.txt 迁入 docs/,归档评估流程、错误诊断、
报告生成与建议有效性验证的设计说明。
2026-07-24 14:36:11 +08:00
caoxiaozhu
663b73af2e feat(frontend): 新增登录页 hero 流程图资源 2026-07-24 14:36:11 +08:00
caoxiaozhu
4a2f1f5dcd build(frontend): 纳管 pnpm 锁文件与 workspace 配置
pnpm-lock.yaml 锁定依赖版本以保证可复现安装;
pnpm-workspace.yaml 声明构建脚本白名单配置。
2026-07-24 14:36:11 +08:00
caoxiaozhu
8ac39cf007 chore: 移除已废弃的 .env.example 示例文件
配置入口已迁移至 backend/config.yaml(本地,已 gitignore)
与各环境实际的 .env 文件,示例文件不再需要维护。
2026-07-24 14:36:11 +08:00
caoxiaozhu
faad88dfcd chore: 忽略本地敏感配置与工具产物
- backend/config.yaml 含数据库账号密码等敏感信息
- .codex-backups/、.pnpm-store/、.zcode/ 为工具产物,不应进版本库
2026-07-24 14:36:11 +08:00
caoxiaozhu
d3a25f3a4b build(frontend): 更新切分缓存版本构建 2026-07-24 14:36:11 +08:00
caoxiaozhu
b801bd314b fix(data-process): 升级预览切分缓存版本 2026-07-24 14:36:11 +08:00
caoxiaozhu
476502fc0d build(frontend): 更新数据预览生产构建 2026-07-24 14:36:11 +08:00
caoxiaozhu
975f55d06c fix(data-process): 优化预览分页与PDF兼容 2026-07-24 14:36:11 +08:00
caoxiaozhu
93373bc61f fix(data-process): 合并文档结构短切片 2026-07-24 14:36:11 +08:00
caoxiaozhu
b79a8e1499 build(frontend): 更新生产构建产物 2026-07-24 14:36:11 +08:00
caoxiaozhu
0124e28d77 docs(data-process): 更新文件处理与切分设计 2026-07-24 14:35:03 +08:00
caoxiaozhu
ad64e44860 feat(data-process): 优化上传切分流程与PDF高亮预览 2026-07-24 14:35:03 +08:00
caoxiaozhu
33d0ed2e01 feat(data-process): 完善文件解析与切分存储链路 2026-07-24 14:35:03 +08:00
wuyongtao
6d4bf85284 feat: 提交 Docker 环境配置文件,更新 .gitignore 规则
- 新增 docker/app/.env 应用环境配置
- 新增 docker/compute/.env 计算节点环境配置
- 更新 .gitignore 允许提交 .env 配置文件

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-24 10:45:10 +08:00
wuyongtao
2b10c013ce feat: 更新后端平台模块、前端组件及构建产物,新增工作计划文档
- 更新 backend 平台 API endpoints 及 platform_store
- 更新前端 ComputeNodesView、DataProcessCreateView、FineTuneCreateView 等组件
- 更新前端 API 模块(compute、fineTune)
- 重构 frontend/dist 构建产物(新 hash)
- 新增 docs/2026-07-24-work-plan.md 工作计划文档

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-24 10:27:52 +08:00
wuyongtao
b28cfbc6fa feat: 更新后端平台模块、Compute引擎、前端组件及构建产物
- 更新 backend 平台 API、platform_store、compute_gateway sync
- 更新 compute agent/engine/adapter 及 API
- 更新 Docker 部署配置(app/compute)
- 新增 frontend/src/utils/ 工具模块
- 新增 scripts/ops_diagnostics.py 运维诊断脚本
- 新增 docs/2026-07-23-development-summary.md 开发总结
- 重构 frontend/dist 构建产物(新 hash)
- 更新前端多个视图组件及 API 模块

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-23 19:32:42 +08:00
caoxiaozhu
f04dc479bb feat: 完成数据处理接口与前端接入 2026-07-23 15:10:13 +08:00
caoxiaozhu
f453234057 feat: 增加前后端一键启动脚本 2026-07-23 11:09:04 +08:00
wuyongtao
a6868ec2e5 fix: 恢复 gitignore 全局 logs/ 规则,通过 !docker/compute/data/yg-ft/logs/ 放行
Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-22 20:50:37 +08:00
wuyongtao
6cd1e46e86 feat: 添加 compute 数据目录结构,配置 gitignore 规则
- docker/compute/data/yg-ft/ 下创建 models、datasets、outputs、logs 目录
- 保留目录结构和 README,忽略子目录实际内容(日志、模型、数据集等)
- 修正 gitignore 全局 logs/ 规则为 /logs/,避免误伤

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-22 20:47:05 +08:00
wuyongtao
836343b29e feat: 新增 compute_gateway、compute_poller、agent 模块,重构前端 dist
- 新增 backend/app/modules/compute_gateway(client/sync)计算网关模块
- 新增 backend/app/workers/compute_poller 计算轮询 worker
- 新增 compute/agent/process_manager 进程管理器
- 新增 scripts/ 脚本目录
- 更新 Docker 部署配置(app/compute/nginx)
- 更新后端平台 API、数据库 SQL、core 配置
- 更新前端多个视图组件及 API 模块
- 重构 frontend/dist 构建产物(新 hash)
- 更新多项文档

Co-Authored-By: Claude <noreply@anthropic.com>
2026-07-22 20:26:16 +08:00
621 changed files with 174335 additions and 8924 deletions

0
.cache/.gitkeep Normal file
View File

View File

0
.cache/tiktoken/.gitkeep Normal file
View File

35
.github/workflows/ci.yml vendored Normal file
View File

@@ -0,0 +1,35 @@
name: yg-ft-ci
on:
push:
branches: [main, master, develop]
pull_request:
jobs:
frontend:
runs-on: ubuntu-latest
defaults:
run:
working-directory: frontend
steps:
- uses: actions/checkout@v4
- uses: actions/setup-node@v4
with:
node-version: 20
cache: npm
cache-dependency-path: frontend/package-lock.json
- run: npm ci
- run: npm run build
backend:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v4
- uses: actions/setup-python@v5
with:
python-version: '3.12'
- run: python -m pip install -r backend/requirements.txt
- run: python -m compileall -q backend/app compute
- run: python -m pytest -q backend/tests/test_storage_security.py
env:
PYTHONPATH: backend

47
.gitignore vendored
View File

@@ -12,8 +12,6 @@ __pycache__/
build/ build/
develop-eggs/ develop-eggs/
dist/ dist/
!frontend/dist/
!frontend/dist/**
node_modules/ node_modules/
*.tsbuildinfo *.tsbuildinfo
downloads/ downloads/
@@ -44,6 +42,7 @@ pip-delete-this-directory.txt
# Runtime data and logs # Runtime data and logs
runtime/ runtime/
backend/runtime/ backend/runtime/
backend/storage/
logs/ logs/
backend/logs/ backend/logs/
*.db *.db
@@ -56,7 +55,6 @@ htmlcov/
.nox/ .nox/
.coverage .coverage
.coverage.* .coverage.*
.cache
nosetests.xml nosetests.xml
coverage.xml coverage.xml
*.cover *.cover
@@ -143,6 +141,7 @@ celerybeat.pid
# Environments # Environments
.env .env
!.env.example
.venv .venv
env/ env/
venv/ venv/
@@ -150,6 +149,16 @@ ENV/
env.bak/ env.bak/
venv.bak/ venv.bak/
# Local backend config (含数据库账号密码等敏感信息,勿提交)
backend/config.yaml
# Agent / IDE 工具产物,不应进版本库
.codex-backups/
.pnpm-store/
.zcode/
.claude/
CLAUDE.md
# Spyder project settings # Spyder project settings
.spyderproject .spyderproject
.spyproject .spyproject
@@ -187,5 +196,33 @@ cython_debug/
# PyPI configuration file # PyPI configuration file
.pypirc .pypirc
# Local backend configuration may contain database credentials. docker/llamafactory-latest.tar.gz
backend/config.yaml # Compute data - 保留目录结构和 README忽略子目录内容日志、模型、数据集等
!docker/compute/data/yg-ft/logs/
docker/compute/data/yg-ft/datasets/*
docker/compute/data/yg-ft/models/*
docker/compute/data/yg-ft/outputs/*
docker/compute/data/yg-ft/trained_models/*
docker/compute/data/yg-ft/logs/**
!docker/compute/data/yg-ft/logs/compute/
!docker/compute/data/yg-ft/logs/training/
!docker/compute/data/yg-ft/**/.gitkeep
!docker/compute/data/yg-ft/**/README.md
# Offline deployment bundle - 离线部署包(镜像、运行时等大文件,不提交)
docker/offline/
# MinIO object storage data - 对象存储运行时数据,勿提交,保留目录结构
docker/minio/data/*
!docker/minio/data/.gitkeep
# nlp-eval-demo - 独立演示项目,不进版本库
nlp-eval-demo/
nlp-eval-demo.zip
# 项目本地缓存HuggingFace / tiktoken 等大文件),保留目录结构与占位文件
.cache/*
!.cache/.gitkeep
!.cache/tiktoken/
!.cache/huggingface/
!.cache/**/.gitkeep

222
README.md
View File

@@ -1,197 +1,59 @@
# YG_FT 模型微调平台 # YG_FT 模型微调平台
YG_FT 是一个面向企业治理场景的模型微调平台,覆盖用户中心、多租户、项目隔离、数据集管理、模型管理、训练任务、评测、推理、审批流、审计留存、算力调度和训练引擎适配 YG_FT 是面向多用户、多租户和多算力节点的模型训练与推理平台,提供数据集、模型、训练、权重合并、推理、评测、算力节点、项目隔离、权限和审计能力
当前前端已有基础页面,后端与算力平台已按多人协作开发方式建立工程骨架,并开始实现正式系统主链路能力。当前代码和 SQL 均作为后续生产演进基线维护,不再以一次性演示或静态 Mock 为开发准则。 ## 架构
## 总体架构
```text ```text
YG_FT/ 浏览器 -> Frontend Nginx:16801 -> Backend API:17861
frontend/ # 前端控制台 |-> PostgreSQL元数据、权限、审计、任务状态
backend/ # FastAPI 应用平台后端 |-> Redis缓存及任务辅助状态
app/ |-> MinIO:19000模型和数据唯一对象源
api/v1/ # 对前端暴露的 REST API |-> Compute API:19100
core/ # 配置、日志、中间件、权限等基础能力 |-> Compute Agent/GPU/LLaMA-Factory
db/ # 数据库连接、迁移、事务工具 `-> File Gateway:19101
modules/ # 业务模块目录
schemas/ # Pydantic 入参/出参模型
services/ # 跨模块应用服务
workers/ # 后台任务入口
requirements.txt # 后端 Python 第三方依赖
compute/ # 算力平台与训练框架适配层
api/ # 内部 Compute API
agent/ # 单机多 GPU 调度与进程管理
engines/llama_factory/ # LLaMA-Factory 适配器
file_gateway/ # 本地文件上传、下载、导入、产物管理
docs/ # 需求、接口、数据库、开发计划和部署文档
docker/ # 容器化配置
``` ```
## 平台分层 Backend、MinIO 和 Compute 节点可以部署在不同服务器,不依赖跨服务器 Docker 网络,通过 IP、DNS 或负载均衡地址通信。Compute 节点只保存按需准备的本地缓存MinIO 是模型、数据集、权重、评测报告和训练产物的唯一数据源。
| 层级 | 职责 | 主要目录 | ## 目录
| --- | --- | --- |
| 前端控制台 | 用户操作入口、任务看板、项目/模型/数据集/训练/审批/审计页面 | `frontend/` |
| 应用平台后端 | 用户中心、多租户、RBAC/ABAC、项目隔离、元数据、审批流、审计、API 编排 | `backend/` |
| 算力平台 | GPU 发现、资源锁定、训练进程管理、日志采集、产物归档、任务状态同步 | `compute/` |
| 训练引擎 | 当前固定接入 LLaMA-Factory预留其他训练平台适配标准 | `compute/engines/` |
| 数据层 | PostgreSQL、Redis、本地文件存储、日志归档 | `docs/postgres-schema.sql` |
## 当前开发基线 | 目录 | 作用 |
| --- | --- |
| `frontend/` | Vue 3、TypeScript、Element Plus 控制台 |
| `backend/app/api/v1/` | 平台 REST API |
| `backend/app/core/` | 配置、认证、权限和日志 |
| `backend/app/db/` | PostgreSQL 访问和初始化 SQL |
| `backend/app/modules/` | 系统、资源、审批、数据处理和存储模块 |
| `backend/app/workers/` | 节点轮询、任务对账和资源同步 |
| `compute/` | Compute API、Agent、GPU 和训练引擎 |
| `docker/` | 应用、MinIO、算力服务部署文件 |
| `docs/` | 架构、权限、部署和测试文档 |
- 使用 FastAPI 提供统一 API 响应结构 `{ code, message, data }` ## 端口
- 本地运行阶段统一使用 PostgreSQL后端启动时会在 PG 中初始化当前运行表和系统内置账号模型、数据集、算力节点、GPU、微调任务等业务数据必须通过页面、接口或正式导入流程产生。
- 支持登录、模型管理、数据集管理、微调任务创建/启动/停止/进度轮询。
- 支持训练日志、loss 指标、checkpoint 和训练产物接口;真实训练执行器接入前,联调状态机必须通过显式环境变量开启。
- 支持多算力节点、GPU、任务队列、资源副本和资源同步状态接口。
- 前端新增 `/compute` 算力节点页面展示节点地址、权重、标签、启用状态、GPU、队列和资源副本。
- `compute/engines/llama_factory/adapter.py` 提供 LLaMA-Factory 参数校验、命令生成和日志解析基础能力。
## 一键启动前端和后端 | 服务 | 主机端口 | 容器端口 |
| --- | ---: | ---: |
| Frontend | 16801 | 80 |
| Backend API | 17861 | 8000 |
| Redis | 16379 | 6379 |
| MinIO API/Console | 19000/19001 | 9000/9001 |
| Compute API/File Gateway | 19100/19101 | 9100 |
## 启动
```bash ```bash
./scripts/start-dev.sh --setup # 首次运行,初始化本地依赖 cd frontend && npm ci && npm run build && cd ..
./scripts/start-dev.sh # 后续直接启动 docker build -f docker/app/Dockerfile.backend -t yg-ft-backend-api:latest .
./scripts/start-dev.sh --check # 仅检查依赖和 PostgreSQL 连接 docker build -f docker/app/Dockerfile.frontend -t yg-ft-frontend-runtime:latest .
docker build -f docker/compute/Dockerfile.compute -t yg-ft-compute-api:latest .
cd docker/minio && docker compose up -d
cd ../app && docker compose up -d
cd ../compute && docker compose up -d
``` ```
首次运行前可将 `backend/config.example.yaml` 复制为 `backend/config.yaml` 并填写 ## 权限与性能
本地数据库凭据;`config.yaml` 已加入 Git 忽略。脚本会检查 PostgreSQL 和两个服务
的健康状态,并在退出时同时停止前端和后端。也可以按下面步骤分别启动服务。
一键启动使用的前后端端口也统一配置在同一个 YAML 中: 系统使用角色权限、资源 ACL、用户/项目/租户归属联合校验;删除为软删除,关键操作写入审计。模型合并前准备 base model 和 adapter结果归档 MinIO推理前按选择节点准备缓存。远程 PostgreSQL 延迟会影响全量列表和看板,页面慢时应检查浏览器 Network、Nginx、Backend 日志、连接池和节点可达性。
```yaml 详细部署见 `docker/README.md`,测试见 `测试用例.md`,本次快照见 `docs/20260812/`
server:
frontend_port: 16801
backend_port: 17861
```
端口被占用时修改这两项后重新启动即可。未显式配置 `cors_allow_origins` 时,后端会
根据 `frontend_port` 自动允许本机前端来源。环境变量 `FRONTEND_PORT`
`BACKEND_PORT` 可以临时覆盖 YAML。
## 后端启动
```bash
cd backend
/opt/miniconda3/bin/python3 -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
uvicorn app.main:app --reload --port 17861
```
默认接口前缀为 `/modelTF`,例如:
```text
GET /modelTF/health
POST /modelTF/login
GET /modelTF/model-manage
GET /modelTF/dataset-manage
GET /modelTF/fine-tune
GET /modelTF/compute/nodes
```
本地运行时默认通过 `backend/config.yaml` 连接 PostgreSQL
```yaml
server:
frontend_port: 16801
backend_port: 17861
database:
url: postgresql+psycopg://localhost:5432/yg_ft
username: yg_ft
password: "change_me"
```
本地启动前需要确保 PostgreSQL 已监听 YAML 配置的地址(默认 `localhost:5432`
并已创建 `yg_ft` 数据库和 `yg_ft` 用户。后端启动后会自动创建当前运行表并写入
内置管理员账号,运行数据统一写入 PostgreSQL。
开发阶段内置登录账号:
| 角色 | 账号 | 密码 | 说明 |
| --- | --- | --- | --- |
| 超级管理员 | `admin` | `admin123` | 拥有当前全部页面权限 |
| 操作员 | `operator` | `operator123` | 拥有业务操作相关页面权限 |
以上账号仅用于本地开发和联调。生产环境初始化后应立即修改密码,或改为企业统一身份认证/管理员初始化流程。
## 前端启动
```bash
cd frontend
npm install
npm run dev
```
通过一键脚本启动时,前端端口和 Vite proxy 的后端端口均来自
`backend/config.yaml`。单独运行 `npm run dev` 时,默认使用前端 `16801`、后端
`17861`,也可以通过 `FRONTEND_PORT``BACKEND_PORT` 覆盖。
## 算力服务启动
```bash
cd compute
uvicorn api.main:app --reload --port 19100
```
默认 `COMPUTE_MODE=real`。真实 GPU 接入时,在每台算力服务器上部署 Compute API、Agent、File Gateway 和 LLaMA-Factory应用平台通过 `compute_nodes.api_base_url``compute_nodes.file_gateway_url` 主动轮询。仅在隔离联调环境可显式设置 `COMPUTE_MODE=simulator``COMPUTE_EXECUTION_MODE=simulator`
## 日志
后端日志模块位于 `backend/app/core/logging.py`,说明文档见:
- `docs/backend-logging.md`
默认输出:
```text
logs/backend-YYYY-MM-DD.log
logs/error-YYYY-MM-DD.log
```
日志格式为 JSON Lines单个文件不超过 20MB只保留最近 10 天。
## 主要文档
- `docs/platform-architecture-requirements.md`:平台需求、功能模块、页面补全建议。
- `docs/menu-functional-requirements.md`:当前菜单、二级路由、规划菜单、功能需求、接口和数据库映射。
- `docs/backend-api-design.md`FastAPI 接口分组、参数定义、权限说明。
- `docs/postgres-schema.sql`PostgreSQL 数据库脚本,包含权限、用户中心、多租户、审批、审计等模型。
- `docs/system-development-plan.md`多人协作开发计划按前端、后端、DB、部署拆分。
- `docs/team-development-plan.md`3-4 人并行开发分工计划,按人员边界标注页面、接口、数据库和交付节奏。
- `docs/first-version-development-plan.md`当前系统主链路开发计划覆盖前端、后端、DB、Compute API、GPU 和 LLaMA-Factory 适配。
- `docs/backend-logging.md`:后端日志模块使用说明。
- `docs/deployment-plan.md`:后期部署方案,覆盖单机算力服务器部署与应用/算力分离部署。
- `docker/README.md`Docker 部署入口,包含应用服务器和算力服务器两套 Compose 使用方式。
## Docker 部署入口
应用服务器:
```bash
cd docker/app
cp .env.example .env
docker compose up -d
```
算力服务器:
```bash
cd docker/compute
cp .env.example .env
docker compose up -d
```
两套 Compose 均采用代码外挂方式运行,镜像只包含运行时环境和第三方依赖。项目根目录不再保留 `Dockerfile``docker-compose.yml`,部署时统一进入 `docker/app``docker/compute` 目录执行。
## 后续开发原则
- 接口实现优先遵循 `docs/backend-api-design.md`
- 数据库实现优先遵循 `docs/postgres-schema.sql`,后续通过 Alembic 迁移管理变更。
- 前端页面与后端接口、数据库表之间的映射以文档中的“对应页面/功能模块”为准。
- 训练引擎适配必须通过 `compute/engines/` 下的标准接口,不在应用平台后端直接拼接训练命令。
- 敏感信息不得写入日志,生产环境密钥通过环境变量或密钥管理系统注入。

View File

@@ -1,18 +0,0 @@
# 默认自动读取 backend/config.yaml。设置此变量可切换到其他 YAML 文件。
BACKEND_CONFIG_FILE=./config.yaml
# 下列环境变量按需启用,并优先于 YAML 中的配置。
# DATABASE_BASE_URL=postgresql+psycopg://localhost:5432/yg_ft
# DATABASE_USERNAME=yg_ft
# DATABASE_PASSWORD=change_me
# DATABASE_URL=postgresql+psycopg://yg_ft:change_me@localhost:5432/yg_ft
# DATABASE_URL 是完整连接串覆盖项,优先级高于上面三个分离项。
# MODELTF_ROUTE_PREFIX=/modelTF
# FRONTEND_PORT=16801
# BACKEND_PORT=17861
# CORS_ALLOW_ORIGINS=http://localhost:16801,http://127.0.0.1:16801
# LOG_LEVEL=INFO
# LOG_DIR=./logs
# COMPUTE_MODE=real
# COMPUTE_STATUS_SYNC_MODE=polling
# COMPUTE_POLL_INTERVAL_SECONDS=3

View File

@@ -32,65 +32,17 @@ backend/
services/ # 跨模块应用服务 services/ # 跨模块应用服务
workers/ # 后台任务入口 workers/ # 后台任务入口
requirements.txt # 后端第三方依赖 requirements.txt # 后端第三方依赖
config.example.yaml # 可提交的脱敏配置模板
config.yaml # 本地配置,已忽略,环境变量可覆盖
logs/ # 本地开发日志目录,生产环境建议挂载到独立日志盘 logs/ # 本地开发日志目录,生产环境建议挂载到独立日志盘
``` ```
## 配置
先复制本地配置并填写数据库凭据:
```bash
cp config.example.yaml config.yaml
```
后端默认读取 `config.yaml`,配置优先级为:
```text
环境变量 > BACKEND_CONFIG_FILE 指定的 YAML > backend/config.yaml > 代码默认值
```
前后端端口、数据库地址、跨域来源、日志路径和计算状态轮询参数均可在 YAML 中
配置。敏感信息或不同环境的差异建议通过环境变量覆盖不要写入仓库。YAML 中的
相对日志路径以该 YAML 文件所在目录为基准。使用其他配置文件时:
```bash
BACKEND_CONFIG_FILE=./config.prod.yaml uvicorn app.main:app --reload --port 17861
```
本机直连 PostgreSQL 时,数据库配置拆分为地址、用户名和密码:
```yaml
server:
frontend_port: 16801
backend_port: 17861
database:
url: postgresql+psycopg://localhost:5432/yg_ft
username: yg_ft
password: "change_me"
```
服务启动时会安全拼装完整连接串用户名或密码中的特殊字符会自动编码。Docker
和生产环境仍可用完整的 `DATABASE_URL` 环境变量覆盖以上三项。
这里的 `5432` 是宿主机 PostgreSQL 默认端口;如果数据库容器映射到 `15432`
将 YAML 中的端口改为 `15432` 即可。
`server.frontend_port``server.backend_port` 由根目录的
`scripts/start-dev.sh` 读取,并分别传给 Vite 和 Uvicorn。未配置
`app.cors_allow_origins` 时,本机 CORS 来源会跟随前端端口。直接手工执行
`uvicorn` 时仍需通过 `--port` 指定监听端口。
## 本地启动 ## 本地启动
```bash ```bash
cd backend cd backend
/opt/miniconda3/bin/python3 -m venv .venv python -m venv .venv
source .venv/bin/activate .venv\Scripts\activate
pip install -r requirements.txt pip install -r requirements.txt
uvicorn app.main:app --reload --port 17861 uvicorn app.main:app --reload
``` ```
健康检查: 健康检查:

View File

@@ -0,0 +1,10 @@
from app.db.platform_store import get_platform_store
store = get_platform_store()
with store.connect() as conn:
rows = conn.execute(
"SELECT id, user_id, login_at, logout_at, duration_seconds FROM sessions ORDER BY login_at DESC LIMIT 10"
).fetchall()
print(f"sessions count: {len(rows)}")
for r in rows:
print(f" user={r['user_id'][:25]}... login={r['login_at']} logout={r['logout_at']} dur={r['duration_seconds']}")

File diff suppressed because it is too large Load Diff

View File

@@ -2,6 +2,8 @@
from app.core.logging import get_logger from app.core.logging import get_logger
from app.db.platform_store import get_platform_store from app.db.platform_store import get_platform_store
from app.core.config import get_settings
from app.modules.storage.minio_store import ObjectStorageError, get_object_storage
router = APIRouter() router = APIRouter()
logger = get_logger(__name__) logger = get_logger(__name__)
@@ -9,6 +11,30 @@ logger = get_logger(__name__)
@router.get("/health") @router.get("/health")
async def health_check() -> dict[str, object]: async def health_check() -> dict[str, object]:
logger.info("health check requested") # Health endpoints are called frequently by Docker and the frontend.
return {"code": 0, "message": "ok", "data": get_platform_store().health_metrics()} # Keep failures visible without emitting one INFO line per probe.
logger.debug("health check requested")
storage_status: dict[str, object] = {"enabled": get_settings().minio_enabled, "status": "disabled"}
if get_settings().minio_enabled:
try:
get_object_storage().ensure_bucket()
storage_status = {
"enabled": True,
"status": "ready",
"endpoint": get_settings().minio_endpoint,
"bucket": get_settings().minio_bucket,
}
except (ObjectStorageError, OSError) as exc:
logger.warning("MinIO health check failed: %s", exc)
storage_status = {
"enabled": True,
"status": "unavailable",
"endpoint": get_settings().minio_endpoint,
"error": str(exc),
}
return {
"code": 0,
"message": "ok",
"data": {**get_platform_store().health_metrics(), "storage": storage_status},
}

File diff suppressed because it is too large Load Diff

View File

@@ -1,9 +1,26 @@
from fastapi import APIRouter from fastapi import APIRouter
from app.api.v1.endpoints.data_process import router as data_process_router
from app.api.v1.endpoints.platform import router as platform_router from app.api.v1.endpoints.platform import router as platform_router
from app.api.v1.endpoints.health import router as health_router from app.api.v1.endpoints.health import router as health_router
from app.modules.tenant.router import router as tenant_router
from app.modules.project.router import router as project_router
from app.modules.approval.router import router as approval_router
from app.modules.system.router import router as system_router
from app.modules.retention.router import router as retention_router
from app.modules.resource.router import router as resource_router
from app.modules.gpu.router import router as gpu_router
from app.modules.data_convert.router import router as data_convert_router
api_router = APIRouter() api_router = APIRouter()
api_router.include_router(health_router, tags=["health"]) api_router.include_router(health_router, tags=["health"])
api_router.include_router(data_process_router, tags=["data-process"])
api_router.include_router(platform_router, tags=["platform"]) api_router.include_router(platform_router, tags=["platform"])
api_router.include_router(system_router, tags=["system"])
api_router.include_router(tenant_router, tags=["tenant"])
api_router.include_router(project_router, tags=["project"])
api_router.include_router(approval_router, tags=["approval"])
api_router.include_router(retention_router, tags=["retention"])
api_router.include_router(resource_router, tags=["resource"])
api_router.include_router(gpu_router, tags=["gpu-assignment"])
api_router.include_router(data_convert_router, tags=["data-convert"])

269
backend/app/core/audit.py Normal file
View File

@@ -0,0 +1,269 @@
"""
审计日志装饰器模块
提供 @audit_log 装饰器,用于自动记录关键业务操作的审计日志。
使用示例:
from app.core.audit import audit_log
@audit_log(action="create_dataset", target_type="dataset")
async def create_dataset(request: Request, ...):
...
"""
from __future__ import annotations
import asyncio
import functools
import time
from datetime import datetime, timezone
from typing import Any, Callable, Optional, TypeVar
from fastapi import Request
from app.core.logging import get_client_ip, get_logger, mask_sensitive_string, request_id_var
logger = get_logger("app.audit")
F = TypeVar("F", bound=Callable[..., Any])
def audit_log(
action: str,
target_type: str = "",
*,
detail_template: str = "",
extract_target_id: Optional[Callable[[Any], str]] = None,
) -> Callable[[F], F]:
"""
审计日志装饰器
Args:
action: 操作类型,如 create_dataset、update_model 等
target_type: 目标资源类型,如 dataset、model 等
detail_template: 日志详情模板(支持 format 参数)
extract_target_id: 从返回值中提取目标 ID 的函数
Returns:
装饰后的函数
"""
def decorator(func: F) -> F:
if asyncio.iscoroutinefunction(func):
@functools.wraps(func)
async def async_wrapper(*args, **kwargs):
started_at = time.perf_counter()
trace_id = request_id_var.get("-")
try:
result = await func(*args, **kwargs)
elapsed_ms = (time.perf_counter() - started_at) * 1000
target_id = _extract_target_id(result, kwargs, extract_target_id)
detail = _build_detail(detail_template, kwargs)
_record_audit(
action=action,
actor_id=_extract_actor_id(kwargs),
target_type=target_type,
target_id=target_id,
detail=detail,
trace_id=trace_id,
duration_ms=elapsed_ms,
kwargs=kwargs,
args=args,
)
return result
except Exception as exc:
_record_audit(
action=action,
actor_id=_extract_actor_id(kwargs),
target_type=target_type,
target_id=_extract_target_id(None, kwargs, extract_target_id),
detail=_build_detail(detail_template, kwargs),
trace_id=trace_id,
duration_ms=(time.perf_counter() - started_at) * 1000,
result="failure",
reason=_safe_exception_reason(exc),
kwargs=kwargs,
args=args,
)
logger.warning("业务操作失败 action=%s reason=%s", action, _safe_exception_reason(exc))
raise
return async_wrapper # type: ignore
else:
@functools.wraps(func)
def sync_wrapper(*args, **kwargs):
started_at = time.perf_counter()
trace_id = request_id_var.get("-")
try:
result = func(*args, **kwargs)
elapsed_ms = (time.perf_counter() - started_at) * 1000
target_id = _extract_target_id(result, kwargs, extract_target_id)
detail = _build_detail(detail_template, kwargs)
_record_audit(
action=action,
actor_id=_extract_actor_id(kwargs),
target_type=target_type,
target_id=target_id,
detail=detail,
trace_id=trace_id,
duration_ms=elapsed_ms,
kwargs=kwargs,
args=args,
)
return result
except Exception as exc:
_record_audit(
action=action,
actor_id=_extract_actor_id(kwargs),
target_type=target_type,
target_id=_extract_target_id(None, kwargs, extract_target_id),
detail=_build_detail(detail_template, kwargs),
trace_id=trace_id,
duration_ms=(time.perf_counter() - started_at) * 1000,
result="failure",
reason=_safe_exception_reason(exc),
kwargs=kwargs,
args=args,
)
logger.warning("业务操作失败 action=%s reason=%s", action, _safe_exception_reason(exc))
raise
return sync_wrapper # type: ignore
return decorator
def _extract_target_id(
result: Any, kwargs: dict, extractor: Optional[Callable[[Any], str]]
) -> Optional[str]:
"""从返回值或 kwargs 中提取目标 ID"""
if extractor:
try:
return extractor(result)
except Exception:
pass
if isinstance(result, dict):
return result.get("id")
# 尝试从路径参数中提取
for key in ("dataset_id", "model_id", "task_id", "resource_id"):
val = kwargs.get(key)
if val:
return str(val)
return None
def _build_detail(template: str, kwargs: dict) -> str:
"""构建审计详情"""
if not template:
return ""
try:
return template.format(**kwargs)
except (KeyError, IndexError):
return template
def _record_audit(
action: str,
actor_id: Optional[str],
target_type: str,
target_id: Optional[str],
detail: str,
trace_id: str,
duration_ms: float,
*,
kwargs: dict[str, Any] | None = None,
args: tuple[Any, ...] = (),
result: str = "success",
reason: str | None = None,
) -> None:
"""通过已有的 record_audit 方法写入审计日志"""
try:
from app.db.platform_store import get_platform_store
store = get_platform_store()
kwargs = kwargs or {}
request = _extract_request(args, kwargs)
current_user = kwargs.get("current_user") or kwargs.get("user") or {}
request_id = request.headers.get("X-Request-ID") if request else None
request_id = request_id or trace_id
client_ip = get_client_ip(request) or None
detail_text = f"{detail} trace_id={trace_id} duration_ms={duration_ms:.1f}" if detail else f"trace_id={trace_id} duration_ms={duration_ms:.1f}"
store.record_audit(
action=action,
actor_id=actor_id,
target_type=target_type or None,
target_id=target_id,
tenant_id=str(current_user.get("tenant_id") or "") or None,
detail=mask_sensitive_string(detail_text),
result=result,
reason=mask_sensitive_string(reason or "") or None,
request_id=request_id,
session_id=str(current_user.get("session_id") or "") or None,
ip=client_ip,
)
except Exception:
logger.error("写入审计日志失败 action=%s", action, exc_info=True)
def _extract_actor_id(kwargs: dict) -> Optional[str]:
"""从 FastAPI 注入的当前用户中提取操作人 ID。"""
for key in ("current_user", "user"):
value = kwargs.get(key)
if isinstance(value, dict) and value.get("id"):
return str(value["id"])
return None
def _extract_request(args: tuple[Any, ...], kwargs: dict[str, Any]) -> Request | None:
for value in tuple(kwargs.values()) + tuple(args):
if isinstance(value, Request):
return value
return None
def _safe_exception_reason(exc: Exception) -> str:
"""Keep audit failures useful without recording credentials or tokens."""
value = getattr(exc, "detail", None) or str(exc) or exc.__class__.__name__
return mask_sensitive_string(str(value))[:500]
# ==================== 预定义的审计操作常量 ====================
class AuditActions:
"""预定义的审计操作类型"""
# 数据集操作
CREATE_DATASET = "create_dataset"
UPDATE_DATASET = "update_dataset"
DELETE_DATASET = "delete_dataset"
# 模型操作
CREATE_MODEL = "create_model"
UPDATE_MODEL = "update_model"
DELETE_MODEL = "delete_model"
# 微调任务
CREATE_FINE_TUNE = "create_fine_tune"
UPDATE_FINE_TUNE = "update_fine_tune"
DELETE_FINE_TUNE = "delete_fine_tune"
# 推理任务
CREATE_INFERENCE = "create_inference"
UPDATE_INFERENCE = "update_inference"
DELETE_INFERENCE = "delete_inference"
# 用户管理
CREATE_USER = "create_user"
UPDATE_USER = "update_user"
DELETE_USER = "delete_user"
# 租户管理
CREATE_TENANT = "create_tenant"
UPDATE_TENANT = "update_tenant"
DELETE_TENANT = "delete_tenant"
# 权限授权
GRANT_ACL = "grant_acl"
REVOKE_ACL = "revoke_acl"
# 系统配置
UPDATE_CONFIG = "update_config"

615
backend/app/core/auth.py Normal file
View File

@@ -0,0 +1,615 @@
"""鉴权依赖:从 Authorization header 解析当前用户,提供权限校验。"""
from __future__ import annotations
from typing import Any
import json
from fastapi import Depends, HTTPException, Query, Request, status
from app.db.platform_store import get_platform_store
from app.core.logging import get_client_ip
# 无需鉴权的路径前缀(健康检查、登录等)
PUBLIC_PATHS = ("/health", "/login", "/system-info")
OWNER_TABLES = {
"dataset": ("datasets", "created_by"), "model": ("models", "created_by"),
"trained_model": ("trained_models", "created_by"), "eval": ("eval_tasks", "created_by"),
"fine-tune": ("fine_tune_tasks", "created_by"), "fine_tune_task": ("fine_tune_tasks", "created_by"),
"compare": ("compare_tasks", "payload"), "inference": ("compare_tasks", "payload"),
"project": ("projects", "create_by"), "data_process": ("data_process_tasks", "created_by"),
"data_convert": ("data_convert_tasks", "created_by"),
}
RESOURCE_TABLES = {
"dataset": "datasets",
"model": "models",
"trained_model": "trained_models",
"eval": "eval_tasks",
"fine-tune": "fine_tune_tasks",
"fine_tune_task": "fine_tune_tasks",
"compare": "compare_tasks",
"inference": "compare_tasks",
"project": "projects",
"data_process": "data_process_tasks",
"data_convert": "data_convert_tasks",
}
MODULE_PERMISSIONS = {
"dashboard": "dashboard",
"fine-tune": "fine-tune",
"model-eval": "model-eval",
"model-compare": "model-inference",
"model-inference": "model-inference",
"model-chat": "model-inference",
"model-manage": "model-manage",
"dataset-manage": "dataset",
"data-process": "data-process",
"data-convert": "data-convert",
"compute": "compute",
"hardware": "hardware",
"users": "user-settings",
}
# These endpoints are the user-facing compute view used by training,
# inference, and evaluation forms. They only return the current user's
# assigned nodes/GPUs in the endpoint implementation, so they must remain
# available after an approval without granting access to the admin compute
# management page.
SELF_SERVICE_COMPUTE_PATHS = {
"/compute/nodes",
"/compute/gpus",
"/compute/my-gpus",
}
# Resource actions are deliberately kept separate from module permissions.
# A user may be allowed to open a module while still lacking the action on a
# specific resource (for example, download or delete).
RESOURCE_ACTIONS = frozenset({
"read",
"write",
"execute",
"download",
"export",
"delete",
"admin",
})
RESOURCE_ACTION_ALIASES = {"export": "download"}
def _extract_token(request: Request) -> str | None:
"""从 Authorization header 提取 token格式: Bearer platform-token-{user_id})。"""
auth = request.headers.get("Authorization", "")
token = auth.replace("Bearer ", "").strip()
if token.startswith("platform-token-"):
return token[len("platform-token-"):]
return None
def _session_token(user_id: str, session_id: str) -> str:
return f"platform-token-{user_id}.{session_id}"
def _record_auth_event(actor_id: str | None, action: str, reason: str, request: Request) -> None:
"""Best-effort security audit for authentication and permission denials."""
try:
get_platform_store().record_audit(
action=action,
actor_id=actor_id,
target_type="auth",
target_id=request.url.path,
detail=reason,
result="denied",
reason=reason,
request_id=request.headers.get("X-Request-ID"),
ip=get_client_ip(request) or None,
)
except Exception:
# An audit failure must never turn an authentication decision into an
# accidental allow or an unrelated 500 response.
pass
def get_current_user(request: Request) -> dict[str, Any]:
"""
FastAPI 依赖:解析当前登录用户。
- 公开路径(/health, /login 等)直接放行,返回匿名用户。
- 无 token 或 token 无效时抛 401。
- admin 用户标记为超级管理员,拥有全部权限。
"""
path = request.url.path
# 去掉路由前缀后判断
for prefix in PUBLIC_PATHS:
if path.endswith(prefix):
return {"id": None, "username": "anonymous", "role": "viewer", "permissions": [], "protected": False}
token_value = _extract_token(request)
if not token_value:
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="missing or invalid token")
store = get_platform_store()
user_id, _, session_id = token_value.partition(".")
if session_id:
with store.connect() as conn:
session = conn.execute(
"SELECT user_id, logout_at, expires_at FROM sessions WHERE id=?", (session_id,)
).fetchone()
if not session or session["user_id"] != user_id or session["logout_at"]:
_record_auth_event(user_id, "auth.session.denied", "session expired or logged out", request)
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="session expired")
if session["expires_at"]:
from datetime import datetime, timezone
try:
if datetime.fromisoformat(str(session["expires_at"]).replace("Z", "+00:00")) <= datetime.now(timezone.utc):
_record_auth_event(user_id, "auth.session.denied", "session expired", request)
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="session expired")
except ValueError:
pass
with store.connect() as conn:
user_row = conn.execute("SELECT * FROM users WHERE id=?", (user_id,)).fetchone()
if user_row:
user = store._user(user_row)
if user.get("status") != "active":
_record_auth_event(user_id, "auth.user.disabled", "user is not active", request)
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="user disabled")
requested_tenant = request.headers.get("X-Tenant-ID", "").strip()
if requested_tenant and (is_admin(user) or requested_tenant in user_tenant_ids(user)):
user["tenant_id"] = requested_tenant
# Keep the session identifier in request context so business audit
# records can be traced back to the exact login session.
if session_id:
user["session_id"] = session_id
path_parts = path.strip("/").split("/")
# The API may be mounted directly at /modelTF or behind /api/v1/modelTF.
# Locate the first known module segment instead of relying on a fixed index.
segment = next((part for part in path_parts if part in MODULE_PERMISSIONS), "")
required = MODULE_PERMISSIONS.get(segment)
relative_path = "/" + "/".join(path_parts[path_parts.index(segment):]) if segment else path
self_service_compute = relative_path.rstrip("/") in SELF_SERVICE_COMPUTE_PATHS
if (
required
and not is_admin(user)
and required not in (user.get("permissions") or [])
and not self_service_compute
):
_record_auth_event(user_id, "auth.permission.denied", f"missing permission: {required}", request)
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail=f"missing permission: {required}")
return user
raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="user not found")
def require_admin(current_user: dict[str, Any] = Depends(get_current_user)) -> dict[str, Any]:
"""FastAPI 依赖:要求当前用户是平台管理员。"""
if is_admin(current_user):
return current_user
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="admin permission required")
def require_tenant_admin(
tenant_id: str,
current_user: dict[str, Any] = Depends(get_current_user),
) -> dict[str, Any]:
"""Allow platform admins and active owner/admin tenant members."""
if is_admin(current_user) or is_tenant_admin(current_user, tenant_id):
return current_user
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="tenant admin permission required")
def is_admin(user: dict[str, Any]) -> bool:
"""判断用户是否为平台管理员;兼容历史 role=admin/protected 数据。"""
return (
user.get("platform_role") == "platform_admin"
or user.get("role") == "admin"
or user.get("protected", False)
)
def user_tenant_ids(user: dict[str, Any]) -> set[str]:
"""Return the tenant scope of a user.
Existing installations keep the primary tenant on ``users.tenant_id``.
``tenant_members`` is optional during the migration and adds memberships
when the permission v2 schema is available.
"""
if is_admin(user):
return {"*"}
primary_tenant = str(user.get("tenant_id") or "default")
result: set[str] = set()
user_id = user.get("id")
if not user_id:
return result
try:
store = get_platform_store()
with store.connect() as conn:
rows = conn.execute(
"SELECT tenant_id FROM tenant_members "
"WHERE user_id=? AND status='active' "
"AND (expires_at IS NULL OR expires_at='' OR expires_at > NOW()::text)",
(user_id,),
).fetchall()
result.update(str(row["tenant_id"]) for row in rows if row.get("tenant_id"))
if not result:
active_tenant = conn.execute(
"SELECT id FROM tenants WHERE id=? "
"AND COALESCE(status, 'active')='active' "
"AND COALESCE(deleted_at, '')=''",
(primary_tenant,),
).fetchone()
if active_tenant:
result.add(primary_tenant)
except Exception:
# Older databases are upgraded lazily. The primary users.tenant_id
# remains a valid fallback until the additive table is available.
result.add(primary_tenant)
return result
def tenant_membership(user: dict[str, Any], tenant_id: str | None = None) -> dict[str, Any] | None:
"""Return the active membership for the selected tenant, if any."""
if is_admin(user):
return {"tenant_id": tenant_id or user.get("tenant_id") or "default", "role": "owner", "status": "active"}
target = str(tenant_id or user.get("tenant_id") or "default")
try:
with get_platform_store().connect() as conn:
row = conn.execute(
"SELECT tenant_id, role, status, expires_at FROM tenant_members "
"WHERE tenant_id=? AND user_id=? AND status='active'",
(target, user.get("id")),
).fetchone()
if not row:
return None
if row.get("expires_at"):
from datetime import datetime, timezone
try:
if datetime.fromisoformat(str(row["expires_at"]).replace("Z", "+00:00")) <= datetime.now(timezone.utc):
return None
except ValueError:
return None
return dict(row)
except Exception:
return None
def is_tenant_admin(user: dict[str, Any], tenant_id: str | None = None) -> bool:
membership = tenant_membership(user, tenant_id)
return bool(membership and membership.get("role") in {"owner", "admin"})
def is_tenant_admin_for_resource(resource_type: str, resource: dict[str, Any], user: dict[str, Any]) -> bool:
if is_admin(user) or resource_type == "model":
return False
tenant_id = resource_tenant_id(resource_type, resource)
return bool(tenant_id and is_tenant_admin(user, tenant_id))
def bind_active_tenant(payload: dict[str, Any], user: dict[str, Any]) -> dict[str, Any]:
"""Bind a new resource to the authenticated tenant context.
Platform administrators may explicitly create a resource in another
active tenant. Ordinary users can only use the tenant selected by the
authenticated session/X-Tenant-ID header, never a client-supplied tenant
id alone.
"""
requested = str(payload.get("tenant_id") or user.get("tenant_id") or "default")
if not is_admin(user) and requested not in user_tenant_ids(user):
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="tenant access denied")
payload["tenant_id"] = requested if is_admin(user) else str(user.get("tenant_id") or requested)
return payload
def resource_record(resource_type: str, resource_id: str) -> dict[str, Any] | None:
"""Load a resource row for authorization without exposing storage details."""
table = RESOURCE_TABLES.get(resource_type)
if not table or not resource_id:
return None
store = get_platform_store()
try:
with store.connect() as conn:
row = conn.execute(f"SELECT * FROM {table} WHERE id=?", (resource_id,)).fetchone()
except Exception:
return None
if not row:
return None
result = dict(row)
if result.get("deleted_at"):
return None
return result
def resource_tenant_id(resource_type: str, resource: dict[str, Any]) -> str | None:
"""Resolve tenant ownership, including legacy JSON-backed task rows."""
# Administrator-created base models are platform shared. Online models
# created by ordinary users retain tenant scope, so their visibility can
# be limited to the owner and members of that tenant.
if resource_type == "model" and str(resource.get("model_source") or "").lower() not in {"api", "online"}:
return None
tenant_id = resource.get("tenant_id")
if tenant_id:
return str(tenant_id)
payload = resource.get("payload")
if payload:
try:
data = json.loads(payload) if isinstance(payload, str) else payload
if isinstance(data, dict) and data.get("tenant_id"):
return str(data["tenant_id"])
except (TypeError, ValueError, json.JSONDecodeError):
pass
return None
def resource_in_user_tenant(resource_type: str, resource: dict[str, Any], user: dict[str, Any]) -> bool:
if is_admin(user):
return True
if resource_type == "model" and str(resource.get("model_source") or "").lower() not in {"api", "online"}:
return True
tenant_id = resource_tenant_id(resource_type, resource)
if not tenant_id:
# A missing tenant is not an implicit shared scope. The migration must
# assign historical rows before ordinary users can access them.
return False
return "*" in user_tenant_ids(user) or tenant_id in user_tenant_ids(user)
def has_resource_access(
resource_type: str,
resource_id: str,
user: dict[str, Any],
permission: str = "read",
) -> bool:
"""
检查用户对某资源是否有指定权限。
- admin/protected 用户直接放行(旁路)。
- 其他用户检查 acls 表中是否有对应授权。
"""
if permission not in RESOURCE_ACTIONS:
return False
permission = RESOURCE_ACTION_ALIASES.get(permission, permission)
if is_admin(user):
return True
store = get_platform_store()
resource = resource_record(resource_type, resource_id)
if not resource:
return False
# Some historical datasets were created before tenant_id was introduced.
# They are not shared by default, but an explicit user ACL granted by an
# administrator is still a valid compatibility path for that legacy row.
legacy_unscoped = resource_type != "model" and not resource_tenant_id(resource_type, resource)
if not resource_in_user_tenant(resource_type, resource, user) and not legacy_unscoped:
return False
if resource_type == "model":
model_source = str(resource.get("model_source") or "").lower()
if permission in {"read", "execute"}:
# Local/registered base models are platform resources. For online
# models, only administrator-created records are global; a normal
# user's online model is shared with the active tenant below.
if model_source not in {"api", "online"}:
return True
creator_id = str(resource.get("created_by") or "")
if creator_id:
with store.connect() as conn:
creator = conn.execute(
"SELECT platform_role, role, protected FROM users WHERE id=?",
(creator_id,),
).fetchone()
if creator and (
creator.get("platform_role") == "platform_admin"
or creator.get("role") == "admin"
or bool(creator.get("protected"))
):
return True
# The tenant check above has already rejected models outside the
# user's active tenant. Same-tenant online models are usable.
if resource_tenant_id(resource_type, resource) in user_tenant_ids(user):
return True
if is_tenant_admin_for_resource(resource_type, resource, user):
return True
acls = store.get_acl(resource_type, resource_id)
user_id = user.get("id")
# ACL role principals are tenant roles, not platform roles. Falling back
# to the platform role keeps compatibility with old ACL records.
membership = tenant_membership(user, resource_tenant_id(resource_type, resource))
user_role = (membership or {}).get("role") or user.get("role")
table_info = OWNER_TABLES.get(resource_type)
if table_info and user_id:
table, column = table_info
with store.connect() as conn:
row = conn.execute(f"SELECT {column} FROM {table} WHERE id=?", (resource_id,)).fetchone()
if row:
owner = row[column]
if column == "payload":
try:
owner = json.loads(owner or "{}").get("created_by")
except (TypeError, ValueError):
owner = None
if owner == user_id:
return True
for entry in acls:
if entry.get("revoked_at"):
continue
if entry.get("expires_at"):
from datetime import datetime, timezone
try:
if datetime.fromisoformat(str(entry["expires_at"]).replace("Z", "+00:00")) <= datetime.now(timezone.utc):
continue
except ValueError:
pass
# 按 user 授权
if entry.get("principal_type") == "user" and entry.get("principal_id") == user_id:
if _permission_covers(entry.get("permission"), permission):
return True
# 按 role 授权
if entry.get("principal_type") == "role" and entry.get("principal_id") == user_role:
if _permission_covers(entry.get("permission"), permission):
return True
return False
def _permission_covers(granted: str | None, required: str) -> bool:
"""权限覆盖判断write/execute 覆盖 readadmin 覆盖一切。"""
if not granted:
return False
if granted == "admin":
return True
if granted == required:
return True
# write 覆盖 read
if required == "read" and granted in ("write", "execute"):
return True
return False
def filter_accessible_resource_ids(
resource_type: str,
all_ids: list[str],
user: dict[str, Any],
) -> list[str]:
"""
从全部资源 ID 中过滤出当前用户可访问的 ID 列表。
- admin 直接返回全部。
- 普通用户查 acls 表取交集。
"""
if is_admin(user):
return all_ids
if not all_ids:
return []
accessible = filter_accessible_resource_ids_batch(resource_type, all_ids, user)
return [rid for rid in all_ids if rid in accessible]
def filter_accessible_resource_ids_batch(
resource_type: str,
resource_ids: list[str],
user: dict[str, Any],
) -> set[str]:
"""Filter a list endpoint with a bounded set of SQL queries.
The previous implementation called ``has_resource_access`` once per
resource. Each call loaded the resource, tenant membership and ACL again,
which made ordinary-user list pages slow on a remote PostgreSQL server.
"""
if is_admin(user):
return set(resource_ids)
ids = list(dict.fromkeys(str(item) for item in resource_ids if item))
if not ids:
return set()
table = RESOURCE_TABLES.get(resource_type)
if not table:
return set()
placeholders = ",".join("?" for _ in ids)
primary_tenant = str(user.get("tenant_id") or "default")
tenant_ids = {primary_tenant}
tenant_roles: dict[str, str] = {primary_tenant: str(user.get("role") or "")}
user_id = str(user.get("id") or "")
store = get_platform_store()
with store.connect() as conn:
memberships = conn.execute(
"SELECT tenant_id, role FROM tenant_members "
"WHERE user_id=? AND status='active' "
"AND (expires_at IS NULL OR expires_at='' OR expires_at > NOW()::text)",
(user_id,),
).fetchall()
for membership in memberships:
tenant_id = str(membership["tenant_id"] or "")
if tenant_id:
tenant_ids.add(tenant_id)
tenant_roles[tenant_id] = str(membership["role"] or "")
rows = conn.execute(
f"SELECT * FROM {table} WHERE id IN ({placeholders})", tuple(ids)
).fetchall()
row_by_id = {str(row["id"]): dict(row) for row in rows}
acl_rows = conn.execute(
"SELECT resource_id, principal_type, principal_id, permission, expires_at "
f"FROM acls WHERE resource_type=? AND resource_id IN ({placeholders}) "
"AND (revoked_at IS NULL OR revoked_at='')",
(resource_type, *ids),
).fetchall()
acl_by_resource: dict[str, list[dict[str, Any]]] = {}
for row in acl_rows:
acl_by_resource.setdefault(str(row["resource_id"]), []).append(dict(row))
allowed: set[str] = set()
for resource_id in ids:
resource = row_by_id.get(resource_id)
if not resource or resource.get("deleted_at"):
continue
if resource_type == "model":
model_source = str(resource.get("model_source") or "").lower()
if model_source not in {"api", "online"}:
allowed.add(resource_id)
continue
creator_id = str(resource.get("created_by") or "")
if creator_id:
with store.connect() as conn:
creator = conn.execute(
"SELECT platform_role, role, protected FROM users WHERE id=?",
(creator_id,),
).fetchone()
if creator and (
creator.get("platform_role") == "platform_admin"
or creator.get("role") == "admin"
or bool(creator.get("protected"))
):
allowed.add(resource_id)
continue
tenant_id = resource_tenant_id(resource_type, resource)
if tenant_id and tenant_id in tenant_ids:
allowed.add(resource_id)
continue
tenant_id = resource_tenant_id(resource_type, resource)
if not tenant_id:
# Legacy rows without a tenant are never implicitly visible. Only
# a direct user ACL can expose one to its explicitly named user;
# role ACLs remain blocked until the row is tenant-migrated.
for entry in acl_by_resource.get(resource_id, []):
expires_at = entry.get("expires_at")
if expires_at:
try:
from datetime import datetime, timezone
if datetime.fromisoformat(str(expires_at).replace("Z", "+00:00")) <= datetime.now(timezone.utc):
continue
except ValueError:
continue
if (
entry.get("principal_type") == "user"
and entry.get("principal_id") == user_id
and _permission_covers(entry.get("permission"), "read")
):
allowed.add(resource_id)
break
continue
if tenant_id not in tenant_ids:
continue
if tenant_roles.get(tenant_id) in {"owner", "admin"}:
allowed.add(resource_id)
continue
owner = resource.get("created_by")
if resource_type in {"eval", "fine-tune", "fine_tune_task", "compare", "inference"} and resource.get("payload"):
try:
payload = json.loads(resource["payload"]) if isinstance(resource["payload"], str) else resource["payload"]
if isinstance(payload, dict) and payload.get("created_by"):
owner = payload["created_by"]
except (TypeError, ValueError, json.JSONDecodeError):
pass
if owner == user_id:
allowed.add(resource_id)
continue
role = tenant_roles.get(tenant_id) or str(user.get("role") or "")
for entry in acl_by_resource.get(resource_id, []):
expires_at = entry.get("expires_at")
if expires_at:
try:
from datetime import datetime, timezone
if datetime.fromisoformat(str(expires_at).replace("Z", "+00:00")) <= datetime.now(timezone.utc):
continue
except ValueError:
continue
user_match = entry.get("principal_type") == "user" and entry.get("principal_id") == user_id
role_match = entry.get("principal_type") == "role" and entry.get("principal_id") == role
if (user_match or role_match) and _permission_covers(entry.get("permission"), "read"):
allowed.add(resource_id)
break
return allowed

View File

@@ -0,0 +1,46 @@
"""集中管理项目本地缓存目录HuggingFace / tiktoken
所有 Python 库通过环境变量引用 ``<repo_root>/.cache/{huggingface,tiktoken}``
避免写入用户家目录,也避免不同部署路径(本地 / Docker下缓存位置不一致。
"""
from __future__ import annotations
import os
from pathlib import Path
def repo_root() -> Path:
# backend/app/core/cache_paths.py -> backend -> 仓库根目录
return Path(__file__).resolve().parents[3]
def cache_root() -> Path:
return repo_root() / ".cache"
def huggingface_cache_dir() -> Path:
return cache_root() / "huggingface"
def tiktoken_cache_dir() -> Path:
return cache_root() / "tiktoken"
def setup_local_caches() -> None:
"""进程启动时统一设置 HF / tiktoken 缓存环境变量,并确保目录存在。
必须在 import ``docling`` / ``tiktoken`` 等依赖之前调用,否则首次使用会
仍然走到默认 ``~/.cache`` 路径。
"""
hf_dir = huggingface_cache_dir()
tiktoken_dir = tiktoken_cache_dir()
hf_dir.mkdir(parents=True, exist_ok=True)
(hf_dir / "hub").mkdir(parents=True, exist_ok=True)
tiktoken_dir.mkdir(parents=True, exist_ok=True)
os.environ["HF_HOME"] = str(hf_dir)
os.environ["HUGGINGFACE_HUB_CACHE"] = str(hf_dir / "hub")
os.environ["HF_HUB_CACHE"] = str(hf_dir / "hub")
os.environ["TIKTOKEN_CACHE_DIR"] = str(tiktoken_dir)

View File

@@ -1,326 +1,114 @@
from __future__ import annotations from dataclasses import dataclass
from dataclasses import dataclass, field
from functools import lru_cache from functools import lru_cache
import os import os
from pathlib import Path
from typing import Any from typing import Any
from urllib.parse import quote, urlsplit, urlunsplit
import yaml try:
from pathlib import Path as _Path
from dotenv import load_dotenv
# 显式指定 backend 目录下的 .env并强制覆盖已有环境变量
# 确保远程数据库配置生效,不被本地默认值或残留环境变量影响。
_env_path = _Path(__file__).resolve().parent.parent.parent / ".env"
load_dotenv(dotenv_path=_env_path, override=True)
except ImportError:
pass
BACKEND_ROOT = Path(__file__).resolve().parents[2] def _int_env(name: str, default: int) -> int:
DEFAULT_CONFIG_PATH = BACKEND_ROOT / "config.yaml" raw = os.getenv(name)
DEFAULT_DATABASE_BASE_URL = "postgresql+psycopg://localhost:5432/yg_ft" if raw is None or raw == "":
DEFAULT_DATABASE_USERNAME = "yg_ft"
DEFAULT_DATABASE_PASSWORD = "change_me"
class ConfigurationError(ValueError):
"""Raised when the backend configuration cannot be parsed or validated."""
def _resolve_config_path(config_path: str | Path | None = None) -> tuple[Path, bool]:
configured_path = config_path or os.getenv("BACKEND_CONFIG_FILE")
is_explicit = configured_path is not None
path = Path(configured_path).expanduser() if configured_path else DEFAULT_CONFIG_PATH
if not path.is_absolute():
path = BACKEND_ROOT / path
return path.resolve(), is_explicit
def _load_yaml(config_path: str | Path | None = None) -> tuple[dict[str, Any], Path]:
path, is_explicit = _resolve_config_path(config_path)
if not path.exists():
if is_explicit:
raise ConfigurationError(f"Backend config file does not exist: {path}")
return {}, path
try:
loaded = yaml.safe_load(path.read_text(encoding="utf-8"))
except yaml.YAMLError as exc:
raise ConfigurationError(f"Invalid YAML in backend config file {path}: {exc}") from exc
if loaded is None:
return {}, path
if not isinstance(loaded, dict):
raise ConfigurationError(f"Backend config root must be a mapping: {path}")
return loaded, path
def _yaml_value(config: dict[str, Any], section: str, key: str, default: Any) -> Any:
section_value = config.get(section, {})
if section_value is None:
return default return default
if not isinstance(section_value, dict): return int(raw)
raise ConfigurationError(f"Config section '{section}' must be a mapping")
return section_value.get(key, default)
def _string_setting( def _list_env(name: str, default: list[str]) -> list[str]:
env_name: str, raw = os.getenv(name)
config: dict[str, Any], if raw is None or raw.strip() == "":
section: str, return default
key: str, return [item.strip() for item in raw.split(",") if item.strip()]
default: str,
) -> str:
env_value = os.getenv(env_name)
if env_value is not None:
return env_value
value = _yaml_value(config, section, key, default)
return default if value is None else str(value)
def _int_setting( def _bool_env(name: str, default: bool) -> bool:
env_name: str, raw = os.getenv(name)
config: dict[str, Any], if raw is None or raw.strip() == "":
section: str, return default
key: str, return raw.strip().lower() in {"1", "true", "yes", "on"}
default: int,
) -> int:
env_value = os.getenv(env_name)
value = (
env_value
if env_value not in (None, "")
else _yaml_value(config, section, key, default)
)
if isinstance(value, (bool, float)):
raise ConfigurationError(f"Config value {section}.{key} must be an integer")
try:
return int(value)
except (TypeError, ValueError) as exc:
raise ConfigurationError(f"Config value {section}.{key} must be an integer") from exc
def _port_setting( def docs_kwargs(enabled: bool) -> dict[str, Any]:
env_name: str, """Swagger UI / ReDoc / OpenAPI schema 路由开关。
config: dict[str, Any],
key: str,
default: int,
) -> int:
port = _int_setting(env_name, config, "server", key, default)
if not 1 <= port <= 65535:
raise ConfigurationError(
f"Config value server.{key} must be between 1 and 65535"
)
return port
关闭时 FastAPI 不注册 /docs、/redoc、/openapi.json访问一律返回 404
def _list_setting( 避免未授权访问泄露 API 结构。
env_name: str, """
config: dict[str, Any], if enabled:
section: str, return {}
key: str, return {"docs_url": None, "redoc_url": None, "openapi_url": None}
default: list[str],
) -> list[str]:
env_value = os.getenv(env_name)
value = (
env_value
if env_value and env_value.strip()
else _yaml_value(config, section, key, default)
)
if isinstance(value, str):
return [item.strip() for item in value.split(",") if item.strip()]
if isinstance(value, list):
return [str(item).strip() for item in value if str(item).strip()]
raise ConfigurationError(
f"Config value {section}.{key} must be a list or comma-separated string"
)
def _path_setting(
env_name: str,
config: dict[str, Any],
section: str,
key: str,
default: str,
config_dir: Path,
) -> str:
env_value = os.getenv(env_name)
if env_value is not None:
return env_value
value = Path(str(_yaml_value(config, section, key, default))).expanduser()
if not value.is_absolute():
value = config_dir / value
return str(value.resolve())
def _build_database_url(base_url: str, username: str, password: str) -> str:
try:
parsed = urlsplit(base_url)
host = parsed.hostname
port = parsed.port
except ValueError as exc:
raise ConfigurationError("Config value database.url is not a valid URL") from exc
if not parsed.scheme or not host:
raise ConfigurationError(
"Config value database.url must include a scheme and host"
)
if parsed.scheme not in {"postgresql", "postgresql+psycopg"}:
raise ConfigurationError(
"Config value database.url must use postgresql or postgresql+psycopg"
)
if parsed.path in {"", "/"}:
raise ConfigurationError("Config value database.url must include a database name")
if not username:
raise ConfigurationError("Config value database.username cannot be empty")
if not password:
raise ConfigurationError("Config value database.password cannot be empty")
formatted_host = f"[{host}]" if ":" in host and not host.startswith("[") else host
host_and_port = f"{formatted_host}:{port}" if port is not None else formatted_host
credentials = f"{quote(username, safe='')}:{quote(password, safe='')}"
parts = (
parsed.scheme,
f"{credentials}@{host_and_port}",
parsed.path,
parsed.query,
parsed.fragment,
)
return urlunsplit(parts)
def _database_url(config: dict[str, Any]) -> str:
# 保留原有完整连接串环境变量,便于 Docker 和生产环境注入密钥。
complete_url_override = os.getenv("DATABASE_URL")
if complete_url_override is not None:
return complete_url_override
base_url = _string_setting(
"DATABASE_BASE_URL",
config,
"database",
"url",
DEFAULT_DATABASE_BASE_URL,
)
username = _string_setting(
"DATABASE_USERNAME",
config,
"database",
"username",
DEFAULT_DATABASE_USERNAME,
)
password = _string_setting(
"DATABASE_PASSWORD",
config,
"database",
"password",
DEFAULT_DATABASE_PASSWORD,
)
return _build_database_url(base_url, username, password)
@dataclass(frozen=True) @dataclass(frozen=True)
class Settings: class Settings:
app_name: str = "YG Fine-Tune Platform API" app_name: str = os.getenv("APP_NAME", "YG Zhilian API")
app_env: str = "local" app_env: str = os.getenv("APP_ENV", "local")
route_prefix: str = "/modelTF" route_prefix: str = os.getenv("MODELTF_ROUTE_PREFIX", "/modelTF")
app_mode: str = "local" app_mode: str = os.getenv("APP_MODE", "local")
frontend_port: int = 16801 database_url: str = os.getenv("DATABASE_URL", "postgresql+psycopg://yg_ft:change_me@localhost:15432/yg_ft")
backend_port: int = 17861 cors_allow_origins: list[str] = None # type: ignore[assignment]
database_url: str = field( compute_mode: str = os.getenv("COMPUTE_MODE", "real")
default="postgresql+psycopg://yg_ft:change_me@localhost:5432/yg_ft", compute_status_sync_mode: str = os.getenv("COMPUTE_STATUS_SYNC_MODE", "polling")
repr=False, compute_poll_interval_seconds: int = _int_env("COMPUTE_POLL_INTERVAL_SECONDS", 3)
) compute_request_timeout_seconds: int = _int_env("COMPUTE_REQUEST_TIMEOUT_SECONDS", 5)
cors_allow_origins: list[str] = field( minio_enabled: bool = _bool_env("MINIO_ENABLED", False)
default_factory=lambda: [ # MinIO is an independent service and may run on another host. The
"http://localhost:16801", # endpoint must therefore be reachable from the Backend container.
"http://127.0.0.1:16801", minio_endpoint: str = os.getenv("MINIO_ENDPOINT", "http://host.docker.internal:19000")
] minio_access_key: str = os.getenv("MINIO_ACCESS_KEY", "minioadmin")
) minio_secret_key: str = os.getenv("MINIO_SECRET_KEY", "minioadmin")
compute_mode: str = "real" minio_bucket: str = os.getenv("MINIO_BUCKET", "yg-ft-resources")
compute_status_sync_mode: str = "polling" minio_secure: bool = _bool_env("MINIO_SECURE", False)
compute_poll_interval_seconds: int = 3 # Small text/data files stay inline in PostgreSQL to avoid unnecessary
log_level: str = "INFO" # MinIO round trips. Larger files remain the shared canonical objects.
log_dir: str = "./logs" minio_inline_max_bytes: int = _int_env("MINIO_INLINE_MAX_BYTES", 256 * 1024)
log_file_prefix: str = "backend" minio_presign_max_bytes: int = _int_env("MINIO_PRESIGN_MAX_BYTES", 1024 * 1024 * 1024 * 1024)
log_error_file_prefix: str = "error" storage_wait_seconds: int = _int_env("STORAGE_WAIT_SECONDS", 300)
log_max_bytes: int = 20 * 1024 * 1024 storage_check_interval_seconds: int = _int_env("STORAGE_CHECK_INTERVAL_SECONDS", 10)
log_retention_days: int = 10 compute_service_token: str = os.getenv("COMPUTE_SERVICE_TOKEN", "")
log_level: str = os.getenv("LOG_LEVEL", "INFO")
log_dir: str = os.getenv("LOG_DIR", "./logs")
log_file_prefix: str = os.getenv("LOG_FILE_PREFIX", "backend")
log_error_file_prefix: str = os.getenv("LOG_ERROR_FILE_PREFIX", "error")
log_max_bytes: int = _int_env("LOG_MAX_BYTES", 20 * 1024 * 1024)
log_retention_days: int = _int_env("LOG_RETENTION_DAYS", 10)
enable_docs: bool = None # type: ignore[assignment]
def __post_init__(self) -> None:
def load_settings(config_path: str | Path | None = None) -> Settings: object.__setattr__(
config, resolved_config_path = _load_yaml(config_path) self,
frontend_port = _port_setting( "cors_allow_origins",
"FRONTEND_PORT", config, "frontend_port", Settings.frontend_port _list_env(
) "CORS_ALLOW_ORIGINS",
backend_port = _port_setting( [
"BACKEND_PORT", config, "backend_port", Settings.backend_port "http://localhost:16801",
) "http://127.0.0.1:16801",
if frontend_port == backend_port: "http://localhost:17861",
raise ConfigurationError( "http://127.0.0.1:17861",
"Config values server.frontend_port and server.backend_port must be different" ],
),
)
# Swagger UI / ReDoc / OpenAPI 文档路由开关:
# 未显式配置 ENABLE_DOCS 时,仅本地/开发环境开放,生产环境默认关闭,
# 避免未授权访问泄露 API 结构。从运行时环境读取 APP_ENV而非类定义时
# 缓存的默认值,保证生产默认关闭始终生效且便于测试。
object.__setattr__(
self,
"enable_docs",
_bool_env("ENABLE_DOCS", os.getenv("APP_ENV", "local") != "prod"),
) )
default_cors = [
f"http://localhost:{frontend_port}",
f"http://127.0.0.1:{frontend_port}",
]
return Settings(
app_name=_string_setting("APP_NAME", config, "app", "name", Settings.app_name),
app_env=_string_setting("APP_ENV", config, "app", "env", Settings.app_env),
route_prefix=_string_setting(
"MODELTF_ROUTE_PREFIX", config, "app", "route_prefix", Settings.route_prefix
),
app_mode=_string_setting("APP_MODE", config, "app", "mode", Settings.app_mode),
frontend_port=frontend_port,
backend_port=backend_port,
database_url=_database_url(config),
cors_allow_origins=_list_setting(
"CORS_ALLOW_ORIGINS", config, "app", "cors_allow_origins", default_cors
),
compute_mode=_string_setting(
"COMPUTE_MODE", config, "compute", "mode", Settings.compute_mode
),
compute_status_sync_mode=_string_setting(
"COMPUTE_STATUS_SYNC_MODE",
config,
"compute",
"status_sync_mode",
Settings.compute_status_sync_mode,
),
compute_poll_interval_seconds=_int_setting(
"COMPUTE_POLL_INTERVAL_SECONDS",
config,
"compute",
"poll_interval_seconds",
Settings.compute_poll_interval_seconds,
),
log_level=_string_setting(
"LOG_LEVEL", config, "logging", "level", Settings.log_level
),
log_dir=_path_setting(
"LOG_DIR",
config,
"logging",
"directory",
Settings.log_dir,
resolved_config_path.parent,
),
log_file_prefix=_string_setting(
"LOG_FILE_PREFIX", config, "logging", "file_prefix", Settings.log_file_prefix
),
log_error_file_prefix=_string_setting(
"LOG_ERROR_FILE_PREFIX",
config,
"logging",
"error_file_prefix",
Settings.log_error_file_prefix,
),
log_max_bytes=_int_setting(
"LOG_MAX_BYTES", config, "logging", "max_bytes", Settings.log_max_bytes
),
log_retention_days=_int_setting(
"LOG_RETENTION_DAYS",
config,
"logging",
"retention_days",
Settings.log_retention_days,
),
)
@lru_cache @lru_cache
def get_settings() -> Settings: def get_settings() -> Settings:
return load_settings() return Settings()

View File

@@ -1,58 +1,278 @@
from __future__ import annotations from __future__ import annotations
from contextvars import ContextVar
from datetime import date, datetime, timedelta
import json import json
import logging import logging
import re
import socket
import sys
import time
from contextvars import ContextVar
from datetime import date, datetime, timedelta
from logging import Handler, LogRecord from logging import Handler, LogRecord
from pathlib import Path from pathlib import Path
import re from typing import Any, Callable
import time
from typing import Any
from uuid import uuid4 from uuid import uuid4
from fastapi import FastAPI, Request from fastapi import FastAPI, Request
from app.core.config import Settings, get_settings from app.core.config import Settings, get_settings
# ==================== 链路追踪 ContextVar ====================
request_id_var: ContextVar[str] = ContextVar("request_id", default="-") request_id_var: ContextVar[str] = ContextVar("request_id", default="-")
user_id_var: ContextVar[str] = ContextVar("user_id", default="")
client_ip_var: ContextVar[str] = ContextVar("client_ip", default="")
class RequestIdFilter(logging.Filter): def get_client_ip(request: Request | None) -> str:
"""获取客户端地址,兼容前置反向代理传递的真实地址。"""
if request is None:
return ""
for header in ("X-Real-IP", "X-Forwarded-For"):
value = request.headers.get(header, "")
if value:
return value.split(",", 1)[0].strip()
return request.client.host if request.client else ""
# ==================== 敏感数据脱敏 ====================
SENSITIVE_KEYS: set[str] = {
"password", "token", "access_token", "refresh_token",
"secret_key", "authorization", "bearer", "api_key",
"private_key", "secret", "cookie",
}
FULL_MASK_KEYS: set[str] = {
"password", "token", "access_token", "refresh_token",
"secret_key", "authorization", "bearer", "api_key",
"private_key", "secret", "cookie",
}
def _mask_phone(value: str) -> str:
"""手机号脱敏138****5678"""
if len(value) >= 11:
return value[:3] + "****" + value[-4:]
return value
def _mask_id_card(value: str) -> str:
"""身份证号脱敏110***********1234"""
if len(value) >= 18:
return value[:3] + "***********" + value[-4:]
return value
def mask_value(key: str, value: Any) -> Any:
"""对单个值进行脱敏处理。"""
if value is None:
return ""
key_lower = key.lower()
if key_lower in FULL_MASK_KEYS:
return "***"
str_val = str(value)
# 手机号模式11位数字1开头
if re.match(r"^1[3-9]\d{9}$", str_val):
return _mask_phone(str_val)
# 身份证模式18位
if re.match(r"^\d{17}[\dXx]$", str_val):
return _mask_id_card(str_val)
return value
def mask_sensitive_dict(data: dict) -> dict:
"""递归脱敏字典中的敏感字段。"""
if not data or not isinstance(data, dict):
return data
result: dict[str, Any] = {}
for key, value in data.items():
if isinstance(value, dict):
result[key] = mask_sensitive_dict(value)
elif isinstance(value, list):
result[key] = [
mask_sensitive_dict(item) if isinstance(item, dict) else item
for item in value
]
else:
result[key] = mask_value(key, value)
return result
def mask_sensitive_string(text: str) -> str:
"""从文本中脱敏常见敏感信息。"""
if not text:
return text
# Mask the value as well as the key. Replacing only ``api_key=`` would
# still leak the credential in audit messages and exception text.
assignment_pattern = (
r"(Bearer\s+|(?:api[-_]?key|access[_-]?token|refresh[_-]?token|"
r"secret[_-]?key|password|private[_-]?key|token)\s*[:=]\s*)"
r"(\"[^\"]*\"|'[^']*'|[^\s,;]+)"
)
try:
text = re.sub(assignment_pattern, r"\1***", text, flags=re.IGNORECASE)
except re.error:
pass
patterns: list[tuple[str, str]] = [
(r"Bearer\s+[A-Za-z0-9\-._]+", "Bearer ***"),
(r"(?i)token\s*[:=]\s*\S+", "token=***"),
(r"(?i)password\s*[:=]\s*\S+", "password=***"),
(r"(?i)secret[_-]?key\s*[:=]\s*\S+", "secret_key=***"),
(r"(?i)api[-_]?key\s*[:=]\s*\S+", "api_key=***"),
(r"(?i)private[_-]?key\s*[:=]\s*\S+", "private_key=***"),
(r"(?i)authorization\s*[:=]\s*\S+", "authorization=***"),
]
for pattern, replacement in patterns:
text = re.sub(pattern, replacement, text)
# 手机号脱敏
text = re.sub(r"\b1[3-9]\d{9}\b", lambda m: _mask_phone(m.group()), text)
return text
# ==================== 大对象截断 ====================
MAX_FIELD_SIZE = 1024 # 超过 1KB 的内容自动截断
def truncate_large_value(value: Any, max_size: int = MAX_FIELD_SIZE) -> Any:
"""超过 max_size 的字符串自动截断(前 500 + 后 500"""
if isinstance(value, str) and len(value) > max_size:
half = max_size // 2
return value[:half] + f"...[truncated {len(value) - max_size} chars]..." + value[-half:]
if isinstance(value, dict):
return {k: truncate_large_value(v, max_size) for k, v in value.items()}
if isinstance(value, list):
return [truncate_large_value(v, max_size) for v in value]
return value
# ==================== TraceId Filter ====================
class TraceIdFilter(logging.Filter):
"""自动注入 traceId / userId / clientIp 到每条日志记录。"""
def filter(self, record: LogRecord) -> bool: def filter(self, record: LogRecord) -> bool:
record.request_id = request_id_var.get() record.traceId = request_id_var.get()
record.userId = user_id_var.get("")
record.clientIp = client_ip_var.get("")
record.host = getattr(self, "_host", None) or socket.gethostname()
record.app = getattr(self, "_app", "yg-ft-platform")
record.env = getattr(self, "_env", "dev")
return True return True
def set_context(self, app: str, env: str, host: str) -> None:
self._app = app
self._env = env
self._host = host
# ==================== JSON Formatter ====================
class JsonLogFormatter(logging.Formatter): class JsonLogFormatter(logging.Formatter):
"""Format one JSON object per line for ELK/Filebeat collection.""" """
生产级 JSON 日志格式化器,符合方案文档 §3.2 字段规范。
输出示例:
{
"@timestamp": "2026-08-19T10:30:45.123+08:00",
"level": "INFO",
"logger": "app.api.v1.endpoints.platform",
"traceId": "abc-123-def-456",
"userId": "u_admin",
"message": "数据集创建成功",
"fields": {"datasetId": "ds_001", "costMs": 23},
"file": "platform.py:156",
"thread": "MainThread",
"host": "pod-7x9k2",
"app": "yg-ft-platform",
"env": "dev"
}
"""
# 标准 LogRecord 属性名集合,用于区分 extra 字段
_STD_ATTRS: set[str] = set(vars(logging.LogRecord("", 0, "", 0, "", None, None)).keys()) | {
"traceId", "userId", "clientIp", "host", "app", "env",
"request_id", "user_id", "client_ip",
"asctime", "message", "module", "function", "process",
"thread", "threadName", "levelname", "levelno", "name",
"pathname", "filename", "lineno", "funcName", "created",
"msecs", "relativeCreated", "exc_info", "exc_text",
"stack_info", "msg", "args", "processName", "process",
}
def format(self, record: LogRecord) -> str: def format(self, record: LogRecord) -> str:
# 时间戳ISO8601 带时区
timestamp = datetime.fromtimestamp(record.created).astimezone().isoformat(
timespec="milliseconds"
)
payload: dict[str, Any] = { payload: dict[str, Any] = {
"@timestamp": datetime.fromtimestamp(record.created).astimezone().isoformat( "@timestamp": timestamp,
timespec="milliseconds"
),
"level": record.levelname, "level": record.levelname,
"logger": record.name, "logger": record.name,
"traceId": getattr(record, "traceId", "-"),
"message": record.getMessage(), "message": record.getMessage(),
"module": record.module, "file": f"{Path(record.pathname).name}:{record.lineno}",
"function": record.funcName, "thread": record.threadName,
"file": record.pathname, "host": getattr(record, "host", ""),
"line": record.lineno, "app": getattr(record, "app", ""),
"process": record.process, "env": getattr(record, "env", ""),
"thread": record.thread,
"thread_name": record.threadName,
"request_id": getattr(record, "request_id", "-"),
} }
if record.exc_info:
# userId业务必填未登录可为空
user_id = getattr(record, "userId", "") or getattr(record, "user_id", "")
if user_id:
payload["userId"] = user_id
# clientIp
client_ip = getattr(record, "clientIp", "") or getattr(record, "client_ip", "")
if client_ip:
payload["clientIp"] = client_ip
# 提取结构化业务字段:只收集通过 extra 传入的非标准属性
fields: dict[str, Any] = {}
for attr in dir(record):
if attr.startswith("_"):
continue
if attr in self._STD_ATTRS:
continue
if attr in ("traceId", "userId", "clientIp", "host", "app", "env"):
continue
val = getattr(record, attr, None)
if val is not None and not callable(val):
fields[attr] = truncate_large_value(val)
if fields:
payload["fields"] = mask_sensitive_dict(fields)
# ERROR 级别额外字段
if record.levelname == "ERROR" or record.exc_info:
error_obj: dict[str, Any] = {
"type": type(record.exc_info[1]).__name__ if record.exc_info and record.exc_info[1] else "Error",
"message": record.getMessage(),
}
if record.exc_info:
error_obj["stack_trace"] = self.formatException(record.exc_info)
if record.stack_info:
error_obj["stack_trace"] = self.formatStack(record.stack_info)
payload["error"] = error_obj
# 兼容旧字段名 exception
if record.exc_info and "error" not in payload:
payload["exception"] = self.formatException(record.exc_info) payload["exception"] = self.formatException(record.exc_info)
if record.stack_info:
payload["stack"] = self.formatStack(record.stack_info)
return json.dumps(payload, ensure_ascii=False, separators=(",", ":")) return json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
# ==================== DateSizeRotatingFileHandler ====================
class DateSizeRotatingFileHandler(Handler): class DateSizeRotatingFileHandler(Handler):
"""Rotate log files by date and size while keeping date in every file name.""" """按日期+大小滚动的文件日志处理器。
- 按天创建文件,文件名包含日期
- 单文件超过 max_bytes 时自动滚动(带序号后缀)
- 自动清理超过 retention_days 的旧日志
"""
def __init__( def __init__(
self, self,
@@ -110,10 +330,8 @@ class DateSizeRotatingFileHandler(Handler):
today = date.today() today = date.today()
if not force and self._stream and self._current_date == today: if not force and self._stream and self._current_date == today:
return return
if self._stream and not self._stream.closed: if self._stream and not self._stream.closed:
self._stream.close() self._stream.close()
self._current_date = today self._current_date = today
self._current_path = self._dated_path(today) self._current_path = self._dated_path(today)
self._stream = self._current_path.open("a", encoding=self.encoding) self._stream = self._current_path.open("a", encoding=self.encoding)
@@ -128,11 +346,9 @@ class DateSizeRotatingFileHandler(Handler):
def _rotate_by_size(self) -> None: def _rotate_by_size(self) -> None:
if not self._current_path or not self._current_path.exists(): if not self._current_path or not self._current_path.exists():
return return
if self._stream and not self._stream.closed: if self._stream and not self._stream.closed:
self._stream.close() self._stream.close()
self._stream = None self._stream = None
stem = self._current_path.stem stem = self._current_path.stem
suffix = self._current_path.suffix suffix = self._current_path.suffix
index = 1 index = 1
@@ -146,7 +362,6 @@ class DateSizeRotatingFileHandler(Handler):
def _cleanup_expired_files(self) -> None: def _cleanup_expired_files(self) -> None:
if self.retention_days <= 0: if self.retention_days <= 0:
return return
cutoff = date.today() - timedelta(days=self.retention_days - 1) cutoff = date.today() - timedelta(days=self.retention_days - 1)
pattern = re.compile( pattern = re.compile(
rf"^{re.escape(self.file_prefix)}-(\d{{4}}-\d{{2}}-\d{{2}})(?:\.\d+)?\.log$" rf"^{re.escape(self.file_prefix)}-(\d{{4}}-\d{{2}}-\d{{2}})(?:\.\d+)?\.log$"
@@ -160,94 +375,308 @@ class DateSizeRotatingFileHandler(Handler):
path.unlink(missing_ok=True) path.unlink(missing_ok=True)
# ==================== StructuredLogger 封装 ====================
class StructuredLogger:
"""
结构化日志记录器,提供符合方案文档 §4.2 的 5W1H 日志接口。
使用方式:
logger = get_structured_logger('app.api.dataset')
logger.info('数据集创建成功', datasetId='ds_001', costMs=23)
"""
def __init__(self, name: str, module: str = ""):
self.logger = logging.getLogger(name)
self.name = name
self.module = module
@property
def trace_id(self) -> str:
return request_id_var.get("-")
def info(self, message: str, **fields: Any) -> None:
self._log(logging.INFO, message, **fields)
def warning(self, message: str, **fields: Any) -> None:
self._log(logging.WARNING, message, **fields)
def error(self, message: str, **fields: Any) -> None:
self._log(logging.ERROR, message, **fields)
def debug(self, message: str, **fields: Any) -> None:
self._log(logging.DEBUG, message, **fields)
def _log(self, level: int, message: str, **fields: Any) -> None:
"""统一日志记录方法,通过 extra 传递结构化字段。"""
extra: dict[str, Any] = {}
if self.module:
extra["module"] = self.module
# 脱敏 + 截断
for k, v in fields.items():
extra[k] = truncate_large_value(v)
self.logger.log(level, message, extra=extra, stack_info=False)
def get_structured_logger(name: str, module: str = "") -> StructuredLogger:
"""获取结构化日志记录器。"""
return StructuredLogger(name, module)
# ==================== 快捷函数 ====================
def get_logger(name: str) -> logging.Logger:
"""获取标准 Python logger。"""
return logging.getLogger(name)
def set_request_id(request_id: str) -> None:
"""设置当前请求的追踪 ID。"""
request_id_var.set(request_id)
def set_user_context(user_id: str = "", client_ip: str = "") -> None:
"""设置当前请求的用户上下文(在鉴权后调用)。"""
if user_id:
user_id_var.set(user_id)
if client_ip:
client_ip_var.set(client_ip)
# ==================== 请求日志中间件 ====================
def setup_request_logging(app: FastAPI) -> None:
"""配置 FastAPI 请求日志中间件,符合方案文档 §五(链路追踪)和 §十(访问日志)。"""
logger = get_logger("app.access")
@app.middleware("http")
async def request_logging_middleware(request: Request, call_next): # type: ignore[no-untyped-def]
# 入口生成 traceId优先使用前端传入的 X-Trace-Id
trace_id = request.headers.get("X-Trace-Id") or request.headers.get("X-Request-ID") or str(uuid4())
token = request_id_var.set(trace_id)
ip_token = client_ip_var.set(get_client_ip(request))
started_at = time.perf_counter()
# 提取客户端 IP
client_ip = "-"
if request.client:
client_ip = request.client.host
# 支持反向代理传递的真实 IP
forwarded_for = request.headers.get("X-Forwarded-For", "")
if forwarded_for:
client_ip = forwarded_for.split(",")[0].strip()
client_ip_var.set(client_ip)
try:
response = await call_next(request)
elapsed_ms = (time.perf_counter() - started_at) * 1000
# 噪声路径降级为 DEBUG健康检查等
noisy_paths = ("/health", "/system-info", "/compute/jobs/", "/model-eval/", "/model-compare/")
log_method = logger.info
if any(request.url.path.endswith(p) or p in request.url.path for p in noisy_paths) and response.status_code < 400:
log_method = logger.debug
if response.status_code >= 500:
log_method = logger.error
elif response.status_code >= 400:
log_method = logger.warning
# 结构化访问日志(中文 message方便直接阅读
log_method(
f"HTTP请求 {request.method} {request.url.path}{response.status_code}(耗时{round(elapsed_ms, 2)}ms",
extra={
"request_method": request.method,
"request_path": request.url.path,
"status_code": response.status_code,
"duration_ms": round(elapsed_ms, 2),
"client_ip": client_ip,
"user_agent": request.headers.get("User-Agent", "")[:200],
},
)
# 5xx 系统错误自动写入操作日志
if response.status_code >= 500:
try:
from app.core.op_log import log_operation, OpModule, OpStatus
log_operation(
module=OpModule.SYSTEM,
action="request",
target_type="api",
target_name=request.url.path,
status=OpStatus.FAILURE,
error_message=f"HTTP {response.status_code} - 系统内部错误",
error_type="HTTPError",
detail=f'{{"method":"{request.method}","path":"{request.url.path}","status":{response.status_code}}}',
func_name="request_logging_middleware",
request=request,
duration_ms=elapsed_ms,
)
except Exception:
pass
response.headers["X-Trace-Id"] = trace_id
response.headers["X-Request-ID"] = trace_id
return response
except Exception:
elapsed_ms = (time.perf_counter() - started_at) * 1000
logger.error(
f"HTTP请求异常 {request.method} {request.url.path}(耗时{round(elapsed_ms, 2)}ms— 服务内部错误",
extra={
"request_method": request.method,
"request_path": request.url.path,
"duration_ms": round(elapsed_ms, 2),
"client_ip": client_ip,
},
exc_info=True,
)
# 未被捕获的异常,写入操作日志
try:
import traceback as _tb
from app.core.op_log import log_operation, OpModule, OpStatus
log_operation(
module=OpModule.SYSTEM,
action="request",
target_type="api",
target_name=request.url.path,
status=OpStatus.FAILURE,
error_message=str(sys.exc_info()[1])[:1000] if sys.exc_info()[1] else "未知异常",
error_type=type(sys.exc_info()[1]).__name__ if sys.exc_info()[1] else "UnknownError",
error_traceback="".join(_tb.format_exception(*sys.exc_info()))[:5000],
func_name="request_logging_middleware",
detail=f'{{"method":"{request.method}","path":"{request.url.path}"}}',
request=request,
duration_ms=elapsed_ms,
)
except Exception:
pass
raise
finally:
request_id_var.reset(token)
client_ip_var.reset(ip_token)
# ==================== 配置函数 ====================
def configure_logging(settings: Settings | None = None) -> None: def configure_logging(settings: Settings | None = None) -> None:
"""
生产级日志配置,符合方案文档 §二(分类分流)和 §六(性能安全)。
日志分类:
- 业务日志 (app-biz): INFO+ 业务流程(保留 7 天)
- 系统日志 (app-sys): 框架/中间件日志(保留 7 天)
- 访问日志 (app-access): HTTP 请求日志(保留 15 天)
- 错误日志 (app-error): ERROR 级别(保留 30 天)
"""
settings = settings or get_settings() settings = settings or get_settings()
root_logger = logging.getLogger() root_logger = logging.getLogger()
root_logger.handlers.clear() root_logger.handlers.clear()
root_logger.setLevel(settings.log_level.upper()) root_logger.setLevel(settings.log_level.upper())
# ---- Formatter ----
console_formatter = logging.Formatter( console_formatter = logging.Formatter(
fmt=( fmt=(
"%(asctime)s | %(levelname)s | pid=%(process)d | %(threadName)s | " "%(asctime)s | %(levelname)s | pid=%(process)d | %(threadName)s | "
"request_id=%(request_id)s | %(name)s | %(pathname)s:%(lineno)d | %(message)s" "traceId=%(traceId)s | %(name)s | %(pathname)s:%(lineno)d | %(message)s"
), ),
datefmt="%Y-%m-%d %H:%M:%S", datefmt="%Y-%m-%d %H:%M:%S",
) )
json_formatter = JsonLogFormatter() json_formatter = JsonLogFormatter()
request_filter = RequestIdFilter()
# ---- TraceIdFilter全局注入 traceId/userId/host/app/env----
trace_filter = TraceIdFilter()
trace_filter.set_context(
app=settings.app_name,
env=settings.app_env,
host=socket.gethostname(),
)
# ---- 控制台 Handler ----
console_handler = logging.StreamHandler() console_handler = logging.StreamHandler()
console_handler.setFormatter(console_formatter) console_handler.setFormatter(console_formatter)
console_handler.addFilter(request_filter) console_handler.addFilter(trace_filter)
file_handler = DateSizeRotatingFileHandler( # ---- 业务日志文件 Handler (app-biz) ----
biz_file_handler = DateSizeRotatingFileHandler(
log_dir=settings.log_dir, log_dir=settings.log_dir,
file_prefix=settings.log_file_prefix, file_prefix="app-biz",
max_bytes=settings.log_max_bytes, max_bytes=settings.log_max_bytes,
retention_days=settings.log_retention_days, retention_days=7,
) )
file_handler.setFormatter(json_formatter) biz_file_handler.setFormatter(json_formatter)
file_handler.addFilter(request_filter) biz_file_handler.addFilter(trace_filter)
# ---- 访问日志文件 Handler (app-access) ----
access_file_handler = DateSizeRotatingFileHandler(
log_dir=settings.log_dir,
file_prefix="app-access",
max_bytes=settings.log_max_bytes,
retention_days=15,
)
access_file_handler.setFormatter(json_formatter)
access_file_handler.addFilter(trace_filter)
# ---- 错误日志文件 Handler (app-error) ----
error_file_handler = DateSizeRotatingFileHandler( error_file_handler = DateSizeRotatingFileHandler(
log_dir=settings.log_dir, log_dir=settings.log_dir,
file_prefix=settings.log_error_file_prefix, file_prefix="app-error",
max_bytes=settings.log_max_bytes, max_bytes=settings.log_max_bytes,
retention_days=settings.log_retention_days, retention_days=30,
) )
error_file_handler.setLevel(logging.ERROR) error_file_handler.setLevel(logging.ERROR)
error_file_handler.setFormatter(json_formatter) error_file_handler.setFormatter(json_formatter)
error_file_handler.addFilter(request_filter) error_file_handler.addFilter(trace_filter)
# ---- 注册 Handler ----
root_logger.addHandler(console_handler) root_logger.addHandler(console_handler)
root_logger.addHandler(file_handler) root_logger.addHandler(biz_file_handler)
root_logger.addHandler(access_file_handler)
root_logger.addHandler(error_file_handler) root_logger.addHandler(error_file_handler)
# ---- 访问日志 Logger 独立路由到访问日志文件 ----
access_logger = logging.getLogger("app.access")
access_logger.propagate = False # 不向 root 传播,避免重复写入业务日志
access_logger.addHandler(console_handler)
access_logger.addHandler(access_file_handler)
# 访问日志中的 ERROR 也要进错误日志
access_logger.addHandler(error_file_handler)
# ---- 框架类 Logger 降级 ----
for logger_name in ("uvicorn", "uvicorn.error", "uvicorn.access"): for logger_name in ("uvicorn", "uvicorn.error", "uvicorn.access"):
logger = logging.getLogger(logger_name) lg = logging.getLogger(logger_name)
logger.handlers.clear() lg.handlers.clear()
logger.propagate = True lg.propagate = True
# 框架类日志归入系统日志,生产环境设为 WARN
logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
logging.getLogger("psycopg.pool").setLevel(logging.ERROR)
logging.getLogger("httpx").setLevel(logging.WARNING)
def get_logger(name: str) -> logging.Logger: # ---- 兼容旧文件前缀(向后兼容)----
return logging.getLogger(name) # 如果配置了旧的 log_file_prefix也创建一个对应的 handler
if settings.log_file_prefix and settings.log_file_prefix != "app-biz":
legacy_file_handler = DateSizeRotatingFileHandler(
log_dir=settings.log_dir,
file_prefix=settings.log_file_prefix,
max_bytes=settings.log_max_bytes,
retention_days=settings.log_retention_days,
)
legacy_file_handler.setFormatter(json_formatter)
legacy_file_handler.addFilter(trace_filter)
root_logger.addHandler(legacy_file_handler)
# 旧错误日志前缀兼容
def set_request_id(request_id: str) -> None: if settings.log_error_file_prefix and settings.log_error_file_prefix != "app-error":
request_id_var.set(request_id) legacy_error_handler = DateSizeRotatingFileHandler(
log_dir=settings.log_dir,
file_prefix=settings.log_error_file_prefix,
def setup_request_logging(app: FastAPI) -> None: max_bytes=settings.log_max_bytes,
logger = get_logger("app.access") retention_days=settings.log_retention_days,
)
@app.middleware("http") legacy_error_handler.setLevel(logging.ERROR)
async def request_logging_middleware(request: Request, call_next): # type: ignore[no-untyped-def] legacy_error_handler.setFormatter(json_formatter)
request_id = request.headers.get("X-Request-ID") or str(uuid4()) legacy_error_handler.addFilter(trace_filter)
token = request_id_var.set(request_id) root_logger.addHandler(legacy_error_handler)
started_at = time.perf_counter()
try:
response = await call_next(request)
elapsed_ms = (time.perf_counter() - started_at) * 1000
logger.info(
"request completed method=%s path=%s status_code=%s duration_ms=%.2f client=%s",
request.method,
request.url.path,
response.status_code,
elapsed_ms,
request.client.host if request.client else "-",
)
response.headers["X-Request-ID"] = request_id
return response
except Exception:
elapsed_ms = (time.perf_counter() - started_at) * 1000
logger.exception(
"request failed method=%s path=%s duration_ms=%.2f client=%s",
request.method,
request.url.path,
elapsed_ms,
request.client.host if request.client else "-",
)
raise
finally:
request_id_var.reset(token)

503
backend/app/core/op_log.py Normal file
View File

@@ -0,0 +1,503 @@
"""
操作日志工具模块
提供 @op_log 装饰器和 log_operation 函数,用于记录用户在各业务模块的详细操作。
自动捕获成功/失败状态、完整报错堆栈、操作耗时等。
核心设计:
- 失败操作必须清晰记录完整异常堆栈traceback
- 记录异常类型(如 RuntimeError / ValueError / ConnectionError
- 记录具体出错的函数名和文件位置,方便定位 bug
- 记录 HTTP 状态码,方便区分用户错误(4xx)和系统错误(5xx)
使用示例:
from app.core.op_log import op_log, OpModule, OpAction
@router.post("/inference/start")
@op_log(module=OpModule.INFERENCE, action=OpAction.START, target_type="inference")
async def start_inference(...):
...
"""
from __future__ import annotations
import asyncio
import functools
import json
import time
import traceback
from datetime import datetime, timezone
from typing import Any, Callable, Optional, TypeVar
from fastapi import Request
from app.core.logging import get_logger, get_structured_logger, request_id_var
from app.db.platform_store import get_platform_store, new_id, utcnow
logger = get_logger("app.op_log")
biz_logger = get_structured_logger("app.biz")
F = TypeVar("F", bound=Callable[..., Any])
class OpModule:
"""业务模块常量"""
FINE_TUNE = "fine-tune" # 模型训练
MODEL_EVAL = "model-eval" # 模型评测
INFERENCE = "model-inference" # 模型推理
MODEL_MANAGE = "model-manage" # 模型管理
DATASET = "dataset" # 数据集
DATA_PROCESS = "data-process" # 数据处理
DATA_CONVERT = "data-convert" # 数据类型转换
COMPUTE = "compute" # 算力节点
SYSTEM = "system" # 系统
class OpAction:
"""操作动作常量"""
CREATE = "create"
UPDATE = "update"
DELETE = "delete"
START = "start"
STOP = "stop"
UPLOAD = "upload"
DOWNLOAD = "download"
CONVERT = "convert"
MERGE = "merge"
IMPORT = "import"
LOGIN = "login"
LOGOUT = "logout"
PUBLISH = "publish"
RETRY = "retry"
class OpStatus:
"""操作状态常量"""
SUCCESS = "success"
FAILURE = "failure"
# ==================== 中文映射表(让日志 message 直接可读)====================
MODULE_CN: dict[str, str] = {
"fine-tune": "模型训练",
"model-eval": "模型评测",
"model-inference": "模型推理",
"model-manage": "模型管理",
"dataset": "数据集",
"data-process": "数据处理",
"data-convert": "数据转换",
"compute": "算力节点",
"system": "系统",
}
ACTION_CN: dict[str, str] = {
"create": "创建",
"update": "更新",
"delete": "删除",
"start": "启动",
"stop": "停止",
"upload": "上传",
"download": "下载",
"convert": "转换",
"merge": "合并",
"import": "导入",
"login": "登录",
"logout": "退出登录",
"publish": "发布",
"retry": "重试",
"request": "请求",
}
TARGET_TYPE_CN: dict[str, str] = {
"fine_tune": "训练任务",
"eval": "评测任务",
"inference": "推理任务",
"model": "模型",
"trained_model": "训练产出模型",
"dataset": "数据集",
"dataset_version": "数据集版本",
"data_process_task": "数据处理任务",
"user": "用户",
"api": "接口",
}
def _build_cn_message(
module: str,
action: str,
target_type: str,
target_name: str | None,
target_id: str | None,
status: str,
username: str | None,
error_type: str,
error_message: str,
) -> str:
"""构建中文人类可读的日志消息,格式:[用户] 对 [模块] 执行了 [动作],结果:成功/失败"""
user_part = f"用户[{username}]" if username else "系统"
module_cn = MODULE_CN.get(module, module)
action_cn = ACTION_CN.get(action, action)
target_cn = TARGET_TYPE_CN.get(target_type, target_type or "")
target_label = target_name or target_id or ""
# 拼接操作对象描述
if target_cn and target_label:
target_part = f"{target_cn}{target_label}"
elif target_cn:
target_part = target_cn
elif target_label:
target_part = f"{target_label}"
else:
target_part = ""
if status == OpStatus.SUCCESS:
result = "成功"
msg = f"{user_part} {action_cn}{module_cn}{target_part},结果:成功"
else:
result = "失败"
err_brief = error_message[:120] if error_message else ""
err_part = f"{error_type}: {err_brief}" if error_type and err_brief else f"{error_type}" if error_type else ""
msg = f"{user_part} {action_cn}{module_cn}{target_part},结果:失败{err_part}"
return msg
def op_log(
module: str,
action: str,
target_type: str = "",
*,
target_name_param: str = "name",
detail_params: Optional[list[str]] = None,
) -> Callable[[F], F]:
"""
操作日志装饰器
自动记录:
- 谁在什么时间操作了什么
- 成功还是失败
- 失败时记录完整异常堆栈(traceback)、异常类型、异常消息
- 出错的函数名和文件位置,方便定位 bug
- 操作耗时(ms)
- 客户端 IP、请求路径
Args:
module: 业务模块OpModule 常量)
action: 操作动作OpAction 常量)
target_type: 资源类型
target_name_param: 从 kwargs 中提取目标名称的参数名
detail_params: 需要记录到 detail 的参数名列表
"""
def decorator(func: F) -> F:
func_name = f"{func.__module__}.{func.__qualname__}"
if asyncio.iscoroutinefunction(func):
@functools.wraps(func)
async def async_wrapper(*args, **kwargs):
started_at = time.perf_counter()
trace_id = request_id_var.get("-")
user = _extract_user(args, kwargs)
request = _extract_request(args)
target_name = _get_param(kwargs, target_name_param, "")
target_id = _get_param(kwargs, "task_id", "") or _get_param(kwargs, "dataset_id", "") or _get_param(kwargs, "model_id", "")
detail_dict = _build_detail(detail_params, kwargs)
detail_str = json.dumps(detail_dict, ensure_ascii=False) if detail_dict else ""
try:
result = await func(*args, **kwargs)
elapsed_ms = (time.perf_counter() - started_at) * 1000
if not target_id and isinstance(result, dict):
target_id = str(result.get("id", ""))
_write_log(
module=module,
action=action,
target_type=target_type,
target_id=str(target_id) if target_id else None,
target_name=str(target_name) if target_name else None,
status=OpStatus.SUCCESS,
error_message="",
error_type="",
error_traceback="",
func_name=func_name,
detail=detail_str,
user=user,
request=request,
trace_id=trace_id,
duration_ms=elapsed_ms,
)
return result
except Exception as exc:
elapsed_ms = (time.perf_counter() - started_at) * 1000
# 捕获完整异常堆栈
tb_lines = traceback.format_exception(type(exc), exc, exc.__traceback__)
full_traceback = "".join(tb_lines)
error_msg = str(exc)[:1000]
error_type = type(exc).__name__
_write_log(
module=module,
action=action,
target_type=target_type,
target_id=str(target_id) if target_id else None,
target_name=str(target_name) if target_name else None,
status=OpStatus.FAILURE,
error_message=error_msg,
error_type=error_type,
error_traceback=full_traceback,
func_name=func_name,
detail=detail_str,
user=user,
request=request,
trace_id=trace_id,
duration_ms=elapsed_ms,
)
raise
return async_wrapper # type: ignore
else:
@functools.wraps(func)
def sync_wrapper(*args, **kwargs):
started_at = time.perf_counter()
trace_id = request_id_var.get("-")
user = _extract_user(args, kwargs)
request = _extract_request(args)
target_name = _get_param(kwargs, target_name_param, "")
target_id = _get_param(kwargs, "task_id", "") or _get_param(kwargs, "dataset_id", "") or _get_param(kwargs, "model_id", "")
detail_dict = _build_detail(detail_params, kwargs)
detail_str = json.dumps(detail_dict, ensure_ascii=False) if detail_dict else ""
try:
result = func(*args, **kwargs)
elapsed_ms = (time.perf_counter() - started_at) * 1000
if not target_id and isinstance(result, dict):
target_id = str(result.get("id", ""))
_write_log(
module=module,
action=action,
target_type=target_type,
target_id=str(target_id) if target_id else None,
target_name=str(target_name) if target_name else None,
status=OpStatus.SUCCESS,
error_message="",
error_type="",
error_traceback="",
func_name=func_name,
detail=detail_str,
user=user,
request=request,
trace_id=trace_id,
duration_ms=elapsed_ms,
)
return result
except Exception as exc:
elapsed_ms = (time.perf_counter() - started_at) * 1000
tb_lines = traceback.format_exception(type(exc), exc, exc.__traceback__)
full_traceback = "".join(tb_lines)
error_msg = str(exc)[:1000]
error_type = type(exc).__name__
_write_log(
module=module,
action=action,
target_type=target_type,
target_id=str(target_id) if target_id else None,
target_name=str(target_name) if target_name else None,
status=OpStatus.FAILURE,
error_message=error_msg,
error_type=error_type,
error_traceback=full_traceback,
func_name=func_name,
detail=detail_str,
user=user,
request=request,
trace_id=trace_id,
duration_ms=elapsed_ms,
)
raise
return sync_wrapper # type: ignore
return decorator
def log_operation(
*,
module: str,
action: str,
target_type: str = "",
target_id: str = "",
target_name: str = "",
status: str = OpStatus.SUCCESS,
error_message: str = "",
error_type: str = "",
error_traceback: str = "",
detail: str = "",
func_name: str = "",
user: Optional[dict] = None,
request: Optional[Request] = None,
duration_ms: float = 0,
) -> None:
"""手动记录操作日志(不方便用装饰器时使用)"""
trace_id = request_id_var.get("-")
_write_log(
module=module,
action=action,
target_type=target_type,
target_id=target_id or None,
target_name=target_name or None,
status=status,
error_message=error_message,
error_type=error_type,
error_traceback=error_traceback,
func_name=func_name,
detail=detail,
user=user,
request=request,
trace_id=trace_id,
duration_ms=duration_ms,
)
def _build_detail(detail_params: Optional[list[str]], kwargs: dict) -> dict:
"""从 kwargs 中提取需要记录的参数"""
detail_dict = {}
if detail_params:
for p in detail_params:
val = kwargs.get(p)
if val is not None:
detail_dict[p] = str(val)[:200]
return detail_dict
def _extract_user(args: tuple, kwargs: dict) -> Optional[dict]:
"""从函数参数中提取 current_user dict"""
for arg in args:
if isinstance(arg, dict) and "id" in arg and "username" in arg:
return arg
for v in kwargs.values():
if isinstance(v, dict) and "id" in v and "username" in v:
return v
return None
def _extract_request(args: tuple) -> Optional[Request]:
"""从函数参数中提取 Request 对象"""
for arg in args:
if isinstance(arg, Request):
return arg
return None
def _get_param(kwargs: dict, key: str, default: str = "") -> str:
"""安全获取参数值"""
val = kwargs.get(key, default)
if val is None:
return default
return str(val)
def _write_log(
module: str,
action: str,
target_type: str,
target_id: Optional[str],
target_name: Optional[str],
status: str,
error_message: str,
error_type: str,
error_traceback: str,
func_name: str,
detail: str,
user: Optional[dict],
request: Optional[Request],
trace_id: str,
duration_ms: float,
) -> None:
"""写入操作日志到数据库 + 文件日志"""
user_id = user.get("id") if user else None
username = user.get("username") if user else None
client_ip = None
req_method = None
req_path = None
if request:
client_ip = request.client.host if request.client else None
req_method = request.method
req_path = request.url.path
# ---- 写文件日志app-biz----
cn_msg = _build_cn_message(
module=module, action=action, target_type=target_type,
target_name=target_name, target_id=target_id, status=status,
username=username, error_type=error_type, error_message=error_message,
)
log_fields = {
"bizModule": module,
"action": action,
"targetType": target_type or "",
"targetId": target_id or "",
"targetName": target_name or "",
"opStatus": status,
"durationMs": round(duration_ms, 2),
"username": username or "",
}
if req_method:
log_fields["requestMethod"] = req_method
if req_path:
log_fields["requestPath"] = req_path
if detail:
log_fields["detail"] = detail[:500]
if error_message:
log_fields["errorMessage"] = error_message[:500]
if error_type:
log_fields["errorType"] = error_type
if status == OpStatus.SUCCESS:
biz_logger.info(cn_msg, **log_fields)
elif status == OpStatus.FAILURE:
biz_logger.error(cn_msg, **log_fields)
# ---- 写数据库 ----
try:
store = get_platform_store()
log_id = new_id("op")
with store.connect() as conn:
conn.execute(
"""
INSERT INTO operation_logs
(id, user_id, username, module, action, target_type, target_id,
target_name, status, error_message, error_type, error_traceback,
func_name, detail, client_ip,
request_method, request_path, trace_id, duration_ms, create_time)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""",
(
log_id, user_id, username, module, action,
target_type or None, target_id, target_name,
status,
error_message[:1000] if error_message else None,
error_type or None,
error_traceback[:5000] if error_traceback else None,
func_name or None,
detail[:2000] if detail else None,
client_ip, req_method, req_path, trace_id,
round(duration_ms, 2), utcnow(),
),
)
except Exception:
logger.error("写入操作日志到数据库失败 module=%s action=%s", module, action, exc_info=True)

File diff suppressed because it is too large Load Diff

View File

@@ -1,15 +1,14 @@
from __future__ import annotations from __future__ import annotations
import os
from collections.abc import Generator from collections.abc import Generator
from contextlib import contextmanager from contextlib import contextmanager
from sqlalchemy import create_engine from sqlalchemy import create_engine
from sqlalchemy.orm import Session, sessionmaker from sqlalchemy.orm import Session, sessionmaker
from app.core.config import get_settings
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql+psycopg://yg_ft:change_me@localhost:15432/yg_ft")
DATABASE_URL = get_settings().database_url
engine = create_engine( engine = create_engine(
DATABASE_URL, DATABASE_URL,
@@ -38,3 +37,4 @@ def session_scope() -> Generator[Session, None, None]:
raise raise
finally: finally:
db.close() db.close()

File diff suppressed because it is too large Load Diff

View File

@@ -33,7 +33,49 @@ CREATE TABLE IF NOT EXISTS trained_models (
create_time TEXT NOT NULL, create_time TEXT NOT NULL,
merged INTEGER NOT NULL DEFAULT 0, merged INTEGER NOT NULL DEFAULT 0,
merging INTEGER NOT NULL DEFAULT 0, merging INTEGER NOT NULL DEFAULT 0,
merged_path TEXT merged_path TEXT,
artifact_dir TEXT,
compute_node_id TEXT,
compute_node_name TEXT
);
CREATE TABLE IF NOT EXISTS model_lineage (
id TEXT PRIMARY KEY,
child_resource_type TEXT NOT NULL,
child_resource_id TEXT NOT NULL,
parent_resource_type TEXT NOT NULL,
parent_resource_id TEXT NOT NULL,
relation_type TEXT NOT NULL,
compute_job_id TEXT,
payload TEXT NOT NULL,
create_time TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS model_artifacts (
id TEXT PRIMARY KEY,
model_id TEXT NOT NULL,
model_kind TEXT NOT NULL,
artifact_type TEXT NOT NULL,
path TEXT NOT NULL,
size_bytes BIGINT NOT NULL DEFAULT 0,
checksum_sha256 TEXT,
metadata TEXT NOT NULL,
compute_job_id TEXT,
create_time TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS model_export_jobs (
id TEXT PRIMARY KEY,
trained_model_id TEXT,
compute_job_id TEXT NOT NULL,
node_id TEXT,
export_type TEXT NOT NULL,
quantization_bit INTEGER NOT NULL DEFAULT 0,
status TEXT NOT NULL,
output_dir TEXT,
payload TEXT NOT NULL,
create_time TEXT NOT NULL,
completed_at TEXT
); );
CREATE TABLE IF NOT EXISTS datasets ( CREATE TABLE IF NOT EXISTS datasets (
@@ -76,6 +118,9 @@ CREATE TABLE IF NOT EXISTS compute_nodes (
data_root TEXT NOT NULL, data_root TEXT NOT NULL,
model_root TEXT NOT NULL, model_root TEXT NOT NULL,
log_root TEXT NOT NULL, log_root TEXT NOT NULL,
api_version TEXT NOT NULL DEFAULT 'v1',
capabilities TEXT NOT NULL DEFAULT '[]',
description TEXT,
last_health_check_at TEXT, last_health_check_at TEXT,
health_detail TEXT NOT NULL health_detail TEXT NOT NULL
); );
@@ -88,7 +133,8 @@ CREATE TABLE IF NOT EXISTS gpus (
name TEXT NOT NULL, name TEXT NOT NULL,
memory_total_gb DOUBLE PRECISION NOT NULL, memory_total_gb DOUBLE PRECISION NOT NULL,
power_limit_w DOUBLE PRECISION NOT NULL, power_limit_w DOUBLE PRECISION NOT NULL,
base_temperature INTEGER NOT NULL base_temperature INTEGER NOT NULL,
last_seen_at TEXT
); );
CREATE TABLE IF NOT EXISTS fine_tune_tasks ( CREATE TABLE IF NOT EXISTS fine_tune_tasks (
@@ -103,7 +149,64 @@ CREATE TABLE IF NOT EXISTS fine_tune_tasks (
completed_at TEXT, completed_at TEXT,
compute_node_id TEXT REFERENCES compute_nodes(id) ON DELETE SET NULL, compute_node_id TEXT REFERENCES compute_nodes(id) ON DELETE SET NULL,
gpus TEXT NOT NULL, gpus TEXT NOT NULL,
sync_job_id TEXT sync_job_id TEXT,
compute_job_id TEXT
);
CREATE TABLE IF NOT EXISTS fine_tune_metrics (
id TEXT PRIMARY KEY,
task_id TEXT NOT NULL REFERENCES fine_tune_tasks(id) ON DELETE CASCADE,
step INTEGER NOT NULL,
epoch DOUBLE PRECISION,
loss DOUBLE PRECISION,
grad_norm DOUBLE PRECISION,
learning_rate DOUBLE PRECISION,
raw TEXT NOT NULL,
create_time TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS fine_tune_checkpoints (
id TEXT PRIMARY KEY,
task_id TEXT NOT NULL REFERENCES fine_tune_tasks(id) ON DELETE CASCADE,
step INTEGER NOT NULL,
name TEXT NOT NULL,
path TEXT NOT NULL,
size_bytes BIGINT NOT NULL DEFAULT 0,
create_time TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS compute_jobs (
id TEXT PRIMARY KEY,
task_id TEXT REFERENCES fine_tune_tasks(id) ON DELETE SET NULL,
node_id TEXT REFERENCES compute_nodes(id) ON DELETE SET NULL,
engine TEXT NOT NULL,
status TEXT NOT NULL,
command TEXT NOT NULL,
output_dir TEXT,
log_file TEXT,
payload TEXT NOT NULL,
create_time TEXT NOT NULL,
update_time TEXT NOT NULL,
completed_at TEXT
);
CREATE TABLE IF NOT EXISTS gpu_allocations (
id TEXT PRIMARY KEY,
task_id TEXT REFERENCES fine_tune_tasks(id) ON DELETE CASCADE,
compute_job_id TEXT,
node_id TEXT REFERENCES compute_nodes(id) ON DELETE CASCADE,
gpu_index INTEGER NOT NULL,
status TEXT NOT NULL,
create_time TEXT NOT NULL,
released_at TEXT
);
CREATE TABLE IF NOT EXISTS scheduler_locks (
lock_key TEXT PRIMARY KEY,
owner TEXT NOT NULL,
expires_at TEXT NOT NULL,
create_time TEXT NOT NULL,
update_time TEXT NOT NULL
); );
CREATE TABLE IF NOT EXISTS resource_replicas ( CREATE TABLE IF NOT EXISTS resource_replicas (
@@ -114,6 +217,10 @@ CREATE TABLE IF NOT EXISTS resource_replicas (
local_path TEXT NOT NULL, local_path TEXT NOT NULL,
status TEXT NOT NULL, status TEXT NOT NULL,
sync_status TEXT NOT NULL, sync_status TEXT NOT NULL,
checksum_sha256 TEXT,
byte_size BIGINT NOT NULL DEFAULT 0,
last_checked_at TEXT,
last_error TEXT,
create_time TEXT NOT NULL create_time TEXT NOT NULL
); );
@@ -127,7 +234,140 @@ CREATE TABLE IF NOT EXISTS resource_sync_jobs (
completed_at TEXT completed_at TEXT
); );
CREATE TABLE IF NOT EXISTS storage_objects (
id TEXT PRIMARY KEY,
resource_type TEXT NOT NULL,
resource_id TEXT NOT NULL,
version_id TEXT NOT NULL,
bucket TEXT NOT NULL,
object_key TEXT NOT NULL,
file_name TEXT,
content_type TEXT,
checksum_sha256 TEXT,
byte_size BIGINT NOT NULL DEFAULT 0,
status TEXT NOT NULL DEFAULT 'pending',
created_by TEXT,
create_time TEXT NOT NULL,
UNIQUE (resource_type, resource_id, version_id, object_key)
);
CREATE TABLE IF NOT EXISTS storage_cache_jobs (
id TEXT PRIMARY KEY,
storage_object_id TEXT NOT NULL REFERENCES storage_objects(id) ON DELETE CASCADE,
node_id TEXT NOT NULL REFERENCES compute_nodes(id) ON DELETE CASCADE,
direction TEXT NOT NULL,
status TEXT NOT NULL,
progress INTEGER NOT NULL DEFAULT 0,
local_path TEXT,
error TEXT,
create_time TEXT NOT NULL,
completed_at TEXT
);
CREATE INDEX IF NOT EXISTS idx_storage_objects_resource ON storage_objects(resource_type, resource_id, version_id);
CREATE INDEX IF NOT EXISTS idx_storage_cache_jobs_node_status ON storage_cache_jobs(node_id, status);
CREATE TABLE IF NOT EXISTS eval_tasks (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
payload TEXT NOT NULL,
status TEXT NOT NULL,
create_time TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS eval_dimensions (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
payload TEXT NOT NULL,
is_active INTEGER NOT NULL DEFAULT 1,
is_default INTEGER NOT NULL DEFAULT 0,
create_time TEXT NOT NULL
);
CREATE TABLE IF NOT EXISTS compare_tasks (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
payload TEXT NOT NULL,
status TEXT NOT NULL,
create_time TEXT NOT NULL
);
CREATE INDEX IF NOT EXISTS idx_fine_tune_status ON fine_tune_tasks(status); CREATE INDEX IF NOT EXISTS idx_fine_tune_status ON fine_tune_tasks(status);
CREATE INDEX IF NOT EXISTS idx_fine_tune_compute_job ON fine_tune_tasks(compute_job_id);
CREATE INDEX IF NOT EXISTS idx_fine_tune_compute_node_status ON fine_tune_tasks(compute_node_id, status);
CREATE INDEX IF NOT EXISTS idx_model_lineage_child ON model_lineage(child_resource_type, child_resource_id);
CREATE INDEX IF NOT EXISTS idx_model_lineage_parent ON model_lineage(parent_resource_type, parent_resource_id);
CREATE INDEX IF NOT EXISTS idx_model_artifacts_model ON model_artifacts(model_kind, model_id, artifact_type);
CREATE INDEX IF NOT EXISTS idx_model_export_jobs_model ON model_export_jobs(trained_model_id, create_time DESC);
CREATE INDEX IF NOT EXISTS idx_model_export_jobs_compute ON model_export_jobs(compute_job_id);
CREATE INDEX IF NOT EXISTS idx_fine_tune_metrics_task_step ON fine_tune_metrics(task_id, step);
CREATE UNIQUE INDEX IF NOT EXISTS uq_fine_tune_metrics_task_step_epoch ON fine_tune_metrics(task_id, step, epoch);
CREATE INDEX IF NOT EXISTS idx_fine_tune_checkpoints_task_step ON fine_tune_checkpoints(task_id, step);
CREATE UNIQUE INDEX IF NOT EXISTS uq_fine_tune_checkpoints_task_path ON fine_tune_checkpoints(task_id, path);
CREATE INDEX IF NOT EXISTS idx_compute_jobs_task ON compute_jobs(task_id);
CREATE INDEX IF NOT EXISTS idx_compute_jobs_node_status ON compute_jobs(node_id, status);
CREATE INDEX IF NOT EXISTS idx_gpu_allocations_node_status ON gpu_allocations(node_id, status);
CREATE UNIQUE INDEX IF NOT EXISTS uq_gpu_allocations_active ON gpu_allocations(node_id, gpu_index) WHERE status IN ('allocated','running');
CREATE INDEX IF NOT EXISTS idx_scheduler_locks_expires ON scheduler_locks(expires_at);
CREATE INDEX IF NOT EXISTS idx_dataset_files_dataset ON dataset_files(dataset_id); CREATE INDEX IF NOT EXISTS idx_dataset_files_dataset ON dataset_files(dataset_id);
CREATE INDEX IF NOT EXISTS idx_gpus_node ON gpus(node_id); CREATE INDEX IF NOT EXISTS idx_gpus_node ON gpus(node_id);
CREATE UNIQUE INDEX IF NOT EXISTS uq_gpus_node_index ON gpus(node_id, gpu_index);
CREATE INDEX IF NOT EXISTS idx_replicas_resource ON resource_replicas(resource_type, resource_id); CREATE INDEX IF NOT EXISTS idx_replicas_resource ON resource_replicas(resource_type, resource_id);
CREATE UNIQUE INDEX IF NOT EXISTS uq_replicas_node_resource ON resource_replicas(node_id, resource_type, resource_id);
CREATE INDEX IF NOT EXISTS idx_sync_jobs_node_status ON resource_sync_jobs(target_node_id, status);
CREATE INDEX IF NOT EXISTS idx_eval_tasks_status ON eval_tasks(status);
CREATE INDEX IF NOT EXISTS idx_eval_dimensions_active ON eval_dimensions(is_active);
CREATE INDEX IF NOT EXISTS idx_compare_tasks_status ON compare_tasks(status);
-- ===================== Project / Tenant =====================
CREATE TABLE IF NOT EXISTS projects (
id TEXT PRIMARY KEY,
tenant_id TEXT NOT NULL DEFAULT 'default',
name TEXT NOT NULL,
code TEXT NOT NULL,
description TEXT,
quota TEXT,
status TEXT NOT NULL DEFAULT 'active',
create_time TEXT NOT NULL,
create_by TEXT,
updated_at TEXT
);
CREATE TABLE IF NOT EXISTS project_members (
project_id TEXT NOT NULL REFERENCES projects(id) ON DELETE CASCADE,
user_id TEXT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
role TEXT NOT NULL DEFAULT 'member',
create_time TEXT NOT NULL,
PRIMARY KEY (project_id, user_id)
);
CREATE TABLE IF NOT EXISTS roles (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
permissions TEXT NOT NULL DEFAULT '[]',
create_time TEXT
);
CREATE TABLE IF NOT EXISTS sessions (
id TEXT PRIMARY KEY,
user_id TEXT NOT NULL,
username TEXT,
login_at TEXT,
logout_at TEXT,
duration_seconds INTEGER,
issued_at TEXT,
expires_at TEXT,
ip TEXT,
create_time TEXT
);
CREATE TABLE IF NOT EXISTS acls (
id TEXT PRIMARY KEY,
resource_type TEXT NOT NULL,
resource_id TEXT NOT NULL,
principal_type TEXT NOT NULL,
principal_id TEXT NOT NULL,
permission TEXT NOT NULL,
create_time TEXT
);

View File

@@ -0,0 +1,352 @@
-- Data processing migration.
--
-- IMPORTANT: This file is intentionally NOT wired into application startup.
-- Apply it explicitly in a controlled deployment, or call
-- DataProcessStore.ensure_schema() from an administrative command.
BEGIN;
-- This migration targets the current runtime schema created by
-- 001_platform_runtime.sql. Refuse the UUID/JSONB target-design schema instead
-- of partially altering it with incompatible TEXT foreign keys.
DO $$
DECLARE
datasets_id_type TEXT;
BEGIN
SELECT format_type(a.atttypid, a.atttypmod)
INTO datasets_id_type
FROM pg_attribute a
JOIN pg_class c ON c.oid = a.attrelid
JOIN pg_namespace n ON n.oid = c.relnamespace
WHERE n.nspname = current_schema()
AND c.relname = 'datasets'
AND a.attname = 'id'
AND a.attnum > 0
AND NOT a.attisdropped;
IF datasets_id_type IS NULL THEN
RAISE EXCEPTION '002_data_process.sql requires 001_platform_runtime.sql first';
END IF;
IF datasets_id_type <> 'text' THEN
RAISE EXCEPTION
'002_data_process.sql supports only the current TEXT runtime schema; found datasets.id type %',
datasets_id_type;
END IF;
END $$;
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS source_task_id TEXT;
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS size_bytes BIGINT NOT NULL DEFAULT 0;
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS record_count BIGINT NOT NULL DEFAULT 0;
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS metadata TEXT NOT NULL DEFAULT '{}';
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS project_id TEXT;
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS owner_id TEXT;
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS created_at TIMESTAMPTZ NOT NULL DEFAULT now();
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS updated_at TIMESTAMPTZ NOT NULL DEFAULT now();
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS deleted_at TIMESTAMPTZ;
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS storage_object_id TEXT;
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS current_version_id TEXT;
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS size_bytes BIGINT NOT NULL DEFAULT 0;
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS record_count BIGINT NOT NULL DEFAULT 0;
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS file_format VARCHAR(40);
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS checksum_sha256 CHAR(64);
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS version_no INTEGER NOT NULL DEFAULT 1;
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS source_task_id TEXT;
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS project_id TEXT;
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS metadata TEXT NOT NULL DEFAULT '{}';
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS created_at TIMESTAMPTZ NOT NULL DEFAULT now();
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS updated_at TIMESTAMPTZ NOT NULL DEFAULT now();
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS deleted_at TIMESTAMPTZ;
CREATE TABLE IF NOT EXISTS data_process_tasks (
id TEXT PRIMARY KEY,
name VARCHAR(150) NOT NULL,
description TEXT,
status VARCHAR(20) NOT NULL DEFAULT 'pending'
CHECK (status IN ('pending', 'running', 'completed', 'failed', 'stopped')),
process_type VARCHAR(20) NOT NULL
CHECK (process_type IN ('structured', 'unstructured', 'external')),
source_dataset_id TEXT REFERENCES datasets(id) ON DELETE SET NULL,
output_dataset_id TEXT REFERENCES datasets(id) ON DELETE SET NULL,
config TEXT NOT NULL DEFAULT '{}',
progress NUMERIC(5,2) NOT NULL DEFAULT 0 CHECK (progress >= 0 AND progress <= 100),
input_count BIGINT NOT NULL DEFAULT 0 CHECK (input_count >= 0),
output_count BIGINT NOT NULL DEFAULT 0 CHECK (output_count >= 0),
filtered_count BIGINT NOT NULL DEFAULT 0 CHECK (filtered_count >= 0),
duplicate_count BIGINT NOT NULL DEFAULT 0 CHECK (duplicate_count >= 0),
error_count BIGINT NOT NULL DEFAULT 0 CHECK (error_count >= 0),
failure_reason TEXT,
generation_run_id TEXT,
results_confirmed BOOLEAN NOT NULL DEFAULT TRUE,
workflow_step VARCHAR(20) NOT NULL DEFAULT 'create'
CHECK (workflow_step IN ('create', 'model', 'upload', 'preview', 'generate', 'results')),
preview_status VARCHAR(20) NOT NULL DEFAULT 'idle'
CHECK (preview_status IN ('idle', 'queued', 'running', 'completed', 'failed', 'cancelled')),
preview_progress NUMERIC(5,2) NOT NULL DEFAULT 0
CHECK (preview_progress >= 0 AND preview_progress <= 100),
preview_run_id TEXT,
preview_failure_reason TEXT,
preview_total_files INTEGER NOT NULL DEFAULT 0 CHECK (preview_total_files >= 0),
preview_completed_files INTEGER NOT NULL DEFAULT 0 CHECK (preview_completed_files >= 0),
tenant_id TEXT,
project_id TEXT,
owner_id TEXT,
approval_status VARCHAR(30) NOT NULL DEFAULT 'not_required',
created_by TEXT,
updated_by TEXT,
deleted_by TEXT,
started_at TIMESTAMPTZ,
completed_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
deleted_at TIMESTAMPTZ
);
ALTER TABLE data_process_tasks ADD COLUMN IF NOT EXISTS generation_run_id TEXT;
-- 历史任务在引入六步确认流程前已经完成审核,默认保留为已确认;
-- 新任务由创建接口显式写入 FALSE并在第六步确认后转为 TRUE。
ALTER TABLE data_process_tasks
ADD COLUMN IF NOT EXISTS results_confirmed BOOLEAN NOT NULL DEFAULT TRUE;
UPDATE data_process_tasks
SET results_confirmed=FALSE
WHERE status <> 'completed' AND results_confirmed=TRUE;
-- 先以可空列接入旧库,才能只回填历史行;随后再收紧默认值与约束。
ALTER TABLE data_process_tasks ADD COLUMN IF NOT EXISTS workflow_step VARCHAR(20);
ALTER TABLE data_process_tasks ADD COLUMN IF NOT EXISTS preview_status VARCHAR(20);
ALTER TABLE data_process_tasks ADD COLUMN IF NOT EXISTS preview_progress NUMERIC(5,2);
ALTER TABLE data_process_tasks ADD COLUMN IF NOT EXISTS preview_run_id TEXT;
ALTER TABLE data_process_tasks ADD COLUMN IF NOT EXISTS preview_failure_reason TEXT;
ALTER TABLE data_process_tasks ADD COLUMN IF NOT EXISTS preview_total_files INTEGER;
ALTER TABLE data_process_tasks ADD COLUMN IF NOT EXISTS preview_completed_files INTEGER;
CREATE TEMP TABLE data_process_workflow_backfill_ids ON COMMIT DROP AS
SELECT id FROM data_process_tasks WHERE workflow_step IS NULL;
UPDATE data_process_tasks task
SET workflow_step = CASE
WHEN task.status IN ('running', 'failed', 'stopped') THEN 'generate'
WHEN task.status = 'completed' AND task.results_confirmed=FALSE THEN 'generate'
WHEN task.status = 'completed' THEN 'results'
ELSE 'create'
END
WHERE task.workflow_step IS NULL;
UPDATE data_process_tasks
SET preview_status='idle', preview_progress=0,
preview_total_files=0, preview_completed_files=0
WHERE preview_status IS NULL OR preview_progress IS NULL
OR preview_total_files IS NULL OR preview_completed_files IS NULL;
ALTER TABLE data_process_tasks ALTER COLUMN workflow_step SET DEFAULT 'create';
ALTER TABLE data_process_tasks ALTER COLUMN workflow_step SET NOT NULL;
ALTER TABLE data_process_tasks ALTER COLUMN preview_status SET DEFAULT 'idle';
ALTER TABLE data_process_tasks ALTER COLUMN preview_status SET NOT NULL;
ALTER TABLE data_process_tasks ALTER COLUMN preview_progress SET DEFAULT 0;
ALTER TABLE data_process_tasks ALTER COLUMN preview_progress SET NOT NULL;
ALTER TABLE data_process_tasks ALTER COLUMN preview_total_files SET DEFAULT 0;
ALTER TABLE data_process_tasks ALTER COLUMN preview_total_files SET NOT NULL;
ALTER TABLE data_process_tasks ALTER COLUMN preview_completed_files SET DEFAULT 0;
ALTER TABLE data_process_tasks ALTER COLUMN preview_completed_files SET NOT NULL;
DO $$
BEGIN
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conrelid='data_process_tasks'::regclass
AND conname='ck_data_process_tasks_workflow_step'
) THEN
ALTER TABLE data_process_tasks ADD CONSTRAINT ck_data_process_tasks_workflow_step
CHECK (workflow_step IN ('create', 'model', 'upload', 'preview', 'generate', 'results'));
END IF;
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conrelid='data_process_tasks'::regclass
AND conname='ck_data_process_tasks_preview_status'
) THEN
ALTER TABLE data_process_tasks ADD CONSTRAINT ck_data_process_tasks_preview_status
CHECK (preview_status IN ('idle', 'queued', 'running', 'completed', 'failed', 'cancelled'));
END IF;
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conrelid='data_process_tasks'::regclass
AND conname='ck_data_process_tasks_preview_progress'
) THEN
ALTER TABLE data_process_tasks ADD CONSTRAINT ck_data_process_tasks_preview_progress
CHECK (preview_progress >= 0 AND preview_progress <= 100);
END IF;
IF NOT EXISTS (
SELECT 1 FROM pg_constraint
WHERE conrelid='data_process_tasks'::regclass
AND conname='ck_data_process_tasks_preview_file_counts'
) THEN
ALTER TABLE data_process_tasks ADD CONSTRAINT ck_data_process_tasks_preview_file_counts
CHECK (preview_total_files >= 0 AND preview_completed_files >= 0
AND preview_completed_files <= preview_total_files);
END IF;
END $$;
CREATE UNIQUE INDEX IF NOT EXISTS uq_data_process_tasks_name_alive
ON data_process_tasks(name) WHERE deleted_at IS NULL;
CREATE INDEX IF NOT EXISTS idx_data_process_tasks_scope_status
ON data_process_tasks(tenant_id, project_id, status, created_at DESC)
WHERE deleted_at IS NULL;
CREATE INDEX IF NOT EXISTS idx_data_process_tasks_creator_created
ON data_process_tasks(created_by, created_at DESC) WHERE deleted_at IS NULL;
CREATE TABLE IF NOT EXISTS data_process_source_files (
id TEXT PRIMARY KEY,
task_id TEXT NOT NULL REFERENCES data_process_tasks(id) ON DELETE CASCADE,
storage_object_id TEXT,
name TEXT NOT NULL,
size_bytes BIGINT NOT NULL DEFAULT 0 CHECK (size_bytes >= 0),
record_count BIGINT NOT NULL DEFAULT 0 CHECK (record_count >= 0),
file_format VARCHAR(40),
checksum_sha256 CHAR(64) NOT NULL,
version_no INTEGER NOT NULL DEFAULT 1 CHECK (version_no > 0),
content TEXT NOT NULL,
content_preview TEXT,
metadata TEXT NOT NULL DEFAULT '{}',
tenant_id TEXT,
project_id TEXT,
created_by TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
deleted_at TIMESTAMPTZ
);
CREATE INDEX IF NOT EXISTS idx_data_process_source_files_task
ON data_process_source_files(task_id, created_at) WHERE deleted_at IS NULL;
CREATE UNIQUE INDEX IF NOT EXISTS uq_data_process_source_checksum_alive
ON data_process_source_files(task_id, checksum_sha256) WHERE deleted_at IS NULL;
CREATE TABLE IF NOT EXISTS data_process_preview_items (
id TEXT PRIMARY KEY,
task_id TEXT NOT NULL REFERENCES data_process_tasks(id) ON DELETE CASCADE,
source_file_id TEXT REFERENCES data_process_source_files(id) ON DELETE CASCADE,
original_content TEXT NOT NULL DEFAULT '',
edited_content TEXT NOT NULL DEFAULT '',
source_start INTEGER CHECK (source_start IS NULL OR source_start >= 0),
source_end INTEGER CHECK (source_end IS NULL OR source_end >= 0),
source_start_line INTEGER CHECK (source_start_line IS NULL OR source_start_line > 0),
source_end_line INTEGER CHECK (source_end_line IS NULL OR source_end_line > 0),
token_count INTEGER NOT NULL DEFAULT 0 CHECK (token_count >= 0),
status VARCHAR(20) NOT NULL DEFAULT 'original'
CHECK (status IN ('original', 'modified', 'manual', 'invalid')),
quality_score TEXT NOT NULL DEFAULT '{}',
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now(),
CHECK (source_start IS NULL OR source_end IS NULL OR source_end >= source_start),
CHECK (source_start_line IS NULL OR source_end_line IS NULL OR source_end_line >= source_start_line)
);
CREATE INDEX IF NOT EXISTS idx_data_process_preview_task_file
ON data_process_preview_items(task_id, source_file_id, created_at);
-- 子表在新库中到这里才存在;只修复本次新增 workflow_step 前的历史任务。
UPDATE data_process_tasks task
SET workflow_step = CASE
WHEN EXISTS (
SELECT 1 FROM data_process_preview_items preview
WHERE preview.task_id=task.id
) THEN 'preview'
WHEN EXISTS (
SELECT 1 FROM data_process_source_files source_file
WHERE source_file.task_id=task.id AND source_file.deleted_at IS NULL
) THEN 'upload'
ELSE task.workflow_step
END
WHERE task.id IN (SELECT id FROM data_process_workflow_backfill_ids)
AND task.status='pending';
CREATE TABLE IF NOT EXISTS data_process_results (
id TEXT PRIMARY KEY,
task_id TEXT NOT NULL REFERENCES data_process_tasks(id) ON DELETE CASCADE,
preview_item_id TEXT REFERENCES data_process_preview_items(id) ON DELETE SET NULL,
instruction TEXT NOT NULL,
input TEXT NOT NULL DEFAULT '',
output TEXT NOT NULL,
chosen TEXT NOT NULL DEFAULT '',
rejected TEXT NOT NULL DEFAULT '',
original_instruction TEXT,
original_input TEXT,
original_output TEXT,
original_chosen TEXT,
original_rejected TEXT,
status VARCHAR(20) NOT NULL DEFAULT 'valid'
CHECK (status IN ('valid', 'modified', 'invalid')),
error TEXT,
split VARCHAR(20) CHECK (split IS NULL OR split IN ('train', 'validation', 'test')),
quality_score TEXT NOT NULL DEFAULT '{}',
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
ALTER TABLE data_process_results ADD COLUMN IF NOT EXISTS chosen TEXT NOT NULL DEFAULT '';
ALTER TABLE data_process_results ADD COLUMN IF NOT EXISTS rejected TEXT NOT NULL DEFAULT '';
ALTER TABLE data_process_results ADD COLUMN IF NOT EXISTS original_chosen TEXT;
ALTER TABLE data_process_results ADD COLUMN IF NOT EXISTS original_rejected TEXT;
CREATE INDEX IF NOT EXISTS idx_data_process_results_task_status
ON data_process_results(task_id, status, id);
CREATE INDEX IF NOT EXISTS idx_data_process_results_task_split
ON data_process_results(task_id, split);
CREATE TABLE IF NOT EXISTS dataset_file_versions (
id TEXT PRIMARY KEY,
dataset_file_id TEXT NOT NULL REFERENCES dataset_files(id) ON DELETE CASCADE,
version_no INTEGER NOT NULL CHECK (version_no > 0),
storage_object_id TEXT NOT NULL,
content_preview TEXT,
description TEXT,
base_version_id TEXT REFERENCES dataset_file_versions(id) ON DELETE SET NULL,
size_bytes BIGINT NOT NULL DEFAULT 0 CHECK (size_bytes >= 0),
record_count BIGINT NOT NULL DEFAULT 0 CHECK (record_count >= 0),
checksum_sha256 CHAR(64) NOT NULL,
source_task_id TEXT REFERENCES data_process_tasks(id) ON DELETE SET NULL,
metadata TEXT NOT NULL DEFAULT '{}',
created_by TEXT,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
ALTER TABLE dataset_file_versions ADD COLUMN IF NOT EXISTS source_task_id TEXT;
ALTER TABLE dataset_file_versions ADD COLUMN IF NOT EXISTS metadata TEXT NOT NULL DEFAULT '{}';
CREATE UNIQUE INDEX IF NOT EXISTS uq_dataset_file_versions_no_002
ON dataset_file_versions(dataset_file_id, version_no);
CREATE INDEX IF NOT EXISTS idx_dataset_file_versions_source_task_002
ON dataset_file_versions(source_task_id) WHERE source_task_id IS NOT NULL;
CREATE TABLE IF NOT EXISTS dataset_records (
id TEXT PRIMARY KEY,
dataset_id TEXT NOT NULL REFERENCES datasets(id) ON DELETE CASCADE,
dataset_file_id TEXT REFERENCES dataset_files(id) ON DELETE CASCADE,
version_id TEXT REFERENCES dataset_file_versions(id) ON DELETE CASCADE,
line_no INTEGER,
split VARCHAR(20) CHECK (split IS NULL OR split IN ('train', 'validation', 'test')),
instruction TEXT,
input TEXT,
output TEXT,
raw TEXT NOT NULL DEFAULT '{}',
status VARCHAR(20) NOT NULL DEFAULT 'valid'
CHECK (status IN ('valid', 'modified', 'invalid')),
source_task_id TEXT REFERENCES data_process_tasks(id) ON DELETE SET NULL,
source_result_id TEXT REFERENCES data_process_results(id) ON DELETE SET NULL,
preview_item_id TEXT REFERENCES data_process_preview_items(id) ON DELETE SET NULL,
created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);
ALTER TABLE dataset_records ADD COLUMN IF NOT EXISTS source_task_id TEXT;
ALTER TABLE dataset_records ADD COLUMN IF NOT EXISTS source_result_id TEXT;
ALTER TABLE dataset_records ADD COLUMN IF NOT EXISTS preview_item_id TEXT;
CREATE INDEX IF NOT EXISTS idx_dataset_records_dataset_002
ON dataset_records(dataset_id, id);
CREATE INDEX IF NOT EXISTS idx_dataset_records_source_task_002
ON dataset_records(source_task_id, source_result_id);
CREATE INDEX IF NOT EXISTS idx_datasets_source_task_002
ON datasets(source_task_id) WHERE source_task_id IS NOT NULL;
CREATE INDEX IF NOT EXISTS idx_dataset_files_source_task_002
ON dataset_files(source_task_id) WHERE source_task_id IS NOT NULL;
COMMIT;

View File

@@ -0,0 +1,111 @@
-- 平台治理:租户 / 审批 / 审计(字段以 platform_store 实际写入为准)
CREATE TABLE IF NOT EXISTS tenants (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
code TEXT,
status TEXT DEFAULT 'active',
owner_user_id TEXT,
quota TEXT,
retention_policy_id TEXT,
create_time TEXT
);
CREATE TABLE IF NOT EXISTS approval_templates (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
steps TEXT,
create_time TEXT
);
CREATE TABLE IF NOT EXISTS approval_instances (
id TEXT PRIMARY KEY,
template_id TEXT,
resource_type TEXT,
resource_id TEXT,
applicant_id TEXT,
status TEXT DEFAULT 'pending',
current_step INTEGER DEFAULT 0,
create_time TEXT
);
CREATE TABLE IF NOT EXISTS approval_steps (
id TEXT PRIMARY KEY,
instance_id TEXT,
step_index INTEGER,
approver_id TEXT,
status TEXT DEFAULT 'pending',
comment TEXT,
time TEXT
);
CREATE TABLE IF NOT EXISTS audit_logs (
id TEXT PRIMARY KEY,
tenant_id TEXT,
project_id TEXT,
actor_id TEXT,
action TEXT,
target_type TEXT,
target_id TEXT,
detail TEXT,
client_ip TEXT,
time TEXT
);
-- 幂等升级 audit_logs 表:新增字段(已存在则跳过)
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS trace_id TEXT;
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS request_method TEXT;
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS request_path TEXT;
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS status_code INTEGER;
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS duration_ms REAL;
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS extra JSONB;
-- 索引(已存在则跳过)
CREATE INDEX IF NOT EXISTS idx_audit_trace ON audit_logs(trace_id);
CREATE INDEX IF NOT EXISTS idx_audit_actor_time ON audit_logs(actor_id, time);
CREATE INDEX IF NOT EXISTS idx_audit_target ON audit_logs(target_type, target_id);
CREATE TABLE IF NOT EXISTS retention_policies (
id TEXT PRIMARY KEY,
name TEXT NOT NULL,
scope TEXT,
rule TEXT,
status TEXT DEFAULT 'active',
create_time TEXT,
create_by TEXT,
updated_at TEXT
);
-- ===================== 操作日志表 =====================
-- 记录用户在每个业务模块的详细操作(成功/失败、报错信息等)
CREATE TABLE IF NOT EXISTS operation_logs (
id TEXT PRIMARY KEY,
user_id TEXT, -- 操作用户 ID
username TEXT, -- 操作用户名(冗余,便于查询)
module TEXT, -- 业务模块fine-tune/model-eval/model-inference/dataset/data-process/data-convert/model-manage
action TEXT, -- 具体动作create/start/stop/delete/upload/convert 等
target_type TEXT, -- 资源类型
target_id TEXT, -- 资源 ID
target_name TEXT, -- 资源名称(便于阅读)
status TEXT NOT NULL, -- success / failure
error_message TEXT, -- 失败时的报错信息
detail TEXT, -- 操作详情 JSON参数摘要
client_ip TEXT, -- 客户端 IP
request_method TEXT, -- HTTP 方法
request_path TEXT, -- 请求路径
trace_id TEXT, -- 链路追踪 ID
duration_ms REAL, -- 耗时(ms)
create_time TEXT -- 操作时间
);
-- 幂等升级 operation_logs 表:新增字段(已存在则跳过)
ALTER TABLE operation_logs ADD COLUMN IF NOT EXISTS error_type TEXT; -- 异常类型RuntimeError / ValueError / ConnectionError
ALTER TABLE operation_logs ADD COLUMN IF NOT EXISTS error_traceback TEXT; -- 完整异常堆栈
ALTER TABLE operation_logs ADD COLUMN IF NOT EXISTS func_name TEXT; -- 出错的函数名
-- 索引
CREATE INDEX IF NOT EXISTS idx_op_user_time ON operation_logs(user_id, create_time DESC);
CREATE INDEX IF NOT EXISTS idx_op_module_time ON operation_logs(module, create_time DESC);
CREATE INDEX IF NOT EXISTS idx_op_status ON operation_logs(status, create_time DESC);
CREATE INDEX IF NOT EXISTS idx_op_action ON operation_logs(action);
CREATE INDEX IF NOT EXISTS idx_op_create_time ON operation_logs(create_time DESC);
CREATE INDEX IF NOT EXISTS idx_op_error_type ON operation_logs(error_type);

View File

@@ -0,0 +1,20 @@
-- 003_model_path_governance
-- 模型路径治理:增加 can_train 标识,区分本地可训练模型与 API / 远程模型。
-- 训练预检阶段依赖该字段拦截不适合 LLaMA-Factory 本地训练的基座模型。
-- 1. models 表增加 can_train默认 0后设搬迁为 1 的规则如下)
ALTER TABLE models ADD COLUMN IF NOT EXISTS can_train INTEGER NOT NULL DEFAULT 0;
-- 2. 将已有模型按规则推定 can_train
-- - path 非空 且 model_source != 'api' → 可训练 (1)
-- - 其余 → 不可训练 (0)
UPDATE models
SET can_train = CASE
WHEN path IS NOT NULL AND path != '' AND model_source IS NOT NULL AND model_source != 'api' THEN 1
ELSE 0
END;
-- 3. 给 trained_models 增加 artifact_dir训练产物目录扫描结果目录
ALTER TABLE trained_models ADD COLUMN IF NOT EXISTS artifact_dir TEXT;
ALTER TABLE trained_models ADD COLUMN IF NOT EXISTS compute_node_id TEXT;
ALTER TABLE trained_models ADD COLUMN IF NOT EXISTS compute_node_name TEXT;

View File

@@ -0,0 +1,3 @@
-- 租户配额与保留策略扩展(如后续治理表需补列,可在此追加)
ALTER TABLE tenants ADD COLUMN IF NOT EXISTS gpu_quota TEXT;
ALTER TABLE tenants ADD COLUMN IF NOT EXISTS storage_quota TEXT;

View File

@@ -0,0 +1,22 @@
-- ============================================================
-- 权限体系扩展GPU 分配表 + 资源所有权字段
-- ============================================================
-- GPU 分配表:管理员指定哪些用户可以使用哪些 GPU 卡
CREATE TABLE IF NOT EXISTS gpu_assignments (
id TEXT PRIMARY KEY,
node_id TEXT NOT NULL REFERENCES compute_nodes(id) ON DELETE CASCADE,
gpu_index INTEGER NOT NULL,
user_id TEXT NOT NULL,
assigned_by TEXT,
assigned_at TEXT NOT NULL,
UNIQUE (node_id, gpu_index, user_id)
);
CREATE INDEX IF NOT EXISTS idx_gpu_assignments_user ON gpu_assignments(user_id);
CREATE INDEX IF NOT EXISTS idx_gpu_assignments_gpu ON gpu_assignments(node_id, gpu_index);
-- 资源所有权字段:用户创建的数据集/模型/训练产物/评测任务
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE models ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE trained_models ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE eval_tasks ADD COLUMN IF NOT EXISTS created_by TEXT;

View File

@@ -0,0 +1,32 @@
-- YG Fine-Tune Platform additive migration: MinIO completion and cache manifest.
-- Execute with psql against an existing database after taking a schema backup.
-- The statements are idempotent and are also included in 000_full_init.sql.
BEGIN;
ALTER TABLE resource_replicas ADD COLUMN IF NOT EXISTS version_id TEXT;
ALTER TABLE resource_replicas ADD COLUMN IF NOT EXISTS storage_object_id TEXT;
ALTER TABLE users ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default';
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS metadata TEXT NOT NULL DEFAULT '{}';
ALTER TABLE model_artifacts ADD COLUMN IF NOT EXISTS storage_object_id TEXT;
ALTER TABLE model_artifacts ADD COLUMN IF NOT EXISTS storage_backend TEXT NOT NULL DEFAULT 'minio';
ALTER TABLE dataset_files ADD COLUMN IF NOT EXISTS storage_object_id TEXT;
ALTER TABLE data_convert_tasks ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE data_convert_tasks ADD COLUMN IF NOT EXISTS storage_backend TEXT NOT NULL DEFAULT 'minio';
ALTER TABLE data_convert_tasks ADD COLUMN IF NOT EXISTS output_storage_object_id TEXT;
ALTER TABLE data_convert_tasks ADD COLUMN IF NOT EXISTS output_content TEXT;
ALTER TABLE eval_tasks ADD COLUMN IF NOT EXISTS deleted_at TEXT;
ALTER TABLE eval_tasks ADD COLUMN IF NOT EXISTS deleted_by TEXT;
ALTER TABLE eval_tasks ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE eval_tasks ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE eval_tasks ADD COLUMN IF NOT EXISTS project_id TEXT;
ALTER TABLE eval_tasks ADD COLUMN IF NOT EXISTS report_storage_object_id TEXT;
CREATE INDEX IF NOT EXISTS idx_resource_replicas_object_005
ON resource_replicas(storage_object_id, node_id);
CREATE INDEX IF NOT EXISTS idx_storage_objects_resource_005
ON storage_objects(resource_type, resource_id, version_id, status);
CREATE INDEX IF NOT EXISTS idx_eval_tasks_active_005
ON eval_tasks(create_time DESC) WHERE deleted_at IS NULL;
COMMIT;

View File

@@ -0,0 +1,21 @@
-- 原子 GPU 预留:评测和推理与训练统一纳入调度占用模型。
BEGIN;
CREATE TABLE IF NOT EXISTS gpu_reservations (
id TEXT PRIMARY KEY,
node_id TEXT NOT NULL REFERENCES compute_nodes(id) ON DELETE CASCADE,
gpu_index INTEGER NOT NULL,
owner_type TEXT NOT NULL CHECK (owner_type IN ('eval', 'inference')),
owner_id TEXT NOT NULL,
status TEXT NOT NULL DEFAULT 'reserved' CHECK (status IN ('reserved', 'released')),
create_time TEXT NOT NULL,
released_at TEXT
);
CREATE INDEX IF NOT EXISTS idx_gpu_reservations_owner
ON gpu_reservations(owner_type, owner_id, status);
CREATE INDEX IF NOT EXISTS idx_gpu_reservations_node_status
ON gpu_reservations(node_id, status);
CREATE UNIQUE INDEX IF NOT EXISTS uq_gpu_reservations_active
ON gpu_reservations(node_id, gpu_index) WHERE status = 'reserved';
COMMIT;

View File

@@ -0,0 +1,33 @@
-- 平台可靠性闭环:导出权限、对象清理、缓存 manifest 元数据。
BEGIN;
ALTER TABLE model_export_jobs ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE model_export_jobs ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE model_export_jobs ADD COLUMN IF NOT EXISTS archive_status TEXT NOT NULL DEFAULT 'pending';
ALTER TABLE model_export_jobs ADD COLUMN IF NOT EXISTS archive_error TEXT;
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS deleted_at TEXT;
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS deleted_by TEXT;
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS cleanup_attempts INTEGER NOT NULL DEFAULT 0;
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS last_cleanup_error TEXT;
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS last_verified_at TEXT;
ALTER TABLE storage_cache_jobs ADD COLUMN IF NOT EXISTS version_id TEXT;
ALTER TABLE storage_cache_jobs ADD COLUMN IF NOT EXISTS checksum_sha256 TEXT;
ALTER TABLE storage_cache_jobs ADD COLUMN IF NOT EXISTS byte_size BIGINT NOT NULL DEFAULT 0;
ALTER TABLE storage_cache_jobs ADD COLUMN IF NOT EXISTS last_accessed_at TEXT;
ALTER TABLE storage_cache_jobs ADD COLUMN IF NOT EXISTS protected_until TEXT;
CREATE TABLE IF NOT EXISTS storage_cleanup_jobs (
id TEXT PRIMARY KEY,
storage_object_id TEXT NOT NULL REFERENCES storage_objects(id) ON DELETE CASCADE,
action TEXT NOT NULL DEFAULT 'delete',
status TEXT NOT NULL DEFAULT 'pending',
attempts INTEGER NOT NULL DEFAULT 0,
error TEXT,
create_time TEXT NOT NULL,
completed_at TEXT
);
CREATE INDEX IF NOT EXISTS idx_storage_cleanup_status ON storage_cleanup_jobs(status, create_time);
COMMIT;

View File

@@ -0,0 +1,53 @@
-- 权限 2.0租户成员、资源申请、ACL 生命周期和审批动作
CREATE TABLE IF NOT EXISTS tenant_members (
tenant_id TEXT NOT NULL,
user_id TEXT NOT NULL,
role TEXT NOT NULL DEFAULT 'member',
status TEXT NOT NULL DEFAULT 'active',
invited_by TEXT,
joined_at TEXT,
expires_at TEXT,
PRIMARY KEY (tenant_id, user_id)
);
CREATE INDEX IF NOT EXISTS idx_tenant_members_user ON tenant_members(user_id, status);
CREATE INDEX IF NOT EXISTS idx_tenant_members_tenant ON tenant_members(tenant_id, status);
CREATE TABLE IF NOT EXISTS resource_access_requests (
id TEXT PRIMARY KEY,
tenant_id TEXT NOT NULL,
resource_type TEXT NOT NULL,
resource_id TEXT NOT NULL,
applicant_id TEXT NOT NULL,
principal_type TEXT NOT NULL DEFAULT 'user',
principal_id TEXT NOT NULL,
requested_permissions TEXT NOT NULL DEFAULT '[]',
reason TEXT,
approval_id TEXT,
status TEXT NOT NULL DEFAULT 'pending',
expires_at TEXT,
created_at TEXT NOT NULL,
decided_at TEXT,
decided_by TEXT
);
CREATE INDEX IF NOT EXISTS idx_resource_access_requests_resource
ON resource_access_requests(resource_type, resource_id, status);
CREATE INDEX IF NOT EXISTS idx_resource_access_requests_applicant
ON resource_access_requests(applicant_id, status);
ALTER TABLE acls ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE acls ADD COLUMN IF NOT EXISTS granted_by TEXT;
ALTER TABLE acls ADD COLUMN IF NOT EXISTS source_request_id TEXT;
ALTER TABLE acls ADD COLUMN IF NOT EXISTS expires_at TEXT;
ALTER TABLE acls ADD COLUMN IF NOT EXISTS revoked_at TEXT;
CREATE INDEX IF NOT EXISTS idx_acls_tenant_active ON acls(tenant_id, revoked_at, expires_at);
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS action TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS requested_permissions TEXT NOT NULL DEFAULT '[]';
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS reason TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS expires_at TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS decided_by TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS decided_at TEXT;
ALTER TABLE approval_steps ADD COLUMN IF NOT EXISTS approver_type TEXT NOT NULL DEFAULT 'user';

View File

@@ -0,0 +1,116 @@
-- 权限 2.0 完整闭环:审批策略、执行状态、结构化审计与历史租户归属
-- 该迁移可独立执行:兼容仅执行过 000_full_init.sql、或未执行 008 的旧数据库。
CREATE TABLE IF NOT EXISTS tenant_members (
tenant_id TEXT NOT NULL,
user_id TEXT NOT NULL,
role TEXT NOT NULL DEFAULT 'member',
status TEXT NOT NULL DEFAULT 'active',
invited_by TEXT,
joined_at TEXT,
expires_at TEXT,
PRIMARY KEY (tenant_id, user_id)
);
CREATE INDEX IF NOT EXISTS idx_tenant_members_user ON tenant_members(user_id, status);
CREATE INDEX IF NOT EXISTS idx_tenant_members_tenant ON tenant_members(tenant_id, status);
INSERT INTO tenant_members (tenant_id, user_id, role, status, joined_at)
SELECT COALESCE(u.tenant_id, 'default'), u.id,
CASE WHEN u.role='admin' OR COALESCE(u.protected, 0)=1 THEN 'owner' ELSE 'member' END,
'active', COALESCE(u.create_time, NOW()::text)
FROM users u
ON CONFLICT (tenant_id, user_id) DO NOTHING;
CREATE TABLE IF NOT EXISTS resource_access_requests (
id TEXT PRIMARY KEY,
tenant_id TEXT NOT NULL DEFAULT 'default',
resource_type TEXT NOT NULL,
resource_id TEXT NOT NULL,
applicant_id TEXT NOT NULL,
principal_type TEXT NOT NULL DEFAULT 'user',
principal_id TEXT NOT NULL,
requested_permissions TEXT NOT NULL DEFAULT '[]',
reason TEXT,
approval_id TEXT,
status TEXT NOT NULL DEFAULT 'pending',
expires_at TEXT,
created_at TEXT NOT NULL DEFAULT to_char(now(), 'YYYY-MM-DD"T"HH24:MI:SS.MS"Z"'),
decided_at TEXT,
decided_by TEXT
);
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS action TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS requested_permissions TEXT NOT NULL DEFAULT '[]';
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS reason TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS expires_at TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS decided_by TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS decided_at TEXT;
ALTER TABLE approval_steps ADD COLUMN IF NOT EXISTS approver_type TEXT NOT NULL DEFAULT 'user';
ALTER TABLE acls ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE acls ADD COLUMN IF NOT EXISTS granted_by TEXT;
ALTER TABLE acls ADD COLUMN IF NOT EXISTS source_request_id TEXT;
ALTER TABLE acls ADD COLUMN IF NOT EXISTS expires_at TEXT;
ALTER TABLE acls ADD COLUMN IF NOT EXISTS revoked_at TEXT;
ALTER TABLE datasets ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE models ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE trained_models ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE eval_tasks ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE approval_templates ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE approval_templates ADD COLUMN IF NOT EXISTS action TEXT;
ALTER TABLE approval_templates ADD COLUMN IF NOT EXISTS resource_type TEXT;
ALTER TABLE approval_templates ADD COLUMN IF NOT EXISTS scope TEXT NOT NULL DEFAULT 'tenant';
ALTER TABLE approval_templates ADD COLUMN IF NOT EXISTS status TEXT NOT NULL DEFAULT 'active';
ALTER TABLE approval_templates ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE approval_templates ADD COLUMN IF NOT EXISTS updated_at TEXT;
CREATE INDEX IF NOT EXISTS idx_approval_templates_match
ON approval_templates(tenant_id, action, resource_type, status);
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS execution_status TEXT NOT NULL DEFAULT 'pending';
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS execution_error TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS executed_at TEXT;
ALTER TABLE approval_instances ADD COLUMN IF NOT EXISTS executed_by TEXT;
CREATE INDEX IF NOT EXISTS idx_approval_instances_execution
ON approval_instances(status, execution_status, action, resource_type, resource_id);
ALTER TABLE resource_access_requests ADD COLUMN IF NOT EXISTS cancelled_at TEXT;
ALTER TABLE resource_access_requests ADD COLUMN IF NOT EXISTS cancelled_by TEXT;
CREATE INDEX IF NOT EXISTS idx_resource_access_requests_expiry
ON resource_access_requests(status, expires_at);
ALTER TABLE projects ADD COLUMN IF NOT EXISTS deleted_at TEXT;
ALTER TABLE projects ADD COLUMN IF NOT EXISTS deleted_by TEXT;
CREATE INDEX IF NOT EXISTS idx_projects_active_tenant ON projects(tenant_id, status, deleted_at);
ALTER TABLE tenants ADD COLUMN IF NOT EXISTS deleted_at TEXT;
ALTER TABLE tenants ADD COLUMN IF NOT EXISTS deleted_by TEXT;
CREATE INDEX IF NOT EXISTS idx_tenants_active ON tenants(status, deleted_at);
ALTER TABLE data_convert_tasks ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE data_convert_tasks ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default';
CREATE INDEX IF NOT EXISTS idx_data_convert_tasks_tenant ON data_convert_tasks(tenant_id, deleted_at);
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS result TEXT NOT NULL DEFAULT 'success';
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS reason TEXT;
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS request_id TEXT;
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS session_id TEXT;
ALTER TABLE audit_logs ADD COLUMN IF NOT EXISTS metadata TEXT;
CREATE INDEX IF NOT EXISTS idx_audit_request ON audit_logs(request_id);
CREATE INDEX IF NOT EXISTS idx_audit_result ON audit_logs(result, time);
-- 历史资源按创建者租户补齐归属。无法识别的资源保留 default后续由管理员复核。
UPDATE datasets d
SET tenant_id = COALESCE(u.tenant_id, 'default')
FROM users u
WHERE d.created_by = u.id AND (d.tenant_id IS NULL OR d.tenant_id = '');
UPDATE models m
SET tenant_id = COALESCE(u.tenant_id, 'default')
FROM users u
WHERE m.created_by = u.id AND (m.tenant_id IS NULL OR m.tenant_id = '');
UPDATE trained_models m
SET tenant_id = COALESCE(u.tenant_id, 'default')
FROM users u
WHERE m.created_by = u.id AND (m.tenant_id IS NULL OR m.tenant_id = '');
UPDATE eval_tasks e
SET tenant_id = COALESCE(u.tenant_id, 'default')
FROM users u
WHERE e.created_by = u.id AND (e.tenant_id IS NULL OR e.tenant_id = '');

View File

@@ -0,0 +1,21 @@
-- Permission 2.0: atomic tenant quota reservations for GPU-backed tasks.
-- This migration is additive and safe to run repeatedly.
CREATE TABLE IF NOT EXISTS tenant_quota_reservations (
id TEXT PRIMARY KEY,
tenant_id TEXT NOT NULL,
owner_type TEXT NOT NULL,
owner_id TEXT NOT NULL,
gpu_count INTEGER NOT NULL DEFAULT 0,
storage_bytes BIGINT NOT NULL DEFAULT 0,
status TEXT NOT NULL DEFAULT 'reserved',
create_time TEXT NOT NULL,
released_at TEXT
);
CREATE INDEX IF NOT EXISTS idx_tenant_quota_reservations_tenant
ON tenant_quota_reservations(tenant_id, status);
CREATE INDEX IF NOT EXISTS idx_tenant_quota_reservations_owner
ON tenant_quota_reservations(owner_type, owner_id, status);
CREATE UNIQUE INDEX IF NOT EXISTS uq_tenant_quota_reservations_active_owner
ON tenant_quota_reservations(owner_type, owner_id) WHERE status = 'reserved';

View File

@@ -0,0 +1,20 @@
-- Permission 2.0: user and inference task lifecycle tombstones.
-- Additive migration, safe to execute repeatedly.
ALTER TABLE users ADD COLUMN IF NOT EXISTS deleted_at TEXT;
ALTER TABLE users ADD COLUMN IF NOT EXISTS deleted_by TEXT;
ALTER TABLE compare_tasks ADD COLUMN IF NOT EXISTS deleted_at TEXT;
ALTER TABLE compare_tasks ADD COLUMN IF NOT EXISTS deleted_by TEXT;
ALTER TABLE eval_tasks ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE eval_tasks ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE compare_tasks ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE compare_tasks ADD COLUMN IF NOT EXISTS tenant_id TEXT;
ALTER TABLE data_convert_tasks ADD COLUMN IF NOT EXISTS deleted_by TEXT;
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS deleted_at TEXT;
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS deleted_by TEXT;
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS cleanup_attempts INTEGER NOT NULL DEFAULT 0;
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS last_cleanup_error TEXT;
ALTER TABLE storage_objects ADD COLUMN IF NOT EXISTS last_verified_at TEXT;
CREATE INDEX IF NOT EXISTS idx_users_active ON users(status, deleted_at);
CREATE INDEX IF NOT EXISTS idx_compare_tasks_active ON compare_tasks(status, deleted_at);

View File

@@ -0,0 +1,63 @@
-- Tenant/user hierarchy alignment. Additive and safe to run repeatedly.
-- New business flows use tenant_members, platform_role, tenant_id and
-- created_by. Legacy role/project fields remain for compatibility only.
ALTER TABLE users ADD COLUMN IF NOT EXISTS platform_role TEXT NOT NULL DEFAULT 'platform_user';
UPDATE users
SET platform_role = CASE
WHEN role = 'admin' OR COALESCE(protected, 0) = 1 THEN 'platform_admin'
ELSE 'platform_user'
END
WHERE platform_role IS NULL OR platform_role NOT IN ('platform_admin', 'platform_user');
CREATE INDEX IF NOT EXISTS idx_users_platform_role ON users(platform_role, status, deleted_at);
ALTER TABLE fine_tune_tasks ADD COLUMN IF NOT EXISTS tenant_id TEXT NOT NULL DEFAULT 'default';
ALTER TABLE fine_tune_tasks ADD COLUMN IF NOT EXISTS created_by TEXT;
ALTER TABLE fine_tune_tasks ADD COLUMN IF NOT EXISTS deleted_at TEXT;
ALTER TABLE fine_tune_tasks ADD COLUMN IF NOT EXISTS deleted_by TEXT;
UPDATE fine_tune_tasks
SET tenant_id = COALESCE(NULLIF(tenant_id, ''), payload::json->>'tenant_id', 'default'),
created_by = COALESCE(NULLIF(created_by, ''), payload::json->>'created_by')
WHERE tenant_id IS NULL OR tenant_id = '' OR created_by IS NULL OR created_by = '';
CREATE INDEX IF NOT EXISTS idx_fine_tune_tasks_tenant_active
ON fine_tune_tasks(tenant_id, status, deleted_at, create_time DESC);
CREATE INDEX IF NOT EXISTS idx_fine_tune_tasks_creator
ON fine_tune_tasks(created_by, deleted_at, create_time DESC);
INSERT INTO tenants (id, name, code, status, quota, create_time)
VALUES ('default', '默认租户', 'default', 'active', '{}', NOW()::text)
ON CONFLICT (id) DO UPDATE
SET status = CASE WHEN COALESCE(tenants.status, 'active') = 'deleted' THEN 'active' ELSE tenants.status END,
deleted_at = CASE WHEN COALESCE(tenants.status, 'active') = 'deleted' THEN NULL ELSE tenants.deleted_at END,
deleted_by = CASE WHEN COALESCE(tenants.status, 'active') = 'deleted' THEN NULL ELSE tenants.deleted_by END;
-- Make the tenant owner relationship explicit and repair older tenant rows.
INSERT INTO tenant_members (tenant_id, user_id, role, status, joined_at)
SELECT t.id, t.owner_user_id, 'owner', 'active', COALESCE(t.create_time, NOW()::text)
FROM tenants t
JOIN users u ON u.id = t.owner_user_id
WHERE t.owner_user_id IS NOT NULL
AND COALESCE(t.status, 'active') = 'active'
ON CONFLICT (tenant_id, user_id) DO UPDATE
SET role = 'owner', status = 'active';
-- Existing account creation used admin as a tenant owner. Platform role and
-- tenant role are separate, so keep only explicit tenant owners as owners.
UPDATE tenant_members tm
SET role = 'member'
FROM users u
WHERE tm.user_id = u.id
AND tm.role = 'owner'
AND (u.role <> 'admin' AND COALESCE(u.protected, 0) = 0)
AND NOT EXISTS (
SELECT 1 FROM tenants t
WHERE t.id = tm.tenant_id AND t.owner_user_id = tm.user_id
);
-- Project is no longer a business isolation boundary. Keep historical columns
-- readable, but make tenant-scoped queries the only supported new path.
CREATE INDEX IF NOT EXISTS idx_datasets_tenant_active ON datasets(tenant_id, deleted_at, create_time DESC);
CREATE INDEX IF NOT EXISTS idx_models_tenant_active ON models(tenant_id, deleted_at, create_time DESC);
CREATE INDEX IF NOT EXISTS idx_trained_models_tenant_active ON trained_models(tenant_id, deleted_at, create_time DESC);
CREATE INDEX IF NOT EXISTS idx_eval_tasks_tenant_active ON eval_tasks(tenant_id, deleted_at, create_time DESC);
CREATE INDEX IF NOT EXISTS idx_compare_tasks_tenant_active ON compare_tasks(tenant_id, deleted_at, create_time DESC);

View File

@@ -1,16 +1,25 @@
import asyncio
from contextlib import suppress
from fastapi import FastAPI from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware from fastapi.middleware.cors import CORSMiddleware
from app.api.v1.router import api_router from app.core.cache_paths import setup_local_caches
from app.core.config import get_settings
from app.core.logging import configure_logging, setup_request_logging # 在任何 docling / tiktoken 模块被实例化之前设置缓存路径,避免首调用走到 ~/.cache。
setup_local_caches()
from app.api.v1.router import api_router # noqa: E402
from app.core.config import docs_kwargs, get_settings # noqa: E402
from app.core.logging import configure_logging, setup_request_logging # noqa: E402
from app.workers.compute_poller import run_compute_poller # noqa: E402
def create_app() -> FastAPI: def create_app() -> FastAPI:
settings = get_settings() settings = get_settings()
configure_logging(settings) configure_logging(settings)
app = FastAPI(title=settings.app_name) app = FastAPI(title=settings.app_name, **docs_kwargs(settings.enable_docs))
app.add_middleware( app.add_middleware(
CORSMiddleware, CORSMiddleware,
allow_origins=settings.cors_allow_origins, allow_origins=settings.cors_allow_origins,
@@ -20,6 +29,19 @@ def create_app() -> FastAPI:
) )
setup_request_logging(app) setup_request_logging(app)
app.include_router(api_router, prefix=settings.route_prefix) app.include_router(api_router, prefix=settings.route_prefix)
@app.on_event("startup")
async def start_workers() -> None:
app.state.compute_poller_task = asyncio.create_task(run_compute_poller())
@app.on_event("shutdown")
async def stop_workers() -> None:
task = getattr(app.state, "compute_poller_task", None)
if task:
task.cancel()
with suppress(asyncio.CancelledError):
await task
return app return app

View File

@@ -0,0 +1,396 @@
from __future__ import annotations
import json
from fastapi import APIRouter, Body, Depends
from typing import Any
from app.api.v1.endpoints.platform import ok, fail
from app.db.platform_store import get_platform_store
from app.core.auth import (
get_current_user,
has_resource_access,
is_tenant_admin,
is_admin,
resource_in_user_tenant,
resource_record,
)
router = APIRouter(prefix="/approvals", tags=["approval"])
# Approval actions are deliberately finite. A caller must not be able to
# create an arbitrary approval that no executor or audit policy understands.
ALLOWED_APPROVAL_ACTIONS = {
"resource.access",
"gpu.assign",
"tenant.quota.update",
"tenant.member.add",
"tenant.member.remove",
"model.use",
"dataset.use",
"dataset.delete",
"trained_model.merge",
"trained_model.export",
"trained_model.delete",
"fine_tune.stop",
"fine_tune.delete",
"eval.delete",
"inference.delete",
"project.archive",
"project.delete",
}
def _available_gpu_options() -> list[dict[str, Any]]:
"""Return only online nodes and currently unassigned, idle GPUs."""
store = get_platform_store()
nodes = store.compute_nodes()
gpus = store.gpus()
assigned = {(str(item.get("node_id")), int(item.get("gpu_index"))) for item in store.gpu_assignments()}
by_node: dict[str, list[dict[str, Any]]] = {}
for gpu in gpus:
node_id = str(gpu.get("node_id") or "")
index = int(gpu.get("id") or 0)
if gpu.get("status") != "idle" or (node_id, index) in assigned:
continue
by_node.setdefault(node_id, []).append({
"index": index,
"name": gpu.get("name") or "GPU",
"memory_total_gb": float(gpu.get("memory_total_gb") or 0),
})
result = []
for node in nodes:
node_id = str(node.get("id") or "")
if not node.get("enabled") or node.get("scheduler_status") != "online" or not by_node.get(node_id):
continue
result.append({
"id": node_id,
"code": node.get("code") or node_id,
"name": node.get("name") or node.get("code") or node_id,
"gpus": sorted(by_node[node_id], key=lambda item: item["index"]),
})
return result
def _validate_gpu_request(assignments: Any) -> list[dict[str, Any]]:
if not isinstance(assignments, list) or not assignments:
raise fail(400, "assignments 不能为空")
available = {
(node["id"], gpu["index"])
for node in _available_gpu_options()
for gpu in node["gpus"]
}
normalized = []
seen: set[tuple[str, int]] = set()
for item in assignments:
if not isinstance(item, dict) or not item.get("node_id") or item.get("gpu_index") is None:
raise fail(400, "每项必须包含 node_id 和 gpu_index")
try:
key = (str(item["node_id"]), int(item["gpu_index"]))
except (TypeError, ValueError):
raise fail(400, "gpu_index 必须是整数")
if key not in available:
raise fail(409, f"GPU {key[0]}:{key[1]} 当前不可申请")
if key not in seen:
seen.add(key)
normalized.append({"node_id": key[0], "gpu_index": key[1]})
return normalized
@router.get("/templates")
def list_templates(current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
if not is_admin(current_user):
raise fail(403, "admin permission required")
return ok(get_platform_store().approval_templates())
@router.post("/templates")
def create_template(payload: dict[str, Any] = Body(...), current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
if not is_admin(current_user): raise fail(403, "admin permission required")
if not payload.get("name"):
raise fail(400, "name 必填")
steps = payload.get("steps") or []
if not isinstance(steps, list) or any(not isinstance(step, dict) for step in steps):
raise fail(400, "steps 格式无效")
if payload.get("action") and payload["action"] not in ALLOWED_APPROVAL_ACTIONS:
raise fail(400, "不支持的审批动作")
payload = {
**payload,
"created_by": current_user.get("id"),
"tenant_id": payload.get("tenant_id") or current_user.get("tenant_id") or "default",
"scope": payload.get("scope") or "tenant",
"status": payload.get("status") or "active",
}
template = get_platform_store().create_approval_template(payload)
get_platform_store().record_audit(
action="approval.template.create", actor_id=current_user.get("id"),
target_type="approval_template", target_id=template["id"],
tenant_id=template.get("tenant_id"),
)
return ok(template)
@router.get("/templates/{template_id}")
def get_template(template_id: str, current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
if not is_admin(current_user):
raise fail(403, "admin permission required")
try:
return ok(get_platform_store().approval_template(template_id))
except KeyError:
raise fail(404, "template not found")
@router.put("/templates/{template_id}")
def update_template(template_id: str, payload: dict[str, Any] = Body(...), current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
if not is_admin(current_user): raise fail(403, "admin permission required")
if payload.get("action") and payload["action"] not in ALLOWED_APPROVAL_ACTIONS:
raise fail(400, "不支持的审批动作")
try:
template = get_platform_store().update_approval_template(template_id, payload)
get_platform_store().record_audit(
action="approval.template.update", actor_id=current_user.get("id"),
target_type="approval_template", target_id=template_id,
tenant_id=template.get("tenant_id"), detail=f"fields={','.join(payload.keys())}",
)
return ok(template)
except KeyError:
raise fail(404, "template not found")
@router.delete("/templates/{template_id}")
def delete_template(template_id: str, current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
if not is_admin(current_user): raise fail(403, "admin permission required")
try:
template = get_platform_store().delete_approval_template(template_id)
get_platform_store().record_audit(
action="approval.template.delete", actor_id=current_user.get("id"),
target_type="approval_template", target_id=template_id,
tenant_id=template.get("tenant_id"),
)
return ok(template)
except KeyError:
raise fail(404, "template not found")
@router.get("")
def list_instances(
status: str | None = None,
mine: bool = False,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
user_id = current_user.get("id")
items = get_platform_store().approval_instances(
status=status,
applicant_id=user_id if mine and not is_admin(current_user) else None,
)
if is_admin(current_user):
return ok(items)
if mine:
return ok(items)
visible = []
for item in items:
if item.get("applicant_id") == user_id:
visible.append(item)
continue
if any(step.get("approver_id") == user_id and step.get("status") == "pending" for step in item.get("steps", [])):
visible.append(item)
continue
if is_tenant_admin(current_user, item.get("tenant_id")) and item.get("status") == "pending":
visible.append(item)
return ok(visible)
@router.get("/gpu-options")
def gpu_request_options(current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
"""Self-service GPU options; does not expose the admin compute page."""
return ok({"nodes": _available_gpu_options()})
@router.post("/gpu-requests")
def create_gpu_request(payload: dict[str, Any] = Body(...), current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
assignments = _validate_gpu_request(payload.get("assignments"))
user_id = str(current_user.get("id") or "")
if is_admin(current_user):
try:
return ok({"approval_required": False, "assignments": get_platform_store().assign_gpus(
[{**item, "user_id": user_id} for item in assignments], assigned_by=user_id,
)})
except ValueError as exc:
raise fail(409, str(exc))
normalized = [{**item, "user_id": user_id} for item in assignments]
instance = get_platform_store().create_approval_instance({
"resource_type": "gpu",
"resource_id": f"batch:{user_id}",
"applicant_id": user_id,
"action": "gpu.assign",
"tenant_id": current_user.get("tenant_id") or "default",
"reason": json.dumps({"assignments": normalized, "reason": payload.get("reason")}, ensure_ascii=False),
})
get_platform_store().record_audit(
action="gpu.assign.request", actor_id=user_id, target_type="gpu",
target_id=instance["id"], tenant_id=current_user.get("tenant_id") or "default",
detail=f"count={len(normalized)}",
)
return ok({"approval_required": True, "approval_id": instance["id"], "approval": instance})
@router.post("")
def create_instance(payload: dict[str, Any] = Body(...), current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
payload["applicant_id"] = current_user.get("id")
for field in ("resource_type", "resource_id"):
if not payload.get(field):
raise fail(400, f"{field} 必填")
resource = resource_record(str(payload["resource_type"]), str(payload["resource_id"]))
if not resource and not is_admin(current_user):
raise fail(404, "resource not found")
action = str(payload.get("action") or "")
if not action or action not in ALLOWED_APPROVAL_ACTIONS:
raise fail(400, "不支持的审批动作")
if action != "resource.access" and not is_admin(current_user) and not has_resource_access(
str(payload["resource_type"]), str(payload["resource_id"]), current_user, "read"
):
raise fail(403, "no permission to request approval for this resource")
if resource and not is_admin(current_user) and not resource_in_user_tenant(str(payload["resource_type"]), resource, current_user):
raise fail(403, "resource belongs to another tenant")
requested = payload.get("requested_permissions") or []
allowed = {"read", "write", "execute", "download"}
if any(permission not in allowed for permission in requested):
raise fail(400, "invalid requested permission")
payload["tenant_id"] = current_user.get("tenant_id") or "default"
payload["requested_permissions"] = requested
try:
return ok(get_platform_store().create_approval_instance(payload))
except KeyError:
raise fail(404, "template not found")
@router.get("/{instance_id}")
def get_instance(instance_id: str, current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
try:
item = get_platform_store().approval_instance(instance_id)
if not is_admin(current_user) and item.get("applicant_id") != current_user.get("id") and not is_tenant_admin(current_user, item.get("tenant_id")):
raise fail(403, "no permission to access approval")
return ok(item)
except KeyError:
raise fail(404, "instance not found")
@router.post("/{instance_id}/steps/{step_index}/decision")
def decide(
instance_id: str,
step_index: int,
payload: dict[str, Any] = Body(...),
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
try:
instance = get_platform_store().approval_instance(instance_id)
if instance.get("applicant_id") == current_user.get("id"):
raise fail(403, "applicant cannot approve own request")
step = next((item for item in instance.get("steps", []) if int(item.get("step_index", -1)) == step_index), None)
if not step and is_admin(current_user) and not instance.get("steps"):
step = {"approver_id": None, "status": "pending"}
if not step:
raise fail(404, "approval step not found")
designated = step.get("approver_id")
if not is_admin(current_user) and designated != current_user.get("id") and not (
step.get("approver_type") == "admin" and is_tenant_admin(current_user, instance.get("tenant_id"))
):
raise fail(403, "current user is not the designated approver")
if instance.get("action") == "tenant.quota.update" and not is_admin(current_user):
raise fail(403, "only platform administrator can approve tenant quota changes")
submitted_approver = payload.get("approver_id")
if submitted_approver and submitted_approver != current_user.get("id"):
raise fail(403, "approver_id must match current session")
store = get_platform_store()
result = store.decide_approval_step(
instance_id,
step_index,
approver_id=str(current_user.get("id")),
approved=bool(payload.get("approved", False)),
comment=payload.get("comment"),
)
if result.get("status") == "approved" and result.get("execution_status") == "ready":
try:
effect = store.apply_approval_effect(result, str(current_user.get("id") or ""))
if effect is not None:
result = store.approval_instance(instance_id)
except Exception as exc:
with store.connect() as conn:
conn.execute(
"UPDATE approval_instances SET execution_status='failed', execution_error=? WHERE id=?",
(str(exc), instance_id),
)
raise fail(409, f"审批已通过,但执行失败:{exc}")
store.record_audit(
action="approval.decision", actor_id=current_user.get("id"),
target_type="approval_instance", target_id=instance_id,
tenant_id=result.get("tenant_id"), result="success" if payload.get("approved") else "rejected",
detail=f"step={step_index}", reason=payload.get("comment"),
)
return ok(result)
except (KeyError, ValueError) as e:
raise fail(400, str(e))
@router.post("/resource-access/requests")
def create_resource_access_request(payload: dict[str, Any] = Body(...), current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
resource_type = str(payload.get("resource_type") or "")
resource_id = str(payload.get("resource_id") or "")
resource = resource_record(resource_type, resource_id)
if not resource:
raise fail(404, "resource not found")
if not is_admin(current_user) and not resource_in_user_tenant(resource_type, resource, current_user):
raise fail(403, "resource belongs to another tenant")
permissions = payload.get("requested_permissions") or ["read"]
allowed = {"read", "write", "execute", "download"}
if not permissions or any(permission not in allowed for permission in permissions):
raise fail(400, "invalid requested permissions")
result = get_platform_store().create_resource_access_request({
"resource_type": resource_type,
"resource_id": resource_id,
"applicant_id": current_user.get("id"),
"principal_type": "user",
"principal_id": current_user.get("id"),
"requested_permissions": permissions,
"reason": payload.get("reason"),
"template_id": payload.get("template_id"),
"tenant_id": current_user.get("tenant_id") or "default",
"expires_at": payload.get("expires_at"),
})
get_platform_store().record_audit(
action="resource.access.request", actor_id=current_user.get("id"),
target_type=resource_type, target_id=resource_id,
tenant_id=current_user.get("tenant_id") or "default",
detail=f"permissions={','.join(permissions)}",
)
return ok(result)
@router.get("/resource-access/requests")
def list_resource_access_requests(status: str | None = None, current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
items = get_platform_store().resource_access_requests(
user_id=None if is_admin(current_user) else current_user.get("id"),
status=status,
)
return ok(items)
@router.post("/resource-access/requests/{request_id}/cancel")
def cancel_resource_access_request(request_id: str, current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
try:
item = get_platform_store().cancel_resource_access_request(
request_id,
str(current_user.get("id") or ""),
is_admin_actor=is_admin(current_user),
)
except KeyError:
raise fail(404, "access request not found")
except PermissionError as exc:
raise fail(403, str(exc))
except ValueError as exc:
raise fail(409, str(exc))
get_platform_store().record_audit(
action="resource.access.cancel", actor_id=current_user.get("id"),
target_type="resource_access_request", target_id=request_id,
tenant_id=item.get("tenant_id"),
)
return ok(item)

View File

@@ -0,0 +1,282 @@
from __future__ import annotations
import time
from typing import Any
from urllib.parse import urljoin
import httpx
from app.core.config import get_settings
def _join_url(base_url: str, path: str) -> str:
return urljoin(base_url.rstrip("/") + "/", path.lstrip("/"))
def _unwrap_items(payload: Any) -> list[dict[str, Any]]:
if isinstance(payload, list):
return [item for item in payload if isinstance(item, dict)]
if isinstance(payload, dict):
data = payload.get("data")
if isinstance(data, dict) and isinstance(data.get("items"), list):
return [item for item in data["items"] if isinstance(item, dict)]
if isinstance(payload.get("items"), list):
return [item for item in payload["items"] if isinstance(item, dict)]
if isinstance(data, list):
return [item for item in data if isinstance(item, dict)]
return []
def _unwrap_dict(payload: Any) -> dict[str, Any]:
if isinstance(payload, dict) and isinstance(payload.get("data"), dict):
return payload["data"]
return payload if isinstance(payload, dict) else {}
# Inference calls are intentionally short-timeout:
# - load dispatch only confirms the compute node accepted the request
# (the actual model load now runs asynchronously on the node).
# - status/unload must never block the platform for long when a node is
# unreachable but still marked online.
INFERENCE_LOAD_TIMEOUT = httpx.Timeout(30, connect=10)
INFERENCE_STATUS_TIMEOUT = httpx.Timeout(30, connect=5)
INFERENCE_UNLOAD_TIMEOUT = httpx.Timeout(30, connect=5)
class ComputeNodeClient:
"""Application-side client for one compute node.
The client accepts both current YG Compute API responses and common
wrapper shapes such as `{code,message,data}` to make future engine/node
adapters less brittle.
"""
def __init__(self, api_base_url: str, token: str | None = None, timeout: float | None = None) -> None:
settings = get_settings()
self.api_base_url = api_base_url.rstrip("/")
self.token = token or settings.compute_service_token
self.timeout = timeout or settings.compute_request_timeout_seconds
self.route_prefix = settings.route_prefix.rstrip("/") or "/modelTF"
def headers(self) -> dict[str, str]:
if not self.token:
return {}
return {"X-Compute-Token": self.token}
async def test_connection(self) -> dict[str, Any]:
started = time.perf_counter()
health = await self.health()
gpus = await self.gpus()
return {
"success": True,
"latency_ms": int((time.perf_counter() - started) * 1000),
"health": health,
"gpus": gpus,
}
async def health(self) -> dict[str, Any]:
paths = [f"{self.route_prefix}/v1/compute/health", f"{self.route_prefix}/health", "/health"]
last_error = ""
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
for path in paths:
try:
response = await client.get(_join_url(self.api_base_url, path))
response.raise_for_status()
return _unwrap_dict(response.json())
except Exception as exc: # noqa: BLE001 - keep endpoint compatibility fallback broad
last_error = str(exc)
raise RuntimeError(last_error or "compute health check failed")
async def gpus(self) -> list[dict[str, Any]]:
paths = [
f"{self.route_prefix}/compute/resources/gpus",
f"{self.route_prefix}/v1/compute/resources/gpus",
"/compute/resources/gpus",
]
last_error = ""
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
for path in paths:
try:
response = await client.get(_join_url(self.api_base_url, path))
response.raise_for_status()
return _unwrap_items(response.json())
except Exception as exc: # noqa: BLE001
last_error = str(exc)
raise RuntimeError(last_error or "compute gpu discovery failed")
async def create_job(self, payload: dict[str, Any]) -> dict[str, Any]:
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
response = await client.post(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs"), json=payload)
response.raise_for_status()
return _unwrap_dict(response.json())
async def preview_job(self, payload: dict[str, Any]) -> dict[str, Any]:
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
response = await client.post(
_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/preview"),
json=payload,
)
response.raise_for_status()
return _unwrap_dict(response.json())
async def validate_job(self, payload: dict[str, Any]) -> dict[str, Any]:
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
response = await client.post(
_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/validate"),
json=payload,
)
response.raise_for_status()
return _unwrap_dict(response.json())
async def check_paths(self, paths: list[dict[str, Any]]) -> dict[str, Any]:
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
response = await client.post(
_join_url(self.api_base_url, f"{self.route_prefix}/compute/files/check-paths"),
json={"paths": paths},
)
response.raise_for_status()
return _unwrap_dict(response.json())
async def list_files(
self,
root: str = "data",
relative_path: str = "",
directories_only: bool = False,
) -> dict[str, Any]:
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
response = await client.get(
_join_url(self.api_base_url, f"{self.route_prefix}/compute/files/list"),
params={"root": root, "relative_path": relative_path, "directories_only": directories_only},
)
response.raise_for_status()
return _unwrap_dict(response.json())
async def get_job(self, job_id: str) -> dict[str, Any]:
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
response = await client.get(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}"))
response.raise_for_status()
return _unwrap_dict(response.json())
async def stop_job(self, job_id: str) -> dict[str, Any]:
# Compute API waits for the child process to exit (up to 10 seconds)
# before returning. The normal polling timeout is intentionally short,
# but is too aggressive for a stop request and used to surface as a
# platform 500 even when the node eventually stopped the job.
stop_timeout = max(float(self.timeout), 30.0)
async with httpx.AsyncClient(timeout=stop_timeout, headers=self.headers()) as client:
response = await client.post(_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}/stop"))
response.raise_for_status()
return _unwrap_dict(response.json())
async def job_logs(
self,
job_id: str,
tail_lines: int | None = None,
offset: int | None = None,
limit: int | None = None,
) -> dict[str, Any]:
params = {
key: value
for key, value in {"tail_lines": tail_lines, "offset": offset, "limit": limit}.items()
if value is not None
}
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
response = await client.get(
_join_url(self.api_base_url, f"{self.route_prefix}/compute/jobs/{job_id}/logs"),
params=params,
)
response.raise_for_status()
return _unwrap_dict(response.json())
async def import_local_file(self, payload: dict[str, Any]) -> dict[str, Any]:
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
response = await client.post(
_join_url(self.api_base_url, f"{self.route_prefix}/compute/files/import-local"),
json=payload,
)
response.raise_for_status()
return _unwrap_dict(response.json())
async def prepare_cache(self, payload: dict[str, Any]) -> dict[str, Any]:
async with httpx.AsyncClient(timeout=httpx.Timeout(900, connect=30), headers=self.headers()) as client:
response = await client.post(
_join_url(self.api_base_url, f"{self.route_prefix}/compute/cache/prepare"),
json=payload,
)
response.raise_for_status()
return _unwrap_dict(response.json())
async def cache_status(self, resource_id: str, version_id: str | None = None) -> dict[str, Any]:
params = {"resource_id": resource_id}
if version_id:
params["version_id"] = version_id
async with httpx.AsyncClient(timeout=self.timeout, headers=self.headers()) as client:
response = await client.get(
_join_url(self.api_base_url, f"{self.route_prefix}/compute/cache/status"),
params=params,
)
response.raise_for_status()
return _unwrap_dict(response.json())
async def _request(
self,
method: str,
path: str,
json_data: dict[str, Any] | None = None,
timeout: float | None = None,
) -> dict[str, Any]:
"""Generic request method for compute API endpoints."""
url = _join_url(self.api_base_url, f"{self.route_prefix}{path}")
async with httpx.AsyncClient(timeout=timeout or 300, headers=self.headers()) as client:
if method.upper() == "GET":
response = await client.get(url)
else:
response = await client.post(url, json=json_data)
response.raise_for_status()
return _unwrap_dict(response.json())
# ── Inference helpers (short timeouts — see module constants) ──────────
async def inference_load(self, payload: dict[str, Any]) -> dict[str, Any]:
"""Dispatch a model load. Returns as soon as the node accepts the
request; the node now loads asynchronously (status goes 'loading')."""
return await self._request("POST", "/inference/load", json_data=payload, timeout=INFERENCE_LOAD_TIMEOUT)
async def inference_status(self) -> dict[str, Any]:
return await self._request("GET", "/inference/status", timeout=INFERENCE_STATUS_TIMEOUT)
async def gpu_resources(self) -> list[dict[str, Any]]:
"""Read live per-GPU metrics from this compute node."""
return await self.gpus()
async def inference_unload(self) -> dict[str, Any]:
return await self._request("POST", "/inference/unload", json_data={}, timeout=INFERENCE_UNLOAD_TIMEOUT)
async def upload_file(
self,
filename: str,
content: bytes,
target_relative_path: str,
resource_type: str | None = None,
resource_id: str | None = None,
) -> dict[str, Any]:
data = {
"target_relative_path": target_relative_path,
"resource_type": resource_type or "",
"resource_id": resource_id or "",
}
files = {"file": (filename, content)}
timeout = httpx.Timeout(max(self.timeout, 60), connect=self.timeout)
async with httpx.AsyncClient(timeout=timeout, headers=self.headers()) as client:
response = await client.post(
_join_url(self.api_base_url, f"{self.route_prefix}/compute/files/upload"),
data=data,
files=files,
)
response.raise_for_status()
return _unwrap_dict(response.json())
async def upload_file_to_url(self, path: str, upload_url: str, object_key: str = "", content_type: str = "application/octet-stream") -> dict[str, Any]:
return await self._request("POST", "/compute/files/upload-to-url", json_data={
"path": path, "upload_url": upload_url, "object_key": object_key, "content_type": content_type,
}, timeout=900)

View File

@@ -0,0 +1,454 @@
from __future__ import annotations
import json
import time
from pathlib import Path
from typing import Any
from app.db.platform_store import get_platform_store
from app.core.config import get_settings
from app.modules.compute_gateway.client import ComputeNodeClient
from app.modules.storage.minio_store import get_object_storage
# starting 状态允许的最大轮询次数(约 40 * 3s ≈ 2 分钟),超过即判定节点不可达
MAX_STARTING_ATTEMPTS = 40
def _extract_job_failure_reason(log_text: str, limit: int = 2000) -> str:
"""Return a concise actionable reason from a failed Compute job log."""
lines = [line.strip() for line in str(log_text or "").splitlines() if line.strip()]
if not lines:
return ""
markers = ("[eval] FAILED", "Traceback", "RuntimeError", "Error:", "ERROR")
for index in range(len(lines) - 1, -1, -1):
if any(marker in lines[index] for marker in markers):
return "\n".join(lines[index : index + 8])[-limit:]
return "\n".join(lines[-8:])[-limit:]
async def _archive_node_directory(
store: Any,
client: ComputeNodeClient,
node: dict[str, Any],
source_path: str,
resource_type: str,
resource_id: str,
version_id: str,
object_prefix: str,
) -> list[dict[str, Any]]:
"""Archive a completed node directory to MinIO, preserving subdirectories."""
data_root = Path(str(node.get("data_root") or "/data/yg-ft")).resolve()
source = Path(source_path).resolve()
if source == data_root:
raise RuntimeError("refuse to archive compute data root; output_dir must be a task subdirectory")
try:
relative_root = source.relative_to(data_root).as_posix()
except ValueError as exc:
raise RuntimeError(f"artifact path is outside compute data root: {source_path}") from exc
queue = [relative_root]
archived: list[dict[str, Any]] = []
max_files = 10000
while queue:
relative = queue.pop(0)
listing = await client.list_files(root="data", relative_path=relative)
for item in listing.get("items") or []:
item_relative = str(item.get("relative_path") or "")
if item.get("type") == "directory":
queue.append(item_relative)
continue
path = str(item.get("path") or "")
if not path:
continue
try:
relative_file = Path(item_relative).relative_to(Path(relative_root)).as_posix()
except ValueError:
relative_file = Path(str(item.get("name") or Path(path).name)).name
object_key = f"{object_prefix}/{version_id}/{relative_file}"
if len(archived) >= max_files:
raise RuntimeError(f"archive file count exceeds limit {max_files}")
upload_url = get_object_storage().presigned_put(object_key)
result = await client.upload_file_to_url(path, upload_url, object_key)
metadata = get_object_storage().stat(object_key)
archived.append(
store.create_storage_object(
{
"resource_type": resource_type,
"resource_id": resource_id,
"version_id": version_id,
"bucket": get_object_storage().bucket,
"object_key": object_key,
"file_name": relative_file,
"content_type": "application/octet-stream",
"byte_size": metadata.get("byte_size") or result.get("byte_size") or 0,
"checksum_sha256": result.get("checksum_sha256") or "",
"status": "available",
}
)
)
return archived
def _node_for_task(task: dict[str, Any]) -> dict[str, Any] | None:
return next((node for node in get_platform_store().compute_nodes() if node["id"] == task.get("compute_node_id")), None)
def _parse_inference_load_status(task: dict[str, Any]) -> tuple[list[dict[str, Any]], dict[str, Any]]:
load_status = task.get("load_status") or {}
if isinstance(load_status, str):
try:
load_status = json.loads(load_status)
except (json.JSONDecodeError, TypeError):
load_status = {}
return load_status.get("loaded_models") or [], load_status
async def reconcile_inference_loads(store: Any) -> list[dict[str, Any]]:
"""推进处于 starting 状态的推理加载。
模型加载已改为异步派发:/model-compare/{id}/load 立即返回,这里在每次
轮询时查询对应计算节点的 /inference/status把任务从 starting 推进到
ready/error。使用短超时单节点不可达不会阻塞整轮轮询。
"""
reconciled: list[dict[str, Any]] = []
now = time.time()
for task in store.compare_tasks():
items, _ = _parse_inference_load_status(task)
if not any(item.get("status") == "starting" for item in items):
continue
# dirty 只要处理过任一 starting 项就置位load_attempts / last_polled_at
# 必须落库,否则节点不可达时计数不会累积,封顶逻辑永远触发不了
dirty = False
for item in items:
if item.get("status") != "starting":
continue
# 节流:同一 item 每 3s 只查询一次
if now - float(item.get("last_polled_at") or 0) < 3:
continue
item["last_polled_at"] = now
item["load_attempts"] = int(item.get("load_attempts") or 0) + 1
dirty = True
node = next((n for n in store.compute_nodes() if n["id"] == item.get("node_id")), None)
if not node:
item["status"] = "error"
item["error"] = "compute node deleted"
store.mark_inference_unloaded(item.get("node_id") or "")
store.release_external_gpus("inference", str(task["id"]), item.get("node_id"))
continue
if not node.get("enabled") or node.get("scheduler_status") != "online":
item["status"] = "error"
item["error"] = "compute node offline"
store.mark_inference_unloaded(node["id"])
store.release_external_gpus("inference", str(task["id"]), node["id"])
continue
try:
status = await ComputeNodeClient(node["api_base_url"]).inference_status()
except Exception as exc: # noqa: BLE001 - node unreachable; keep retrying until cap
if int(item.get("load_attempts") or 0) >= MAX_STARTING_ATTEMPTS:
item["status"] = "error"
item["error"] = f"compute node unreachable: {exc}"
store.mark_inference_unloaded(node["id"])
store.release_external_gpus("inference", str(task["id"]), node["id"])
continue
node_status = status.get("status")
if node_status == "ready":
item["status"] = "ready"
item.pop("error", None)
selected_gpus = item.get("gpu_indices") or item.get("gpus")
store.mark_inference_loaded(node["id"], selected_gpus)
elif node_status == "error":
item["status"] = "error"
item["error"] = status.get("error") or "model load failed on compute node"
store.mark_inference_unloaded(node["id"])
store.release_external_gpus("inference", str(task["id"]), node["id"])
elif node_status == "idle":
# 节点重启导致已加载模型丢失
item["status"] = "error"
item["error"] = "model disappeared from compute node (node may have restarted)"
store.mark_inference_unloaded(node["id"])
store.release_external_gpus("inference", str(task["id"]), node["id"])
# node_status == "loading" -> 保持 starting下轮再查
if dirty:
if any(i.get("status") in {"ready", "running"} for i in items):
new_status = "loaded"
elif any(i.get("status") == "starting" for i in items):
new_status = "starting" # 仍在加载中,保持 starting
else:
new_status = "failed"
store.update_compare_task(task["id"], {"status": new_status, "load_status": {"loaded_models": items}})
reconciled.append({"task_id": task["id"], "status": new_status})
return reconciled
async def fetch_eval_result_content(
client: ComputeNodeClient,
node: dict[str, Any],
job: dict[str, Any],
file_name: str = "eval_results.json",
) -> dict[str, Any] | None:
output_dir = job.get("output_dir")
if not output_dir:
return None
full_path = f"{str(output_dir).rstrip('/')}/{file_name}"
data_root = "/data/yg-ft/"
if full_path.startswith(data_root):
full_path = full_path[len(data_root):]
rel_path = full_path.lstrip("/")
import httpx
url = f"{node['api_base_url'].rstrip('/')}/modelTF/compute/files/read"
async with httpx.AsyncClient(timeout=30, headers=client.headers()) as http:
response = await http.get(url, params={"path": rel_path})
response.raise_for_status()
payload = response.json()
return payload if isinstance(payload, dict) else None
async def fetch_eval_progress_content(
client: ComputeNodeClient,
node: dict[str, Any],
job: dict[str, Any],
) -> dict[str, Any] | None:
return await fetch_eval_result_content(client, node, job, "eval_progress.json")
async def poll_compute_jobs_once(store: Any | None = None) -> dict[str, Any]:
store = store or get_platform_store()
synced: list[dict[str, Any]] = []
failed: list[dict[str, str]] = []
online_nodes = {
str(node.get("id"))
for node in store.compute_nodes()
if node.get("enabled") and node.get("scheduler_status") in {"online", "draining"}
}
training_tasks = {str(task["id"]): task for task in store.running_compute_tasks()}
# Completed tasks whose MinIO archive was interrupted remain eligible for
# reconciliation after a Backend restart or a transient node failure.
if get_settings().minio_enabled:
for task in store.tasks():
if task.get("status") != "completed" or not task.get("compute_job_id"):
continue
if (
str(task.get("archive_status") or "") != "completed"
and str(task.get("compute_node_id")) in online_nodes
):
training_tasks.setdefault(str(task["id"]), task)
for task in training_tasks.values():
node = _node_for_task(task)
if not node:
failed.append({"task_id": task["id"], "error": "compute node not found"})
continue
try:
client = ComputeNodeClient(node["api_base_url"])
job = await client.get_job(task["compute_job_id"])
try:
logs = await client.job_logs(task["compute_job_id"], tail_lines=5000)
store.record_training_log_metrics(task["id"], str(logs.get("content") or ""))
except Exception:
pass
# P0-4: Force-fetch last log snippet when job reaches terminal state
if job.get("status") in {"failed", "stopped"}:
try:
last_logs = await client.job_logs(task["compute_job_id"], tail_lines=200)
job["log_snippet"] = str(last_logs.get("content") or "")[:8192]
except Exception:
pass
updated_task = store.apply_compute_job(task["id"], job)
if (
get_settings().minio_enabled
and job.get("status") == "completed"
and job.get("output_dir")
):
trained_model = next(
(
item
for item in store.trained_models()
if item.get("name")
== (task.get("output_model_name") or f"{task.get('name')}-lora")
),
None,
)
if trained_model:
try:
archived = await _archive_node_directory(
store,
client,
node,
str(job["output_dir"]),
"trained_model",
str(trained_model["id"]),
str(job.get("id") or task.get("compute_job_id") or task["id"]),
f"trained_models/{trained_model['id']}",
)
artifacts = store.model_artifacts(str(trained_model["id"]))
if archived and artifacts:
store.link_model_artifact_storage_object(
str(artifacts[0]["id"]), str(archived[0]["id"])
)
store.update_task(task["id"], {
"archive_status": "completed",
"archive_object_ids": [str(item["id"]) for item in archived],
"archive_error": "",
})
except Exception as archive_exc:
store.update_task(task["id"], {
"archive_status": "pending",
"archive_error": str(archive_exc)[:2000],
})
raise
synced.append(updated_task)
except Exception as exc: # noqa: BLE001 - keep polling other jobs
failed.append({"task_id": task["id"], "error": str(exc)})
standalone_synced: list[dict[str, Any]] = []
standalone_jobs = {
str(record["id"]): record
for record in store.active_standalone_compute_jobs()
if str(record.get("node_id")) in online_nodes
}
if get_settings().minio_enabled:
for record in store.standalone_compute_jobs_pending_archive():
if str(record.get("node_id")) in online_nodes:
standalone_jobs.setdefault(str(record["id"]), record)
for record in standalone_jobs.values():
node = next((item for item in store.compute_nodes() if item["id"] == record.get("node_id")), None)
if not node:
failed.append({"job_id": record["id"], "error": "compute node not found"})
continue
try:
job = await ComputeNodeClient(node["api_base_url"]).get_job(record["id"])
standalone_synced.append(store.sync_model_merge_job(record["id"], job))
if get_settings().minio_enabled and job.get("status") == "completed" and job.get("output_dir"):
payload = (store.compute_job(record["id"]).get("payload") or {})
trained_model_id = str(payload.get("trained_model_id") or payload.get("model_name") or "")
if trained_model_id:
trained_model = next(
(item for item in store.trained_models() if item.get("id") == trained_model_id or item.get("name") == trained_model_id),
None,
)
if trained_model:
archived = await _archive_node_directory(
store,
ComputeNodeClient(node["api_base_url"], timeout=900),
node,
str(job["output_dir"]),
"trained_model",
str(trained_model["id"]),
str(job.get("id") or record["id"]),
f"trained_models/{trained_model['id']}",
)
artifacts = store.model_artifacts(str(trained_model["id"]))
if archived and artifacts:
store.link_model_artifact_storage_object(
str(artifacts[0]["id"]), str(archived[0]["id"])
)
store.update_compute_job_archive(record["id"], "completed", [str(item["id"]) for item in archived])
except Exception as exc: # noqa: BLE001 - keep polling other jobs
try:
store.update_compute_job_archive(record["id"], "pending", [], str(exc)[:2000])
except Exception:
pass
failed.append({"job_id": record["id"], "error": str(exc)})
# ── Eval job sync ────────────────────────────────────────────────
eval_synced = 0
eval_tasks = {str(task["id"]): task for task in store.running_eval_tasks()}
if get_settings().minio_enabled:
# A completed evaluation can win the race with the poller: its status
# is persisted before the report archive finishes. Keep such tasks in
# the reconciliation set until the report object is available.
for task in store.eval_tasks():
if (
task.get("status") == "completed"
and task.get("compute_job_id")
and str(task.get("archive_status") or "") != "completed"
and str(task.get("compute_node_id")) in online_nodes
):
eval_tasks.setdefault(str(task["id"]), task)
for eval_task in eval_tasks.values():
if str(eval_task.get("compute_node_id")) not in online_nodes:
continue
node = next(
(item for item in store.compute_nodes() if item["id"] == eval_task.get("compute_node_id")),
None,
)
if not node:
failed.append({"eval_task_id": eval_task["id"], "error": "compute node not found"})
continue
try:
client = ComputeNodeClient(node["api_base_url"])
job = await client.get_job(eval_task["compute_job_id"])
result_content = None
# Read live progress and partial results while the evaluator is running.
if job.get("status") in {"queued", "running"} and job.get("output_dir"):
try:
progress_content = await fetch_eval_progress_content(client, node, job)
if progress_content:
store.update_eval_task(
eval_task["id"],
{
"progress_detail": progress_content,
"progress": progress_content.get("percentage", eval_task.get("progress", 0)),
},
)
except Exception:
pass
try:
result_content = await fetch_eval_result_content(client, node, job)
except Exception:
result_content = None
# Try to read eval_results.json from the job output directory on completion.
if job.get("status") == "completed" and job.get("output_dir"):
try:
result_content = await fetch_eval_result_content(client, node, job)
except Exception:
pass
if job.get("status") in {"failed", "stopped"} and not job.get("error"):
try:
failure_logs = await client.job_logs(eval_task["compute_job_id"], tail_lines=120)
job["error"] = _extract_job_failure_reason(str(failure_logs.get("content") or ""))
except Exception:
pass
store.apply_eval_job_result(eval_task["id"], job, result_content)
if get_settings().minio_enabled and job.get("status") == "completed" and job.get("output_dir"):
archived = await _archive_node_directory(
store,
client,
node,
str(job["output_dir"]),
"eval",
str(eval_task["id"]),
str(job.get("id") or eval_task.get("compute_job_id") or eval_task["id"]),
f"evaluations/{eval_task['id']}",
)
report_object = next(
(item for item in archived if Path(str(item.get("file_name") or "")).name == "eval_results.json"),
archived[0] if archived else None,
)
store.update_eval_task(eval_task["id"], {
"report_storage_object_id": str(report_object["id"]) if report_object else "",
"archive_status": "completed",
"archive_object_ids": [str(item["id"]) for item in archived],
"archive_error": "",
})
# 评测 GPU 占用由 eval_tasks 状态派生,无需维护推理内存标记
eval_synced += 1
except Exception as exc: # noqa: BLE001
try:
current_eval = store.eval_task(eval_task["id"])
except Exception:
current_eval = eval_task
if current_eval.get("status") == "completed":
try:
store.update_eval_task(eval_task["id"], {"archive_status": "pending", "archive_error": str(exc)[:2000]})
except Exception:
pass
failed.append({"eval_task_id": eval_task["id"], "error": str(exc)})
# ── Inference load reconciliation ─────────────────────────────────────
try:
inference_reconciled = await reconcile_inference_loads(store)
except Exception as exc: # noqa: BLE001 - keep polling alive
failed.append({"inference_reconcile": str(exc)})
inference_reconciled = []
return {"synced": len(synced) + len(standalone_synced) + eval_synced, "failed": failed,
"items": synced, "standalone": standalone_synced, "eval_synced": eval_synced,
"inference_reconciled": inference_reconciled}

View File

@@ -0,0 +1,3 @@
from .router import router
__all__ = ["router"]

View File

@@ -0,0 +1,570 @@
from __future__ import annotations
import json
import hashlib
import os
from pathlib import Path
from typing import Any
from fastapi import APIRouter, Body, Depends, File, Request, UploadFile
from fastapi.responses import FileResponse, Response
from app.api.v1.endpoints.platform import ok, fail
from app.core.auth import get_current_user, has_resource_access, is_admin
from app.core.config import get_settings
from app.core.op_log import op_log, OpModule, OpAction
from app.db.platform_store import get_platform_store, new_id
from app.modules.storage.minio_store import get_object_storage
from app.modules.storage.policy import should_store_in_minio
def _authorize_data_convert_request(
request: Request,
task_id: str | None = None,
current_user: dict[str, Any] = Depends(get_current_user),
) -> None:
"""Protect every task-scoped conversion endpoint with resource ACL."""
if not task_id or is_admin(current_user):
return
permission = "read" if request.method in {"GET", "HEAD"} else "write"
if request.url.path.endswith("/run") or request.url.path.endswith("/import-as-dataset"):
permission = "execute"
if not has_resource_access("data_convert", task_id, current_user, permission):
raise fail(403, "no permission to access this data convert task")
router = APIRouter(
prefix="/data-convert",
tags=["data-convert"],
dependencies=[Depends(_authorize_data_convert_request)],
)
# 存储根目录
STORAGE_ROOT = Path(__file__).resolve().parents[3] / "storage" / "data-convert"
def _safe_output_filename(value: Any) -> str:
"""输出文件名白名单校验:仅允许普通文件名,阻断 ``../``、``/``、``\\`` 等路径穿越。
转换结果始终写入 ``STORAGE_ROOT/<task_id>/output/<output_filename>``
若文件名可被注入路径分隔符,将导致任意文件读写/删除。
"""
name = str(value or "converted-data.jsonl").strip()
if (
not name
or name in {".", ".."}
or name != Path(name).name
or "/" in name
or "\\" in name
or any(ord(character) < 32 or ord(character) == 127 for character in name)
):
raise fail(400, "output filename must be a plain file name")
return name
def _task_output_path(task: dict[str, Any]) -> Path:
"""返回经过白名单校验的转换输出文件路径(始终位于任务 output 目录内)。"""
return _output_dir(task["id"]) / _safe_output_filename(task.get("output_filename"))
def _task_dir(task_id: str) -> Path:
return STORAGE_ROOT / task_id
def _input_dir(task_id: str) -> Path:
return _task_dir(task_id) / "input"
def _output_dir(task_id: str) -> Path:
return _task_dir(task_id) / "output"
def _minio_enabled() -> bool:
return bool(get_settings().minio_enabled)
def _input_object_key(task_id: str, name: str) -> str:
return f"data-convert/{task_id}/input/{Path(name).name}"
def _output_object_key(task_id: str, name: str) -> str:
return f"data-convert/{task_id}/output/{Path(name).name}"
def _task_objects(task_id: str) -> list[dict[str, Any]]:
return get_platform_store().storage_objects_for_resource("data_convert", task_id)
def _register_object(
task_id: str,
*,
version_id: str,
object_key: str,
file_name: str,
content_type: str,
content: bytes,
created_by: str | None,
) -> dict[str, Any]:
storage = get_object_storage()
uploaded = storage.put_bytes(object_key, content, content_type)
return get_platform_store().create_storage_object(
{
"resource_type": "data_convert",
"resource_id": task_id,
"version_id": version_id,
"bucket": uploaded["bucket"],
"object_key": object_key,
"file_name": file_name,
"content_type": content_type,
"byte_size": len(content),
"checksum_sha256": hashlib.sha256(content).hexdigest(),
"status": "available",
"created_by": created_by,
}
)
def _input_objects(task_id: str) -> list[dict[str, Any]]:
prefix = f"data-convert/{task_id}/input/"
return sorted(
[item for item in _task_objects(task_id) if str(item.get("object_key") or "").startswith(prefix)],
key=lambda item: str(item.get("file_name") or item.get("object_key") or ""),
)
def _output_object(task: dict[str, Any]) -> dict[str, Any] | None:
key = _output_object_key(task["id"], _safe_output_filename(task.get("output_filename")))
return next((item for item in _task_objects(task["id"]) if item.get("object_key") == key), None)
def _read_output(task: dict[str, Any]) -> bytes | None:
if _minio_enabled():
item = _output_object(task)
if item:
return get_object_storage().get_bytes(item["object_key"])
inline = task.get("output_content")
return str(inline).encode("utf-8") if inline is not None else None
path = _task_output_path(task)
return path.read_bytes() if path.exists() else None
def _convert_from_minio(task: dict[str, Any], created_by: str | None) -> tuple[int, int, bytes]:
output_name = _safe_output_filename(task.get("output_filename"))
output_lines: list[str] = []
input_count = 0
output_count = 0
for item in _input_objects(task["id"]):
input_count += 1
raw = get_object_storage().get_bytes(item["object_key"])
data = json.loads(raw.decode("utf-8"))
if isinstance(data, list):
records = data
elif isinstance(data, dict):
records = [data]
else:
raise ValueError(f"JSON must be object or array: {item.get('file_name')}")
for record in records:
output_lines.append(json.dumps(record, ensure_ascii=False) + "\n")
output_count += 1
output = "".join(output_lines).encode("utf-8")
store = get_platform_store()
with store.connect() as conn:
if should_store_in_minio(len(output)):
output_object = _register_object(
task["id"],
version_id="output",
object_key=_output_object_key(task["id"], output_name),
file_name=output_name,
content_type="application/jsonl",
content=output,
created_by=created_by,
)
conn.execute(
"UPDATE data_convert_tasks SET output_storage_object_id=%s, output_content=NULL, storage_backend='minio' WHERE id=%s",
(output_object["id"], task["id"]),
)
else:
conn.execute(
"UPDATE data_convert_tasks SET output_storage_object_id=NULL, output_content=%s, storage_backend='database' WHERE id=%s",
(output.decode("utf-8"), task["id"]),
)
return input_count, output_count, output
def _convert_from_local(task: dict[str, Any]) -> tuple[int, int, bytes]:
input_dir = _input_dir(task["id"])
output_dir = _output_dir(task["id"])
input_dir.mkdir(parents=True, exist_ok=True)
output_dir.mkdir(parents=True, exist_ok=True)
output_path = _task_output_path(task)
output_path.unlink(missing_ok=True)
input_count = 0
output_count = 0
with output_path.open("w", encoding="utf-8") as output_file:
for json_file in sorted(input_dir.iterdir()):
if not json_file.is_file() or not json_file.name.lower().endswith(".json"):
continue
input_count += 1
data = json.loads(json_file.read_text(encoding="utf-8"))
records = data if isinstance(data, list) else [data] if isinstance(data, dict) else None
if records is None:
raise ValueError(f"JSON must be object or array: {json_file.name}")
for record in records:
output_file.write(json.dumps(record, ensure_ascii=False) + "\n")
output_count += 1
return input_count, output_count, output_path.read_bytes()
@router.get("")
def list_tasks(
page: int = 1,
page_size: int = 20,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
store = get_platform_store()
with store.connect() as conn:
if is_admin(current_user):
# 管理员可见全部
rows = conn.execute(
"SELECT task.*, creator.display_name AS creator_name, processor.display_name AS processor_name "
"FROM data_convert_tasks task "
"LEFT JOIN users creator ON creator.id=task.created_by "
"LEFT JOIN users processor ON processor.id=task.processed_by "
"WHERE task.deleted_at IS NULL "
"ORDER BY task.create_time DESC LIMIT %s OFFSET %s",
(page_size, (page - 1) * page_size),
).fetchall()
total = conn.execute(
"SELECT COUNT(*) FROM data_convert_tasks WHERE deleted_at IS NULL"
).fetchone()[0]
else:
# 普通用户只能看到本租户且由自己创建的任务;跨租户 ACL 通过任务级依赖访问。
user_id = current_user.get("id")
rows = conn.execute(
"SELECT task.*, creator.display_name AS creator_name, processor.display_name AS processor_name "
"FROM data_convert_tasks task "
"LEFT JOIN users creator ON creator.id=task.created_by "
"LEFT JOIN users processor ON processor.id=task.processed_by "
"WHERE task.deleted_at IS NULL AND task.tenant_id=%s AND task.created_by=%s "
"ORDER BY task.create_time DESC LIMIT %s OFFSET %s",
(current_user.get("tenant_id") or "default", user_id, page_size, (page - 1) * page_size),
).fetchall()
total = conn.execute(
"SELECT COUNT(*) FROM data_convert_tasks WHERE deleted_at IS NULL AND tenant_id=%s AND created_by=%s",
(current_user.get("tenant_id") or "default", user_id,)
).fetchone()[0]
return ok({"items": [dict(r) for r in rows], "total": total})
@router.post("")
@op_log(module=OpModule.DATA_CONVERT, action=OpAction.CREATE, target_type="convert_task", target_name_param="name")
def create_task(
payload: dict[str, Any] = Body(...),
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
name = str(payload.get("name") or "").strip()
if not name:
raise fail(400, "name is required")
task_id = new_id("dct")
output_filename = _safe_output_filename(payload.get("output_filename"))
description = str(payload.get("description") or "").strip()
user_id = current_user.get("id")
store = get_platform_store()
tenant_id = current_user.get("tenant_id") or "default"
try:
store.assert_active_tenant(tenant_id)
except ValueError as exc:
raise fail(400, str(exc))
with store.connect() as conn:
conn.execute(
"INSERT INTO data_convert_tasks (id, name, description, output_filename, created_by, tenant_id) "
"VALUES (%s, %s, %s, %s, %s, %s)",
(task_id, name, description, output_filename, user_id, tenant_id),
)
# MinIO 是正式存储;本地目录只在关闭 MinIO 的旧兼容模式下创建。
if not _minio_enabled():
_input_dir(task_id).mkdir(parents=True, exist_ok=True)
_output_dir(task_id).mkdir(parents=True, exist_ok=True)
return ok(_get_task(task_id))
@router.get("/{task_id}")
def get_task(
task_id: str,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
task = _get_task(task_id)
if not task:
raise fail(404, "task not found")
# 附加输入文件列表;旧任务没有对象记录时继续读取本地兼容目录。
files = []
if _minio_enabled():
files = [
{"name": item.get("file_name") or Path(item["object_key"]).name, "size": item.get("byte_size") or 0}
for item in _input_objects(task_id)
]
else:
input_dir = _input_dir(task_id)
if input_dir.exists():
for f in sorted(input_dir.iterdir()):
if f.is_file():
files.append({"name": f.name, "size": f.stat().st_size})
task["input_files"] = files
return ok(task)
@router.post("/{task_id}/source-files")
@op_log(module=OpModule.DATA_CONVERT, action=OpAction.UPLOAD, target_type="convert_task", target_name_param="task_id")
async def upload_source_files(
task_id: str,
files: list[UploadFile] = File(...),
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
task = _get_task(task_id)
if not task:
raise fail(404, "task not found")
if task["status"] not in ("pending", "uploaded"):
raise fail(400, "task is not editable")
staged = []
for upload in files:
name = Path(upload.filename or "input.json").name
if not name.lower().endswith(".json"):
raise fail(415, f"only JSON files are supported: {name}")
content = await upload.read()
if _minio_enabled():
_register_object(
task_id,
version_id=f"input-{hashlib.sha256(name.encode('utf-8')).hexdigest()[:16]}",
object_key=_input_object_key(task_id, name),
file_name=name,
content_type=upload.content_type or "application/json",
content=content,
created_by=task.get("created_by") or current_user.get("id"),
)
else:
input_dir = _input_dir(task_id)
input_dir.mkdir(parents=True, exist_ok=True)
(input_dir / name).write_bytes(content)
staged.append({"name": name, "size": len(content)})
store = get_platform_store()
# 标记上传完成
with store.connect() as conn:
conn.execute(
"UPDATE data_convert_tasks SET status='uploaded', processed_by=%s, processed_at=NOW(), update_time=NOW() WHERE id=%s",
(current_user.get("id"), task_id),
)
# 自动转换并导入数据集
try:
if _minio_enabled():
input_count, output_count, output = _convert_from_minio(
task, task.get("created_by") or current_user.get("id")
)
else:
input_count, output_count, output = _convert_from_local(task)
with store.connect() as conn:
conn.execute(
"UPDATE data_convert_tasks SET status='completed', "
"input_count=%s, output_count=%s, processed_by=%s, processed_at=NOW(), update_time=NOW() WHERE id=%s",
(input_count, output_count, current_user.get("id"), task_id),
)
# 自动导入数据集
content = output.decode("utf-8")
size_bytes = len(content.encode("utf-8"))
dataset = store.create_dataset({
"name": task["name"],
"type": "train",
"storage_type": "minio" if should_store_in_minio(size_bytes) else ("database" if _minio_enabled() else "local"),
"source": "upload",
"task_id": task_id,
"size": f"{size_bytes} B",
"count": output_count,
"description": f"由数据类型转换任务 {task_id} 自动导入",
"created_by": task.get("created_by") or current_user.get("id"),
"tenant_id": task.get("tenant_id") or current_user.get("tenant_id") or "default",
})
dataset_id = dataset["id"]
with store.connect() as conn:
dataset_file = store.add_dataset_file(conn, dataset_id, _safe_output_filename(task.get("output_filename")), content)
if should_store_in_minio(len(output)):
output_name = _safe_output_filename(task.get("output_filename"))
object_key = f"datasets/{dataset_id}/versions/{dataset_file.get('active_version_id') or dataset_file['id']}/{output_name}"
uploaded = get_object_storage().put_bytes(object_key, output, "application/jsonl")
storage_object = store.create_storage_object({
"resource_type": "dataset", "resource_id": dataset_id,
"version_id": dataset_file.get("active_version_id") or dataset_file["id"],
"bucket": uploaded["bucket"], "object_key": object_key,
"file_name": output_name, "content_type": "application/jsonl",
"byte_size": len(output), "checksum_sha256": hashlib.sha256(output).hexdigest(),
"status": "available", "created_by": task.get("created_by") or current_user.get("id"),
})
store.link_dataset_file_storage_object(dataset_file["id"], storage_object["id"])
return ok({
"staged_files": staged,
"auto_converted": True,
"dataset_id": dataset_id,
"input_count": input_count,
"output_count": output_count,
})
except Exception as exc:
with store.connect() as conn:
conn.execute(
"UPDATE data_convert_tasks SET status='failed', error_message=%s, processed_by=%s, processed_at=NOW(), update_time=NOW() WHERE id=%s",
(str(exc)[:500], current_user.get("id"), task_id),
)
return ok({"staged_files": staged, "auto_converted": False, "error": str(exc)[:500]})
@router.post("/{task_id}/run")
@op_log(module=OpModule.DATA_CONVERT, action=OpAction.CONVERT, target_type="convert_task", target_name_param="task_id")
def run_convert(
task_id: str,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
task = _get_task(task_id)
if not task:
raise fail(404, "task not found")
if task["status"] not in ("uploaded", "completed", "failed"):
raise fail(400, "please upload source files first")
# 标记运行中
store = get_platform_store()
with store.connect() as conn:
conn.execute(
"UPDATE data_convert_tasks SET status='running', error_message='', processed_by=%s, processed_at=NOW(), update_time=NOW() WHERE id=%s",
(current_user.get("id"), task_id),
)
try:
if _minio_enabled():
input_count, output_count, _ = _convert_from_minio(task, task.get("created_by") or current_user.get("id"))
else:
input_count, output_count, _ = _convert_from_local(task)
# 更新任务状态
with store.connect() as conn:
conn.execute(
"UPDATE data_convert_tasks SET status='completed', "
"input_count=%s, output_count=%s, processed_by=%s, processed_at=NOW(), update_time=NOW() WHERE id=%s",
(input_count, output_count, current_user.get("id"), task_id),
)
except Exception as exc:
with store.connect() as conn:
conn.execute(
"UPDATE data_convert_tasks SET status='failed', error_message=%s, processed_by=%s, processed_at=NOW(), update_time=NOW() WHERE id=%s",
(str(exc)[:500], current_user.get("id"), task_id),
)
raise fail(500, f"convert failed: {exc}")
return ok(_get_task(task_id))
@router.get("/{task_id}/download")
def download_result(
task_id: str,
current_user: dict = Depends(get_current_user),
):
task = _get_task(task_id)
if not task:
raise fail(404, "task not found")
if task["status"] != "completed":
raise fail(400, "task is not completed")
output = _read_output(task)
if output is None:
raise fail(404, "output file not found")
if _minio_enabled():
return Response(content=output, media_type="application/octet-stream", headers={
"Content-Disposition": f"attachment; filename={_safe_output_filename(task.get('output_filename'))}"
})
return FileResponse(
str(_task_output_path(task)),
media_type="application/octet-stream",
filename=_safe_output_filename(task.get("output_filename")),
)
@router.post("/{task_id}/import-as-dataset")
def import_as_dataset(
task_id: str,
payload: dict[str, Any] = Body(default={}),
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
"""把已转换的 JSONL 文件导入为数据集管理中的上传任务记录source='task')。"""
task = _get_task(task_id)
if not task:
raise fail(404, "task not found")
if task["status"] != "completed":
raise fail(400, "task is not completed")
output = _read_output(task)
if output is None:
raise fail(404, "output file not found")
content = output.decode("utf-8")
dataset_name = str(payload.get("name") or task["name"]).strip()
description = str(payload.get("description") or f"由数据类型转换任务 {task_id} 导入").strip()
size_bytes = len(content.encode("utf-8"))
store = get_platform_store()
# 用 store 提供的接口创建数据集与文件
dataset = store.create_dataset({
"name": dataset_name,
"type": "train",
"storage_type": "minio" if should_store_in_minio(size_bytes) else ("database" if _minio_enabled() else "local"),
"source": "upload",
"task_id": task_id,
"size": f"{size_bytes} B",
"count": task["output_count"],
"description": description,
"created_by": task.get("created_by") or (current_user.get("id") if current_user else None),
})
dataset_id = dataset["id"]
with store.connect() as conn:
dataset_file = store.add_dataset_file(conn, dataset_id, _safe_output_filename(task.get("output_filename")), content)
if should_store_in_minio(len(output)):
output_name = _safe_output_filename(task.get("output_filename"))
object_key = f"datasets/{dataset_id}/versions/{dataset_file.get('active_version_id') or dataset_file['id']}/{output_name}"
uploaded = get_object_storage().put_bytes(object_key, output, "application/jsonl")
storage_object = store.create_storage_object({
"resource_type": "dataset", "resource_id": dataset_id,
"version_id": dataset_file.get("active_version_id") or dataset_file["id"],
"bucket": uploaded["bucket"], "object_key": object_key,
"file_name": output_name, "content_type": "application/jsonl",
"byte_size": len(output), "checksum_sha256": hashlib.sha256(output).hexdigest(),
"status": "available", "created_by": task.get("created_by") or (current_user.get("id") if current_user else None),
})
store.link_dataset_file_storage_object(dataset_file["id"], storage_object["id"])
return ok({"dataset_id": dataset_id, "name": dataset_name})
@router.delete("/{task_id}")
@op_log(module=OpModule.DATA_CONVERT, action=OpAction.DELETE, target_type="convert_task", target_name_param="task_id")
def delete_task(
task_id: str,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
task = _get_task(task_id)
if not task:
raise fail(404, "task not found")
store = get_platform_store()
with store.connect() as conn:
conn.execute(
"UPDATE data_convert_tasks SET deleted_at=NOW() WHERE id=%s",
(task_id,),
)
if _minio_enabled():
for item in _task_objects(task_id):
try:
get_object_storage().delete(item["object_key"])
store.update_storage_object(item["id"], {"status": "deleted"})
except Exception:
pass
else:
# 旧兼容数据仍清理本地目录。
import shutil
task_dir = _task_dir(task_id)
if task_dir.exists():
shutil.rmtree(task_dir, ignore_errors=True)
return ok({"deleted": task_id})
def _get_task(task_id: str) -> dict[str, Any] | None:
store = get_platform_store()
with store.connect() as conn:
row = conn.execute(
"SELECT * FROM data_convert_tasks WHERE id=%s AND deleted_at IS NULL",
(task_id,),
).fetchone()
return dict(row) if row else None

View File

@@ -1 +1,46 @@
"""Data processing module.""" """数据处理模块:从原始文件接入到生成标准训练数据的全流程。
整体分层
--------
- ``algorithms/`` 纯算法层(无副作用:不访问 DB / 文件系统 / 网络)。
解析、格式检测、质量评分、去重、数据集切分、结构化预处理。
API、后台任务与测试共同复用。
- ``store/`` 持久层PostgreSQL。按 Mixin 拆分:
tasks / source_files / preview / generation / results / datasets。
构造时不连库、不迁移;部署方须显式执行 002 迁移(见 schema_cli
- 其余顶层文件 围绕上述两层的“服务 / 适配器”模块,由 endpoints 编排。
顶层模块速查
------------
constants.py 共享常量MAX_QA_PAIRS_PER_ITEM、MODEL_GENERATION_BATCH_SIZE
storage.py 原始源文件的本地对象存储(受控暂存于 storage/data-process
office_preview.py Word/Excel 原文件的安全受限预览(仅返回绘制所需的结构化数据)。
document_chunking.py 基于 Docling / LlamaIndex 的文档切分。
dataset_format.py Alpaca/ShareGPT/DPO/CPT 数据集格式校验(训练提交前预检)。
generation.py 大模型生成适配器,把预览内容转为标准 instruction/output 记录。
schema_cli.py data_process 运行表的显式检查 / 安装命令(运维工具,应用启动不自动调用)。
调用关系
--------
``app/api/v1/endpoints/data_process.py`` 是编排入口,组合调用以上模块与两个子包;
``dataset_format.py`` 另被 ``db/platform_store.py`` 用于训练预检。
典型链路
--------
上传源文件 → storage 暂存 → algorithms.parse_text_content 解析
→ office_preview / document_chunking 处理 → store 落库预览
→ generation 调模型生成 → store 写 results → dataset_format 校验后发布。
导入约定
--------
为避免循环导入,本 ``__init__`` 不统一再导出;请按需从子模块直接导入:
from app.modules.data_process.store import DataProcessStore
from app.modules.data_process.algorithms import estimate_token_count
from app.modules.data_process.generation import generate_model_records
"""
# 注意:为了避免循环导入,不在此处导入所有内容
# 请直接从子模块导入所需功能
__all__ = ["store", "algorithms"]

View File

@@ -0,0 +1,151 @@
"""数据处理算法模块。
重构为多个按职责拆分的子模块:
- types.py: 类型定义、常量与 dataclass
- text_utils.py: 文本解码、归一化与格式检测
- format_detection.py: 文档结构检测
- parsers/: PDF / Office / JSON / CSV 解析器
- quality.py: 质量评分与去重
- transforms.py: 数据集分割
- structured_processing.py: 结构化数据预处理
使用方式:
from app.modules.data_process.algorithms import estimate_token_count
子模块之间存在导入分层,顶层按依赖顺序导入以避免循环导入。
"""
from __future__ import annotations
# Layer 0: 类型与常量(叶子,不依赖内部模块)
from .types import (
DatasetSplit,
DocumentHeading,
DocumentNoiseSpan,
DocumentStructure,
MAX_QA_PAIRS_PER_ITEM,
ParsedText,
PdfPageText,
ProcessedStructuredRecord,
QualityScore,
SUPPORTED_TEXT_FORMATS,
StructuredPreprocessOption,
TextFormat,
_MAX_WORKBOOK_COLUMNS,
_MAX_WORKBOOK_HEADER_SCAN_ROWS,
)
# Layer 2: 文本工具(仅依赖 types对 parsers 的依赖在函数体内延迟导入)
from .text_utils import (
_normalize_spreadsheet_value,
decode_utf8,
detect_text_format,
normalize_text,
parse_text_content,
parse_utf8_text,
structured_json_dumps,
)
# Layer 1: 解析器(依赖 types 与 text_utils
from .parsers import (
LayoutRepeatedBlock,
_infer_xlsx_header_region,
_rewrite_xlsx_workbook_relationships,
_validate_office_archive,
_xlsx_sheet_merge_ranges,
detect_layout_repeated_blocks,
detect_pdf_document_noise,
extract_pdf_page_texts,
remove_document_noise,
remove_layout_repeated_blocks,
)
# Layer 3: 数据转换与质量评分
from .transforms import stable_split, stable_split_assignments
from .quality import (
content_quality_flags,
deduplicate_structured_records,
estimate_token_count,
fingerprints_are_near_duplicate,
is_low_quality_content,
is_near_duplicate,
near_duplicate_fingerprint,
record_fingerprint,
score_quality,
)
# Layer 4: 结构化数据处理(依赖 text_utils / quality / transforms / parsers
from .structured_processing import (
canonical_record_json,
desensitize_pii,
desensitize_structured_record,
expand_to_context_boundaries,
extract_structured_records,
filter_anomalous_structured_records,
flatten_structured_record,
generate_standard_records,
merge_short_blocks,
normalize_structured_record,
preprocess_structured_records,
preprocess_structured_records_with_lineage,
protected_context_ranges,
)
# Layer 5: 文档结构检测(依赖 structured_processing
from .format_detection import detect_document_structure
__all__ = [
"DatasetSplit",
"DocumentHeading",
"DocumentNoiseSpan",
"DocumentStructure",
"MAX_QA_PAIRS_PER_ITEM",
"ParsedText",
"PdfPageText",
"ProcessedStructuredRecord",
"QualityScore",
"SUPPORTED_TEXT_FORMATS",
"StructuredPreprocessOption",
"TextFormat",
"_MAX_WORKBOOK_COLUMNS",
"_MAX_WORKBOOK_HEADER_SCAN_ROWS",
"_normalize_spreadsheet_value",
"canonical_record_json",
"content_quality_flags",
"decode_utf8",
"deduplicate_structured_records",
"desensitize_pii",
"desensitize_structured_record",
"detect_document_structure",
"detect_layout_repeated_blocks",
"detect_pdf_document_noise",
"detect_text_format",
"estimate_token_count",
"expand_to_context_boundaries",
"extract_pdf_page_texts",
"extract_structured_records",
"filter_anomalous_structured_records",
"fingerprints_are_near_duplicate",
"flatten_structured_record",
"generate_standard_records",
"is_low_quality_content",
"is_near_duplicate",
"merge_short_blocks",
"near_duplicate_fingerprint",
"normalize_structured_record",
"normalize_text",
"parse_text_content",
"parse_utf8_text",
"preprocess_structured_records",
"preprocess_structured_records_with_lineage",
"protected_context_ranges",
"record_fingerprint",
"LayoutRepeatedBlock",
"remove_document_noise",
"remove_layout_repeated_blocks",
"score_quality",
"stable_split",
"stable_split_assignments",
"structured_json_dumps",
]

View File

@@ -0,0 +1,24 @@
"""数据处理算法 - 本地语义嵌入模型共享单例。"""
from __future__ import annotations
import os
from functools import lru_cache
from typing import Any
@lru_cache(maxsize=1)
def semantic_embedding_model() -> Any:
"""加载本地嵌入模型,供语义分块与语义质量评分共用。
模型可在部署环境覆盖;默认模型体积较小且适合中英文语义判断。
返回 LlamaIndex BaseEmbedding通过 ``get_text_embedding`` 使用。
"""
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
return HuggingFaceEmbedding(
model_name=os.getenv("DATA_PROCESS_EMBEDDING_MODEL", "BAAI/bge-small-zh-v1.5"),
device=os.getenv("DATA_PROCESS_EMBEDDING_DEVICE", "cpu"),
trust_remote_code=False,
)

View File

@@ -0,0 +1,119 @@
"""数据处理算法 - 格式检测。"""
from __future__ import annotations
import re
from .structured_processing import protected_context_ranges
from .text_utils import normalize_text
from .types import DocumentHeading, DocumentStructure
def detect_document_structure(text: str) -> DocumentStructure:
"""识别 Markdown、中文章节和数字编号标题及不可拆分块。"""
normalized = normalize_text(text)
if not normalized:
return DocumentStructure(
line_count=0,
paragraph_count=0,
headings=(),
code_block_count=0,
table_block_count=0,
list_block_count=0,
)
code_ranges = protected_context_ranges(
normalized,
preserve_code_blocks=True,
preserve_tables=False,
preserve_lists=False,
)
table_candidates = protected_context_ranges(
normalized,
preserve_code_blocks=False,
preserve_tables=True,
preserve_lists=False,
)
list_candidates = protected_context_ranges(
normalized,
preserve_code_blocks=False,
preserve_tables=False,
preserve_lists=True,
)
table_ranges = tuple(
item
for item in table_candidates
if not any(
item[0] < code_end and item[1] > code_start
for code_start, code_end in code_ranges
)
)
list_ranges = tuple(
item
for item in list_candidates
if not any(
item[0] < code_end and item[1] > code_start
for code_start, code_end in code_ranges
)
)
markdown_heading = re.compile(r"^\s*(?P<marks>#{1,6})\s+(?P<title>.+?)\s*#*\s*$")
chinese_heading = re.compile(
r"^\s*(?P<title>第[一二三四五六七八九十百千万0-9]+[章节篇部分].*)$"
)
numbered_heading = re.compile(
r"^\s*(?P<number>\d+(?:\.\d+)*)(?:[、.]|\s+)\s*(?P<title>\S.*)$"
)
headings: list[DocumentHeading] = []
cursor = 0
for line_number, raw_line in enumerate(normalized.splitlines(keepends=True), start=1):
line = raw_line.rstrip("\n")
line_end = cursor + len(line)
if not any(range_start <= cursor < range_end for range_start, range_end in code_ranges):
match = markdown_heading.match(line)
if match:
headings.append(
DocumentHeading(
level=len(match.group("marks")),
title=normalize_text(match.group("title")),
line_number=line_number,
start=cursor,
end=line_end,
)
)
else:
match = chinese_heading.match(line)
if match:
headings.append(
DocumentHeading(
level=1,
title=normalize_text(match.group("title")),
line_number=line_number,
start=cursor,
end=line_end,
)
)
else:
match = numbered_heading.match(line)
if match:
headings.append(
DocumentHeading(
level=min(6, match.group("number").count(".") + 1),
title=normalize_text(match.group("title")),
line_number=line_number,
start=cursor,
end=line_end,
)
)
cursor += len(raw_line)
paragraphs = [part for part in re.split(r"\n\s*\n", normalized) if part.strip()]
return DocumentStructure(
line_count=len(normalized.splitlines()),
paragraph_count=len(paragraphs),
headings=tuple(headings),
code_block_count=len(code_ranges),
table_block_count=len(table_ranges),
list_block_count=len(list_ranges),
)

View File

@@ -0,0 +1,27 @@
"""文档解析器模块。"""
from .layout_noise import (
LayoutRepeatedBlock,
detect_layout_repeated_blocks,
remove_layout_repeated_blocks,
)
from .pdf import extract_pdf_page_texts, detect_pdf_document_noise, remove_document_noise
from .office import (
_validate_office_archive,
_rewrite_xlsx_workbook_relationships,
_xlsx_sheet_merge_ranges,
_infer_xlsx_header_region,
)
__all__ = [
'extract_pdf_page_texts',
'detect_pdf_document_noise',
'remove_document_noise',
'LayoutRepeatedBlock',
'detect_layout_repeated_blocks',
'remove_layout_repeated_blocks',
'_validate_office_archive',
'_rewrite_xlsx_workbook_relationships',
'_xlsx_sheet_merge_ranges',
'_infer_xlsx_header_region',
]

View File

@@ -0,0 +1,89 @@
"""数据处理算法 - CSV 解析。"""
from __future__ import annotations
import csv
import io
import re
from typing import Any
from ..text_utils import normalize_text
def _source_line_offsets(
text: str,
start_line: int,
end_line: int,
) -> tuple[int, int]:
"""把 1-based 物理行范围转换为左闭右开的字符范围。"""
line_starts = [0]
line_starts.extend(match.end() for match in re.finditer("\n", text))
if start_line < 1 or end_line < start_line or end_line > len(line_starts):
raise ValueError("source line range is outside normalized text")
source_start = line_starts[start_line - 1]
source_end = (
line_starts[end_line] - 1
if end_line < len(line_starts)
else len(text)
)
return source_start, source_end
def _extract_csv_records_with_locators(
text: str,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""从 CSV 文本中提取记录及其行级定位信息。"""
normalized_text = normalize_text(text)
if not normalized_text:
return [], []
try:
dialect = csv.Sniffer().sniff(normalized_text[:8192], delimiters=",\t;")
except csv.Error:
dialect = csv.excel
reader = csv.DictReader(io.StringIO(normalized_text), dialect=dialect)
if not reader.fieldnames:
raise ValueError("CSV header is required")
headers = [normalize_text(header or "") for header in reader.fieldnames]
if any(not header for header in headers):
raise ValueError("CSV header cannot be empty")
if len(set(headers)) != len(headers):
raise ValueError("CSV headers must be unique")
reader.fieldnames = headers
records: list[dict[str, Any]] = []
locators: list[dict[str, Any]] = []
source_lines = normalized_text.splitlines()
previous_end_line = reader.line_num
for row in reader:
end_line = reader.line_num
start_line = previous_end_line + 1
previous_end_line = end_line
while start_line < end_line and not source_lines[start_line - 1].strip():
start_line += 1
if None in row:
raise ValueError("CSV row has more fields than the header")
normalized_row = {
key: normalize_text(value or "")
for key, value in row.items()
}
if any(value for value in normalized_row.values()):
records.append(normalized_row)
source_start, source_end = _source_line_offsets(
normalized_text,
start_line,
end_line,
)
locators.append(
{
"kind": "csv",
"record_index": len(records),
"start_line": start_line,
"end_line": end_line,
"source_start": source_start,
"source_end": source_end,
}
)
return records, locators

View File

@@ -0,0 +1,395 @@
"""数据处理算法 - JSON/JSONL 解析。"""
from __future__ import annotations
import json
import re
from collections.abc import Mapping, Sequence
from decimal import Decimal
from typing import Any
from ..text_utils import _normalize_format, _normalize_value, normalize_text
from ..types import (
_DuplicateJsonKeyError,
_JSON_ENVELOPE_KEYS,
_JSON_RECORD_ARRAY_KEYS,
_JSON_RESPONSE_METADATA_KEYS,
_JSON_WRAPPER_METADATA_KEYS,
_MAX_JSON_DEPTH,
)
def _record_from_value(value: Any, *, normalize: bool = True) -> dict[str, Any]:
if isinstance(value, Mapping):
return dict(_normalize_value(value)) if normalize else dict(value)
return {"value": _normalize_value(value) if normalize else value}
def _json_pointer_segment(value: Any) -> str:
return str(value).replace("~", "~0").replace("/", "~1")
def _reject_duplicate_json_keys(pairs: list[tuple[str, Any]]) -> dict[str, Any]:
result: dict[str, Any] = {}
for key, value in pairs:
if key in result:
raise _DuplicateJsonKeyError(f"duplicate JSON object key: {key!r}")
result[key] = value
return result
def _reject_json_constant(value: str) -> Any:
raise ValueError(f"non-finite JSON number is not allowed: {value}")
def _skip_json_whitespace(text: str, offset: int) -> int:
while offset < len(text) and text[offset] in " \t\r\n":
offset += 1
return offset
def _validate_json_nesting(text: str) -> None:
"""在构造 Python 对象前限制容器深度,避免依赖解释器递归阈值。"""
depth = 0
in_string = False
escaped = False
for char in text:
if in_string:
if escaped:
escaped = False
elif char == "\\":
escaped = True
elif char == '"':
in_string = False
continue
if char == '"':
in_string = True
elif char in "[{":
depth += 1
if depth > _MAX_JSON_DEPTH:
raise ValueError(
f"JSON nesting exceeds the supported depth of {_MAX_JSON_DEPTH}"
)
elif char in "]}":
depth = max(0, depth - 1)
def _strict_json_loads(text: str) -> tuple[Any, int, int]:
"""严格解析单个 JSON 值并返回其左闭右开源码区间。"""
start = _skip_json_whitespace(text, 0)
if start >= len(text):
raise ValueError("JSON content is empty")
_validate_json_nesting(text)
decoder = json.JSONDecoder(
object_pairs_hook=_reject_duplicate_json_keys,
parse_float=Decimal,
parse_int=int,
parse_constant=_reject_json_constant,
strict=True,
)
try:
payload, end = decoder.raw_decode(text, start)
except json.JSONDecodeError as exc:
raise ValueError(
f"invalid JSON at line {exc.lineno}, column {exc.colno}: {exc.msg}"
) from exc
except RecursionError as exc:
raise ValueError("JSON nesting exceeds the supported depth") from exc
except _DuplicateJsonKeyError as exc:
raise ValueError(str(exc)) from exc
except ValueError as exc:
# parse_int/parse_float/parse_constant 的异常也必须稳定映射为客户端错误。
raise ValueError(f"invalid JSON number: {exc}") from exc
trailing = _skip_json_whitespace(text, end)
if trailing != len(text):
line = text.count("\n", 0, trailing) + 1
line_start = text.rfind("\n", 0, trailing) + 1
column = trailing - line_start + 1
raise ValueError(
f"invalid JSON at line {line}, column {column}: extra data"
)
return payload, start, end
def _json_value_end(text: str, start: int) -> int:
"""在已验证 JSON 中定位一个值的结束偏移,不对数值做二次解析。"""
if start >= len(text):
raise ValueError("invalid JSON source span")
first = text[start]
if first == '"':
escaped = False
for offset in range(start + 1, len(text)):
char = text[offset]
if escaped:
escaped = False
elif char == "\\":
escaped = True
elif char == '"':
return offset + 1
raise ValueError("invalid JSON source span")
if first in "[{":
stack = [first]
in_string = False
escaped = False
for offset in range(start + 1, len(text)):
char = text[offset]
if in_string:
if escaped:
escaped = False
elif char == "\\":
escaped = True
elif char == '"':
in_string = False
continue
if char == '"':
in_string = True
elif char in "[{":
stack.append(char)
elif char in "]}":
expected = "[" if char == "]" else "{"
if not stack or stack[-1] != expected:
raise ValueError("invalid JSON source span")
stack.pop()
if not stack:
return offset + 1
raise ValueError("invalid JSON source span")
end = start
while end < len(text) and text[end] not in " \t\r\n,]}":
end += 1
if end == start:
raise ValueError("invalid JSON source span")
return end
def _json_object_value_spans(
text: str,
start: int,
end: int,
) -> dict[str, tuple[int, int]]:
"""返回已验证 JSON 对象直接子字段的值区间。"""
if text[start] != "{" or text[end - 1] != "}":
raise ValueError("JSON source value is not an object")
result: dict[str, tuple[int, int]] = {}
offset = _skip_json_whitespace(text, start + 1)
key_decoder = json.JSONDecoder()
while offset < end - 1:
key, key_end = key_decoder.raw_decode(text, offset)
if not isinstance(key, str):
raise ValueError("invalid JSON object key")
offset = _skip_json_whitespace(text, key_end)
if offset >= end or text[offset] != ":":
raise ValueError("invalid JSON object member")
value_start = _skip_json_whitespace(text, offset + 1)
value_end = _json_value_end(text, value_start)
result[key] = (value_start, value_end)
offset = _skip_json_whitespace(text, value_end)
if offset >= end - 1:
break
if text[offset] != ",":
raise ValueError("invalid JSON object member")
offset = _skip_json_whitespace(text, offset + 1)
return result
def _json_array_item_spans(
text: str,
start: int,
end: int,
) -> list[tuple[int, int]]:
"""返回已验证 JSON 数组中每个直接元素的源码区间。"""
if text[start] != "[" or text[end - 1] != "]":
raise ValueError("JSON source value is not an array")
result: list[tuple[int, int]] = []
offset = _skip_json_whitespace(text, start + 1)
while offset < end - 1:
item_end = _json_value_end(text, offset)
result.append((offset, item_end))
offset = _skip_json_whitespace(text, item_end)
if offset >= end - 1:
break
if text[offset] != ",":
raise ValueError("invalid JSON array item")
offset = _skip_json_whitespace(text, offset + 1)
return result
def _json_span_at_path(
text: str,
root_span: tuple[int, int],
path: Sequence[str],
) -> tuple[int, int]:
span = root_span
for key in path:
try:
span = _json_object_value_spans(text, *span)[key]
except KeyError as exc:
raise ValueError(f"JSON source path cannot be located: {key}") from exc
return span
def _pure_json_record_wrapper(
payload: Any,
) -> tuple[list[Any], tuple[str, ...]] | None:
"""识别不会与业务字段冲突的纯记录包装对象。"""
if not isinstance(payload, Mapping):
return None
def direct_wrapper(
value: Mapping[str, Any],
metadata_keys: frozenset[str] | set[str] = _JSON_WRAPPER_METADATA_KEYS,
) -> tuple[list[Any], tuple[str, ...]] | None:
candidates = [
key
for key in _JSON_RECORD_ARRAY_KEYS
if isinstance(value.get(key), list)
]
if len(candidates) != 1:
return None
record_key = candidates[0]
records = value[record_key]
if any(not isinstance(record, Mapping) for record in records):
return None
if any(
key != record_key and key not in metadata_keys
for key in value
):
return None
return records, (record_key,)
direct = direct_wrapper(payload)
if direct is not None:
return direct
envelope_keys = [
key
for key in _JSON_ENVELOPE_KEYS
if isinstance(payload.get(key), Mapping)
]
if len(envelope_keys) != 1:
return None
envelope_key = envelope_keys[0]
if any(
key != envelope_key and key not in _JSON_RESPONSE_METADATA_KEYS
for key in payload
):
return None
nested = direct_wrapper(
payload[envelope_key],
_JSON_RESPONSE_METADATA_KEYS,
)
if nested is None:
return None
records, nested_path = nested
return records, (envelope_key, *nested_path)
def _json_record_locator(
text: str,
*,
record_index: int,
json_pointer: str,
span: tuple[int, int],
) -> dict[str, Any]:
source_start, source_end = span
start_line = text.count("\n", 0, source_start) + 1
last_character = max(source_start, source_end - 1)
end_line = text.count("\n", 0, last_character) + 1
return {
"kind": "json",
"record_index": record_index,
"json_pointer": json_pointer,
"source_start": source_start,
"source_end": source_end,
"start_line": start_line,
"end_line": end_line,
}
def _extract_structured_records_with_locators(
text: str,
file_format: str,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
"""从 JSON、JSONL 或 CSV 中提取记录。
JSON 根数组始终表示多条记录;对象仅在满足纯包装契约时展开,其他
对象均视为一条业务记录。JSON 字段和值在解析阶段保持原样,只有用户
明确选择 ``normalize_format`` 后才会规范化。
"""
normalized_format = _normalize_format(file_format)
if normalized_format not in {"json", "jsonl", "csv"}:
raise ValueError("structured record extraction only supports JSON, JSONL and CSV")
if normalized_format == "json":
if _skip_json_whitespace(text, 0) == len(text):
return [], []
payload, root_start, root_end = _strict_json_loads(text)
values: Sequence[Any]
pointer_path: tuple[str, ...] = ()
record_spans: list[tuple[int, int]]
if isinstance(payload, list):
values = payload
record_spans = _json_array_item_spans(text, root_start, root_end)
else:
wrapper = _pure_json_record_wrapper(payload)
if wrapper is None:
values = [payload]
record_spans = [(root_start, root_end)]
else:
values, pointer_path = wrapper
array_span = _json_span_at_path(
text,
(root_start, root_end),
pointer_path,
)
record_spans = _json_array_item_spans(text, *array_span)
if len(record_spans) != len(values):
raise ValueError("JSON record source spans do not match parsed records")
records = [_record_from_value(value, normalize=False) for value in values]
pointer_prefix = "".join(
f"/{_json_pointer_segment(segment)}" for segment in pointer_path
)
locators = [
_json_record_locator(
text,
record_index=index + 1,
json_pointer=(
f"{pointer_prefix}/{index}"
if pointer_path or isinstance(payload, list)
else ""
),
span=record_spans[index],
)
for index in range(len(records))
]
return records, locators
if normalized_format == "jsonl":
records: list[dict[str, Any]] = []
locators: list[dict[str, Any]] = []
source_offset = 0
for line_number, line in enumerate(text.split("\n"), start=1):
line_content_end = len(line)
if _skip_json_whitespace(line, 0) == line_content_end:
source_offset += len(line) + 1
continue
try:
value, value_start, value_end = _strict_json_loads(line)
except ValueError as exc:
raise ValueError(
f"invalid JSONL at line {line_number}: {exc}"
) from exc
records.append(_record_from_value(value, normalize=False))
locators.append(
{
"kind": "jsonl",
"record_index": len(records),
"start_line": line_number,
"end_line": line_number,
"source_start": source_offset + value_start,
"source_end": source_offset + value_end,
}
)
source_offset += len(line) + 1
return records, locators
# CSV 分支独立放在 csv_parser 中,避免与 JSON 机制耦合。
from .csv_parser import _extract_csv_records_with_locators
return _extract_csv_records_with_locators(text)

View File

@@ -0,0 +1,174 @@
"""基于 Docling 输出的版面噪声检测与剔除。
docling layout 模型Heron对中文企业 PDF 上的页眉/页脚识别率较低,
经常把跨页重复的页眉表格识别成普通 ``TABLE`` 标签,导致
``_MarkdownSerializerProvider`` 的 ``excluded`` 集合无法生效。
本模块提供第二层启发式:扫描 docling 输出的所有 ``TableItem``
对每个表按"首列标签序列"聚合。如果同一组标签在文档中多页重复出现,
则判定为页眉/页脚类重复块,并在最终 chunk 文本中按行剔除。
"""
from __future__ import annotations
import math
import re
from collections.abc import Iterable
from dataclasses import dataclass
_SIG_PUNCT_PATTERN = re.compile(r"[\s\W_]+", re.UNICODE)
_SIG_DIGIT_PATTERN = re.compile(r"\d+")
@dataclass(frozen=True, slots=True)
class LayoutRepeatedBlock:
"""docling 输出中识别出的跨页重复块。"""
labels: tuple[str, ...]
occurrences: int
@property
def signature(self) -> str:
"""拼接签名(用于日志与向后兼容)。"""
return "".join(self.labels)
def _normalize_signature(text: str) -> str:
"""归一化:删除所有数字、去除空白/标点、转小写。"""
stripped = _SIG_DIGIT_PATTERN.sub("", text)
return _SIG_PUNCT_PATTERN.sub("", stripped).casefold()
def _extract_first_column_labels(table_text: str) -> tuple[str, ...]:
"""提取 docling TableItem markdown 表示中的"首列标签"序列。"""
labels: list[str] = []
seen: set[str] = set()
for raw_line in table_text.splitlines():
line = raw_line.strip()
if "|" not in line:
continue
parts = [cell.strip() for cell in line.strip("|").split("|")]
if not parts or not parts[0]:
continue
# 过滤掉分隔行(如 "| - | - |"
if all(re.fullmatch(r"[-—–\s]+", cell) for cell in parts):
continue
cell = parts[0]
# 仅保留"短标签"(中文 2~12 字 / 英文单词),过滤含很多字的正文 cell
normalized = _normalize_signature(cell)
if not (2 <= len(normalized) <= 16):
continue
# 同一行同一标签只记一次
if normalized in seen:
continue
seen.add(normalized)
labels.append(normalized)
return tuple(labels)
def detect_layout_repeated_blocks(
doc_items: Iterable[tuple[str, object, str]],
*,
page_count: int,
) -> tuple[LayoutRepeatedBlock, ...]:
"""扫描 docling 输出,识别跨页重复出现的标签组。
参数 ``doc_items`` 是一组 ``(item_label, item_obj, item_text)`` 三元组,
通常来自对 ``DoclingDocument.iterate_items()`` 的遍历。
判定条件(与 ``detect_pdf_document_noise`` 保持一致):
- 同一组首列标签至少在 ``max(3, ceil(page_count * 0.3))`` 个不同 item 中出现;
- 标签序列长度在 ``[1, 8]`` 之间。
"""
if page_count < 3:
return ()
label_groups: dict[tuple[str, ...], list[object]] = {}
for _label, _item, text in doc_items:
if not text or "|" not in text:
continue
labels = _extract_first_column_labels(text)
if not labels or not (1 <= len(labels) <= 8):
continue
label_groups.setdefault(labels, []).append(_item)
minimum_occurrences = max(3, math.ceil(page_count * 0.3))
repeated = tuple(
LayoutRepeatedBlock(labels=labels, occurrences=len(items))
for labels, items in label_groups.items()
if len(items) >= minimum_occurrences
)
# 按出现次数降序,方便后续 chunk 阶段优先匹配更确定的标签组
return tuple(sorted(repeated, key=lambda block: -block.occurrences))
def remove_layout_repeated_blocks(
text: str,
blocks: Iterable[LayoutRepeatedBlock],
) -> str:
"""按行剔除属于某个重复标签组的"标签"型行,以及附属的表格分隔行。
仅剔除整行的首列归一化结果命中某个 block 的标签集(子集判定);
含正文的长行不会因子串匹配被误删。
紧接着被剔除的标签行的分隔行(如 ``| - | - | - |``)与紧随其后的空行也会被删除,
避免残留"裸表格"格式。
"""
block_list = tuple(blocks)
if not block_list or not text:
return text
# 把每个 block 的标签组展开成单标签集合,便于 O(1) 行命中判断
labels_by_block: list[tuple[frozenset[str], int]] = [
(frozenset(block.labels), block.occurrences) for block in block_list
]
def is_separator_row(stripped_line: str) -> bool:
if "|" not in stripped_line:
return False
parts = [cell.strip() for cell in stripped_line.strip("|").split("|")]
if not parts:
return False
return all(re.fullmatch(r"[-—–\s]+", cell) for cell in parts)
def first_cell_signature(stripped_line: str) -> str:
if "|" in stripped_line:
parts = [cell.strip() for cell in stripped_line.strip("|").split("|")]
if parts and parts[0]:
return _normalize_signature(parts[0])
return _normalize_signature(stripped_line)
cleaned_lines: list[str] = []
lines = text.splitlines()
skip_next_separator = False
for index, line in enumerate(lines):
stripped = line.strip()
if not stripped:
cleaned_lines.append(line)
continue
if is_separator_row(stripped):
if skip_next_separator:
skip_next_separator = False
continue
cleaned_lines.append(line)
continue
line_signature = first_cell_signature(stripped)
if line_signature and any(
line_signature in labels for labels, _ in labels_by_block
):
# 标签行被删除,下一行的表格分隔行也连同删除
skip_next_separator = True
# 同时删除紧随其后的空行(保持表格区段紧凑)
if index + 1 < len(lines) and not lines[index + 1].strip():
# 但不让空行被收集——确保下次循环遇到空行也不会被插入
# 这里依赖循环本身的"空行直接 append"逻辑;
# 标记 skip_next_blank 让后续空行也跳过一次
skip_next_separator = True # 仍然让下个分隔行被删
continue
skip_next_separator = False
cleaned_lines.append(line)
return "\n".join(cleaned_lines)

View File

@@ -0,0 +1,716 @@
"""数据处理算法 - Office 文档解析。"""
from __future__ import annotations
import io
import re
import unicodedata
import zipfile
import xml.etree.ElementTree as ET
from collections.abc import Iterator, Mapping, Sequence
from pathlib import PurePosixPath
from typing import Any
from urllib.parse import unquote, urlsplit
from docx import Document
from docx.oxml.ns import qn
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table
from docx.text.paragraph import Paragraph
from openpyxl import load_workbook
from openpyxl.utils.cell import range_boundaries
from pptx import Presentation
from ..text_utils import _append_bounded_text, _normalize_spreadsheet_value, normalize_text
from ..types import (
_MAX_ARCHIVE_COMPRESSION_RATIO,
_MAX_ARCHIVE_ENTRIES,
_MAX_ARCHIVE_ENTRY_BYTES,
_MAX_ARCHIVE_UNCOMPRESSED_BYTES,
_MAX_PRESENTATION_SLIDES,
_MAX_WORKBOOK_CELLS,
_MAX_WORKBOOK_COLUMNS,
_MAX_WORKBOOK_HEADER_ROWS,
_MAX_WORKBOOK_HEADER_SCAN_ROWS,
_MAX_WORKBOOK_MERGED_RANGES,
_MAX_WORKBOOK_ROWS,
_MAX_WORKBOOK_SCANNED_ROWS,
_MAX_WORKBOOK_SHEETS,
TextFormat,
)
_XLSX_REPORT_METADATA_PATTERN = re.compile(
r"^(?:报表|报告|标题|说明|备注|制表|统计|日期|时间|期间|"
r"report|title|note|remark|date|time|period)\b",
flags=re.IGNORECASE,
)
def _validate_office_archive(raw: bytes, file_format: TextFormat) -> None:
"""在交给 Office 解析库前限制 ZIP 包规模并拒绝活动 XML。"""
required_members = {
"docx": {"[Content_Types].xml", "word/document.xml"},
"xlsx": {"[Content_Types].xml", "xl/workbook.xml"},
"pptx": {"[Content_Types].xml", "ppt/presentation.xml"},
}
try:
with zipfile.ZipFile(io.BytesIO(raw)) as archive:
members = archive.infolist()
if len(members) > _MAX_ARCHIVE_ENTRIES:
raise ValueError(
f"{file_format.upper()} archive contains too many entries "
f"(limit {_MAX_ARCHIVE_ENTRIES})"
)
names: set[str] = set()
total_size = 0
for member in members:
path = PurePosixPath(member.filename.replace("\\", "/"))
if path.is_absolute() or ".." in path.parts:
raise ValueError(
f"{file_format.upper()} archive contains an unsafe member path"
)
if member.filename in names:
raise ValueError(
f"{file_format.upper()} archive contains duplicate member names"
)
names.add(member.filename)
if member.flag_bits & 0x1:
raise ValueError(f"encrypted {file_format.upper()} files are not supported")
if member.is_dir():
continue
if member.file_size > _MAX_ARCHIVE_ENTRY_BYTES:
raise ValueError(
f"{file_format.upper()} archive entry exceeds "
f"{_MAX_ARCHIVE_ENTRY_BYTES} bytes"
)
total_size += member.file_size
if total_size > _MAX_ARCHIVE_UNCOMPRESSED_BYTES:
raise ValueError(
f"{file_format.upper()} archive expands beyond "
f"{_MAX_ARCHIVE_UNCOMPRESSED_BYTES} bytes"
)
if member.file_size >= 1024 * 1024:
if member.compress_size <= 0:
raise ValueError(
f"{file_format.upper()} archive has an unsafe compression ratio"
)
ratio = member.file_size / member.compress_size
if ratio > _MAX_ARCHIVE_COMPRESSION_RATIO:
raise ValueError(
f"{file_format.upper()} archive has an unsafe compression ratio"
)
missing = required_members[file_format] - names
if missing:
raise ValueError(
f"invalid {file_format.upper()} package: missing "
f"{', '.join(sorted(missing))}"
)
for member in members:
if member.is_dir() or not member.filename.lower().endswith((".xml", ".rels")):
continue
with archive.open(member) as stream:
prefix = stream.read(min(member.file_size, 1024 * 1024)).upper()
if b"<!DOCTYPE" in prefix or b"<!ENTITY" in prefix:
raise ValueError(
f"{file_format.upper()} archive contains unsupported active XML"
)
except zipfile.BadZipFile as exc:
raise ValueError(f"invalid {file_format.upper()} file: not an Office ZIP package") from exc
def iter_document_blocks(parent: Any) -> Iterator[Any]:
"""按文档顺序产出正文段落与表格,并下钻 SDT 内容控件。
Word 的目录、复选框等内容控件包在 ``w:sdt`` 元素里,只遍历 body
直接子级会把这些段落整段丢掉。
"""
for child in parent.iterchildren():
if isinstance(child, (CT_P, CT_Tbl)):
yield child
elif child.tag == qn("w:sdt"):
content = child.find(qn("w:sdtContent"))
if content is not None:
yield from iter_document_blocks(content)
def _extract_docx_text(raw: bytes) -> str:
_validate_office_archive(raw, "docx")
try:
document = Document(io.BytesIO(raw))
except Exception as exc:
raise ValueError(f"invalid DOCX file: {exc}") from exc
parts: list[str] = []
total = 0
for child in iter_document_blocks(document.element.body):
if isinstance(child, CT_P):
total = _append_bounded_text(parts, Paragraph(child, document).text, total)
continue
if isinstance(child, CT_Tbl):
table = Table(child, document)
for row in table.rows:
cells = [normalize_text(cell.text) for cell in row.cells]
total = _append_bounded_text(parts, "\t".join(cells), total)
return normalize_text("\n\n".join(parts))
def _presentation_shape_text(shape: Any) -> list[str]:
if getattr(shape, "has_table", False):
return [
"\t".join(normalize_text(cell.text) for cell in row.cells)
for row in shape.table.rows
]
if getattr(shape, "has_text_frame", False):
return [shape.text]
child_shapes = getattr(shape, "shapes", None)
if child_shapes is not None:
values: list[str] = []
for child in child_shapes:
values.extend(_presentation_shape_text(child))
return values
return []
def _extract_pptx_text(raw: bytes) -> str:
_validate_office_archive(raw, "pptx")
try:
presentation = Presentation(io.BytesIO(raw))
except Exception as exc:
raise ValueError(f"invalid PPTX file: {exc}") from exc
if len(presentation.slides) > _MAX_PRESENTATION_SLIDES:
raise ValueError(
f"PPTX contains too many slides (limit {_MAX_PRESENTATION_SLIDES})"
)
parts: list[str] = []
total = 0
for slide in presentation.slides:
slide_parts: list[str] = []
for shape in slide.shapes:
slide_parts.extend(_presentation_shape_text(shape))
total = _append_bounded_text(parts, "\n".join(slide_parts), total)
return normalize_text("\n\n".join(parts))
def _xml_local_name(tag: str) -> str:
return tag.rsplit("}", 1)[-1]
def _resolve_xlsx_relationship_target(
archive: zipfile.ZipFile,
target: str,
*,
source_part: str = "xl/workbook.xml",
) -> str:
"""按 OPC URI 规则解析内部关系目标,并保证结果仍位于 ZIP 根内。"""
raw_target = target.strip()
if (
raw_target != target
or not raw_target
or "\\" in raw_target
or any(unicodedata.category(char).startswith("C") for char in raw_target)
or re.search(r"%(?![0-9A-Fa-f]{2})", raw_target)
):
raise ValueError("XLSX workbook contains an unsafe worksheet path")
try:
parsed = urlsplit(raw_target)
except ValueError as exc:
raise ValueError("XLSX workbook contains an unsafe worksheet path") from exc
if parsed.scheme or parsed.netloc or parsed.query or parsed.fragment:
raise ValueError("XLSX workbook contains an unsafe worksheet path")
if re.search(r"%(?:2[fF]|5[cC]|0{2})", parsed.path):
raise ValueError("XLSX workbook contains an unsafe worksheet path")
try:
decoded_path = unquote(parsed.path, encoding="utf-8", errors="strict")
except UnicodeDecodeError as exc:
raise ValueError("XLSX workbook contains an unsafe worksheet path") from exc
if not decoded_path or "\\" in decoded_path or "\x00" in decoded_path or "%" in decoded_path:
raise ValueError("XLSX workbook contains an unsafe worksheet path")
parts = [] if decoded_path.startswith("/") else list(PurePosixPath(source_part).parent.parts)
for part in decoded_path.lstrip("/").split("/"):
if part in {"", "."}:
continue
if part == "..":
if not parts:
raise ValueError("XLSX workbook contains an unsafe worksheet path")
parts.pop()
continue
if unicodedata.category(part[0]).startswith("C") or any(
unicodedata.category(char).startswith("C") for char in part
):
raise ValueError("XLSX workbook contains an unsafe worksheet path")
parts.append(part)
if not parts:
raise ValueError("XLSX workbook contains an unsafe worksheet path")
member_name = "/".join(parts)
try:
member = archive.getinfo(member_name)
except KeyError as exc:
raise ValueError(
f"XLSX worksheet relationship target does not exist: {member_name}"
) from exc
if member.is_dir():
raise ValueError("XLSX worksheet relationship target must be a file")
return member_name
def _rewrite_xlsx_workbook_relationships(
raw: bytes,
replacements: Mapping[str, str],
) -> bytes:
"""把已验证的 worksheet Target 改为解析库稳定支持的包内绝对路径。"""
relationships_member = "xl/_rels/workbook.xml.rels"
output = io.BytesIO()
with zipfile.ZipFile(io.BytesIO(raw)) as source, zipfile.ZipFile(output, "w") as target:
target.comment = source.comment
for member in source.infolist():
if member.filename == relationships_member:
root = ET.fromstring(source.read(member))
pending = dict(replacements)
for element in root:
if _xml_local_name(element.tag) != "Relationship":
continue
relationship_id = element.attrib.get("Id")
if relationship_id in pending:
element.set("Target", pending.pop(relationship_id))
if pending:
raise ValueError(
"XLSX workbook relationship changed during normalization"
)
content = ET.tostring(root, encoding="utf-8", xml_declaration=True)
target.writestr(member, content)
continue
if member.is_dir():
target.writestr(member, b"")
continue
with source.open(member) as source_stream, target.open(
member,
"w",
force_zip64=True,
) as target_stream:
while chunk := source_stream.read(1024 * 1024):
target_stream.write(chunk)
return output.getvalue()
def _xlsx_sheet_merge_ranges(
raw: bytes,
) -> tuple[
dict[str, tuple[tuple[int, int, int, int], ...]],
dict[str, str],
]:
"""流式读取 XLSX 合并单元格,不把工作表 XML 整体载入内存。"""
relationship_namespace = (
"http://schemas.openxmlformats.org/officeDocument/2006/relationships"
)
try:
with zipfile.ZipFile(io.BytesIO(raw)) as archive:
relationships: dict[str, tuple[str, str, str]] = {}
with archive.open("xl/_rels/workbook.xml.rels") as stream:
for _, element in ET.iterparse(stream, events=("end",)):
if _xml_local_name(element.tag) != "Relationship":
element.clear()
continue
relationship_id = element.attrib.get("Id")
target = element.attrib.get("Target")
target_mode = element.attrib.get("TargetMode", "Internal")
relationship_type = element.attrib.get("Type", "")
if relationship_id:
if relationship_id in relationships:
raise ValueError(
"XLSX workbook contains duplicate relationship identifiers"
)
relationships[relationship_id] = (
target or "",
target_mode,
relationship_type,
)
element.clear()
sheet_paths: dict[str, str] = {}
normalized_targets: dict[str, str] = {}
with archive.open("xl/workbook.xml") as stream:
for _, element in ET.iterparse(stream, events=("end",)):
if _xml_local_name(element.tag) != "sheet":
element.clear()
continue
title = element.attrib.get("name")
relationship_id = element.attrib.get(
f"{{{relationship_namespace}}}id"
)
if title and relationship_id:
relationship = relationships.get(relationship_id)
if relationship is None:
raise ValueError(
f"XLSX worksheet {title!r} references a missing relationship"
)
target, target_mode, relationship_type = relationship
if target_mode.strip().lower() != "internal":
raise ValueError(
f"XLSX worksheet {title!r} uses an external relationship"
)
if not relationship_type.endswith("/worksheet"):
raise ValueError(
f"XLSX worksheet {title!r} has an invalid relationship type"
)
member_name = _resolve_xlsx_relationship_target(
archive,
target,
)
sheet_paths[title] = member_name
canonical_target = f"/{member_name}"
if target != canonical_target:
normalized_targets[relationship_id] = canonical_target
element.clear()
result: dict[str, tuple[tuple[int, int, int, int], ...]] = {}
total_ranges = 0
for title, member_name in sheet_paths.items():
ranges: list[tuple[int, int, int, int]] = []
with archive.open(member_name) as stream:
for _, element in ET.iterparse(stream, events=("end",)):
if _xml_local_name(element.tag) != "mergeCell":
element.clear()
continue
reference = element.attrib.get("ref")
if reference:
try:
boundaries = range_boundaries(reference)
except ValueError as exc:
raise ValueError(
f"XLSX worksheet {title!r} has an invalid merged range"
) from exc
ranges.append(boundaries)
total_ranges += 1
if total_ranges > _MAX_WORKBOOK_MERGED_RANGES:
raise ValueError(
"XLSX workbook contains too many merged ranges "
f"(limit {_MAX_WORKBOOK_MERGED_RANGES})"
)
element.clear()
result[title] = tuple(ranges)
return result, normalized_targets
except (KeyError, ET.ParseError, zipfile.BadZipFile) as exc:
raise ValueError(f"invalid XLSX workbook structure: {exc}") from exc
def _xlsx_header_end_row(
first_row: int,
rows: Mapping[int, Sequence[Any]],
merged_ranges: Sequence[tuple[int, int, int, int]],
) -> int:
"""在固定深度内闭包表头合并关系,忽略越界或跨空行的可疑级联。"""
header_end = first_row
maximum_end = first_row + _MAX_WORKBOOK_HEADER_ROWS - 1
has_header_hierarchy = any(
max_column > min_column and min_row == first_row
for min_column, min_row, max_column, _ in merged_ranges
)
changed = True
while changed:
changed = False
for min_column, min_row, max_column, max_row in merged_ranges:
if min_row < first_row or min_row > header_end or max_row < first_row:
continue
if max_column == min_column and not has_header_hierarchy:
continue
candidate = max_row + 1 if max_column > min_column else max_row
if candidate <= header_end or candidate > maximum_end:
continue
if (
max_column > min_column
and len(_xlsx_nonempty_values(rows.get(min_row, ()))) < 2
and len(_xlsx_nonempty_values(rows.get(candidate, ()))) < 2
):
continue
if any(
not rows.get(row_number)
for row_number in range(header_end + 1, candidate + 1)
):
continue
header_end = candidate
changed = True
return header_end
def _xlsx_headers(
title: str,
rows: Mapping[int, Sequence[Any]],
first_row: int,
header_end: int,
merged_ranges: Sequence[tuple[int, int, int, int]],
) -> list[str]:
width = max((len(row) for row in rows.values()), default=0)
for min_column, min_row, max_column, max_row in merged_ranges:
if min_row <= header_end and max_row >= first_row:
width = max(width, max_column)
if width > _MAX_WORKBOOK_COLUMNS:
raise ValueError(
f"XLSX worksheet {title!r} exceeds {_MAX_WORKBOOK_COLUMNS} columns"
)
matrix = [
list(rows.get(row_number, ())) + [None] * (width - len(rows.get(row_number, ())))
for row_number in range(first_row, header_end + 1)
]
for min_column, min_row, max_column, max_row in merged_ranges:
if min_row > header_end or max_row < first_row:
continue
source_row = max(first_row, min_row) - first_row
source_column = min_column - 1
source = matrix[source_row][source_column]
for row_number in range(max(first_row, min_row), min(header_end, max_row) + 1):
for column_number in range(min_column, max_column + 1):
matrix[row_number - first_row][column_number - 1] = source
headers: list[str] = []
for column in range(width):
components: list[str] = []
for row in matrix:
component = normalize_text(str(row[column] or ""))
if component and (not components or component != components[-1]):
components.append(component)
header = ".".join(components)
if not header:
raise ValueError(f"XLSX worksheet {title!r} contains an empty header")
headers.append(header)
if len(set(headers)) != len(headers):
raise ValueError(f"XLSX worksheet {title!r} contains duplicate headers")
return headers
def _xlsx_nonempty_values(row: Sequence[Any]) -> list[Any]:
return [value for value in row if value not in {None, ""}]
def _infer_xlsx_header_region(
title: str,
rows: Mapping[int, Sequence[Any]],
merged_ranges: Sequence[tuple[int, int, int, int]],
) -> tuple[int, int, list[str]]:
"""从有界前缀中推断表头,区分报表说明、多级表头和正文合并。"""
ordered_rows = sorted(rows)
if not ordered_rows:
return 0, 0, []
first_nonempty_row = ordered_rows[0]
horizontal_merge_rows = {
min_row
for min_column, min_row, max_column, _ in merged_ranges
if max_column > min_column
}
candidates: list[tuple[float, int, int, list[str]]] = []
for first_row in ordered_rows:
raw_values = _xlsx_nonempty_values(rows[first_row])
if not raw_values:
continue
if len(raw_values) < 2 and first_row in horizontal_merge_rows:
continue
header_end = _xlsx_header_end_row(first_row, rows, merged_ranges)
header_rows = {
row_number: rows[row_number]
for row_number in range(first_row, header_end + 1)
if row_number in rows
}
try:
headers = _xlsx_headers(
title,
header_rows,
first_row,
header_end,
merged_ranges,
)
except ValueError:
if header_end == first_row:
continue
header_end = first_row
try:
headers = _xlsx_headers(
title,
{first_row: rows[first_row]},
first_row,
first_row,
merged_ranges,
)
except ValueError:
continue
data_rows = [rows[row_number] for row_number in ordered_rows if row_number > header_end]
if any(len(row) > len(headers) for row in data_rows):
continue
text_ratio = sum(isinstance(value, str) for value in raw_values) / len(raw_values)
score = text_ratio * 6 + min(len(raw_values), 4)
if data_rows:
first_data_values = _xlsx_nonempty_values(data_rows[0])
score += 4 * min(len(first_data_values), len(headers)) / len(headers)
if first_data_values:
score += (
2
* sum(
not isinstance(value, str)
for value in first_data_values
)
/ len(first_data_values)
)
first_label = normalize_text(str(raw_values[0]))
if len(raw_values) <= 2 and _XLSX_REPORT_METADATA_PATTERN.match(first_label):
score -= 8
candidates.append((score, first_row, header_end, headers))
if not candidates:
first_row = first_nonempty_row
headers = _xlsx_headers(
title,
{first_row: rows[first_row]},
first_row,
first_row,
(),
)
return first_row, first_row, headers
_, first_row, header_end, headers = max(
candidates,
key=lambda candidate: (candidate[0], -candidate[1]),
)
return first_row, header_end, headers
def _extract_xlsx_records(
raw: bytes,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
_validate_office_archive(raw, "xlsx")
merged_by_sheet, normalized_targets = _xlsx_sheet_merge_ranges(raw)
workbook_raw = (
_rewrite_xlsx_workbook_relationships(raw, normalized_targets)
if normalized_targets
else raw
)
try:
workbook = load_workbook(
io.BytesIO(workbook_raw),
read_only=True,
data_only=True,
keep_links=False,
)
except Exception as exc:
raise ValueError(f"invalid XLSX file: {exc}") from exc
records: list[dict[str, Any]] = []
locators: list[dict[str, Any]] = []
total_cells = 0
try:
if len(workbook.worksheets) > _MAX_WORKBOOK_SHEETS:
raise ValueError(
f"XLSX contains too many worksheets (limit {_MAX_WORKBOOK_SHEETS})"
)
for sheet_index, worksheet in enumerate(workbook.worksheets):
reset_dimensions = getattr(worksheet, "reset_dimensions", None)
if callable(reset_dimensions):
reset_dimensions()
merged_ranges = merged_by_sheet.get(worksheet.title, ())
sheet_rows = 0
scanned_rows = 0
row_iterator = enumerate(
worksheet.iter_rows(values_only=True),
start=1,
)
buffered_rows: dict[int, Sequence[Any]] = {}
def normalized_row_values(
row: Sequence[Any],
sheet_title: str = worksheet.title,
) -> list[Any]:
values = list(row)
while values and values[-1] in {None, ""}:
values.pop()
if len(values) > _MAX_WORKBOOK_COLUMNS:
raise ValueError(
f"XLSX worksheet {sheet_title!r} exceeds "
f"{_MAX_WORKBOOK_COLUMNS} columns"
)
return values
for row_number, row in row_iterator:
scanned_rows += 1
if scanned_rows > _MAX_WORKBOOK_SCANNED_ROWS:
raise ValueError(
f"XLSX worksheet {worksheet.title!r} exceeds "
f"{_MAX_WORKBOOK_SCANNED_ROWS} scanned rows"
)
values = normalized_row_values(row)
if not values or all(value in {None, ""} for value in values):
continue
buffered_rows[row_number] = tuple(values)
if len(buffered_rows) >= _MAX_WORKBOOK_HEADER_SCAN_ROWS:
break
if not buffered_rows:
continue
_, header_end_row, headers = _infer_xlsx_header_region(
worksheet.title,
buffered_rows,
merged_ranges,
)
def append_record(
row_number: int,
values: Sequence[Any],
record_headers: Sequence[str] = tuple(headers),
locator_sheet_index: int = sheet_index,
sheet_title: str = worksheet.title,
) -> None:
nonlocal total_cells, sheet_rows
row_values = list(values)
if len(row_values) > len(record_headers):
raise ValueError(
f"XLSX worksheet {sheet_title!r} has a row wider than its header"
)
row_values.extend([None] * (len(record_headers) - len(row_values)))
record = {
header: _normalize_spreadsheet_value(value)
for header, value in zip(record_headers, row_values, strict=True)
}
if not any(value not in {"", None} for value in record.values()):
return
sheet_record_index = sheet_rows
sheet_rows += 1
total_cells += len(record_headers)
if sheet_rows > _MAX_WORKBOOK_ROWS:
raise ValueError(
f"XLSX worksheet {sheet_title!r} exceeds "
f"{_MAX_WORKBOOK_ROWS} data rows"
)
if total_cells > _MAX_WORKBOOK_CELLS:
raise ValueError(
f"XLSX workbook exceeds {_MAX_WORKBOOK_CELLS} populated cells"
)
records.append(record)
locators.append(
{
"kind": "xlsx",
"record_index": len(records),
"sheet_index": locator_sheet_index,
"sheet_name": sheet_title,
"row_number": row_number,
"sheet_record_index": sheet_record_index,
}
)
for row_number, values in buffered_rows.items():
if row_number > header_end_row:
append_record(row_number, values)
for row_number, row in row_iterator:
scanned_rows += 1
if scanned_rows > _MAX_WORKBOOK_SCANNED_ROWS:
raise ValueError(
f"XLSX worksheet {worksheet.title!r} exceeds "
f"{_MAX_WORKBOOK_SCANNED_ROWS} scanned rows"
)
values = normalized_row_values(row)
if not values or all(value in {None, ""} for value in values):
continue
append_record(row_number, values)
finally:
workbook.close()
return records, locators

View File

@@ -0,0 +1,299 @@
"""数据处理算法 - PDF 文档解析。"""
from __future__ import annotations
import io
import math
import re
import unicodedata
from collections.abc import Sequence
from typing import Literal
from pypdf import PdfReader
from ..text_utils import normalize_text
from ..types import (
DocumentNoiseSpan,
PdfPageText,
_MAX_EXTRACTED_TEXT_CHARS,
_MAX_PDF_PAGES,
_PdfLine,
)
_PDF_PAGE_NUMBER_LINE_PATTERN = re.compile(
r"^(?:页次\s*)?(?:第\s*)?(?P<page>\d+)\s*页\s*"
r"(?:(?:[/]\s*)?共\s*(?P<total>\d+)\s*页)?$"
)
_PDF_FRACTION_PAGE_LINE_PATTERN = re.compile(
r"^[—–-]?\s*(?P<page>\d+)\s*[/]\s*(?P<total>\d+)\s*[—–-]?$"
)
_PDF_CLASSIFICATION_LABEL_PATTERN = re.compile(
r"^(?:(?:秘密等级|密级)\s*)?(?:商密|秘密|机密|绝密)"
r"\s*(?:[【\[(][^】\])]{1,8}[】\])])?$"
)
_TOC_TITLE_PATTERN = re.compile(r"^(?:目\s*录|contents)$", re.IGNORECASE)
_TOC_LEADER_ENTRY_PATTERN = re.compile(
r"(?:[..…·•]\s*){3,}\s*\d{1,4}\s*$"
)
_TOC_NUMBERED_ENTRY_PATTERN = re.compile(
r"^(?:第[\u3400-\u4dbf\u4e00-\u9fff]{1,12}章|附表\s*\d+|\d+(?:\.\d+)+)"
r"\s+.+\s+\d{1,4}\s*$",
re.IGNORECASE,
)
_MARGIN_TEMPLATE_KEYWORDS = (
"",
"页次",
"版本",
"文件编码",
"秘密等级",
"密级",
"商密",
"confidential",
)
def extract_pdf_page_texts(raw: bytes) -> tuple[PdfPageText, ...]:
"""提取 PDF 各页文本,并保留与切片字符偏移一致的页范围。"""
if b"%PDF-" not in raw[:1024]:
raise ValueError("invalid PDF file: missing PDF header")
try:
reader = PdfReader(io.BytesIO(raw), strict=True)
except Exception as exc:
raise ValueError(f"invalid PDF file: {exc}") from exc
if reader.is_encrypted and not reader.decrypt(""):
raise ValueError("password-protected PDF files are not supported")
if len(reader.pages) > _MAX_PDF_PAGES:
raise ValueError(f"PDF contains too many pages (limit {_MAX_PDF_PAGES})")
pages: list[PdfPageText] = []
total = 0
has_text = False
for page_number, page in enumerate(reader.pages, start=1):
try:
text = normalize_text(page.extract_text() or "")
except Exception as exc:
raise ValueError(
f"failed to extract text from PDF page {page_number}: {exc}"
) from exc
if not text:
pages.append(
PdfPageText(
page_number=page_number,
text="",
source_start=total,
source_end=total,
)
)
continue
if has_text:
total += 2
start = total
total += len(text)
if total > _MAX_EXTRACTED_TEXT_CHARS:
raise ValueError(
f"extracted document text exceeds {_MAX_EXTRACTED_TEXT_CHARS} characters"
)
pages.append(
PdfPageText(
page_number=page_number,
text=text,
source_start=start,
source_end=total,
)
)
has_text = True
if not has_text:
raise ValueError(
"PDF contains no extractable text; scanned or image-only PDF files are not supported"
)
return tuple(pages)
def _pdf_page_lines(page: PdfPageText) -> tuple[_PdfLine, ...]:
lines: list[_PdfLine] = []
local_offset = 0
for raw_line in page.text.splitlines(keepends=True):
content = raw_line.rstrip("\r\n")
leading = len(content) - len(content.lstrip())
trailing = len(content.rstrip())
text = content.strip()
if text:
lines.append(
_PdfLine(
text=text,
start=page.source_start + local_offset + leading,
end=page.source_start + local_offset + trailing,
)
)
local_offset += len(raw_line)
return tuple(lines)
def _is_standalone_page_number(
text: str,
*,
physical_page: int,
page_count: int,
) -> bool:
normalized = unicodedata.normalize("NFKC", text).strip()
match = _PDF_PAGE_NUMBER_LINE_PATTERN.fullmatch(normalized)
if match is None:
match = _PDF_FRACTION_PAGE_LINE_PATTERN.fullmatch(normalized)
if match is None or int(match.group("page")) != physical_page:
return False
total = match.groupdict().get("total")
return total is None or int(total) == page_count
def _margin_signature(text: str) -> str:
normalized = unicodedata.normalize("NFKC", text).casefold()
normalized = re.sub(r"\s+", " ", normalized).strip()
if any(keyword in normalized for keyword in _MARGIN_TEMPLATE_KEYWORDS):
normalized = re.sub(r"\d+", "#", normalized)
return normalized
def _has_margin_metadata_keyword(text: str) -> bool:
normalized = unicodedata.normalize("NFKC", text).casefold()
return any(keyword in normalized for keyword in _MARGIN_TEMPLATE_KEYWORDS)
def _has_meaningful_margin_signature(signature: str) -> bool:
return len(re.sub(r"[#\W_]+", "", signature, flags=re.UNICODE)) >= 2
def _is_toc_leader_entry(text: str) -> bool:
return bool(_TOC_LEADER_ENTRY_PATTERN.search(text))
def _is_toc_numbered_entry(text: str) -> bool:
return bool(_TOC_NUMBERED_ENTRY_PATTERN.fullmatch(text))
def detect_pdf_document_noise(
pages: Sequence[PdfPageText],
) -> tuple[DocumentNoiseSpan, ...]:
"""识别 PDF 中的独立页码、重复页边内容和高置信目录。
规则只查看每页顶部 5 行和底部 3 行来推断页眉页脚;目录必须有
明显的点引导线密度,避免仅因正文中出现“目录”或章节标题而误删。
"""
page_lines = tuple(_pdf_page_lines(page) for page in pages)
detected: dict[tuple[int, int], DocumentNoiseSpan] = {}
def mark(
line: _PdfLine,
kind: Literal["page_number", "repeated_margin", "table_of_contents"],
) -> None:
detected.setdefault(
(line.start, line.end),
DocumentNoiseSpan(
start=line.start,
end=line.end,
kind=kind,
),
)
for page, lines in zip(pages, page_lines, strict=True):
for line in lines:
if _is_standalone_page_number(
line.text,
physical_page=page.page_number,
page_count=len(pages),
):
mark(line, "page_number")
outer_margin_lines = (*lines[:2], *lines[-2:])
for line in outer_margin_lines:
if _PDF_CLASSIFICATION_LABEL_PATTERN.fullmatch(line.text):
mark(line, "repeated_margin")
# 只在三页及以上文档中推断通用页眉页脚,避免短文档误删。
if len(pages) >= 3:
signature_pages: dict[str, set[int]] = {}
candidate_lines: list[tuple[int, _PdfLine, str]] = []
for page_index, lines in enumerate(page_lines):
boundary_lines = (
*((line, index < 2) for index, line in enumerate(lines[:5])),
*((line, index < 2) for index, line in enumerate(reversed(lines[-3:]))),
)
seen_ranges: set[tuple[int, int]] = set()
for line, is_outer_margin in boundary_lines:
line_range = (line.start, line.end)
if (
line_range in seen_ranges
or line_range in detected
or len(line.text) > 160
):
continue
seen_ranges.add(line_range)
if not is_outer_margin and not _has_margin_metadata_keyword(line.text):
continue
signature = _margin_signature(line.text)
if not _has_meaningful_margin_signature(signature):
continue
signature_pages.setdefault(signature, set()).add(page_index)
candidate_lines.append((page_index, line, signature))
minimum_pages = max(3, math.ceil(len(pages) * 0.3))
repeated_signatures = {
signature
for signature, matching_pages in signature_pages.items()
if len(matching_pages) >= minimum_pages
}
for _, line, signature in candidate_lines:
if signature in repeated_signatures:
mark(line, "repeated_margin")
# 先依据强证据判定目录页,再补充删除少量不带点引导线的编号目录项。
toc_active = False
for lines in page_lines:
content_lines = [
line for line in lines if (line.start, line.end) not in detected
]
leader_entries = [line for line in content_lines if _is_toc_leader_entry(line.text)]
titles = [line for line in content_lines if _TOC_TITLE_PATTERN.fullmatch(line.text)]
starts_toc = bool(titles and len(leader_entries) >= 2)
is_toc_dense = bool(
len(leader_entries) >= 3
and len(leader_entries) / max(1, len(content_lines)) >= 0.5
)
if not (starts_toc or (toc_active and is_toc_dense)):
toc_active = False
continue
toc_active = True
for line in content_lines:
if (
line in titles
or _is_toc_leader_entry(line.text)
or _is_toc_numbered_entry(line.text)
):
mark(line, "table_of_contents")
return tuple(sorted(detected.values(), key=lambda span: (span.start, span.end)))
def remove_document_noise(
text: str,
spans: Sequence[DocumentNoiseSpan],
*,
source_offset: int = 0,
) -> str:
"""按原文绝对偏移移除噪声,不改动调用方保留的原文及偏移。"""
text_end = source_offset + len(text)
intersections = sorted(
(
max(0, span.start - source_offset),
min(len(text), span.end - source_offset),
)
for span in spans
if span.start < text_end and span.end > source_offset
)
if not intersections:
return text
parts: list[str] = []
cursor = 0
for start, end in intersections:
if end <= cursor:
continue
if start > cursor:
parts.append(text[cursor:start])
cursor = end
parts.append(text[cursor:])
cleaned = normalize_text("".join(parts))
return re.sub(r"\n{3,}", "\n\n", cleaned)
def _extract_pdf_text(raw: bytes) -> str:
return "\n\n".join(page.text for page in extract_pdf_page_texts(raw) if page.text)

View File

@@ -0,0 +1,428 @@
"""数据处理算法 - 质量评分和去重。"""
from __future__ import annotations
import hashlib
import json
import math
import re
import unicodedata
from collections import Counter
from collections.abc import Iterable, Mapping, Sequence
from copy import deepcopy
from typing import Any
from .text_utils import normalize_text
from .types import (
_MAX_ANOMALY_TEXT_CHARS,
_MOJIBAKE_MARKERS,
_TOKEN_PATTERN,
ProcessedStructuredRecord,
QualityScore,
)
def estimate_token_count(text: str) -> int:
"""粗略估计文本的 token 数量。"""
return len(_TOKEN_PATTERN.findall(text))
def content_quality_flags(
text: str,
*,
min_chars: int = 20,
min_tokens: int = 5,
max_chars: int = _MAX_ANOMALY_TEXT_CHARS,
) -> tuple[str, ...]:
"""返回非结构化内容的确定性低质量原因。"""
if min_chars < 0 or min_tokens < 0 or max_chars <= 0:
raise ValueError("content quality limits must be non-negative")
normalized = normalize_text(text)
if not normalized:
return ("empty_content",)
flags: list[str] = []
if len(normalized) < min_chars or estimate_token_count(normalized) < min_tokens:
flags.append("content_too_short")
if len(normalized) > max_chars:
flags.append("content_too_long")
if any(marker in normalized for marker in _MOJIBAKE_MARKERS):
flags.append("mojibake")
nonspace = [char for char in normalized if not char.isspace()]
if nonspace:
readable_ratio = sum(
char.isprintable()
and unicodedata.category(char) not in {"Co", "Cs", "Cn"}
for char in nonspace
) / len(nonspace)
if readable_ratio < 0.85:
flags.append("low_printable_ratio")
if len(nonspace) >= 100:
most_common = Counter(nonspace).most_common(1)[0][1]
if most_common / len(nonspace) > 0.9:
flags.append("repetitive_content")
return tuple(dict.fromkeys(flags))
def is_low_quality_content(
text: str,
*,
min_chars: int = 20,
min_tokens: int = 5,
max_chars: int = _MAX_ANOMALY_TEXT_CHARS,
) -> bool:
"""判断内容是否命中任一低质量规则。"""
return bool(
content_quality_flags(
text,
min_chars=min_chars,
min_tokens=min_tokens,
max_chars=max_chars,
)
)
def _deduplicate_structured_entries(
entries: Sequence[ProcessedStructuredRecord],
) -> list[ProcessedStructuredRecord]:
"""仅按整条 canonical JSON 稳定去重,避免误删同 ID 的更新记录。"""
# canonical_record_json 位于 structured_processing延迟导入以断开循环依赖。
from .structured_processing import canonical_record_json
exact_seen: set[str] = set()
unique: list[ProcessedStructuredRecord] = []
for entry in entries:
record = entry.record
fingerprint = hashlib.sha256(canonical_record_json(record).encode("utf-8")).hexdigest()
if fingerprint in exact_seen:
continue
exact_seen.add(fingerprint)
unique.append(
ProcessedStructuredRecord(
entry.source_index,
deepcopy(dict(record)),
)
)
return unique
def deduplicate_structured_records(
records: Sequence[Mapping[str, Any]],
) -> list[dict[str, Any]]:
"""仅按整条 canonical JSON 稳定去重。"""
entries = [
ProcessedStructuredRecord(index, deepcopy(dict(record)))
for index, record in enumerate(records)
]
return [entry.record for entry in _deduplicate_structured_entries(entries)]
def _near_duplicate_features(text: str, shingle_size: int) -> tuple[str, ...]:
if isinstance(shingle_size, bool) or not isinstance(shingle_size, int):
raise TypeError("shingle_size must be an integer")
if shingle_size <= 0:
raise ValueError("shingle_size must be greater than 0")
tokens = re.findall(
r"[\u3400-\u4dbf\u4e00-\u9fff]|[A-Za-z0-9_]+",
normalize_text(text).casefold(),
)
if not tokens:
return ()
if len(tokens) < shingle_size:
return ("\x1f".join(tokens),)
return tuple(
"\x1f".join(tokens[index : index + shingle_size])
for index in range(len(tokens) - shingle_size + 1)
)
def near_duplicate_fingerprint(text: str, *, shingle_size: int = 3) -> str:
"""生成 64 位 SimHash 指纹,用于低成本近重复候选筛选。"""
if isinstance(shingle_size, bool) or not isinstance(shingle_size, int):
raise TypeError("shingle_size must be an integer")
if shingle_size <= 0:
raise ValueError("shingle_size must be greater than 0")
features = Counter(_near_duplicate_features(text, shingle_size))
if not features:
return "0" * 16
vector = [0] * 64
for feature, weight in features.items():
digest = int.from_bytes(hashlib.sha256(feature.encode("utf-8")).digest()[:8], "big")
for bit in range(64):
vector[bit] += weight if digest & (1 << bit) else -weight
fingerprint = sum(1 << bit for bit, value in enumerate(vector) if value >= 0)
return f"{fingerprint:016x}"
def fingerprints_are_near_duplicate(
left: str,
right: str,
*,
max_hamming_distance: int = 3,
) -> bool:
"""比较两个 64 位十六进制 SimHash 指纹。"""
if isinstance(max_hamming_distance, bool) or not isinstance(max_hamming_distance, int):
raise TypeError("max_hamming_distance must be an integer")
if not 0 <= max_hamming_distance <= 64:
raise ValueError("max_hamming_distance must be in [0, 64]")
if not re.fullmatch(r"[0-9a-fA-F]{16}", left) or not re.fullmatch(
r"[0-9a-fA-F]{16}", right
):
raise ValueError("fingerprints must be 16-character hexadecimal strings")
distance = (int(left, 16) ^ int(right, 16)).bit_count()
return distance <= max_hamming_distance
def is_near_duplicate(
left: str,
right: str,
*,
shingle_size: int = 3,
similarity_threshold: float = 0.9,
max_hamming_distance: int = 3,
) -> bool:
"""结合词片 Jaccard 和 SimHash 判断两段内容是否近重复。"""
if isinstance(similarity_threshold, bool) or not isinstance(
similarity_threshold, (int, float)
):
raise TypeError("similarity_threshold must be a number")
if not 0 <= similarity_threshold <= 1:
raise ValueError("similarity_threshold must be in [0, 1]")
if isinstance(max_hamming_distance, bool) or not isinstance(max_hamming_distance, int):
raise TypeError("max_hamming_distance must be an integer")
if not 0 <= max_hamming_distance <= 64:
raise ValueError("max_hamming_distance must be in [0, 64]")
left_normalized = normalize_text(left)
right_normalized = normalize_text(right)
if not left_normalized or not right_normalized:
return left_normalized == right_normalized
if left_normalized.casefold() == right_normalized.casefold():
return True
left_features = set(_near_duplicate_features(left_normalized, shingle_size))
right_features = set(_near_duplicate_features(right_normalized, shingle_size))
union = left_features | right_features
similarity = len(left_features & right_features) / len(union) if union else 1.0
if similarity >= similarity_threshold:
return True
return fingerprints_are_near_duplicate(
near_duplicate_fingerprint(left_normalized, shingle_size=shingle_size),
near_duplicate_fingerprint(right_normalized, shingle_size=shingle_size),
max_hamming_distance=max_hamming_distance,
)
def record_fingerprint(record: Mapping[str, Any]) -> str:
"""计算与字典键顺序无关的稳定记录指纹。"""
canonical = {
"instruction": normalize_text(str(record.get("instruction") or "")),
"input": normalize_text(str(record.get("input") or "")),
"output": normalize_text(str(record.get("output") or "")),
}
raw = json.dumps(canonical, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
def _readability_score(text: str) -> float:
if not text:
return 0.0
nonspace = [char for char in text if not char.isspace()]
if not nonspace:
return 0.0
printable_ratio = sum(char.isprintable() for char in nonspace) / len(nonspace)
useful_ratio = sum(
char.isalnum() or "\u3400" <= char <= "\u9fff" or unicodedata.category(char).startswith("P")
for char in nonspace
) / len(nonspace)
return round(100 * (0.65 * printable_ratio + 0.35 * useful_ratio), 2)
def _internal_duplicate_score(text: str) -> float:
units = [unit.strip().lower() for unit in re.split(r"[\n。!?;]+", text) if unit.strip()]
if len(units) <= 1:
return 100.0
return round(100 * len(set(units)) / len(units), 2)
def _source_relevance_score(record: Mapping[str, Any], source_content: str) -> float:
"""估算结果与来源文本的词元覆盖率。
这是无外部模型依赖、可重复的首版评分。没有来源文本(例如人工新增结果)
时不扣分;存在来源时,以结果中的有效词元被来源覆盖的比例计分。
"""
source = normalize_text(source_content)
if not source:
return 100.0
candidate = normalize_text(
"\n".join(
str(record.get(field) or "") for field in ("instruction", "input", "output")
)
)
def semantic_tokens(text: str) -> set[str]:
return {
token.lower()
for token in _TOKEN_PATTERN.findall(text)
if token.isalnum() or "\u3400" <= token <= "\u9fff"
}
source_tokens = semantic_tokens(source)
candidate_tokens = semantic_tokens(candidate)
if not candidate_tokens:
return 0.0
if not source_tokens:
return 0.0
return round(100 * len(candidate_tokens & source_tokens) / len(candidate_tokens), 2)
def score_quality(
record: Mapping[str, Any],
*,
min_output_length: int = 20,
source_content: str = "",
known_fingerprints: Iterable[str] = (),
threshold: float = 60.0,
) -> QualityScore:
"""按完整性、长度、可读性、来源相关性和重复度计算质量分。"""
if min_output_length <= 0:
raise ValueError("min_output_length must be greater than 0")
if not 0 <= threshold <= 100:
raise ValueError("threshold must be in [0, 100]")
instruction = normalize_text(str(record.get("instruction") or ""))
input_text = normalize_text(str(record.get("input") or ""))
output = normalize_text(str(record.get("output") or ""))
flags: list[str] = []
completeness = 100.0
if not instruction:
completeness -= 50
flags.append("missing_instruction")
if not output:
completeness -= 50
flags.append("missing_output")
output_length = len(output)
length_score = round(min(100.0, output_length / min_output_length * 100), 2)
if output_length < min_output_length:
flags.append("output_too_short")
readability = _readability_score("\n".join((instruction, input_text, output)))
if readability < 70:
flags.append("low_readability")
relevance = _source_relevance_score(record, source_content)
if source_content and relevance < 30:
flags.append("low_source_relevance")
fingerprint = record_fingerprint(record)
known = set(known_fingerprints)
duplicate = 0.0 if fingerprint in known else _internal_duplicate_score(output)
if duplicate == 0:
flags.append("duplicate_record")
elif duplicate < 70:
flags.append("repetitive_output")
overall = round(
completeness * 0.35
+ length_score * 0.20
+ readability * 0.20
+ relevance * 0.15
+ duplicate * 0.10,
2,
)
hard_valid = bool(instruction and output)
return QualityScore(
overall=overall,
completeness=completeness,
length=length_score,
readability=readability,
relevance=relevance,
duplicate=duplicate,
is_valid=hard_valid and overall >= threshold,
flags=tuple(flags),
fingerprint=fingerprint,
)
def _cosine_similarity(left: Sequence[float], right: Sequence[float]) -> float:
if not left or not right or len(left) != len(right):
return 0.0
dot = math.fsum(a * b for a, b in zip(left, right))
norm_left = math.sqrt(math.fsum(a * a for a in left))
norm_right = math.sqrt(math.fsum(b * b for b in right))
if not norm_left or not norm_right:
return 0.0
return dot / (norm_left * norm_right)
def semantic_quality_scores(
record: Mapping[str, Any],
*,
source_content: str = "",
embed_model: Any = None,
) -> dict[str, Any] | None:
"""用本地嵌入向量计算语义相关性0-100
返回 ``question_answer``(问题↔答案)、``answer_source``(答案↔来源,
无来源时缺省)与 ``overall``;嵌入模型不可用时返回 None 降级,不阻断流程。
"""
try:
if embed_model is None:
from .embedding import semantic_embedding_model
embed_model = semantic_embedding_model()
if embed_model is None:
return None
question = normalize_text(
" ".join(
str(record.get(field) or "")
for field in ("instruction", "input")
)
)
answer = normalize_text(
str(record.get("output") or "") or str(record.get("chosen") or "")
)
source = normalize_text(source_content)
texts = [text for text in {question, answer, source} if text]
if not texts:
return None
vectors = {text: embed_model.get_text_embedding(text) for text in texts}
except Exception:
return None
scores: dict[str, Any] = {}
if question and answer:
scores["question_answer"] = round(
100 * max(0.0, _cosine_similarity(vectors[question], vectors[answer])), 2
)
if answer and source:
scores["answer_source"] = round(
100 * max(0.0, _cosine_similarity(vectors[answer], vectors[source])), 2
)
if not scores:
return None
scores["overall"] = round(sum(scores.values()) / len(scores), 2)
return scores
def composite_overall(
*,
rule: float | None,
semantic: float | None = None,
judge: float | None = None,
) -> float:
"""三层加权组合:规则 35% + 语义 20% + 评审 45%,缺失层自动重归一。"""
if rule is None:
rule = 0.0
if judge is not None and semantic is not None:
overall = rule * 0.35 + semantic * 0.20 + judge * 0.45
elif semantic is not None:
overall = rule * 0.60 + semantic * 0.40
elif judge is not None:
overall = rule * 0.55 + judge * 0.45
else:
overall = rule
return round(max(0.0, min(100.0, overall)), 2)

View File

@@ -0,0 +1,809 @@
"""数据处理算法 - 结构化数据处理。"""
from __future__ import annotations
import hashlib
import json
import math
import re
from collections.abc import Iterable, Mapping, Sequence
from copy import deepcopy
from datetime import date, datetime, time
from decimal import Decimal
from typing import Any
from .parsers.json_parser import _extract_structured_records_with_locators
from .quality import (
_deduplicate_structured_entries,
content_quality_flags,
estimate_token_count,
)
from .text_utils import _normalize_field_name, normalize_text, structured_json_dumps
from .transforms import stable_split_assignments
from .types import (
_CHINESE_NAME_CONTEXT_PATTERN,
_EMAIL_PATTERN,
_ENGLISH_NAME_CONTEXT_PATTERN,
_ID_CARD_PATTERN,
_IDENTITY_FIELD_PATTERN,
_MAX_STRUCTURED_DEPTH,
_MAX_STRUCTURED_FIELDS,
_NAME_FIELD_NAMES,
_PHONE_PATTERN,
_STRUCTURED_OPTIONS,
MAX_QA_PAIRS_PER_ITEM,
ProcessedStructuredRecord,
StructuredPreprocessOption,
)
def _canonical_value(value: Any) -> Any:
if value is None or isinstance(value, (bool, int)):
return value
if isinstance(value, Decimal):
if not value.is_finite():
raise ValueError("non-finite JSON number is not allowed")
return value
if isinstance(value, float):
if not math.isfinite(value):
raise ValueError("non-finite JSON number is not allowed")
return value
if isinstance(value, str):
return normalize_text(value)
if isinstance(value, (datetime, date, time)):
return value.isoformat()
if isinstance(value, Mapping):
normalized: dict[str, Any] = {}
for key, item in sorted(value.items(), key=lambda pair: str(pair[0])):
normalized_key = normalize_text(str(key))
if not normalized_key:
raise ValueError("structured record contains an empty field name")
if normalized_key in normalized:
raise ValueError(
f"structured record fields collide after normalization: {normalized_key}"
)
normalized[normalized_key] = _canonical_value(item)
return normalized
if isinstance(value, (list, tuple)):
return [_canonical_value(item) for item in value]
if isinstance(value, (set, frozenset)):
items = [_canonical_value(item) for item in value]
return sorted(
items,
key=lambda item: structured_json_dumps(item, sort_keys=True),
)
if isinstance(value, (bytes, bytearray, memoryview)):
return bytes(value).hex()
return normalize_text(str(value))
def _is_empty_value(value: Any) -> bool:
if value is None:
return True
if isinstance(value, str):
return not normalize_text(value)
if isinstance(value, Mapping):
return not value or all(_is_empty_value(item) for item in value.values())
if isinstance(value, (list, tuple, set, frozenset)):
return not value or all(_is_empty_value(item) for item in value)
return False
def _is_name_field(field: Any) -> bool:
raw_field = normalize_text(str(field))
if not raw_field:
return False
# 只匹配明确表示自然人姓名的字段,避免将 table_name、product_name、
# chinese_name 等业务名称或元数据字段误判为个人敏感信息。
if _normalize_field_name(raw_field, "snake_case") in _NAME_FIELD_NAMES:
return True
# detect_structure 会使用点号生成扁平化路径(例如 profile.name此时仅
# 判断最后一个路径段,不能退回到宽泛的 ``*_name`` 后缀匹配。
if "." not in raw_field:
return False
leaf_field = raw_field.rsplit(".", 1)[-1]
return _normalize_field_name(leaf_field, "snake_case") in _NAME_FIELD_NAMES
def _embedded_structure(value: Any) -> Any:
if not isinstance(value, str):
return value
candidate = value.strip()
if not candidate or candidate[0] not in "[{":
return value
try:
parsed = json.loads(candidate)
except json.JSONDecodeError:
return value
return parsed if isinstance(parsed, (Mapping, list)) else value
def _structured_options(options: Iterable[str] | Mapping[str, Any]) -> set[str]:
if isinstance(options, str):
raise TypeError("options must be an iterable or mapping of option names")
if isinstance(options, Mapping):
enabled = {str(key) for key, value in options.items() if bool(value)}
else:
enabled = {str(option) for option in options}
unknown = enabled - _STRUCTURED_OPTIONS
if unknown:
raise ValueError(f"unsupported structured preprocess options: {', '.join(sorted(unknown))}")
return enabled
def _clean_invalid_structured_entries(
entries: Sequence[ProcessedStructuredRecord],
) -> list[ProcessedStructuredRecord]:
if not entries:
return []
fields: list[str] = []
for entry in entries:
record = entry.record
for field in record:
if field not in fields:
fields.append(field)
active_fields = [
field
for field in fields
if any(not _is_empty_value(entry.record.get(field)) for entry in entries)
]
if not active_fields:
return []
cleaned: list[ProcessedStructuredRecord] = []
for entry in entries:
record = entry.record
values = {field: deepcopy(record.get(field)) for field in active_fields}
# 清洗只依据整行是否为空。外键、父级 ID 等字段天然允许为空,不能
# 因为字段名以 *_id 结尾就把它们全部提升为联合必填项。
if all(_is_empty_value(value) for value in values.values()):
continue
cleaned.append(ProcessedStructuredRecord(entry.source_index, values))
return cleaned
def _percentile(values: Sequence[float], fraction: float) -> float:
if not values:
raise ValueError("cannot calculate a percentile of an empty sequence")
ordered = sorted(values)
position = (len(ordered) - 1) * fraction
lower = math.floor(position)
upper = math.ceil(position)
if lower == upper:
return ordered[lower]
weight = position - lower
return ordered[lower] * (1 - weight) + ordered[upper] * weight
def _filter_anomalous_structured_entries(
entries: Sequence[ProcessedStructuredRecord],
*,
iqr_multiplier: float = 1.5,
) -> list[ProcessedStructuredRecord]:
"""按字段级数值 IQR、乱码和极端文本长度过滤异常记录。"""
if iqr_multiplier <= 0:
raise ValueError("iqr_multiplier must be greater than 0")
numeric_values: dict[str, list[float]] = {}
text_lengths: dict[str, list[float]] = {}
for entry in entries:
record = entry.record
for field, value in record.items():
if (
isinstance(value, (int, float))
and not isinstance(value, bool)
and not _IDENTITY_FIELD_PATTERN.search(
_normalize_field_name(field, "snake_case")
)
):
number = float(value)
if math.isfinite(number):
numeric_values.setdefault(field, []).append(number)
elif isinstance(value, str) and value:
text_lengths.setdefault(field, []).append(float(len(value)))
numeric_bounds: dict[str, tuple[float, float]] = {}
for field, values in numeric_values.items():
# 小样本不做统计异常判断,避免把合法长尾值误删。
if len(values) < 8:
continue
first_quartile = _percentile(values, 0.25)
third_quartile = _percentile(values, 0.75)
spread = third_quartile - first_quartile
numeric_bounds[field] = (
first_quartile - iqr_multiplier * spread,
third_quartile + iqr_multiplier * spread,
)
text_upper_bounds: dict[str, float] = {}
for field, lengths in text_lengths.items():
if len(lengths) < 8:
continue
first_quartile = _percentile(lengths, 0.25)
third_quartile = _percentile(lengths, 0.75)
spread = third_quartile - first_quartile
text_upper_bounds[field] = max(512.0, third_quartile + 3 * spread)
accepted: list[ProcessedStructuredRecord] = []
for entry in entries:
record = entry.record
anomalous = False
for field, value in record.items():
if (
isinstance(value, (int, float))
and not isinstance(value, bool)
and not _IDENTITY_FIELD_PATTERN.search(
_normalize_field_name(field, "snake_case")
)
):
number = float(value)
if not math.isfinite(number):
anomalous = True
break
bounds = numeric_bounds.get(field)
if bounds and not bounds[0] <= number <= bounds[1]:
anomalous = True
break
if isinstance(value, str):
flags = content_quality_flags(value, min_chars=0, min_tokens=0)
if {"content_too_long", "mojibake", "low_printable_ratio"} & set(flags):
anomalous = True
break
upper_bound = text_upper_bounds.get(field)
if upper_bound is not None and len(value) > upper_bound:
anomalous = True
break
if not anomalous:
accepted.append(
ProcessedStructuredRecord(
entry.source_index,
deepcopy(dict(record)),
)
)
return accepted
def _preview_content(item: Mapping[str, Any]) -> str:
for field in ("edited_content", "editedContent", "original_content", "originalContent", "content"):
value = item.get(field)
if value is not None:
return normalize_text(str(value))
return ""
def _standard_fields(content: str) -> tuple[str, str, str]:
if not content:
return "", "", ""
try:
payload = json.loads(content)
except json.JSONDecodeError:
payload = None
if isinstance(payload, Mapping):
instruction = next(
(
str(payload[key])
for key in ("instruction", "question", "prompt")
if payload.get(key) is not None
),
"",
)
input_text = next(
(str(payload[key]) for key in ("input", "context") if payload.get(key) is not None),
"",
)
output = next(
(str(payload[key]) for key in ("output", "answer", "response") if payload.get(key) is not None),
"",
)
if instruction or output:
return normalize_text(instruction), normalize_text(input_text), normalize_text(output)
question_answer = re.match(
r"^\s*(?:问|question)\s*[:]\s*(.+?)(?:\n|\r\n?)\s*(?:答|answer)\s*[:]\s*(.+)\s*$",
content,
flags=re.IGNORECASE | re.DOTALL,
)
if question_answer:
return normalize_text(question_answer.group(1)), "", normalize_text(question_answer.group(2))
lines = [line.strip() for line in content.splitlines() if line.strip()]
first_line = re.sub(r"^(?:问|question)\s*[:]\s*", "", lines[0], flags=re.IGNORECASE)
output = normalize_text("\n".join(lines[1:])) if len(lines) > 1 else normalize_text(content)
return normalize_text(first_line), "", output
def _protected_markdown_ranges(
text: str,
*,
preserve_code_blocks: bool,
preserve_tables: bool,
preserve_lists: bool,
) -> list[tuple[int, int]]:
"""找出不应从中间切开的 Markdown 代码块、表格和连续列表。"""
lines: list[tuple[int, int, str]] = []
cursor = 0
for raw_line in text.splitlines(keepends=True):
end = cursor + len(raw_line)
lines.append((cursor, end, raw_line.rstrip("\r\n")))
cursor = end
if cursor < len(text) or not lines:
lines.append((cursor, len(text), text[cursor:]))
ranges: list[tuple[int, int]] = []
code_line_indexes: set[int] = set()
if preserve_code_blocks:
open_block: tuple[int, str, int] | None = None
for index, (start, end, content) in enumerate(lines):
fence = re.match(r"^\s*(`{3,}|~{3,})", content)
if not fence:
continue
marker = fence.group(1)[0]
length = len(fence.group(1))
if open_block is None:
open_block = (index, marker, length)
continue
first_index, open_marker, open_length = open_block
if marker == open_marker and length >= open_length:
ranges.append((lines[first_index][0], end))
code_line_indexes.update(range(first_index, index + 1))
open_block = None
if open_block is not None:
first_index = open_block[0]
ranges.append((lines[first_index][0], len(text)))
code_line_indexes.update(range(first_index, len(lines)))
if preserve_tables:
index = 0
while index + 1 < len(lines):
if index in code_line_indexes:
index += 1
continue
header = lines[index][2].strip()
separator = lines[index + 1][2].strip().strip("|")
cells = [cell.strip() for cell in separator.split("|")]
if (
"|" not in header
or len(cells) < 2
or not all(re.fullmatch(r":?-{3,}:?", cell) for cell in cells)
):
index += 1
continue
end_index = index + 1
while (
end_index + 1 < len(lines)
and end_index + 1 not in code_line_indexes
and lines[end_index + 1][2].strip()
and "|" in lines[end_index + 1][2]
):
end_index += 1
ranges.append((lines[index][0], lines[end_index][1]))
index = end_index + 1
if preserve_lists:
list_pattern = re.compile(r"^\s*(?:[-+*]|\d+[.)])\s+\S")
continuation_pattern = re.compile(r"^\s{2,}\S")
index = 0
while index < len(lines):
if index in code_line_indexes or not list_pattern.match(lines[index][2]):
index += 1
continue
end_index = index
item_count = 1
while end_index + 1 < len(lines) and end_index + 1 not in code_line_indexes:
next_line = lines[end_index + 1][2]
if list_pattern.match(next_line):
item_count += 1
end_index += 1
elif continuation_pattern.match(next_line):
end_index += 1
else:
break
if item_count >= 2:
ranges.append((lines[index][0], lines[end_index][1]))
index = end_index + 1
merged: list[tuple[int, int]] = []
for start, end in sorted(ranges):
if merged and start < merged[-1][1]:
merged[-1] = (merged[-1][0], max(merged[-1][1], end))
else:
merged.append((start, end))
return merged
def extract_structured_records(text: str, file_format: str) -> list[dict[str, Any]]:
"""从 JSON、JSONL 或 CSV 中提取规范化记录。"""
records, _ = _extract_structured_records_with_locators(text, file_format)
return records
def canonical_record_json(record: Mapping[str, Any]) -> str:
"""生成与字段顺序无关、可用于比较和落库的 canonical JSON。"""
if not isinstance(record, Mapping):
raise TypeError("record must be a mapping")
return structured_json_dumps(_canonical_value(record), sort_keys=True)
def normalize_structured_record(
record: Mapping[str, Any],
*,
field_name_style: str = "snake_case",
) -> dict[str, Any]:
"""规范字段名、Unicode/空白、容器类型和不可 JSON 化的标量。"""
if not isinstance(record, Mapping):
raise TypeError("record must be a mapping")
normalized: dict[str, Any] = {}
for key, value in record.items():
normalized_key = _normalize_field_name(key, field_name_style)
if not normalized_key:
raise ValueError("structured record contains an empty field name")
if normalized_key in normalized:
raise ValueError(
f"structured record fields collide after normalization: {normalized_key}"
)
normalized[normalized_key] = _canonical_value(value)
return dict(sorted(normalized.items()))
def flatten_structured_record(
record: Mapping[str, Any],
*,
separator: str = ".",
) -> dict[str, Any]:
"""把嵌套对象展平;数组保留为 canonical JSON 兼容值。"""
if not isinstance(record, Mapping):
raise TypeError("record must be a mapping")
if not separator:
raise ValueError("separator cannot be empty")
flattened: dict[str, Any] = {}
def visit(value: Any, path: tuple[str, ...], depth: int) -> None:
if depth > _MAX_STRUCTURED_DEPTH:
raise ValueError(
f"structured record nesting exceeds {_MAX_STRUCTURED_DEPTH} levels"
)
value = _embedded_structure(value)
if isinstance(value, Mapping):
if not value and path:
key = separator.join(path)
flattened[key] = {}
return
for child_key, child_value in value.items():
normalized_key = normalize_text(str(child_key))
if not normalized_key:
raise ValueError("structured record contains an empty field name")
visit(child_value, (*path, normalized_key), depth + 1)
return
key = separator.join(path)
if key in flattened:
raise ValueError(f"structured record fields collide while flattening: {key}")
flattened[key] = _canonical_value(value)
if len(flattened) > _MAX_STRUCTURED_FIELDS:
raise ValueError(
f"structured record exceeds {_MAX_STRUCTURED_FIELDS} flattened fields"
)
for field, value in record.items():
field_name = normalize_text(str(field))
if not field_name:
raise ValueError("structured record contains an empty field name")
visit(value, (field_name,), 1)
return flattened
def filter_anomalous_structured_records(
records: Sequence[Mapping[str, Any]],
*,
iqr_multiplier: float = 1.5,
) -> list[dict[str, Any]]:
"""按字段级数值 IQR、乱码和极端文本长度过滤异常记录。"""
entries = [
ProcessedStructuredRecord(index, deepcopy(dict(record)))
for index, record in enumerate(records)
]
return [
entry.record
for entry in _filter_anomalous_structured_entries(
entries,
iqr_multiplier=iqr_multiplier,
)
]
def desensitize_pii(text: str) -> tuple[str, dict[str, int]]:
"""掩码邮箱、手机号、身份证号及有明确上下文的姓名。"""
if not isinstance(text, str):
raise TypeError("text must be str")
counts: dict[str, int] = {"email": 0, "phone": 0, "id_card": 0}
def replace(pattern: re.Pattern[str], replacement: str, kind: str, value: str) -> str:
def replacer(_: re.Match[str]) -> str:
counts[kind] += 1
return replacement
return pattern.sub(replacer, value)
masked = replace(_EMAIL_PATTERN, "[EMAIL]", "email", text)
masked = replace(_ID_CARD_PATTERN, "[ID_CARD]", "id_card", masked)
masked = replace(_PHONE_PATTERN, "[PHONE]", "phone", masked)
def replace_context_name(match: re.Match[str]) -> str:
counts["name"] = counts.get("name", 0) + 1
return f"{match.group('label')}{match.group('separator')}[NAME]"
masked = _CHINESE_NAME_CONTEXT_PATTERN.sub(replace_context_name, masked)
masked = _ENGLISH_NAME_CONTEXT_PATTERN.sub(replace_context_name, masked)
counts["total"] = sum(counts.values())
return masked, counts
def desensitize_structured_record(
record: Mapping[str, Any],
) -> tuple[dict[str, Any], dict[str, int]]:
"""递归脱敏结构化姓名字段及任意文本中的手机号、邮箱、身份证号。"""
if not isinstance(record, Mapping):
raise TypeError("record must be a mapping")
counts: dict[str, int] = {"email": 0, "phone": 0, "id_card": 0}
def add_counts(values: Mapping[str, int]) -> None:
for kind, count in values.items():
if kind != "total" and count:
counts[kind] = counts.get(kind, 0) + count
def visit(value: Any, field: Any = "") -> Any:
if isinstance(value, Mapping):
return {key: visit(item, key) for key, item in value.items()}
if isinstance(value, list):
return [visit(item, field) for item in value]
if isinstance(value, tuple):
return [visit(item, field) for item in value]
if _is_name_field(field) and not _is_empty_value(value):
counts["name"] = counts.get("name", 0) + 1
return "[NAME]"
if isinstance(value, str):
masked, found = desensitize_pii(value)
add_counts(found)
return masked
return deepcopy(value)
masked = {key: visit(value, key) for key, value in record.items()}
counts["total"] = sum(counts.values())
return masked, counts
def preprocess_structured_records_with_lineage(
records: Iterable[Mapping[str, Any]],
options: Iterable[str] | Mapping[str, Any],
) -> list[ProcessedStructuredRecord]:
"""执行结构化预处理,并保留每条结果在原始输入中的稳定索引。"""
enabled = _structured_options(options)
current: list[ProcessedStructuredRecord] = []
for source_index, record in enumerate(records):
if not isinstance(record, Mapping):
if "clean_invalid" in enabled:
continue
raise TypeError("structured records must contain mappings")
value = deepcopy(dict(record))
if "detect_structure" in enabled:
value = flatten_structured_record(value)
if "normalize_format" in enabled:
value = normalize_structured_record(value)
current.append(ProcessedStructuredRecord(source_index, value))
if "clean_invalid" in enabled:
current = _clean_invalid_structured_entries(current)
if "filter_anomaly" in enabled:
current = _filter_anomalous_structured_entries(current)
if "deduplicate" in enabled:
current = _deduplicate_structured_entries(current)
if "desensitize" in enabled:
current = [
ProcessedStructuredRecord(
entry.source_index,
desensitize_structured_record(entry.record)[0],
)
for entry in current
]
return current
def preprocess_structured_records(
records: Iterable[Mapping[str, Any]],
options: Iterable[str] | Mapping[str, Any],
) -> list[dict[str, Any]]:
"""按界面选项执行确定性、无副作用的结构化数据预处理。"""
return [
entry.record
for entry in preprocess_structured_records_with_lineage(records, options)
]
def protected_context_ranges(
text: str,
*,
preserve_code_blocks: bool = True,
preserve_tables: bool = True,
preserve_lists: bool = True,
) -> tuple[tuple[int, int], ...]:
"""返回代码块、表格和列表的不可拆分区间。
返回的偏移量基于规范化后的文本;调用方应在同一份 ``normalize_text``
结果上使用这些区间。
"""
normalized = normalize_text(text)
return tuple(
_protected_markdown_ranges(
normalized,
preserve_code_blocks=preserve_code_blocks,
preserve_tables=preserve_tables,
preserve_lists=preserve_lists,
)
)
def expand_to_context_boundaries(
text: str,
start: int,
end: int,
*,
preserve_paragraph: bool = True,
preserve_code_blocks: bool = True,
preserve_tables: bool = True,
preserve_lists: bool = True,
) -> tuple[int, int]:
"""将一个文本区间扩展到段落及受保护 Markdown 结构边界。"""
normalized = normalize_text(text)
if isinstance(start, bool) or isinstance(end, bool):
raise TypeError("start and end must be integers")
if not isinstance(start, int) or not isinstance(end, int):
raise TypeError("start and end must be integers")
if not 0 <= start <= end <= len(normalized):
raise ValueError("start and end must define a valid normalized text range")
expanded_start = start
expanded_end = end
if preserve_paragraph and normalized:
paragraph_start = normalized.rfind("\n\n", 0, start)
expanded_start = 0 if paragraph_start < 0 else paragraph_start + 2
paragraph_end = normalized.find("\n\n", end)
expanded_end = len(normalized) if paragraph_end < 0 else paragraph_end
ranges = protected_context_ranges(
normalized,
preserve_code_blocks=preserve_code_blocks,
preserve_tables=preserve_tables,
preserve_lists=preserve_lists,
)
changed = True
while changed:
changed = False
for range_start, range_end in ranges:
overlaps = range_start < expanded_end and range_end > expanded_start
contains_boundary = (
range_start <= expanded_start < range_end
or range_start < expanded_end <= range_end
)
if not overlaps and not contains_boundary:
continue
next_start = min(expanded_start, range_start)
next_end = max(expanded_end, range_end)
if (next_start, next_end) != (expanded_start, expanded_end):
expanded_start, expanded_end = next_start, next_end
changed = True
return expanded_start, expanded_end
def merge_short_blocks(
blocks: Iterable[str],
*,
min_token_count: int = 100,
separator: str = "\n\n",
) -> list[str]:
"""按原顺序合并短内容块,并把末尾残块归入前一块。"""
if isinstance(min_token_count, bool) or not isinstance(min_token_count, int):
raise TypeError("min_token_count must be an integer")
if min_token_count <= 0:
raise ValueError("min_token_count must be greater than 0")
if not isinstance(separator, str):
raise TypeError("separator must be str")
merged: list[str] = []
pending: list[str] = []
pending_tokens = 0
for block in blocks:
if not isinstance(block, str):
raise TypeError("blocks must contain strings")
normalized = normalize_text(block)
if not normalized:
continue
token_count = estimate_token_count(normalized)
if not pending and token_count >= min_token_count:
merged.append(normalized)
continue
pending.append(normalized)
pending_tokens += token_count
if pending_tokens >= min_token_count:
merged.append(separator.join(pending))
pending = []
pending_tokens = 0
if pending:
tail = separator.join(pending)
if merged:
merged[-1] = separator.join((merged[-1], tail))
else:
merged.append(tail)
return merged
def generate_standard_records(
preview_items: Iterable[Mapping[str, Any]],
*,
qa_pairs_per_item: int = 1,
semantic_enrichment: bool = False,
split: Mapping[str, int] | None = None,
split_seed: str = "",
) -> list[dict[str, Any]]:
"""把预览内容确定性转换为标准 instruction/input/output 记录。
该函数只负责本地标准化,不冒充 LLM服务层可将其作为无模型模式或
LLM 响应解析后的统一落库步骤。
"""
if not 1 <= qa_pairs_per_item <= MAX_QA_PAIRS_PER_ITEM:
raise ValueError(
f"qa_pairs_per_item must be in [1, {MAX_QA_PAIRS_PER_ITEM}]"
)
prefixes = (
"请结合实际情况说明:",
"请用通俗易懂的方式说明:",
"请从实际应用角度说明:",
"请简洁自然地说明:",
"请详细解答:",
)
results: list[dict[str, Any]] = []
for item_index, item in enumerate(preview_items):
content = _preview_content(item)
instruction, input_text, output = _standard_fields(content)
preview_id = str(item.get("id") or f"preview-{item_index + 1}")
for variant_index in range(qa_pairs_per_item):
variant_instruction = instruction
if variant_index:
if semantic_enrichment:
prefix = prefixes[variant_index % len(prefixes)]
if variant_index >= len(prefixes):
prefix = (
f"{prefix.removesuffix('')}"
f"(问法 {variant_index + 1}"
)
variant_instruction = f"{prefix}{instruction}"
else:
variant_instruction = f"{instruction}(问法 {variant_index + 1})"
raw_id = f"{preview_id}:{variant_index + 1}"
result_id = f"result_{hashlib.sha256(raw_id.encode('utf-8')).hexdigest()[:16]}"
status = "valid" if variant_instruction and output else "invalid"
results.append(
{
"id": result_id,
"preview_item_id": preview_id,
"instruction": variant_instruction,
"input": input_text,
"output": output,
"original_instruction": variant_instruction,
"original_input": input_text,
"original_output": output,
"status": status,
"split": "train",
}
)
assignments = stable_split_assignments(
[str(result["id"]) for result in results],
split,
seed=split_seed,
)
for result, assignment in zip(results, assignments, strict=True):
result["split"] = assignment
return results

View File

@@ -0,0 +1,328 @@
"""数据处理算法 - 文本处理工具。"""
from __future__ import annotations
import csv
import io
import json
import math
import re
import unicodedata
from datetime import date, datetime, time
from decimal import Decimal
from pathlib import Path
from typing import Any, Literal
from collections.abc import Mapping
from .types import ParsedText, TextFormat
# 格式别名映射
_FORMAT_ALIASES: dict[str, TextFormat] = {
"json": "json",
"jsonl": "jsonl",
"ndjson": "jsonl",
"csv": "csv",
"tsv": "csv",
"md": "markdown",
"markdown": "markdown",
"txt": "txt",
"text": "txt",
"pdf": "pdf",
"docx": "docx",
"xlsx": "xlsx",
"pptx": "pptx",
}
# 旧版 Office 格式映射
_LEGACY_OFFICE_FORMATS: dict[str, str] = {
"doc": "docx",
"xls": "xlsx",
"ppt": "pptx",
}
# Office Open XML 格式集合
_OFFICE_OPEN_XML_FORMATS = {"docx", "xlsx", "pptx"}
# 文本提取限制
_MAX_EXTRACTED_TEXT_CHARS = 20_000_000
def decode_utf8(raw: bytes | bytearray | memoryview | str) -> str:
"""严格解码 UTF-8 文本,并移除可选 BOM。
不使用 ``errors='replace'``,避免上传内容损坏后仍被静默接收。
"""
if isinstance(raw, str):
return raw.removeprefix("\ufeff")
if not isinstance(raw, (bytes, bytearray, memoryview)):
raise TypeError("raw must be bytes-like or str")
try:
return bytes(raw).decode("utf-8-sig")
except UnicodeDecodeError as exc:
raise ValueError(f"content is not valid UTF-8 at byte {exc.start}") from exc
def parse_utf8_text(raw: bytes | bytearray | memoryview | str) -> str:
"""``decode_utf8`` 的语义化别名,供上传服务直接调用。"""
return decode_utf8(raw)
def normalize_text(text: str) -> str:
"""规范 Unicode、换行和行尾空白同时保留段落结构。"""
if not isinstance(text, str):
raise TypeError("text must be str")
normalized = unicodedata.normalize("NFKC", text.removeprefix("\ufeff"))
normalized = normalized.replace("\r\n", "\n").replace("\r", "\n")
normalized = "".join(
char
for char in normalized
if char in {"\n", "\t"} or not unicodedata.category(char).startswith("C")
)
lines = [re.sub(r"[\t \f\v]+$", "", line) for line in normalized.split("\n")]
return "\n".join(lines).strip()
def _normalize_format(value: str | None) -> TextFormat | None:
if value is None:
return None
normalized = value.strip().lower().removeprefix(".")
if normalized in _LEGACY_OFFICE_FORMATS:
replacement = _LEGACY_OFFICE_FORMATS[normalized]
raise ValueError(
f"legacy .{normalized} format is not supported; "
f"convert the file to .{replacement} and upload again"
)
try:
return _FORMAT_ALIASES[normalized]
except KeyError as exc:
raise ValueError(f"unsupported text format: {value}") from exc
def detect_text_format(
*,
filename: str | None = None,
text: str = "",
file_format: str | None = None,
) -> TextFormat:
"""按显式格式、扩展名和内容特征依次识别文本格式。"""
explicit = _normalize_format(file_format)
if explicit:
return explicit
if filename:
suffix = Path(filename).suffix.lower().removeprefix(".")
if suffix in _LEGACY_OFFICE_FORMATS:
replacement = _LEGACY_OFFICE_FORMATS[suffix]
raise ValueError(
f"legacy .{suffix} format is not supported; "
f"convert the file to .{replacement} and upload again"
)
detected = _FORMAT_ALIASES.get(suffix)
if detected:
return detected
stripped = text.strip()
if stripped:
if stripped[0] in "[{":
try:
json.loads(stripped)
except json.JSONDecodeError:
pass
else:
return "json"
nonempty_lines = [line for line in stripped.splitlines() if line.strip()]
if len(nonempty_lines) > 1:
try:
for line in nonempty_lines:
json.loads(line)
except json.JSONDecodeError:
pass
else:
return "jsonl"
if re.search(r"(?m)^(?:#{1,6}\s+|```|~~~)", stripped) or re.search(
r"(?m)^\s*\|.+\|\s*$", stripped
):
return "markdown"
sample = stripped[:8192]
try:
dialect = csv.Sniffer().sniff(sample, delimiters=",\t;")
rows = list(csv.reader(io.StringIO(sample), dialect))
if len(rows) >= 2 and len(rows[0]) >= 2:
return "csv"
except csv.Error:
pass
return "txt"
def _append_bounded_text(parts: list[str], value: Any, total: int) -> int:
text = normalize_text(str(value or ""))
if not text:
return total
total += len(text)
if total > _MAX_EXTRACTED_TEXT_CHARS:
raise ValueError(
f"extracted document text exceeds {_MAX_EXTRACTED_TEXT_CHARS} characters"
)
parts.append(text)
return total
def _normalize_spreadsheet_value(value: Any) -> Any:
if value is None:
return ""
if isinstance(value, str):
return normalize_text(value)
if isinstance(value, (datetime, date, time)):
return value.isoformat()
if isinstance(value, (bool, int, float)):
return value
return normalize_text(str(value))
def _normalize_value(value: Any) -> Any:
if isinstance(value, str):
return normalize_text(value)
if isinstance(value, Mapping):
return {normalize_text(str(key)): _normalize_value(item) for key, item in value.items()}
if isinstance(value, list):
return [_normalize_value(item) for item in value]
return value
def parse_text_content(
raw: bytes | bytearray | memoryview | str,
*,
filename: str | None = None,
file_format: str | None = None,
) -> ParsedText:
"""安全解析 UTF-8 文本、文本型 PDF 和现代 Office 文件。"""
# 解析器依赖 text_utilsnormalize_text 等),这里延迟导入以断开循环依赖。
from .parsers.json_parser import _extract_structured_records_with_locators
from .parsers.office import _extract_docx_text, _extract_pptx_text, _extract_xlsx_records
from .parsers.pdf import _extract_pdf_text
detected_format = detect_text_format(
filename=filename,
text="",
file_format=file_format,
)
if detected_format in _OFFICE_OPEN_XML_FORMATS or detected_format == "pdf":
binary = _binary_bytes(raw, detected_format)
if detected_format == "pdf":
text = _extract_pdf_text(binary)
return ParsedText(format=detected_format, text=text, records=())
if detected_format == "docx":
text = _extract_docx_text(binary)
return ParsedText(format=detected_format, text=text, records=())
if detected_format == "pptx":
text = _extract_pptx_text(binary)
return ParsedText(format=detected_format, text=text, records=())
records, record_locators = _extract_xlsx_records(binary)
text = "\n".join(
json.dumps(record, ensure_ascii=False, separators=(",", ":"))
for record in records
)
return ParsedText(
format=detected_format,
text=normalize_text(text),
records=tuple(records),
record_locators=tuple(record_locators),
)
decoded_text = decode_utf8(raw)
detected_format = detect_text_format(
filename=filename,
text=decoded_text,
file_format=file_format,
)
# JSON/JSONL 是有损规范化的禁区NFKC、控制字符删除或 trim 都可能改变字段值、
# 掩盖非法输入,甚至把原本合法的字符串变成语法错误。其他格式保持历史行为。
text = (
decoded_text
if detected_format in {"json", "jsonl"}
else normalize_text(decoded_text)
)
records: list[dict[str, Any]] = []
record_locators: list[dict[str, Any]] = []
if detected_format in {"json", "jsonl", "csv"}:
records, record_locators = _extract_structured_records_with_locators(
text,
detected_format,
)
return ParsedText(
format=detected_format,
text=text,
records=tuple(records),
record_locators=tuple(record_locators),
)
def _normalize_field_name(value: Any, style: str) -> str:
name = normalize_text(str(value))
if style == "preserve":
return name
if style == "lower":
return name.lower()
if style != "snake_case":
raise ValueError("field_name_style must be snake_case, lower or preserve")
name = re.sub(r"(?<=[a-z0-9])(?=[A-Z])", "_", name)
name = re.sub(r"[^\w\u3400-\u4dbf\u4e00-\u9fff]+", "_", name, flags=re.UNICODE)
return re.sub(r"_+", "_", name).strip("_").lower()
def structured_json_dumps(value: Any, *, sort_keys: bool = False) -> str:
"""序列化紧凑 JSON并把 ``Decimal`` 保持为原值对应的 JSON 数字。
标准库会要求先把 ``Decimal`` 转成 float 或字符串;前者可能静默舍入,
后者会改变 JSON 类型。这里直接输出有限 Decimal 的十进制表示。
"""
def serialize(item: Any) -> str:
if item is None:
return "null"
if item is True:
return "true"
if item is False:
return "false"
if isinstance(item, int):
return str(item)
if isinstance(item, Decimal):
if not item.is_finite():
raise ValueError("non-finite JSON number is not allowed")
return str(item)
if isinstance(item, float):
if not math.isfinite(item):
raise ValueError("non-finite JSON number is not allowed")
return json.dumps(item, allow_nan=False)
if isinstance(item, str):
return json.dumps(item, ensure_ascii=False)
if isinstance(item, Mapping):
pairs: list[tuple[str, Any]] = []
seen_keys: set[str] = set()
for key, child in item.items():
if not isinstance(key, str):
raise TypeError("JSON object keys must be strings")
if key in seen_keys:
raise ValueError(f"duplicate JSON object key: {key!r}")
seen_keys.add(key)
pairs.append((key, child))
if sort_keys:
pairs.sort(key=lambda pair: pair[0])
return "{" + ",".join(
f"{json.dumps(key, ensure_ascii=False)}:{serialize(child)}"
for key, child in pairs
) + "}"
if isinstance(item, (list, tuple)):
return "[" + ",".join(serialize(child) for child in item) + "]"
raise TypeError(f"value of type {type(item).__name__} is not JSON serializable")
return serialize(value)
def _binary_bytes(
raw: bytes | bytearray | memoryview | str,
file_format: TextFormat,
) -> bytes:
if isinstance(raw, str):
raise ValueError(f"{file_format.upper()} content must be uploaded as binary data")
if not isinstance(raw, (bytes, bytearray, memoryview)):
raise TypeError("raw must be bytes-like or str")
return bytes(raw)

View File

@@ -0,0 +1,81 @@
"""数据处理算法 - 数据集转换和分割。"""
from __future__ import annotations
import hashlib
import math
from collections.abc import Mapping, Sequence
from .types import DatasetSplit
def stable_split(
value: str | int,
split: Mapping[str, int] | None = None,
*,
seed: str = "",
) -> DatasetSplit:
"""按稳定哈希将记录划分到 train/validation/test。"""
ratios = dict(split or {"train": 80, "validation": 10, "test": 10})
required = {"train", "validation", "test"}
if set(ratios) != required:
raise ValueError("split must contain exactly train, validation and test")
if any(isinstance(value, bool) or not isinstance(value, int) or value < 0 for value in ratios.values()):
raise ValueError("split ratios must be non-negative integers")
if sum(ratios.values()) != 100:
raise ValueError("split ratios must sum to 100")
digest = hashlib.sha256(f"{seed}:{value}".encode("utf-8")).digest()
bucket = int.from_bytes(digest[:8], "big") % 10_000
train_boundary = ratios["train"] * 100
validation_boundary = train_boundary + ratios["validation"] * 100
if bucket < train_boundary:
return "train"
if bucket < validation_boundary:
return "validation"
return "test"
def stable_split_assignments(
values: Sequence[str | int],
split: Mapping[str, int] | None = None,
*,
seed: str = "",
) -> list[DatasetSplit]:
"""按稳定顺序和精确配额批量划分数据集。
单条哈希分桶只能在大样本下近似比例。这里先按哈希稳定排序,再用
最大余数法计算各切分配额,确保小数据集也严格遵循配置比例。
"""
ratios = dict(split or {"train": 80, "validation": 10, "test": 10})
# 复用单条划分的参数校验,避免两套规则逐渐漂移。
stable_split("validation", ratios, seed=seed)
if not values:
return []
split_order: tuple[DatasetSplit, ...] = ("train", "validation", "test")
exact = {name: len(values) * ratios[name] / 100 for name in split_order}
quotas = {name: math.floor(exact[name]) for name in split_order}
remaining = len(values) - sum(quotas.values())
remainder_order = sorted(
split_order,
key=lambda name: (-(exact[name] - quotas[name]), split_order.index(name)),
)
for name in remainder_order[:remaining]:
quotas[name] += 1
ranked_indices = sorted(
range(len(values)),
key=lambda index: (
hashlib.sha256(f"{seed}:{values[index]}".encode("utf-8")).digest(),
index,
),
)
assignments: list[DatasetSplit] = ["train"] * len(values)
cursor = 0
for name in split_order:
for index in ranked_indices[cursor : cursor + quotas[name]]:
assignments[index] = name
cursor += quotas[name]
return assignments

View File

@@ -0,0 +1,266 @@
"""数据处理模块使用的无副作用算法。
本模块不访问数据库、文件系统或网络,便于 API、后台任务和测试共同复用。
所有偏移量均为 Python 字符串偏移量。
"""
from __future__ import annotations
import csv
import hashlib
import io
import json
import math
import re
import unicodedata
import xml.etree.ElementTree as ET
import zipfile
from collections import Counter
from collections.abc import Iterable, Mapping, Sequence
from copy import deepcopy
from dataclasses import dataclass
from datetime import date, datetime, time
from decimal import Decimal
from pathlib import Path, PurePosixPath
from typing import Any, Literal
from urllib.parse import unquote, urlsplit
from docx import Document
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table
from docx.text.paragraph import Paragraph
from openpyxl import load_workbook
from openpyxl.utils.cell import range_boundaries
from pptx import Presentation
from pypdf import PdfReader
# 避免循环导入:直接定义常量而不是从 constants 导入
MAX_QA_PAIRS_PER_ITEM = 50
TextFormat = Literal[
"json",
"jsonl",
"csv",
"markdown",
"txt",
"pdf",
"docx",
"xlsx",
"pptx",
]
DatasetSplit = Literal["train", "validation", "test"]
StructuredPreprocessOption = Literal[
"clean_invalid",
"detect_structure",
"deduplicate",
"normalize_format",
"filter_anomaly",
"desensitize",
]
SUPPORTED_TEXT_FORMATS: tuple[TextFormat, ...] = (
"json",
"jsonl",
"csv",
"markdown",
"txt",
"pdf",
"docx",
"xlsx",
"pptx",
)
_FORMAT_ALIASES: dict[str, TextFormat] = {
"json": "json",
"jsonl": "jsonl",
"ndjson": "jsonl",
"csv": "csv",
"tsv": "csv",
"md": "markdown",
"markdown": "markdown",
"txt": "txt",
"text": "txt",
"pdf": "pdf",
"docx": "docx",
"xlsx": "xlsx",
"pptx": "pptx",
}
_LEGACY_OFFICE_FORMATS: dict[str, str] = {
"doc": "docx",
"xls": "xlsx",
"ppt": "pptx",
}
_OFFICE_OPEN_XML_FORMATS = {"docx", "xlsx", "pptx"}
_MAX_ARCHIVE_ENTRIES = 10_000
_MAX_ARCHIVE_UNCOMPRESSED_BYTES = 512 * 1024 * 1024
_MAX_ARCHIVE_ENTRY_BYTES = 128 * 1024 * 1024
_MAX_ARCHIVE_COMPRESSION_RATIO = 200
_MAX_EXTRACTED_TEXT_CHARS = 20_000_000
_MAX_PDF_PAGES = 2_000
_MAX_PRESENTATION_SLIDES = 2_000
_MAX_WORKBOOK_SHEETS = 100
_MAX_WORKBOOK_ROWS = 100_000
_MAX_WORKBOOK_SCANNED_ROWS = 200_000
_MAX_WORKBOOK_COLUMNS = 256
_MAX_WORKBOOK_CELLS = 2_000_000
_MAX_WORKBOOK_HEADER_ROWS = 8
_MAX_WORKBOOK_HEADER_SCAN_ROWS = 64
_MAX_WORKBOOK_MERGED_RANGES = 100_000
_MAX_STRUCTURED_FIELDS = 1_024
_MAX_STRUCTURED_DEPTH = 16
_MAX_JSON_DEPTH = 64
_MAX_ANOMALY_TEXT_CHARS = 1_000_000
_STRUCTURED_OPTIONS = {
"clean_invalid",
"detect_structure",
"deduplicate",
"normalize_format",
"filter_anomaly",
"desensitize",
}
_IDENTITY_FIELD_PATTERN = re.compile(r"(?:^|[._])(?:id|uuid|key|code)$|(?:^|[._]).+_id$")
_MOJIBAKE_MARKERS = ("\ufffd", "锟斤拷", "烫烫烫", "屯屯屯", "Ã", "Â", "â€")
_JSON_RECORD_ARRAY_KEYS = ("records", "data", "items", "rows")
_JSON_ENVELOPE_KEYS = ("response", "payload")
_JSON_WRAPPER_METADATA_KEYS = frozenset(
{
"page",
"page_size",
"pageSize",
"per_page",
"perPage",
"total",
"total_count",
"totalCount",
"count",
"offset",
"limit",
"cursor",
"next_cursor",
"nextCursor",
"has_more",
"hasMore",
}
)
_JSON_RESPONSE_METADATA_KEYS = _JSON_WRAPPER_METADATA_KEYS | {
"success",
"status",
"code",
"message",
"error",
}
_NAME_FIELD_NAMES = {
"name",
"full_name",
"fullname",
"real_name",
"contact_name",
"customer_name",
"recipient_name",
"姓名",
"中文姓名",
"真实姓名",
"联系人",
"联系人姓名",
"客户姓名",
"收件人",
"收件人姓名",
}
_EMAIL_PATTERN = re.compile(
r"(?<![\w.+-])[A-Za-z0-9.!#$%&'*+/=?^_`{|}~-]+"
r"@[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?"
r"(?:\.[A-Za-z0-9](?:[A-Za-z0-9-]{0,61}[A-Za-z0-9])?)+(?![\w.-])"
)
_PHONE_PATTERN = re.compile(r"(?<!\d)(?:(?:\+|00)?86[-\s]?)?1[3-9]\d{9}(?!\d)")
_ID_CARD_PATTERN = re.compile(r"(?<!\d)(?:\d{17}[\dXx]|\d{15})(?!\d)")
_CHINESE_NAME_CONTEXT_PATTERN = re.compile(
r"(?P<label>姓名|真实姓名|联系人(?:姓名)?|收件人)"
r"(?P<separator>\s*(?:[:=]|为)\s*|\s+)"
r"(?P<name>[\u3400-\u4dbf\u4e00-\u9fff·]{2,8})"
)
_ENGLISH_NAME_CONTEXT_PATTERN = re.compile(
r"(?im)(?P<label>full\s+name|contact\s+name|name)"
r"(?P<separator>\s*[:=]\s*)"
r"(?P<name>[A-Za-z][A-Za-z'-]*(?:[ \t]+[A-Za-z][A-Za-z'-]*){0,3})"
)
_TOKEN_PATTERN = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff]|[A-Za-z0-9_]+|[^\s]")
@dataclass(frozen=True, slots=True)
class ParsedText:
"""文本、文档或工作簿的统一解析结果。"""
format: TextFormat
text: str
records: tuple[dict[str, Any], ...]
record_locators: tuple[dict[str, Any], ...] = ()
@dataclass(frozen=True, slots=True)
class ProcessedStructuredRecord:
"""保留原始记录索引的结构化预处理结果。"""
source_index: int
record: dict[str, Any]
@dataclass(frozen=True, slots=True)
class PdfPageText:
"""PDF 物理页在统一提取文本中的字符范围。"""
page_number: int
text: str
source_start: int
source_end: int
@dataclass(frozen=True, slots=True)
class DocumentNoiseSpan:
"""PDF 中可安全从展示内容移除的文本范围。"""
start: int
end: int
kind: Literal["page_number", "repeated_margin", "table_of_contents"]
@dataclass(frozen=True, slots=True)
class QualityScore:
"""标准 instruction/input/output 记录的可解释质量分。"""
overall: float
completeness: float
length: float
readability: float
relevance: float
duplicate: float
is_valid: bool
flags: tuple[str, ...]
fingerprint: str
@dataclass(frozen=True, slots=True)
class DocumentHeading:
"""文档标题的位置和层级。"""
level: int
title: str
line_number: int
start: int
end: int
@dataclass(frozen=True, slots=True)
class DocumentStructure:
"""无需外部模型即可复现的文档结构摘要。"""
line_count: int
paragraph_count: int
headings: tuple[DocumentHeading, ...]
code_block_count: int
table_block_count: int
list_block_count: int
@dataclass(frozen=True, slots=True)
class _PdfLine:
text: str
start: int
end: int
class _DuplicateJsonKeyError(ValueError):
"""严格 JSON 解析时发现同一对象内的重复键。"""

View File

@@ -0,0 +1,7 @@
"""数据处理模块的共享限制。"""
MAX_QA_PAIRS_PER_ITEM = 50
MODEL_GENERATION_BATCH_SIZE = 10
__all__ = ["MAX_QA_PAIRS_PER_ITEM", "MODEL_GENERATION_BATCH_SIZE"]

View File

@@ -0,0 +1,158 @@
"""Dataset format validation for Alpaca, ShareGPT, DPO, CPT formats.
Used by the training preflight flow to validate that uploaded dataset files
conform to the declared format before submitting to the compute node.
"""
from __future__ import annotations
import json
from typing import Any
def _load_sample(path: str | None, content: str | None = None, max_samples: int = 20) -> list[dict[str, Any]]:
"""Load up to max_samples records from JSONL file path or raw content string."""
try:
if content is not None:
text = content.strip()
elif path:
with open(path, "r", encoding="utf-8") as fh:
text = fh.read().strip()
else:
return []
except Exception:
return []
if not text:
return []
# 先按整文件 JSON数组/单对象)解析,兼容 .json失败再按 jsonl 逐行解析
try:
value = json.loads(text)
except json.JSONDecodeError:
value = None
if isinstance(value, list):
return [item for item in value[:max_samples] if isinstance(item, dict)]
if isinstance(value, dict):
return [value]
lines = text.splitlines()[:max_samples]
records: list[dict[str, Any]] = []
for line in lines:
line = line.strip()
if not line:
continue
try:
record = json.loads(line)
except json.JSONDecodeError:
continue
if isinstance(record, dict):
records.append(record)
return records
def _check_alpaca(records: list[dict[str, Any]]) -> list[str]:
"""Validate Alpaca format: requires 'instruction' field."""
errors: list[str] = []
if not records:
errors.append("Alpaca 格式数据集无有效记录")
return errors
missing_instruction = sum(1 for r in records if not r.get("instruction"))
if missing_instruction:
errors.append(
f"Alpaca 格式要求每条记录包含 instruction 字段,"
f"{len(records)}条中有{missing_instruction}条缺失"
)
return errors
def _check_sharegpt(records: list[dict[str, Any]]) -> list[str]:
"""Validate ShareGPT format: requires 'messages' (list of dicts with role/content)."""
errors: list[str] = []
if not records:
errors.append("ShareGPT 格式数据集无有效记录")
return errors
bad = 0
for r in records:
messages = r.get("messages")
if not isinstance(messages, list) or not messages:
bad += 1
continue
for msg in messages:
if not isinstance(msg, dict) or "role" not in msg or "content" not in msg:
bad += 1
break
if bad:
errors.append(
f"ShareGPT 格式要求每条记录包含 messages 列表,"
f"每条消息需有 role 和 content 字段,前{len(records)}条中有{bad}条不符合"
)
return errors
def _check_dpo(records: list[dict[str, Any]]) -> list[str]:
"""Validate DPO format: requires 'chosen' and 'rejected' fields."""
errors: list[str] = []
if not records:
errors.append("DPO 格式数据集无有效记录")
return errors
missing_chosen = sum(1 for r in records if not r.get("chosen"))
missing_rejected = sum(1 for r in records if not r.get("rejected"))
if missing_chosen:
errors.append(f"DPO 格式要求 chosen 字段,前{len(records)}条中有{missing_chosen}条缺失")
if missing_rejected:
errors.append(f"DPO 格式要求 rejected 字段,前{len(records)}条中有{missing_rejected}条缺失")
return errors
def _check_cpt(records: list[dict[str, Any]]) -> list[str]:
"""Validate CPT format: requires 'text' field, should NOT have instruction/output."""
errors: list[str] = []
if not records:
errors.append("CPT 格式数据集无有效记录")
return errors
missing_text = sum(1 for r in records if not r.get("text"))
has_instruction = sum(1 for r in records if r.get("instruction") or r.get("output"))
if missing_text:
errors.append(f"CPT 格式要求 text 字段,前{len(records)}条中有{missing_text}条缺失")
if has_instruction:
errors.append(
f"CPT 格式不应包含 instruction/output 字段(疑似 Alpaca 格式),"
f"{len(records)}条中有{has_instruction}条包含此类字段"
)
return errors
FORMAT_VALIDATORS = {
"alpaca": _check_alpaca,
"alpaca_jsonl": _check_alpaca,
"sharegpt": _check_sharegpt,
"dpo": _check_dpo,
"cpt": _check_cpt,
"pt": _check_cpt,
}
def validate_dataset_format(
dataset_format: str,
content: str | None = None,
path: str | None = None,
max_samples: int = 20,
) -> list[str]:
"""Validate dataset content against expected format.
Args:
dataset_format: One of 'alpaca', 'sharegpt', 'dpo', 'cpt'.
content: Raw file content (JSONL text). Mutually exclusive with path.
path: File path to read content from.
max_samples: Maximum records to sample for validation.
Returns:
List of error messages (empty if valid).
"""
fmt = str(dataset_format).lower().strip()
validator = FORMAT_VALIDATORS.get(fmt)
if not validator:
return [f"不支持的数据集格式: {dataset_format},支持的格式: {', '.join(sorted(FORMAT_VALIDATORS))}"]
records = _load_sample(path=path, content=content, max_samples=max_samples)
return validator(records)

View File

@@ -0,0 +1,693 @@
"""基于 Docling 与 LlamaIndex 的文档切分实现。"""
from __future__ import annotations
import logging
import re
import threading
import time
import unicodedata
from dataclasses import dataclass
from functools import lru_cache
from io import BytesIO
from typing import Any, Literal
import tiktoken
from docling_core.transforms.chunker.hierarchical_chunker import ChunkingSerializerProvider
from llama_index.core import Document
from llama_index.core.base.embeddings.base import BaseEmbedding
from llama_index.core.node_parser import SemanticSplitterNodeParser, SentenceSplitter
from app.modules.data_process.algorithms import normalize_text
from app.modules.data_process.algorithms.embedding import semantic_embedding_model
ChunkMethod = Literal["layout_hybrid", "semantic", "fixed"]
_PAGE_FURNITURE = re.compile(
r"(?m)^\s*(?:第\s*\d+\s*页\s*共\s*\d+\s*页|[-—–]?\s*\d+\s*[/]\s*\d+\s*[-—–]?)\s*$"
)
# Docling 的 markdown 序列化会给列表项补上自动编号,而 Word 的编号存放在
# numbering.xml 中python-docx 抽取的正文不含这些编号;紧凑匹配前剥掉
# 行首编号,否则带列表的切片会整体定位失败。
_LIST_MARKER_PREFIX = re.compile(
r"(?m)^[ \t>]*(?:(?:\d{1,3}[.)])+|\([a-zA-Z0-9]{1,3}\)|[a-zA-Z][.)]|[-*+•·])[ \t]+"
)
_COMPACT_CHARACTER = re.compile(r"[\w\u3400-\u4dbf\u4e00-\u9fff]", re.UNICODE)
_CONVERTER_LOCK = threading.Lock()
logger = logging.getLogger(__name__)
@dataclass(frozen=True, slots=True)
class DocumentChunk:
"""切片正文及其在原文件中的可追溯信息。"""
original_content: str
contextualized_content: str
source_start: int | None
source_end: int | None
source_start_line: int | None
source_end_line: int | None
token_count: int
heading_path: tuple[str, ...] = ()
source_pages: tuple[int, ...] = ()
doc_item_refs: tuple[str, ...] = ()
source_bboxes: tuple[dict[str, Any], ...] = ()
def _sentence_chunks(text: str) -> list[str]:
"""提供稳定的中英文句界,避免 LlamaIndex 默认分词器下载额外资源。"""
boundary = re.compile(
r".*?(?:\n\s*\n|[。!?!?;](?:[\"'”’)】》]*)|\.(?:\s+|$)|$)",
re.DOTALL,
)
return [part for part in boundary.findall(text) if part]
@lru_cache(maxsize=1)
def _tokenizer() -> tiktoken.Encoding:
"""加载 cl100k_base 编码器,优先在线下载,失败时使用本地缓存以支持离线环境。"""
import base64
from app.core.cache_paths import tiktoken_cache_dir
# 缓存目录已在 app.core.cache_paths.setup_local_caches 中统一指向 <repo>/.cache/tiktoken
# 此处直接读取TIKTOKEN_CACHE_DIR 已在启动阶段写入。
offline_cache = tiktoken_cache_dir()
try:
# 尝试标准方式加载(环境变量 TIKTOKEN_CACHE_DIR 已被统一设置)
return tiktoken.get_encoding("cl100k_base")
except Exception:
# 如果失败,尝试手动从本地文件构造
try:
local_file = offline_cache / "9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
if not local_file.exists():
# 尝试另一个可能的文件名
local_file = offline_cache / "cl100k_base.tiktoken"
if local_file.exists():
# 读取 BPE 文件内容
with open(local_file, "rb") as f:
contents = f.read()
# 解析 BPE 文件
mergeable_ranks = {}
for line in contents.splitlines():
if line:
token, rank = line.split()
mergeable_ranks[base64.b64decode(token)] = int(rank)
# 构造 Encoding 对象(模块顶部已 import tiktoken
# 此处不能再 import tiktoken.core否则会把 tiktoken
# 变成局部变量,使函数开头的 tiktoken.get_encoding 抛
# UnboundLocalError
return tiktoken.core.Encoding(
name="cl100k_base",
pat_str=r"""'(?i:[sdmt]|ll|ve|re)|[^\r\n\p{L}\p{N}]?+\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]++[\r\n]*|\s*[\r\n]|\s+(?!\S)|\s+""",
mergeable_ranks=mergeable_ranks,
special_tokens={
"": 100257,
"<|fim_prefix|>": 100258,
"<|fim_middle|>": 100259,
"<|fim_suffix|>": 100260,
"<|endofprompt|>": 100276,
},
)
except Exception:
pass
raise RuntimeError(
f"无法加载 cl100k_base 编码器\n"
f"请确保以下任一条件满足:\n"
f"1. 服务器可以访问网络\n"
f"2. 本地存在缓存文件: {offline_cache}/9b5ad71b2ce5302211f9c61530b329a4922fc6a4"
)
def _text_chunks(
text: str,
*,
chunk_size: int,
chunk_overlap: int,
) -> list[DocumentChunk]:
normalized = normalize_text(text)
if not normalized:
return []
splitter = SentenceSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
tokenizer=_tokenizer().encode,
chunking_tokenizer_fn=_sentence_chunks,
include_metadata=False,
include_prev_next_rel=False,
)
nodes = splitter.get_nodes_from_documents([Document(text=normalized)])
return _nodes_to_chunks(nodes, normalized)
def chunk_fixed_text(
text: str,
*,
chunk_size: int,
chunk_overlap: int,
) -> list[DocumentChunk]:
"""使用 LlamaIndex SentenceSplitter 按句界控制固定 Token 长度。"""
return _text_chunks(text, chunk_size=chunk_size, chunk_overlap=chunk_overlap)
def chunk_semantic_text(
text: str,
*,
chunk_size: int,
chunk_overlap: int,
breakpoint_percentile_threshold: int,
embed_model: BaseEmbedding | None = None,
) -> list[DocumentChunk]:
"""使用 LlamaIndex SemanticSplitter 识别主题跳变,再限制最大长度。"""
normalized = normalize_text(text)
if not normalized:
return []
splitter = SemanticSplitterNodeParser.from_defaults(
embed_model=embed_model or semantic_embedding_model(),
breakpoint_percentile_threshold=breakpoint_percentile_threshold,
buffer_size=1,
sentence_splitter=_sentence_chunks,
include_metadata=False,
include_prev_next_rel=False,
)
semantic_nodes = splitter.get_nodes_from_documents([Document(text=normalized)])
result: list[DocumentChunk] = []
search_from = 0
for node in semantic_nodes:
content = node.get_content().strip()
if not content:
continue
start = _locate_text(normalized, content, search_from)
if start is None:
start = _locate_text(normalized, content, 0)
if start is None:
result.append(_unlocated_chunk(content))
continue
if len(_tokenizer().encode(content)) <= chunk_size:
result.append(_make_text_chunk(normalized, start, start + len(content)))
else:
for child in _text_chunks(
content,
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
):
if child.source_start is None or child.source_end is None:
result.append(_unlocated_chunk(child.original_content))
continue
result.append(
_make_text_chunk(
normalized,
start + child.source_start,
start + child.source_end,
)
)
search_from = start + len(content)
return result
def _nodes_to_chunks(nodes: list[Any], source_text: str) -> list[DocumentChunk]:
chunks: list[DocumentChunk] = []
search_from = 0
for node in nodes:
content = node.get_content().strip()
if not content:
continue
raw_start = getattr(node, "start_char_idx", None)
raw_end = getattr(node, "end_char_idx", None)
if (
isinstance(raw_start, int)
and isinstance(raw_end, int)
and source_text[raw_start:raw_end].strip() == content
):
start = raw_start + len(source_text[raw_start:raw_end]) - len(source_text[raw_start:raw_end].lstrip())
else:
start = _locate_text(source_text, content, search_from)
if start is None:
start = _locate_text(source_text, content, 0)
if start is None:
chunks.append(_unlocated_chunk(content))
continue
end = start + len(content)
chunks.append(_make_text_chunk(source_text, start, end))
search_from = max(search_from, end)
return chunks
def _locate_text(source: str, content: str, start: int) -> int | None:
position = source.find(content, start)
return position if position >= 0 else None
def _unlocated_chunk(content: str) -> DocumentChunk:
"""正文在源文本中定位失败时保底保留切片,只放弃行号信息。"""
return DocumentChunk(
original_content=content,
contextualized_content=content,
source_start=None,
source_end=None,
source_start_line=None,
source_end_line=None,
token_count=len(_tokenizer().encode(content)),
)
def _make_text_chunk(source: str, start: int, end: int) -> DocumentChunk:
content = source[start:end]
return DocumentChunk(
original_content=content,
contextualized_content=content,
source_start=start,
source_end=end,
source_start_line=source.count("\n", 0, start) + 1,
source_end_line=source.count("\n", 0, max(start, end - 1)) + 1,
token_count=len(_tokenizer().encode(content)),
)
@lru_cache(maxsize=1)
def _document_converter():
from docling.datamodel.base_models import InputFormat
from docling.datamodel.pipeline_options import PdfPipelineOptions
from docling.document_converter import DocumentConverter, PdfFormatOption
pipeline_options = PdfPipelineOptions()
pipeline_options.do_ocr = False
return DocumentConverter(
format_options={
InputFormat.PDF: PdfFormatOption(pipeline_options=pipeline_options),
}
)
class _MarkdownSerializerProvider(ChunkingSerializerProvider):
def get_serializer(self, doc: Any):
from docling_core.transforms.chunker.hierarchical_chunker import ChunkingDocSerializer
from docling_core.transforms.serializer.markdown import (
MarkdownParams,
MarkdownTableSerializer,
)
from docling_core.types.doc import DocItemLabel
excluded = {
DocItemLabel.DOCUMENT_INDEX,
DocItemLabel.PAGE_HEADER,
DocItemLabel.PAGE_FOOTER,
}
return ChunkingDocSerializer(
doc=doc,
table_serializer=MarkdownTableSerializer(),
params=MarkdownParams(
labels=set(DocItemLabel) - excluded,
compact_tables=True,
image_placeholder="",
escape_html=False,
escape_underscores=False,
),
)
def _clean_layout_text(value: str) -> str:
return normalize_text(_PAGE_FURNITURE.sub("", value)).strip()
def _compact_with_offsets(value: str) -> tuple[str, list[int]]:
compact: list[str] = []
offsets: list[int] = []
for index, character in enumerate(unicodedata.normalize("NFKC", value)):
if _COMPACT_CHARACTER.fullmatch(character):
compact.append(character.casefold())
offsets.append(index)
return "".join(compact), offsets
def _expand_to_line_boundaries(source_text: str, start: int, end: int) -> tuple[int, int]:
while start > 0 and source_text[start - 1] not in "\r\n":
start -= 1
while end < len(source_text) and source_text[end] not in "\r\n":
end += 1
return start, end
def _project_layout_span(
source_text: str,
content: str,
*,
compact_source: str,
source_offsets: list[int],
compact_start: int,
) -> tuple[int | None, int | None, int]:
for candidate in (content, _LIST_MARKER_PREFIX.sub("", content)):
compact_content, _ = _compact_with_offsets(candidate)
if len(compact_content) < 4:
continue
position = compact_source.find(compact_content, compact_start)
if position < 0:
position = compact_source.find(compact_content)
if position < 0:
continue
start, end = _expand_to_line_boundaries(
source_text,
source_offsets[position],
source_offsets[position + len(compact_content) - 1] + 1,
)
# 重复内容回退匹配可能命中已消费的更早位置,游标只进不退,
# 避免后续切片跟着错位。
return start, end, max(compact_start, position + len(compact_content))
return _project_layout_span_by_anchors(
source_text,
content,
compact_source=compact_source,
source_offsets=source_offsets,
compact_start=compact_start,
)
def _project_layout_span_by_anchors(
source_text: str,
content: str,
*,
compact_source: str,
source_offsets: list[int],
compact_start: int,
) -> tuple[int | None, int | None, int]:
"""按行锚点顺序匹配,容忍切片里插入的重复表头等非连续内容。"""
segments = [
compact
for compact in (
_compact_with_offsets(line)[0]
for line in _LIST_MARKER_PREFIX.sub("", content).split("\n")
)
if len(compact) >= 6
]
if not segments:
return None, None, compact_start
total = sum(len(segment) for segment in segments)
def match_from(cursor: int) -> tuple[list[tuple[int, int]], int]:
matched: list[tuple[int, int]] = []
position = cursor
for segment in segments:
found = compact_source.find(segment, position)
if found < 0:
continue
matched.append((found, found + len(segment)))
position = found + len(segment)
return matched, sum(end - start for start, end in matched)
matched, covered = match_from(compact_start)
if covered * 2 < total:
retried, retry_covered = match_from(0)
if retry_covered > covered:
matched, covered = retried, retry_covered
# 覆盖不足一半时宁可不定位,也不能给出错误的行号。
if not matched or covered * 2 < total:
return None, None, compact_start
start, end = _expand_to_line_boundaries(
source_text,
source_offsets[matched[0][0]],
source_offsets[matched[-1][1] - 1] + 1,
)
return start, end, max(compact_start, matched[-1][1])
def chunk_layout_document(
raw: bytes,
*,
filename: str,
source_text: str,
chunk_size: int,
) -> list[DocumentChunk]:
"""使用 Docling HybridChunker 按版面层级、列表与表格边界切分。"""
from docling.chunking import HybridChunker
from docling.datamodel.base_models import DocumentStream
from docling.exceptions import BaseError as DoclingError
from docling_core.transforms.chunker.tokenizer.openai import OpenAITokenizer
from docling_core.types.doc import DocItemLabel
from app.modules.data_process.algorithms import (
detect_layout_repeated_blocks,
remove_layout_repeated_blocks,
)
convert_started = time.perf_counter()
try:
with _CONVERTER_LOCK:
conversion = _document_converter().convert(
DocumentStream(name=filename, stream=BytesIO(raw))
)
except DoclingError as exc:
raise ValueError(f"文档版面解析失败: {exc}") from exc
logger.info(
"layout chunking convert done file=%s elapsed=%.2fs",
filename,
time.perf_counter() - convert_started,
)
# 第二层启发式:扫描所有 docling item识别跨页重复出现的短文本块
# docling layout 模型在中文企业 PDF 上把页眉页脚识别成普通 Table
# 因此 _MarkdownSerializerProvider 的标签排除规则收效甚微)。
page_count = len(getattr(conversion.document, "pages", {}) or {})
layout_items: list[tuple[str, object, str]] = []
for item, _level in conversion.document.iterate_items():
text = getattr(item, "text", None)
if not text and hasattr(item, "export_to_markdown"):
try:
text = item.export_to_markdown(doc=conversion.document) or ""
except TypeError:
# 旧版 docling_core 无 doc 参数
text = item.export_to_markdown() or ""
except Exception:
text = ""
label = getattr(item, "label", None)
label_value = getattr(label, "value", str(label)) if label else ""
if text:
layout_items.append((label_value, item, text))
repeated_blocks = detect_layout_repeated_blocks(
layout_items, page_count=page_count
)
chunker = HybridChunker(
tokenizer=OpenAITokenizer(tokenizer=_tokenizer(), max_tokens=chunk_size),
serializer_provider=_MarkdownSerializerProvider(),
merge_peers=True,
repeat_table_header=True,
)
compact_source, source_offsets = _compact_with_offsets(source_text)
compact_start = 0
result: list[DocumentChunk] = []
covered_refs: set[str] = set()
excluded = {
DocItemLabel.DOCUMENT_INDEX,
DocItemLabel.PAGE_HEADER,
DocItemLabel.PAGE_FOOTER,
}
for raw_chunk in chunker.chunk(conversion.document):
doc_items = tuple(raw_chunk.meta.doc_items or ())
if doc_items and all(item.label in excluded for item in doc_items):
continue
content = _clean_layout_text(raw_chunk.text)
if not content:
continue
contextualized = _clean_layout_text(chunker.contextualize(raw_chunk)) or content
if repeated_blocks:
content = remove_layout_repeated_blocks(content, repeated_blocks)
contextualized = remove_layout_repeated_blocks(
contextualized, repeated_blocks
)
if not content:
continue
start, end, compact_start = _project_layout_span(
source_text,
content,
compact_source=compact_source,
source_offsets=source_offsets,
compact_start=compact_start,
)
original = source_text[start:end] if start is not None and end is not None else content
pages: set[int] = set()
refs: list[str] = []
bboxes: list[dict[str, Any]] = []
for item in doc_items:
refs.append(str(item.self_ref))
covered_refs.add(str(item.self_ref))
for provenance in item.prov or ():
pages.add(int(provenance.page_no))
bbox = provenance.bbox
bboxes.append(
{
"page": int(provenance.page_no),
"left": float(bbox.l),
"top": float(bbox.t),
"right": float(bbox.r),
"bottom": float(bbox.b),
"origin": str(bbox.coord_origin.value),
}
)
result.append(
DocumentChunk(
original_content=original,
contextualized_content=contextualized,
source_start=start,
source_end=end,
source_start_line=(source_text.count("\n", 0, start) + 1 if start is not None else None),
source_end_line=(
source_text.count("\n", 0, max(start or 0, (end or 1) - 1)) + 1
if end is not None
else None
),
token_count=len(_tokenizer().encode(contextualized)),
heading_path=tuple(str(item) for item in (raw_chunk.meta.headings or ())),
source_pages=tuple(sorted(pages)),
doc_item_refs=tuple(refs),
source_bboxes=tuple(bboxes),
)
)
# HybridChunker(merge_peers=True) 会丢弃"末尾无正文的孤立标题"。
# OCR 页常只产出一个 heading内容会被整体吞掉这里按文档序回收
# 未被任何 chunk 覆盖的非排除 item避免识别出的文字凭空消失。
# 注意 heading 会进入 meta.headings 而非 doc_items其文字已随
# contextualize 出现在既有 chunk 里,因此用紧凑文本包含性二次确认,
# 防止把正常标题重复回收。
chunk_haystack = _compact_with_offsets(
"\n".join(chunk.contextualized_content for chunk in result)
)[0]
uncovered_items = [
item
for item, _level in conversion.document.iterate_items()
if item.label not in excluded
and str(item.self_ref) not in covered_refs
and (getattr(item, "text", None) or "").strip()
and _compact_with_offsets(str(item.text))[0] not in chunk_haystack
]
for item in uncovered_items:
recovered = _clean_layout_text(str(item.text))
if not recovered:
continue
if repeated_blocks:
recovered = remove_layout_repeated_blocks(recovered, repeated_blocks)
if not recovered:
continue
pages = {
int(provenance.page_no) for provenance in item.prov or ()
}
bboxes = [
{
"page": int(provenance.page_no),
"left": float(provenance.bbox.l),
"top": float(provenance.bbox.t),
"right": float(provenance.bbox.r),
"bottom": float(provenance.bbox.b),
"origin": str(provenance.bbox.coord_origin.value),
}
for provenance in item.prov or ()
]
logger.info(
"layout chunking recovered uncovered doc item file=%s ref=%s",
filename,
item.self_ref,
)
result.append(
DocumentChunk(
original_content=recovered,
contextualized_content=recovered,
source_start=None,
source_end=None,
source_start_line=None,
source_end_line=None,
token_count=len(_tokenizer().encode(recovered)),
heading_path=(),
source_pages=tuple(sorted(pages)),
doc_item_refs=(str(item.self_ref),),
source_bboxes=tuple(bboxes),
)
)
return result
def merge_short_chunks(
chunks: list[DocumentChunk],
*,
source_text: str,
min_token_count: int,
max_token_count: int,
) -> list[DocumentChunk]:
"""在不突破长度上限的前提下,把过短块并入相邻内容。"""
result: list[DocumentChunk] = []
index = 0
while index < len(chunks):
current = chunks[index]
if current.token_count >= min_token_count:
result.append(current)
index += 1
continue
if index + 1 < len(chunks):
combined = _combine_chunks(current, chunks[index + 1], source_text)
if combined.token_count <= max_token_count:
result.append(combined)
index += 2
continue
if result:
combined = _combine_chunks(result[-1], current, source_text)
if combined.token_count <= max_token_count:
result[-1] = combined
index += 1
continue
result.append(current)
index += 1
return result
def _combine_chunks(
left: DocumentChunk,
right: DocumentChunk,
source_text: str,
) -> DocumentChunk:
contextualized = "\n\n".join(
part for part in (left.contextualized_content, right.contextualized_content) if part
)
start = left.source_start
end = right.source_end
has_contiguous_source = (
start is not None
and left.source_end is not None
and right.source_start is not None
and end is not None
and left.source_end <= right.source_start
)
original = (
source_text[start:end]
if has_contiguous_source and start is not None and end is not None
else "\n\n".join(
part for part in (left.original_content, right.original_content) if part
)
)
if not has_contiguous_source:
start = None
end = None
return DocumentChunk(
original_content=original,
contextualized_content=contextualized,
source_start=start,
source_end=end,
source_start_line=left.source_start_line if start is not None else None,
source_end_line=right.source_end_line if end is not None else None,
token_count=len(_tokenizer().encode(contextualized)),
heading_path=left.heading_path or right.heading_path,
source_pages=tuple(sorted(set(left.source_pages) | set(right.source_pages))),
doc_item_refs=left.doc_item_refs + right.doc_item_refs,
source_bboxes=left.source_bboxes + right.source_bboxes,
)

View File

@@ -0,0 +1,313 @@
"""数据处理 - 生成结果的多层质量评测。
三层体系:规则层(确定性规则分)+ 语义层(本地嵌入向量)+ 评审层
(复用生成模型按 rubric 打分的 LLM-as-judge。任一层失败自动降级
评测永远返回可用结果,不阻断调用方流程。
"""
from __future__ import annotations
import logging
from collections.abc import Mapping
from dataclasses import asdict
from datetime import UTC, datetime
from typing import Any
import httpx
from .algorithms import normalize_text, score_quality
from .algorithms.quality import composite_overall, semantic_quality_scores
from .generation import (
ModelGenerationError,
_is_retryable_generation_error,
_json_payload,
_message_content,
chat_completions_url,
)
logger = logging.getLogger(__name__)
# 送入评审提示词的来源正文上限,避免超长切片挤占评分输出空间。
_MAX_JUDGE_SOURCE_CHARS = 6000
_JUDGE_DIMENSIONS: dict[str, tuple[str, ...]] = {
"standard": (
"faithfulness",
"correctness",
"clarity",
"completeness",
"alignment",
),
"reasoning": (
"faithfulness",
"correctness",
"clarity",
"completeness",
"alignment",
"reasoning_validity",
),
"dpo": (
"clarity",
"chosen_quality",
"rejected_quality",
"preference_reasonableness",
"faithfulness",
),
}
_DIMENSION_LABELS: dict[str, str] = {
"faithfulness": "忠实度",
"correctness": "正确性",
"clarity": "问题清晰度",
"completeness": "回答完整性",
"alignment": "指令对齐",
"reasoning_validity": "推理有效性",
"chosen_quality": "chosen 回答质量",
"rejected_quality": "rejected 回答质量",
"preference_reasonableness": "偏好区分合理性",
}
_DIMENSION_RULES: dict[str, str] = {
"faithfulness": "忠实度:答案的全部陈述是否被参考资料支持,没有编造、没有引入资料之外的信息;未提供参考资料时按答案内部自洽性评估",
"correctness": "正确性:答案中的事实、概念与计算是否正确",
"clarity": "问题清晰度:问题是否清晰、自包含、无歧义,脱离上下文也能理解",
"completeness": "回答完整性:答案是否充分、直接地回应了问题的全部要点",
"alignment": "指令对齐:答案的形式与范围是否符合问题的要求(如格式、语言、范围限定)",
"reasoning_validity": "推理有效性:思维链步骤是否逻辑连贯、无跳步或循环论证,结论是否由推理过程自然得出",
"chosen_quality": "chosen 回答质量:更优回答的正确性、完整性与表述质量",
"rejected_quality": "rejected 回答质量:较差回答是否仍具备基本可读性,使对比训练有意义",
"preference_reasonableness": "偏好区分合理性chosen 是否明显优于 rejected且优劣差异与问题直接相关",
}
def _judge_system_prompt(output_type: str) -> str:
dimensions = _JUDGE_DIMENSIONS[output_type]
rules = "\n".join(f"- {_DIMENSION_RULES[name]}" for name in dimensions)
scores_schema = ", ".join(f'"{name}": 1-5' for name in dimensions)
return (
"你是大模型训练数据质量评审员。严格依据用户消息中的【参考资料】评审这条训练数据,逐维度按 1-5 分打分:\n"
f"{rules}\n"
"评分锚点5 分=完全符合维度描述3 分=基本符合但有明显不足1 分=严重不符合。\n"
"忠实度只依据参考资料与公认常识判断,无法得到支持的陈述必须扣分;不要因为答案冗长而加分。\n"
"只输出一个 JSON 对象,不要输出 JSON 之外的任何文字。\n"
'输出格式:{"scores": {' + scores_schema + '}, "reason": "一句话总评", "issues": ["具体问题,没有则为空数组"]}'
)
def _judge_user_prompt(record: Mapping[str, Any], source_content: str) -> str:
source = normalize_text(source_content)[:_MAX_JUDGE_SOURCE_CHARS] or "(无参考资料)"
instruction = normalize_text(str(record.get("instruction") or "")) or "(空)"
input_text = normalize_text(str(record.get("input") or ""))
sections = [f"【参考资料】\n{source}", f"【问题】\n{instruction}"]
if input_text:
sections.append(f"【输入】\n{input_text}")
if record.get("chosen") or record.get("rejected"):
sections.append(f"【更优回答 chosen】\n{normalize_text(str(record.get('chosen') or '')) or '(空)'}")
sections.append(f"【较差回答 rejected】\n{normalize_text(str(record.get('rejected') or '')) or '(空)'}")
else:
output = normalize_text(str(record.get("output") or ""))
sections.append(f"【回答】\n{output or '(空)'}")
return "\n\n".join(sections)
def _validated_judge_payload(payload: Any, output_type: str) -> dict[str, Any]:
if not isinstance(payload, Mapping):
raise ModelGenerationError("评审响应不是 JSON 对象")
raw_scores = payload.get("scores")
if not isinstance(raw_scores, Mapping):
raise ModelGenerationError("评审响应缺少 scores 对象")
expected = _JUDGE_DIMENSIONS[output_type]
scores: dict[str, float] = {}
for name in expected:
value = raw_scores.get(name)
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ModelGenerationError(f"评审响应缺少维度 {name} 的有效分数")
scores[name] = round(max(1.0, min(5.0, float(value))), 1)
issues = payload.get("issues")
if not isinstance(issues, list):
issues = []
issues = [str(item)[:200] for item in issues if str(item).strip()][:8]
reason = normalize_text(str(payload.get("reason") or ""))[:300]
return {
"scores": scores,
"overall": round(sum(scores.values()) / len(scores) * 20, 2),
"reason": reason,
"issues": issues,
}
def _judge_record(
record: Mapping[str, Any],
source_content: str,
*,
model: Mapping[str, Any],
config: Mapping[str, Any],
client: httpx.Client | None,
) -> dict[str, Any] | None:
output_type = str(config.get("output_type") or "standard").strip().lower()
if output_type not in _JUDGE_DIMENSIONS:
output_type = "standard"
endpoint = chat_completions_url(str(model.get("api_url") or ""))
model_name = str(model.get("online_model_name") or model.get("name") or "").strip()
if not model_name:
raise ModelGenerationError("generation model name is required")
temperature = 0.1
max_tokens = max(256, min(2048, int(config.get("max_tokens", 1024) or 1024)))
timeout = max(1.0, min(120.0, float(config.get("request_timeout_seconds", 60) or 60)))
retries = max(0, min(5, int(config.get("generation_retries", 2) or 2)))
headers = {"Content-Type": "application/json"}
api_key = str(model.get("api_key") or "").strip()
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
request_payload: dict[str, Any] = {
"model": model_name,
"messages": [
{"role": "system", "content": _judge_system_prompt(output_type)},
{"role": "user", "content": _judge_user_prompt(record, source_content)},
],
"temperature": temperature,
"max_tokens": max_tokens,
}
if bool(config.get("json_mode", False)):
request_payload["response_format"] = {"type": "json_object"}
owns_client = client is None
http_client = client or httpx.Client(timeout=timeout)
try:
last_error: Exception | None = None
for _ in range(retries + 1):
try:
response = http_client.post(endpoint, headers=headers, json=request_payload)
response.raise_for_status()
body = response.json()
if not isinstance(body, Mapping):
raise ModelGenerationError("model response body must be a JSON object")
judged = _validated_judge_payload(
_json_payload(_message_content(body)),
output_type,
)
judged["model"] = model_name
judged["output_type"] = output_type
return judged
except Exception as exc:
last_error = exc
if not _is_retryable_generation_error(exc):
break
raise ModelGenerationError(f"质量评审调用失败: {last_error}")
finally:
if owns_client:
http_client.close()
def evaluate_result_record(
record: Mapping[str, Any],
*,
source_content: str = "",
model: Mapping[str, Any] | None = None,
config: Mapping[str, Any] | None = None,
client: httpx.Client | None = None,
embed_model: Any = None,
min_output_length: int = 20,
) -> dict[str, Any]:
"""对一条生成结果执行三层评测,返回可直接落库的 quality_score 字典。
规则层字段保持原样平铺(向后兼容既有读取方);新增 ``semantic``、
``judge``、``layers``、``evaluated`` 与组合 ``overall``。
"""
config_dict = dict(config or {})
rule = score_quality(
record,
min_output_length=min_output_length,
source_content=source_content,
)
quality: dict[str, Any] = asdict(rule)
semantic = semantic_quality_scores(
record,
source_content=source_content,
embed_model=embed_model,
)
judge: dict[str, Any] | None = None
if model is not None:
try:
judge = _judge_record(
record,
source_content,
model=model,
config=config_dict,
client=client,
)
except Exception as exc:
logger.warning(
"data process judge evaluation degraded: %s",
exc,
)
layers = {
"rule": rule.overall,
"semantic": semantic.get("overall") if semantic else None,
"judge": judge.get("overall") if judge else None,
}
quality.update(
semantic=semantic,
judge=judge,
layers=layers,
evaluated=True,
evaluated_at=datetime.now(UTC).isoformat(),
overall=composite_overall(
rule=layers["rule"],
semantic=layers["semantic"],
judge=layers["judge"],
),
)
return quality
def reevaluate_edited_record(
record: Mapping[str, Any],
*,
source_content: str = "",
previous_quality: Mapping[str, Any] | None = None,
embed_model: Any = None,
min_output_length: int = 20,
) -> dict[str, Any]:
"""手动编辑/恢复后重算规则与语义层,丢弃已过期的评审层。
编辑会改变内容,旧的评审分不再可信;规则与语义层本地重算零成本。
``evaluated`` 标记沿用原值,保证已评测过的结果编辑后仍有可用分数。
"""
rule = score_quality(
record,
min_output_length=min_output_length,
source_content=source_content,
)
quality: dict[str, Any] = asdict(rule)
semantic = semantic_quality_scores(
record,
source_content=source_content,
embed_model=embed_model,
)
previous = dict(previous_quality or {})
evaluated = bool(previous.get("evaluated"))
layers = {
"rule": rule.overall,
"semantic": semantic.get("overall") if semantic else None,
"judge": None,
}
quality.update(
semantic=semantic,
judge=None,
layers=layers,
evaluated=evaluated,
evaluated_at=(
datetime.now(UTC).isoformat() if evaluated else None
),
overall=composite_overall(
rule=layers["rule"],
semantic=layers["semantic"],
),
)
return quality

View File

@@ -0,0 +1,641 @@
"""数据处理任务的大模型生成适配器。"""
from __future__ import annotations
import hashlib
import json
import logging
import re
from collections.abc import Callable, Iterable, Mapping
from typing import Any
from urllib.parse import urlsplit, urlunsplit
import httpx
from app.modules.data_process.algorithms import normalize_text, stable_split_assignments
from app.modules.data_process.constants import (
MAX_QA_PAIRS_PER_ITEM,
MODEL_GENERATION_BATCH_SIZE,
)
class ModelGenerationError(ValueError):
"""模型配置、响应或调用失败。"""
class _TerminalModelGenerationError(ModelGenerationError):
"""使用相同参数重试也无法恢复的模型响应错误。"""
OUTPUT_TYPE_STANDARD = "standard"
OUTPUT_TYPE_REASONING = "reasoning"
OUTPUT_TYPE_DPO = "dpo"
SUPPORTED_OUTPUT_TYPES = {
OUTPUT_TYPE_STANDARD,
OUTPUT_TYPE_REASONING,
OUTPUT_TYPE_DPO,
}
REASONING_DETAIL_NORMAL = "normal"
REASONING_DETAIL_DETAILED = "detailed"
SUPPORTED_REASONING_DETAILS = {
REASONING_DETAIL_NORMAL,
REASONING_DETAIL_DETAILED,
}
MINIMAX_M3_API_HOSTS = {"api.minimax.io", "api.minimaxi.com"}
MINIMAX_M3_MIN_COMPLETION_TOKENS = 4096
logger = logging.getLogger(__name__)
# 问题表述风格规则:防止模型产出“请描述/请说明”式模板化问句。
_QUESTION_STYLE_RULE = (
"各条问题必须覆盖不同的信息点并使用不同的句式,只替换关键词套用同一句式视为重复。"
"问题表述要像真实用户自然提出的问题:具体、口语化、直奔信息点,"
"避免“请描述”“请说明”“根据文档”等模板化开头,"
"也不要把原文句子直接改成问句;多条问题时交替使用直接疑问、场景式提问、追问式等句式。"
"表述示例(仅示意风格,不要照搬内容):"
"避免——“请描述系统的权限控制机制”;"
"推荐——“不同角色能看到的菜单不一样,平台是怎么控制的?”"
)
# 任务配置未提供提示语时的兜底,与前端内置默认提示语保持同等信息量。
_DEFAULT_GENERATION_PROMPT = (
"你是一名专业的数据生成专家。请基于来源内容生成高质量、"
"可直接用于监督微调的问答数据:问题聚焦核心信息点、"
"表述像真实用户自然提出的问题,具体、口语化,多条问题使用不同句式;"
"答案严格依据来源内容,准确、完整、语言自然,不引入来源之外的信息。"
)
def _is_retryable_generation_error(exc: Exception) -> bool:
if isinstance(exc, _TerminalModelGenerationError):
return False
if isinstance(exc, httpx.HTTPStatusError):
status_code = exc.response.status_code
return status_code in {408, 425, 429} or status_code >= 500
if isinstance(exc, httpx.RequestError):
return True
return isinstance(exc, (json.JSONDecodeError, ModelGenerationError))
def _is_official_minimax_m3(endpoint: str, model_name: str) -> bool:
host = (urlsplit(endpoint).hostname or "").casefold()
return host in MINIMAX_M3_API_HOSTS and model_name.casefold() == "minimax-m3"
def chat_completions_url(value: str) -> str:
"""把域名、基础 URL 或完整地址统一为 chat completions 地址。"""
raw = (value or "").strip()
if not raw:
raise ModelGenerationError("generation model api_url is required")
if "://" not in raw:
raw = f"https://{raw}"
parsed = urlsplit(raw)
if parsed.scheme not in {"http", "https"} or not parsed.hostname:
raise ModelGenerationError("generation model api_url must be an HTTP(S) host or URL")
if parsed.username or parsed.password:
raise ModelGenerationError("generation model api_url must not contain credentials")
path = parsed.path.rstrip("/")
if path.endswith("/chat/completions"):
target_path = path
elif path.endswith("/v1"):
target_path = f"{path}/chat/completions"
elif not path:
target_path = "/v1/chat/completions"
else:
target_path = f"{path}/v1/chat/completions"
return urlunsplit((parsed.scheme, parsed.netloc, target_path, "", ""))
def _response_choice(payload: Mapping[str, Any]) -> Mapping[str, Any]:
try:
choice = payload["choices"][0]
except (KeyError, IndexError, TypeError) as exc:
raise ModelGenerationError("模型响应缺少 choices[0]") from exc
if not isinstance(choice, Mapping):
raise ModelGenerationError("模型响应 choices[0] 不是对象")
return choice
def _response_finish_reason(payload: Mapping[str, Any]) -> str:
try:
return str(_response_choice(payload).get("finish_reason") or "").strip().lower()
except ModelGenerationError:
return ""
def _response_content_length(payload: Mapping[str, Any]) -> int:
try:
message = _response_choice(payload).get("message")
if not isinstance(message, Mapping):
return 0
content = message.get("content")
if isinstance(content, str):
return len(content)
if isinstance(content, list):
return sum(
len(str(item.get("text") or ""))
for item in content
if isinstance(item, Mapping)
)
except ModelGenerationError:
pass
return 0
def _raise_for_terminal_response(payload: Mapping[str, Any]) -> Mapping[str, Any]:
choice = _response_choice(payload)
base_response = payload.get("base_resp")
status_code: Any = None
status_message = ""
if isinstance(base_response, Mapping):
status_code = base_response.get("status_code")
status_message = re.sub(
r"\s+", " ", str(base_response.get("status_msg") or "")
).strip()[:200]
if bool(payload.get("input_sensitive")) or status_code in {1026, "1026"}:
raise _TerminalModelGenerationError(
f"模型输入触发内容安全拦截code={status_code or 1026}"
)
if bool(payload.get("output_sensitive")) or status_code in {1027, "1027"}:
raise _TerminalModelGenerationError(
f"模型输出触发内容安全拦截code={status_code or 1027}"
)
finish_reason = str(choice.get("finish_reason") or "").strip().lower()
if finish_reason == "length":
raise _TerminalModelGenerationError(
"模型输出因达到 Token 上限被截断finish_reason=length"
"请提高最大输出长度后重试"
)
if finish_reason == "content_filter":
raise _TerminalModelGenerationError(
"模型输出被内容安全策略拦截finish_reason=content_filter"
)
if finish_reason in {"tool_calls", "function_call"}:
raise _TerminalModelGenerationError(
f"模型返回了当前生成任务不支持的工具调用finish_reason={finish_reason}"
)
if status_code not in {None, "", 0, "0"}:
detail = f"{status_message}" if status_message else ""
raise _TerminalModelGenerationError(
f"模型服务返回业务错误code={status_code}{detail}"
)
return choice
def _message_content(payload: Mapping[str, Any]) -> str:
choice = _raise_for_terminal_response(payload)
message = choice.get("message")
if not isinstance(message, Mapping):
raise ModelGenerationError("模型响应缺少 choices[0].message")
content = message.get("content")
if isinstance(content, str):
result = content
elif isinstance(content, list):
parts = [
str(item.get("text") or "")
for item in content
if isinstance(item, Mapping) and item.get("type") in {None, "text", "output_text"}
]
result = "".join(parts)
elif content is None:
result = ""
else:
raise ModelGenerationError("模型响应 content 必须是文本")
if not result.strip():
raise ModelGenerationError("模型返回的最终内容为空,未生成可解析的 JSON")
return result
def _json_documents(content: str) -> list[Any]:
decoder = json.JSONDecoder()
documents: list[Any] = []
cursor = 0
while cursor < len(content):
match = re.search(r"[\[{]", content[cursor:])
if not match:
break
start = cursor + match.start()
try:
value, end = decoder.raw_decode(content[start:])
except json.JSONDecodeError:
cursor = start + 1
continue
if isinstance(value, (Mapping, list)):
documents.append(value)
cursor = start + max(end, 1)
return documents
def _json_payload(content: str) -> Any:
# 只移除模型在 JSON 之前自行输出的思考过程,不能破坏 JSON 字段中的训练内容。
cleaned = content.strip()
if re.match(r"^\s*<think>", cleaned, flags=re.IGNORECASE) and not re.match(
r"^\s*<think>[\s\S]*?</think>", cleaned, flags=re.IGNORECASE
):
raise ModelGenerationError("模型思考内容未闭合,响应可能已被截断")
cleaned = re.sub(
r"^\s*(?:<think>[\s\S]*?</think>\s*)+",
"",
cleaned,
count=1,
flags=re.IGNORECASE,
).strip()
fenced = re.fullmatch(r"```(?:json)?\s*([\s\S]*?)\s*```", cleaned, flags=re.IGNORECASE)
if fenced:
cleaned = fenced.group(1).strip()
try:
return json.loads(cleaned)
except json.JSONDecodeError as direct_error:
documents = _json_documents(cleaned)
if len(documents) == 1:
return documents[0]
if len(documents) > 1:
raise ModelGenerationError("模型响应包含多个 JSON 对象,无法确定应使用哪一个")
raise ModelGenerationError(
"模型响应中没有找到唯一且完整的 JSON 对象"
f"(第 {direct_error.lineno} 行,第 {direct_error.colno} 列)"
) from direct_error
def _result_items(payload: Any) -> list[Mapping[str, Any]]:
if isinstance(payload, list):
values = payload
elif isinstance(payload, Mapping):
nested = next(
(
payload[key]
for key in ("items", "results", "data", "records")
if isinstance(payload.get(key), list)
),
None,
)
values = nested if isinstance(nested, list) else [payload]
else:
raise ModelGenerationError("model JSON must be an object or array")
items = [item for item in values if isinstance(item, Mapping)]
if not items:
raise ModelGenerationError("model JSON does not contain result objects")
return items
def _prompt_messages(
prompt: str,
content: str,
count: int,
*,
start_index: int,
total_count: int,
output_type: str,
reasoning_detail: str,
) -> list[dict[str, str]]:
end_index = start_index + count - 1
if output_type == OUTPUT_TYPE_REASONING:
schema = '{"items":[{"instruction":"...","input":"...","reasoning":"...","answer":"..."}]}'
detail_rule = (
"推理详细程度为“详细”:完整展开问题条件、来源依据、中间计算或推导,"
"并在得出答案前核对结论;每一步都必须能从来源内容中验证。"
if reasoning_detail == REASONING_DETAIL_DETAILED
else
"推理详细程度为“普通”:只保留得出答案所需的关键依据和必要步骤,"
"避免冗长复述、套话和无依据扩展。"
)
output_rule = (
"你正在生成用于训练推理模型的思维链数据,而不是普通问答数据。"
"instruction、reasoning 和 answer 均不得为空reasoning 必须是基于来源内容、"
f"可核对的推理过程answer 只写最终答案。{detail_rule}"
"这是思维链输出模式,即使其他提示语要求省略分析,也不得省略 reasoning。"
"不要自行添加 <think> 标签,系统会在保存时统一组装。"
)
elif output_type == OUTPUT_TYPE_DPO:
schema = (
'{"items":[{"instruction":"...","input":"...",'
'"chosen":"...","rejected":"..."}]}'
)
output_rule = (
"你正在生成用于直接偏好优化DPO的成对偏好数据。"
"instruction、chosen 和 rejected 均不得为空chosen 必须是忠于来源、"
"准确完整的优选回答rejected 必须是表面合理但存在明确质量缺陷的拒选回答。"
"两者不得相同rejected 不得包含违法危险内容,也不得用空白、乱码或无关文本凑数。"
"不要输出分析过程或 <think> 标签。"
)
else:
schema = '{"items":[{"instruction":"...","input":"...","output":"..."}]}'
output_rule = (
"你正在生成标准监督微调问答数据。instruction 和 output 不得为空;"
"output 只写最终答案,禁止输出分析、推理过程或 <think> 标签。"
)
schema_instruction = (
f"必须只返回 JSON 对象,格式为 {schema}items 必须包含 {count} 条。"
f"这是总计 {total_count} 条中的第 {start_index}-{end_index} 条。"
f"{_QUESTION_STYLE_RULE}{output_rule}"
"不要输出 Markdown 代码围栏或 JSON 之外的说明。"
)
base_prompt = normalize_text(prompt) or _DEFAULT_GENERATION_PROMPT
if "{{ content }}" in base_prompt:
user_prompt = base_prompt.replace("{{ content }}", content)
return [
{"role": "system", "content": schema_instruction},
{"role": "user", "content": user_prompt},
]
return [
{"role": "system", "content": f"{base_prompt}\n{schema_instruction}"},
{"role": "user", "content": f"来源内容:\n{content}"},
]
def generate_model_records(
preview_items: Iterable[Mapping[str, Any]],
*,
model: Mapping[str, Any],
config: Mapping[str, Any],
task_id: str,
split: Mapping[str, int],
qa_pairs_per_item: int,
client: httpx.Client | None = None,
on_progress: Callable[[int, int], None] | None = None,
) -> list[dict[str, Any]]:
"""调用 OpenAI 兼容接口,将预览切片生成标准训练记录。
每个切片按安全批次调用模型;失败批次会产生一条可人工修复的
invalid 结果,已经成功的批次不会丢失。
"""
if not 1 <= qa_pairs_per_item <= MAX_QA_PAIRS_PER_ITEM:
raise ModelGenerationError(f"qa_pairs_per_item must be in [1, {MAX_QA_PAIRS_PER_ITEM}]")
output_type = str(config.get("output_type") or OUTPUT_TYPE_STANDARD).strip().lower()
if output_type not in SUPPORTED_OUTPUT_TYPES:
raise ModelGenerationError(f"output_type must be one of {sorted(SUPPORTED_OUTPUT_TYPES)}")
reasoning_detail = str(
config.get("reasoning_detail") or REASONING_DETAIL_NORMAL
).strip().lower()
if reasoning_detail not in SUPPORTED_REASONING_DETAILS:
raise ModelGenerationError(
f"reasoning_detail must be one of {sorted(SUPPORTED_REASONING_DETAILS)}"
)
endpoint = chat_completions_url(str(model.get("api_url") or ""))
model_name = str(model.get("online_model_name") or model.get("name") or "").strip()
if not model_name:
raise ModelGenerationError("generation model name is required")
is_minimax_m3 = _is_official_minimax_m3(endpoint, model_name)
temperature = float(config.get("temperature", 0.7))
max_tokens = int(config.get("max_tokens", 1024))
timeout = max(1.0, min(120.0, float(config.get("request_timeout_seconds", 60))))
retries = max(0, min(5, int(config.get("generation_retries", 2))))
headers = {"Content-Type": "application/json"}
api_key = str(model.get("api_key") or "").strip()
if api_key:
headers["Authorization"] = f"Bearer {api_key}"
owns_client = client is None
http_client = client or httpx.Client(timeout=timeout)
results: list[dict[str, Any]] = []
try:
preview_list = list(preview_items)
total_items = len(preview_list)
for item_index, item in enumerate(preview_list):
preview_id = str(item.get("id") or f"preview-{item_index + 1}")
content = normalize_text(
str(item.get("edited_content") or item.get("original_content") or "")
)
for batch_offset in range(0, qa_pairs_per_item, MODEL_GENERATION_BATCH_SIZE):
batch_count = min(
MODEL_GENERATION_BATCH_SIZE,
qa_pairs_per_item - batch_offset,
)
batch_start = batch_offset + 1
batch_end = batch_offset + batch_count
request_payload: dict[str, Any] = {
"model": model_name,
"messages": _prompt_messages(
str(config.get("generation_prompt") or ""),
content,
batch_count,
start_index=batch_start,
total_count=qa_pairs_per_item,
output_type=output_type,
reasoning_detail=reasoning_detail,
),
"temperature": temperature,
}
if is_minimax_m3:
request_payload.update(
reasoning_split=True,
max_completion_tokens=max(
max_tokens,
MINIMAX_M3_MIN_COMPLETION_TOKENS,
),
)
else:
request_payload["max_tokens"] = max_tokens
if bool(config.get("json_mode", False)) and not is_minimax_m3:
request_payload["response_format"] = {"type": "json_object"}
last_error: Exception | None = None
generated_items: list[Mapping[str, Any]] | None = None
for _ in range(retries + 1):
try:
response = http_client.post(
endpoint,
headers=headers,
json=request_payload,
)
response.raise_for_status()
body = response.json()
if not isinstance(body, Mapping):
raise ModelGenerationError("model response body must be a JSON object")
try:
candidate_items = _result_items(
_json_payload(_message_content(body))
)
except ModelGenerationError as exc:
logger.warning(
"data process model response rejected task_id=%s model=%s "
"finish_reason=%s response_chars=%s input_sensitive=%s "
"output_sensitive=%s reason=%s",
task_id,
model_name,
_response_finish_reason(body) or "missing",
_response_content_length(body),
bool(body.get("input_sensitive")),
bool(body.get("output_sensitive")),
str(exc),
)
raise
if len(candidate_items) < batch_count:
raise ModelGenerationError(
"model response contains fewer result objects than requested: "
f"expected {batch_count}, got {len(candidate_items)}"
)
generated_items = candidate_items
break
except (
httpx.HTTPError,
json.JSONDecodeError,
ModelGenerationError,
) as exc:
last_error = exc
if not _is_retryable_generation_error(exc):
break
if generated_items is None:
error_message = str(last_error or "model generation failed")[:2000]
failure_instruction = (
f"模型生成失败,请人工补充(第 {batch_start}-{batch_end} 条)"
)
result_id = (
"result_"
f"{hashlib.sha256(f'{preview_id}:error:{batch_start}'.encode()).hexdigest()[:16]}"
)
results.append(
{
"id": result_id,
"preview_item_id": preview_id,
"instruction": failure_instruction,
"input": content,
"output": "",
"chosen": "",
"rejected": "",
"original_instruction": failure_instruction,
"original_input": content,
"original_output": "",
"original_chosen": "",
"original_rejected": "",
"status": "invalid",
"error": error_message,
"split": "train",
}
)
continue
for batch_index, value in enumerate(generated_items[:batch_count]):
variant_index = batch_offset + batch_index
instruction = normalize_text(
str(value.get("instruction") or value.get("question") or "")
)
input_text = normalize_text(
str(value.get("input") or value.get("context") or "")
)
chosen = ""
rejected = ""
if output_type == OUTPUT_TYPE_REASONING:
reasoning = normalize_text(
re.sub(
r"</?think>",
"",
str(value.get("reasoning") or value.get("analysis") or ""),
flags=re.IGNORECASE,
)
)
answer = normalize_text(
re.sub(
r"</?think>",
"",
str(
value.get("answer")
or value.get("final_answer")
or value.get("output")
or ""
),
flags=re.IGNORECASE,
)
)
output = (
f"<think>\n{reasoning}\n</think>\n{answer}"
if reasoning and answer
else answer or (f"<think>\n{reasoning}\n</think>" if reasoning else "")
)
valid = bool(instruction and reasoning and answer)
missing_error = "model result is missing instruction, reasoning or answer"
elif output_type == OUTPUT_TYPE_DPO:
chosen = normalize_text(str(value.get("chosen") or ""))
rejected = normalize_text(str(value.get("rejected") or ""))
chosen = normalize_text(
re.sub(
r"<think>[\s\S]*?(?:</think>|$)",
"",
chosen,
flags=re.IGNORECASE,
)
)
rejected = normalize_text(
re.sub(
r"<think>[\s\S]*?(?:</think>|$)",
"",
rejected,
flags=re.IGNORECASE,
)
)
output = chosen
valid = bool(
instruction
and chosen
and rejected
and chosen.strip() != rejected.strip()
)
missing_error = (
"model result is missing instruction, chosen or rejected, "
"or chosen equals rejected"
)
else:
output = normalize_text(
str(
value.get("output")
or value.get("answer")
or value.get("response")
or ""
)
)
output = normalize_text(
re.sub(
r"<think>[\s\S]*?(?:</think>|$)",
"",
output,
flags=re.IGNORECASE,
)
)
valid = bool(instruction and output)
missing_error = "model result is missing instruction or output"
raw_id = (
f"{preview_id}:{variant_index + 1}:{instruction}:"
f"{output}:{rejected}"
)
result_id = f"result_{hashlib.sha256(raw_id.encode()).hexdigest()[:16]}"
results.append(
{
"id": result_id,
"preview_item_id": preview_id,
"instruction": instruction,
"input": input_text,
"output": output,
"chosen": chosen,
"rejected": rejected,
"original_instruction": instruction,
"original_input": input_text,
"original_output": output,
"original_chosen": chosen,
"original_rejected": rejected,
"status": "valid" if valid else "invalid",
"error": (None if valid else missing_error),
"split": "train",
}
)
if on_progress:
on_progress(item_index + 1, total_items)
finally:
if owns_client:
http_client.close()
assignments = stable_split_assignments(
[str(result["id"]) for result in results],
split,
seed=task_id,
)
for result, assignment in zip(results, assignments, strict=True):
result["split"] = assignment
return results
__all__ = ["ModelGenerationError", "chat_completions_url", "generate_model_records"]

View File

@@ -0,0 +1,319 @@
"""Word 与 Excel 原文件的安全、受限预览模型。
预览只返回浏览器绘制所需的结构化数据,不返回或执行 Office 包中的活动内容。
DOCX 的字符偏移与上传时的正文抽取规则保持一致,供前端定位当前切片。
"""
from __future__ import annotations
import io
import re
from typing import Any
from docx import Document
from docx.oxml.ns import qn
from docx.oxml.table import CT_Tbl
from docx.oxml.text.paragraph import CT_P
from docx.table import Table
from docx.text.paragraph import Paragraph
from openpyxl import load_workbook
from app.modules.data_process.algorithms import (
_MAX_WORKBOOK_COLUMNS,
_MAX_WORKBOOK_HEADER_SCAN_ROWS,
_infer_xlsx_header_region,
_normalize_spreadsheet_value,
_rewrite_xlsx_workbook_relationships,
_validate_office_archive,
_xlsx_sheet_merge_ranges,
normalize_text,
)
from app.modules.data_process.algorithms.parsers.office import iter_document_blocks
MAX_DOCX_PREVIEW_BLOCKS = 2_000
MAX_XLSX_PREVIEW_ROWS = 200
def _docx_alignment(paragraph: Paragraph) -> str:
value = paragraph.alignment
return {
0: "left",
1: "center",
2: "right",
3: "justify",
4: "distribute",
5: "justify",
7: "justify",
8: "distribute",
9: "distribute",
}.get(int(value) if value is not None else -1, "left")
def _docx_heading_level(paragraph: Paragraph) -> int | None:
style = paragraph.style
style_name = str(style.name or "") if style is not None else ""
style_id = str(style.style_id or "") if style is not None else ""
match = re.search(r"(?:heading|标题)\s*([1-6])", f"{style_name} {style_id}", re.IGNORECASE)
if match:
return int(match.group(1))
# Word 的目录和导航窗格依据大纲级别识别标题;未套标题样式但带
# outlineLvl 的段落(如手工排版的编号小节)同样是标题。
outline = paragraph._p.find(f"{qn('w:pPr')}/{qn('w:outlineLvl')}")
if outline is not None:
value = outline.get(qn("w:val"))
if value is not None and value.isdigit():
level = int(value)
if 0 <= level <= 5:
return level + 1
return None
def build_docx_preview(raw: bytes) -> dict[str, Any]:
"""把 DOCX 转为保留标题、段落和表格顺序的浏览器预览模型。"""
_validate_office_archive(raw, "docx")
try:
document = Document(io.BytesIO(raw))
except Exception as exc:
raise ValueError(f"invalid DOCX file: {exc}") from exc
blocks: list[dict[str, Any]] = []
source_cursor = 0
has_source_content = False
rendered_blocks = 0
truncated = False
def source_range(value: str) -> tuple[str, int, int] | None:
nonlocal source_cursor, has_source_content
text = normalize_text(value)
if not text:
return None
if has_source_content:
source_cursor += 2
start = source_cursor
source_cursor += len(text)
has_source_content = True
return text, start, source_cursor
for child in iter_document_blocks(document.element.body):
if rendered_blocks >= MAX_DOCX_PREVIEW_BLOCKS:
truncated = True
break
if isinstance(child, CT_P):
paragraph = Paragraph(child, document)
located = source_range(paragraph.text)
if located is None:
continue
text, start, end = located
style_name = str(paragraph.style.name or "") if paragraph.style else ""
blocks.append(
{
"type": "paragraph",
"text": text,
"style": style_name,
"heading_level": _docx_heading_level(paragraph),
"alignment": _docx_alignment(paragraph),
"is_list": "list" in style_name.casefold() or "列表" in style_name,
"source_start": start,
"source_end": end,
}
)
rendered_blocks += 1
continue
if not isinstance(child, CT_Tbl):
continue
table = Table(child, document)
preview_rows: list[dict[str, Any]] = []
for row in table.rows:
if rendered_blocks >= MAX_DOCX_PREVIEW_BLOCKS:
truncated = True
break
cell_values = [normalize_text(cell.text) for cell in row.cells]
located = source_range("\t".join(cell_values))
if located is None:
continue
_, start, end = located
preview_rows.append(
{
"cells": cell_values,
"source_start": start,
"source_end": end,
}
)
rendered_blocks += 1
if preview_rows:
blocks.append({"type": "table", "rows": preview_rows})
if truncated:
break
return {
"format": "docx",
"blocks": blocks,
"truncated": truncated,
}
def build_xlsx_preview(
raw: bytes,
*,
sheet_index: int = 0,
offset: int = 0,
limit: int = 100,
) -> dict[str, Any]:
"""按工作表分页返回 XLSX 的表头和记录网格。"""
if sheet_index < 0 or offset < 0:
raise ValueError("sheet_index and offset must be non-negative")
if limit < 1 or limit > MAX_XLSX_PREVIEW_ROWS:
raise ValueError(
f"XLSX preview limit must be in [1, {MAX_XLSX_PREVIEW_ROWS}]"
)
_validate_office_archive(raw, "xlsx")
merged_by_sheet, normalized_targets = _xlsx_sheet_merge_ranges(raw)
workbook_raw = (
_rewrite_xlsx_workbook_relationships(raw, normalized_targets)
if normalized_targets
else raw
)
try:
workbook = load_workbook(
io.BytesIO(workbook_raw),
read_only=True,
data_only=True,
keep_links=False,
)
except Exception as exc:
raise ValueError(f"invalid XLSX file: {exc}") from exc
try:
sheets = [
{
"index": index,
"name": worksheet.title,
"state": worksheet.sheet_state,
}
for index, worksheet in enumerate(workbook.worksheets)
]
if not sheets:
raise ValueError("XLSX workbook contains no worksheets")
if sheet_index >= len(sheets):
raise ValueError("XLSX worksheet index is out of range")
worksheet = workbook.worksheets[sheet_index]
reset_dimensions = getattr(worksheet, "reset_dimensions", None)
if callable(reset_dimensions):
reset_dimensions()
row_iterator = enumerate(worksheet.iter_rows(values_only=True), start=1)
buffered_rows: dict[int, tuple[Any, ...]] = {}
def normalized_values(row: tuple[Any, ...]) -> list[Any]:
values = list(row)
while values and values[-1] in {None, ""}:
values.pop()
if len(values) > _MAX_WORKBOOK_COLUMNS:
raise ValueError(
f"XLSX worksheet {worksheet.title!r} exceeds "
f"{_MAX_WORKBOOK_COLUMNS} columns"
)
return values
for row_number, row in row_iterator:
values = normalized_values(row)
if not values or all(value in {None, ""} for value in values):
continue
buffered_rows[row_number] = tuple(values)
if len(buffered_rows) >= _MAX_WORKBOOK_HEADER_SCAN_ROWS:
break
if not buffered_rows:
return {
"format": "xlsx",
"sheets": sheets,
"active_sheet": {
"index": sheet_index,
"name": worksheet.title,
"columns": [],
"rows": [],
"offset": offset,
"limit": limit,
"has_more": False,
},
}
_, header_end_row, headers = _infer_xlsx_header_region(
worksheet.title,
buffered_rows,
merged_by_sheet.get(worksheet.title, ()),
)
preview_rows: list[dict[str, Any]] = []
record_index = 0
has_more = False
def append_row(row_number: int, values: tuple[Any, ...] | list[Any]) -> bool:
nonlocal record_index, has_more
row_values = list(values)
if len(row_values) > len(headers):
raise ValueError(
f"XLSX worksheet {worksheet.title!r} has a row wider than its header"
)
row_values.extend([None] * (len(headers) - len(row_values)))
record = {
header: _normalize_spreadsheet_value(value)
for header, value in zip(headers, row_values, strict=True)
}
if not any(value not in {"", None} for value in record.values()):
return False
current_index = record_index
record_index += 1
if current_index < offset:
return False
if len(preview_rows) >= limit:
has_more = True
return True
preview_rows.append(
{
"row_number": row_number,
"record_index": current_index,
"values": [record[header] for header in headers],
"record": record,
}
)
return False
for row_number, values in buffered_rows.items():
if row_number > header_end_row and append_row(row_number, values):
break
else:
for row_number, row in row_iterator:
values = normalized_values(row)
if not values or all(value in {None, ""} for value in values):
continue
if append_row(row_number, values):
break
return {
"format": "xlsx",
"sheets": sheets,
"active_sheet": {
"index": sheet_index,
"name": worksheet.title,
"columns": headers,
"rows": preview_rows,
"offset": offset,
"limit": limit,
"has_more": has_more,
},
}
finally:
workbook.close()
__all__ = [
"MAX_DOCX_PREVIEW_BLOCKS",
"MAX_XLSX_PREVIEW_ROWS",
"build_docx_preview",
"build_xlsx_preview",
]

View File

@@ -0,0 +1,76 @@
"""数据处理运行表的显式检查与安装命令。"""
from __future__ import annotations
import argparse
from urllib.parse import urlsplit
from app.modules.data_process.store import DataProcessStore
REQUIRED_TASK_COLUMNS = (
"generation_run_id",
"results_confirmed",
"workflow_step",
"preview_status",
"preview_progress",
"preview_run_id",
"preview_failure_reason",
"preview_total_files",
"preview_completed_files",
)
def _target_label(database_url: str) -> str:
parsed = urlsplit(database_url)
database = parsed.path.strip("/") or "(unknown)"
return f"{parsed.hostname or '(unknown)'}:{parsed.port or 5432}/{database}"
def _schema_ready(store: DataProcessStore) -> bool:
with store.connect() as conn:
row = conn.execute(
"""
SELECT COUNT(*) = %s AS ready
FROM information_schema.columns
WHERE table_schema=current_schema()
AND table_name='data_process_tasks'
AND column_name = ANY(%s)
""",
(len(REQUIRED_TASK_COLUMNS), list(REQUIRED_TASK_COLUMNS)),
).fetchone()
return bool(row and row["ready"])
def main() -> int:
parser = argparse.ArgumentParser(
description="检查或显式安装数据处理运行表(不会由应用启动自动执行)"
)
action = parser.add_mutually_exclusive_group(required=True)
action.add_argument("--check", action="store_true", help="只读检查迁移是否已安装")
action.add_argument("--apply", action="store_true", help="执行 002 数据处理迁移")
parser.add_argument(
"--yes",
action="store_true",
help="确认允许修改 DATABASE_URL 指向的数据库;与 --apply 同时使用",
)
args = parser.parse_args()
store = DataProcessStore()
target = _target_label(store.database_url)
if args.check:
ready = _schema_ready(store)
print(f"数据处理 schema{'已安装' if ready else '未安装'};目标:{target}")
return 0 if ready else 1
if not args.yes:
parser.error("--apply 必须同时提供 --yes确认修改目标数据库")
print(f"正在安装数据处理 schema目标{target}")
store.ensure_schema()
if not _schema_ready(store):
raise RuntimeError("迁移执行后仍未检测到 generation_run_id")
print("数据处理 schema 安装完成")
return 0
if __name__ == "__main__":
raise SystemExit(main())

View File

@@ -0,0 +1,669 @@
"""数据处理源文件的受控暂存与分层对象存储。"""
from __future__ import annotations
import os
import re
import stat
import unicodedata
import uuid
from dataclasses import dataclass
from functools import lru_cache
from pathlib import Path, PurePosixPath
from typing import Iterable, Iterator
from urllib.parse import quote, unquote, urlsplit
from app.core.config import get_settings
from app.modules.storage.minio_store import get_object_storage
class DataProcessStorageError(ValueError):
"""本地对象引用或文件系统状态不安全。"""
@dataclass(frozen=True, slots=True)
class StagedSourceObject:
"""尚未发布的原始文件;绝对路径仅在存储模块内部流转。"""
reference: str
_temporary_path: Path
_relative_path: PurePosixPath
def _default_storage_root() -> Path:
data_root = os.getenv("YG_FT_DATA_ROOT", "").strip()
if data_root:
return Path(data_root).expanduser() / "data-process"
return Path(__file__).resolve().parents[3] / "storage" / "data-process"
def _configured_storage_root() -> Path:
configured = os.getenv("DATA_PROCESS_STORAGE_DIR", "").strip()
if not configured:
return _default_storage_root()
path = Path(configured).expanduser()
# 相对配置固定以 backend 目录为基准,
# 避免从不同 cwd 启动时写入不同位置。
return path if path.is_absolute() else Path(__file__).resolve().parents[3] / path
def _safe_component(value: str, label: str) -> str:
if not value or value in {".", ".."} or len(value) > 128:
raise DataProcessStorageError(f"invalid {label}")
if not value[0].isalnum() or any(
not (character.isalnum() or character in {"-", "_", "."})
for character in value
):
raise DataProcessStorageError(f"invalid {label}")
return value
def _safe_basename(value: str) -> str:
if not value or len(value.encode("utf-8")) > 255:
raise DataProcessStorageError("invalid source file name")
if value != Path(value).name or "/" in value or "\\" in value or "\x00" in value:
raise DataProcessStorageError("invalid source file name")
if value in {".", ".."} or any(
unicodedata.category(character).startswith("C") for character in value
):
raise DataProcessStorageError("invalid source file name")
return value
class LocalDataProcessStorage:
"""Stage locally, but publish and read authoritative source files from MinIO."""
def __init__(self, root: str | os.PathLike[str] | Path | None = None) -> None:
configured = Path(root) if root is not None else _configured_storage_root()
configured = configured.expanduser()
if configured.exists() and configured.is_symlink():
raise DataProcessStorageError("data process storage root must not be a symlink")
configured.mkdir(parents=True, exist_ok=True, mode=0o700)
self._root = configured.resolve(strict=True)
# StagedSourceObject 本身是普通 dataclass不能只依赖其中的路径字段判断
# 来源;只接受由当前存储实例实际签发的对象,
# 避免调用方伪造暂存路径。
self._issued_staged_objects: dict[Path, StagedSourceObject] = {}
self._ensure_directory(self._root / ".staging")
@property
def root(self) -> Path:
"""仅供运维和测试检查API 响应不得序列化该属性。"""
return self._root
def new_batch_id(self) -> str:
return f"batch-{uuid.uuid4().hex}"
def stage_bytes(
self,
*,
batch_id: str,
task_id: str,
source_file_id: str,
version: int,
name: str,
content: bytes,
) -> StagedSourceObject:
batch_id = _safe_component(batch_id, "batch id")
task_id = _safe_component(task_id, "task id")
source_file_id = _safe_component(source_file_id, "source file id")
if isinstance(version, bool) or not isinstance(version, int) or version < 1:
raise DataProcessStorageError("invalid source file version")
basename = _safe_basename(name)
if not isinstance(content, bytes):
raise TypeError("content must be bytes")
batch_directory = self._ensure_directory(self._root / ".staging" / batch_id)
temporary_path = batch_directory / f"{source_file_id}-{uuid.uuid4().hex}.tmp"
flags = os.O_CREAT | os.O_EXCL | os.O_WRONLY
if hasattr(os, "O_NOFOLLOW"):
flags |= os.O_NOFOLLOW
descriptor = os.open(temporary_path, flags, 0o600)
try:
with os.fdopen(descriptor, "wb", closefd=True) as stream:
stream.write(content)
stream.flush()
os.fsync(stream.fileno())
except Exception:
temporary_path.unlink(missing_ok=True)
raise
relative_path = PurePosixPath(
task_id,
source_file_id,
f"v{version}",
basename,
)
reference = self._reference(task_id, source_file_id, version, basename)
staged = StagedSourceObject(reference, temporary_path, relative_path)
self._issued_staged_objects[temporary_path] = staged
return staged
def stage_copy(
self,
*,
batch_id: str,
source_reference: str,
expected_source_task_id: str,
expected_source_file_id: str,
task_id: str,
source_file_id: str,
version: int,
name: str,
) -> StagedSourceObject:
"""为不可变源对象创建独立目录项,不把大文件重新读入内存。"""
batch_id = _safe_component(batch_id, "batch id")
task_id = _safe_component(task_id, "task id")
source_file_id = _safe_component(source_file_id, "source file id")
if isinstance(version, bool) or not isinstance(version, int) or version < 1:
raise DataProcessStorageError("invalid source file version")
basename = _safe_basename(name)
source_relative = self._relative_from_reference(source_reference)
if source_relative is None:
raise DataProcessStorageError("original source object is not available")
self._assert_expected_owner(
source_relative,
expected_task_id=expected_source_task_id,
expected_source_file_id=expected_source_file_id,
)
if self._is_minio_reference(source_reference):
content = self.read(source_reference)
if content is None:
raise DataProcessStorageError("original source object is not available")
return self.stage_bytes(
batch_id=batch_id,
task_id=task_id,
source_file_id=source_file_id,
version=version,
name=basename,
content=content,
)
descriptor, source_info = self._open_read_descriptor(source_relative)
os.close(descriptor)
batch_directory = self._ensure_directory(self._root / ".staging" / batch_id)
temporary_path = batch_directory / f"{source_file_id}-{uuid.uuid4().hex}.tmp"
source_path = self._path_for_relative(source_relative)
try:
os.link(source_path, temporary_path, follow_symlinks=False)
copy_info = temporary_path.lstat()
if (
not stat.S_ISREG(copy_info.st_mode)
or source_info.st_dev != copy_info.st_dev
or source_info.st_ino != copy_info.st_ino
):
raise DataProcessStorageError("source storage object changed while copying")
except Exception:
temporary_path.unlink(missing_ok=True)
raise
relative_path = PurePosixPath(
task_id,
source_file_id,
f"v{version}",
basename,
)
reference = self._reference(task_id, source_file_id, version, basename)
staged = StagedSourceObject(reference, temporary_path, relative_path)
self._issued_staged_objects[temporary_path] = staged
return staged
def publish(self, objects: Iterable[StagedSourceObject]) -> None:
staged = list(objects)
published: list[StagedSourceObject] = []
try:
seen_temporary_paths: set[Path] = set()
for item in staged:
self._validate_staged_object(item, require_file=True)
if item._temporary_path in seen_temporary_paths:
raise DataProcessStorageError("duplicate staged source object")
seen_temporary_paths.add(item._temporary_path)
for item in staged:
if self._is_minio_reference(item.reference):
content = item._temporary_path.read_bytes()
get_object_storage().put_bytes(
self.object_key(item.reference),
content,
"application/octet-stream",
)
elif not item.reference.startswith("db://data-process/"):
final_path = self._path_for_relative(item._relative_path)
self._ensure_directory(final_path.parent)
if final_path.exists() or final_path.is_symlink():
raise DataProcessStorageError("source storage object already exists")
os.link(item._temporary_path, final_path, follow_symlinks=False)
self._fsync_directory(final_path.parent)
published.append(item)
item._temporary_path.unlink()
except Exception:
for item in reversed(published):
try:
self.delete(item.reference)
except Exception:
# 回滚必须尽量处理其余对象,并保留真正的发布异常。
pass
for item in staged:
try:
self.discard([item])
except Exception:
pass
raise
self.discard(staged)
def discard(self, objects: Iterable[StagedSourceObject]) -> None:
staged = list(objects)
for item in staged:
self._validate_staged_object(item, require_file=False)
batch_directories: set[Path] = set()
first_error: Exception | None = None
for item in staged:
temporary_path = item._temporary_path
try:
temporary_path.unlink(missing_ok=True)
except Exception as exc:
if first_error is None:
first_error = exc
else:
self._issued_staged_objects.pop(temporary_path, None)
batch_directories.add(temporary_path.parent)
for directory in batch_directories:
self._remove_empty_directory(directory)
if first_error is not None:
raise first_error
def read(self, reference: str) -> bytes | None:
"""Read a MinIO object or legacy local reference."""
if self._is_minio_reference(reference):
try:
return get_object_storage().get_bytes(self.object_key(reference))
except Exception as exc: # noqa: BLE001 - normalize object-not-found for callers
raise DataProcessStorageError("source storage object does not exist") from exc
relative_path = self._relative_from_reference(reference)
if relative_path is None:
return None
descriptor, _ = self._open_read_descriptor(relative_path)
with os.fdopen(descriptor, "rb", closefd=True) as stream:
return stream.read()
def file_size(
self,
reference: str,
*,
expected_task_id: str,
expected_source_file_id: str,
) -> int | None:
"""返回受控 local 对象大小;旧 ``db://`` 对象没有原始文件。"""
if self._is_minio_reference(reference):
try:
return int(get_object_storage().stat(self.object_key(reference)).get("byte_size") or 0)
except Exception as exc: # noqa: BLE001
raise DataProcessStorageError("source storage object does not exist") from exc
relative_path = self._relative_from_reference(reference)
if relative_path is None:
return None
self._assert_expected_owner(
relative_path,
expected_task_id=expected_task_id,
expected_source_file_id=expected_source_file_id,
)
descriptor, info = self._open_read_descriptor(relative_path)
os.close(descriptor)
return info.st_size
def iter_bytes(
self,
reference: str,
*,
expected_task_id: str,
expected_source_file_id: str,
expected_size: int,
start: int = 0,
length: int | None = None,
chunk_size: int = 256 * 1024,
) -> Iterator[bytes]:
"""按范围流式读取原始文件,避免 PDF 预览把大文件整体载入内存。"""
if self._is_minio_reference(reference):
content = self.read(reference) or b""
if start < 0 or expected_size != len(content) or start > expected_size:
raise DataProcessStorageError("source object size does not match metadata")
remaining = expected_size - start if length is None else length
if remaining < 0 or start + remaining > expected_size:
raise DataProcessStorageError("invalid source byte range")
yield content[start : start + remaining]
return
relative_path = self._relative_from_reference(reference)
if relative_path is None:
raise DataProcessStorageError("original source object is not available")
self._assert_expected_owner(
relative_path,
expected_task_id=expected_task_id,
expected_source_file_id=expected_source_file_id,
)
if start < 0 or expected_size < 0 or chunk_size < 1:
raise DataProcessStorageError("invalid source byte range")
descriptor, info = self._open_read_descriptor(relative_path)
if info.st_size != expected_size:
os.close(descriptor)
raise DataProcessStorageError("source object size does not match metadata")
remaining = expected_size - start if length is None else length
if remaining < 0 or start + remaining > expected_size:
os.close(descriptor)
raise DataProcessStorageError("invalid source byte range")
with os.fdopen(descriptor, "rb", closefd=True) as stream:
stream.seek(start)
while remaining:
chunk = stream.read(min(chunk_size, remaining))
if not chunk:
raise DataProcessStorageError("source object ended unexpectedly")
remaining -= len(chunk)
yield chunk
def validate_owner(
self,
reference: str,
*,
expected_task_id: str,
expected_source_file_id: str,
) -> bool:
"""校验 local 引用归属;旧 ``db://`` 引用无需文件系统处理。"""
if self._is_minio_reference(reference):
self._assert_minio_owner(reference, expected_task_id, expected_source_file_id)
return True
if str(reference or "").startswith("db://data-process/"):
self._assert_database_owner(reference, expected_task_id, expected_source_file_id)
return True
relative_path = self._relative_from_reference(reference)
if relative_path is None:
return False
self._assert_expected_owner(
relative_path,
expected_task_id=expected_task_id,
expected_source_file_id=expected_source_file_id,
)
return True
def _open_read_descriptor(
self,
relative_path: PurePosixPath,
) -> tuple[int, os.stat_result]:
path = self._path_for_relative(relative_path)
self._assert_controlled_parent(path)
try:
before_open = path.lstat()
except FileNotFoundError as exc:
raise DataProcessStorageError("source storage object does not exist") from exc
if stat.S_ISLNK(before_open.st_mode) or not stat.S_ISREG(before_open.st_mode):
raise DataProcessStorageError("source storage object is not a regular file")
flags = os.O_RDONLY
if hasattr(os, "O_NOFOLLOW"):
flags |= os.O_NOFOLLOW
descriptor = os.open(path, flags)
after_open = os.fstat(descriptor)
if (
not stat.S_ISREG(after_open.st_mode)
or before_open.st_dev != after_open.st_dev
or before_open.st_ino != after_open.st_ino
):
os.close(descriptor)
raise DataProcessStorageError("source storage object changed while opening")
return descriptor, after_open
def delete(
self,
reference: str,
*,
expected_task_id: str | None = None,
expected_source_file_id: str | None = None,
) -> bool:
"""删除受控 local 对象;旧 ``db://`` 引用保持不变。"""
if self._is_minio_reference(reference):
if (expected_task_id is None) != (expected_source_file_id is None):
raise DataProcessStorageError("both expected storage owner fields are required")
if expected_task_id is not None and expected_source_file_id is not None:
self._assert_minio_owner(reference, expected_task_id, expected_source_file_id)
get_object_storage().delete(self.object_key(reference))
return True
if str(reference or "").startswith("db://data-process/"):
if (expected_task_id is None) != (expected_source_file_id is None):
raise DataProcessStorageError("both expected storage owner fields are required")
if expected_task_id is not None and expected_source_file_id is not None:
self._assert_database_owner(reference, expected_task_id, expected_source_file_id)
return False
relative_path = self._relative_from_reference(reference)
if relative_path is None:
return False
if (expected_task_id is None) != (expected_source_file_id is None):
raise DataProcessStorageError("both expected storage owner fields are required")
if expected_task_id is not None and expected_source_file_id is not None:
self._assert_expected_owner(
relative_path,
expected_task_id=expected_task_id,
expected_source_file_id=expected_source_file_id,
)
path = self._path_for_relative(relative_path)
self._assert_controlled_parent(path)
try:
info = path.lstat()
except FileNotFoundError:
return False
if stat.S_ISLNK(info.st_mode) or not stat.S_ISREG(info.st_mode):
raise DataProcessStorageError("refusing to delete a non-regular storage object")
path.unlink()
self._fsync_directory(path.parent)
for directory in (path.parent, path.parent.parent, path.parent.parent.parent):
self._remove_empty_directory(directory)
return True
@staticmethod
def _assert_expected_owner(
relative_path: PurePosixPath,
*,
expected_task_id: str,
expected_source_file_id: str,
) -> None:
task_id = _safe_component(expected_task_id, "expected task id")
source_file_id = _safe_component(
expected_source_file_id,
"expected source file id",
)
if relative_path.parts[:2] != (task_id, source_file_id):
raise DataProcessStorageError("source storage object owner mismatch")
def _relative_from_reference(self, reference: str) -> PurePosixPath | None:
if reference.startswith("db://"):
return None
parsed = urlsplit(reference)
if parsed.scheme not in {"local", "minio"} or parsed.netloc != "data-process":
raise DataProcessStorageError("unsupported source storage reference")
if parsed.query or parsed.fragment or "\\" in parsed.path:
raise DataProcessStorageError("unsafe source storage reference")
raw_parts = parsed.path.lstrip("/").split("/")
if len(raw_parts) != 4:
raise DataProcessStorageError("unsafe source storage reference")
if any(re.search(r"%(?![0-9A-Fa-f]{2})", part) for part in raw_parts):
raise DataProcessStorageError("unsafe source storage reference")
try:
decoded = [unquote(part, encoding="utf-8", errors="strict") for part in raw_parts]
except UnicodeDecodeError as exc:
raise DataProcessStorageError("unsafe source storage reference") from exc
if any("/" in part or "\\" in part for part in decoded):
raise DataProcessStorageError("unsafe source storage reference")
canonical_parts = [
quote(decoded[0], safe="-_."),
quote(decoded[1], safe="-_."),
quote(decoded[2], safe="-_."),
quote(decoded[3], safe=""),
]
if canonical_parts != raw_parts:
raise DataProcessStorageError("source storage reference is not canonical")
task_id = _safe_component(decoded[0], "task id")
source_file_id = _safe_component(decoded[1], "source file id")
version_component = decoded[2]
if not version_component.startswith("v") or not version_component[1:].isdigit():
raise DataProcessStorageError("invalid source file version")
version = int(version_component[1:])
if version < 1:
raise DataProcessStorageError("invalid source file version")
basename = _safe_basename(decoded[3])
return PurePosixPath(task_id, source_file_id, f"v{version}", basename)
@staticmethod
def _is_minio_reference(reference: str) -> bool:
return str(reference or "").startswith("minio://data-process/")
@staticmethod
def _reference(task_id: str, source_file_id: str, version: int, basename: str) -> str:
scheme = "minio" if get_settings().minio_enabled else "local"
return f"{scheme}://data-process/{task_id}/{source_file_id}/v{version}/{quote(basename, safe='')}"
@staticmethod
def object_key(reference: str) -> str:
parsed = urlsplit(reference)
if parsed.scheme != "minio" or parsed.netloc != "data-process":
raise DataProcessStorageError("reference is not a MinIO source object")
return "data-process/" + parsed.path.lstrip("/")
def _assert_minio_owner(self, reference: str, task_id: str, source_file_id: str) -> None:
relative = self._relative_from_reference(reference)
if relative is None:
raise DataProcessStorageError("invalid MinIO source reference")
self._assert_expected_owner(relative, expected_task_id=task_id, expected_source_file_id=source_file_id)
@staticmethod
def _assert_database_owner(reference: str, task_id: str, source_file_id: str) -> None:
parsed = urlsplit(reference)
parts = parsed.path.lstrip("/").split("/")
if parsed.netloc != "data-process" or len(parts) != 3:
raise DataProcessStorageError("invalid database source reference")
expected_task_id = _safe_component(task_id, "expected task id")
expected_source_file_id = _safe_component(source_file_id, "expected source file id")
if tuple(parts[:2]) != (expected_task_id, expected_source_file_id) or parts[2] != "v1":
raise DataProcessStorageError("source storage object owner mismatch")
def _path_for_relative(self, relative_path: PurePosixPath) -> Path:
if relative_path.is_absolute() or any(
part in {"", ".", ".."} for part in relative_path.parts
):
raise DataProcessStorageError("storage path escapes the configured root")
path = self._root.joinpath(*relative_path.parts)
self._assert_controlled_parent(path)
return path
def _validate_staged_object(
self,
item: StagedSourceObject,
*,
require_file: bool,
) -> None:
if not isinstance(item, StagedSourceObject):
raise DataProcessStorageError("invalid staged source object")
if self._issued_staged_objects.get(item._temporary_path) is not item:
raise DataProcessStorageError("staged source object was not issued by this storage")
if item.reference.startswith("db://data-process/"):
parsed = urlsplit(item.reference)
parts = parsed.path.lstrip("/").split("/")
expected = item._relative_path.parts[:3]
if (
parsed.netloc != "data-process"
or parsed.query
or parsed.fragment
or len(parts) != 3
or tuple(parts) != expected
):
raise DataProcessStorageError("staged source object reference mismatch")
else:
expected_relative = self._relative_from_reference(item.reference)
if expected_relative is None or expected_relative != item._relative_path:
raise DataProcessStorageError("staged source object reference mismatch")
staging_root = self._root / ".staging"
try:
relative_temporary = item._temporary_path.relative_to(staging_root)
except ValueError as exc:
raise DataProcessStorageError("staged source object escapes staging") from exc
if len(relative_temporary.parts) != 2:
raise DataProcessStorageError("invalid staged source object path")
_safe_component(relative_temporary.parts[0], "batch id")
_safe_basename(relative_temporary.parts[1])
self._assert_controlled_parent(item._temporary_path)
try:
info = item._temporary_path.lstat()
except FileNotFoundError:
if require_file:
raise DataProcessStorageError("staged source object does not exist") from None
return
if stat.S_ISLNK(info.st_mode) or not stat.S_ISREG(info.st_mode):
raise DataProcessStorageError("staged source object is not a regular file")
def _ensure_directory(self, directory: Path) -> Path:
try:
relative = directory.relative_to(self._root)
except ValueError as exc:
raise DataProcessStorageError("storage path escapes the configured root") from exc
current = self._root
for component in relative.parts:
current = current / component
try:
current.mkdir(mode=0o700)
except FileExistsError:
pass
info = current.lstat()
if stat.S_ISLNK(info.st_mode) or not stat.S_ISDIR(info.st_mode):
raise DataProcessStorageError("storage path contains a symlink or non-directory")
return directory
def _assert_controlled_parent(self, path: Path) -> None:
try:
relative_parent = path.parent.relative_to(self._root)
except ValueError as exc:
raise DataProcessStorageError("storage path escapes the configured root") from exc
current = self._root
for component in relative_parent.parts:
current = current / component
if not current.exists():
continue
info = current.lstat()
if stat.S_ISLNK(info.st_mode) or not stat.S_ISDIR(info.st_mode):
raise DataProcessStorageError("storage path contains a symlink or non-directory")
@staticmethod
def _fsync_directory(directory: Path) -> None:
# Windows 不支持以 O_RDONLY 打开目录做 fsync跳过即可。
# 数据完整性在 Linux 生产环境保障Windows 开发环境忽略。
if os.name == "nt":
return
descriptor = os.open(directory, os.O_RDONLY)
try:
os.fsync(descriptor)
finally:
os.close(descriptor)
def _remove_empty_directory(self, directory: Path) -> None:
if directory in {self._root, self._root / ".staging"}:
return
self._assert_controlled_parent(directory / "placeholder")
try:
directory.rmdir()
except (FileNotFoundError, OSError):
return
@lru_cache
def get_data_process_storage() -> LocalDataProcessStorage:
return LocalDataProcessStorage()
__all__ = [
"DataProcessStorageError",
"LocalDataProcessStorage",
"StagedSourceObject",
"get_data_process_storage",
]

View File

@@ -0,0 +1,69 @@
"""数据处理存储层。"""
from .base import (
StoreBase,
DataProcessStoreError,
NotFoundError,
ConflictError,
InvalidStateError,
utcnow,
new_id,
repeat_task_id,
json_dumps,
TASK_STATUSES,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_decode_row,
_preview_config_changed,
_reasoning_output_is_valid,
_source_storage_descriptor,
)
from .tasks import TasksMixin
from .source_files import SourceFilesMixin
from .preview import PreviewMixin
from .generation import GenerationMixin
from .results import ResultsMixin
from .datasets import DatasetsMixin
class DataProcessStore(
StoreBase,
TasksMixin,
SourceFilesMixin,
PreviewMixin,
GenerationMixin,
ResultsMixin,
DatasetsMixin,
):
"""数据处理持久层。
构造函数不会连接数据库或执行迁移。部署方必须显式执行 002 SQL
或在受控的管理命令中调用 :meth:`ensure_schema`,避免应用启动时
修改远程数据库。
"""
pass
def get_data_process_store() -> DataProcessStore:
"""获取数据处理存储实例。"""
return DataProcessStore()
__all__ = [
"DataProcessStore",
"DataProcessStoreError",
"NotFoundError",
"ConflictError",
"InvalidStateError",
"get_data_process_store",
"utcnow",
"new_id",
"repeat_task_id",
"TASK_STATUSES",
"EDITABLE_STATUSES",
"_decode_row",
"_preview_config_changed",
"_reasoning_output_is_valid",
"_source_storage_descriptor",
]

View File

@@ -0,0 +1,310 @@
"""数据处理存储层 - 基础设施。
包含:异常类、工具函数、常量定义、基类。
"""
from __future__ import annotations
import hashlib
import json
import re
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from datetime import UTC, date, datetime
from decimal import Decimal
from pathlib import Path
from typing import Any
import psycopg
from psycopg.rows import dict_row
from app.core.config import get_settings
# 常量定义
TASK_STATUSES = {"pending", "running", "completed", "failed", "stopped"}
EDITABLE_STATUSES = {"pending", "failed", "stopped", "completed"}
ACTIVE_PREVIEW_STATUSES = {"queued", "running"}
WORKFLOW_STEPS = {"create", "model", "upload", "preview", "generate", "results"}
_PREVIEW_CONFIG_ALIASES = {
"preprocess_options": "preprocessOptions",
"chunk_method": "chunkMethod",
"chunk_size": "chunkSize",
"chunk_overlap": "chunkOverlap",
"min_chunk_size": "minChunkSize",
"semantic_breakpoint_percentile": "semanticBreakpointPercentile",
"preserve_tables": "preserveTables",
"preserve_code_blocks": "preserveCodeBlocks",
"preserve_lists": "preserveLists",
}
_UNSTRUCTURED_PREVIEW_DEFAULTS: dict[str, Any] = {
"chunk_method": "layout_hybrid",
"chunk_size": 800,
"chunk_overlap": 100,
"min_chunk_size": 100,
"semantic_breakpoint_percentile": 95,
"preserve_tables": True,
"preserve_code_blocks": True,
"preserve_lists": True,
}
_REGENERATION_MARKER_KEY = "_regeneration_prepared"
_REPEAT_SOURCE_TASK_KEY = "_repeat_source_task_id"
_REPEAT_REQUEST_KEY = "_repeat_request_id"
_INTERNAL_CONFIG_KEYS = {
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
}
# 异常类
class DataProcessStoreError(RuntimeError):
pass
class NotFoundError(DataProcessStoreError):
pass
class ConflictError(DataProcessStoreError):
pass
class InvalidStateError(DataProcessStoreError):
pass
# 工具函数
def utcnow() -> str:
return datetime.now(UTC).replace(microsecond=0).isoformat().replace("+00:00", "Z")
def new_id(prefix: str) -> str:
return f"{prefix}_{uuid.uuid4().hex[:20]}"
def repeat_task_id(source_task_id: str, request_id: str) -> str:
"""按源任务和请求幂等键生成稳定的新任务 ID。"""
digest = hashlib.sha256(f"{source_task_id}:{request_id}".encode()).hexdigest()
return f"dpt_{digest[:20]}"
def json_dumps(value: Any) -> str:
return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
def _database_url(value: str) -> str:
return value.replace("postgresql+psycopg://", "postgresql://")
def _json_value(value: Any, default: Any) -> Any:
if value is None or value == "":
return default
if isinstance(value, (dict, list)):
return value
try:
return json.loads(value)
except (TypeError, json.JSONDecodeError):
return default
def _task_output_type(task: dict[str, Any]) -> str:
config = _json_value(task.get("config"), {})
if not isinstance(config, dict):
return "standard"
return str(config.get("output_type") or config.get("outputType") or "standard")
def _task_reasoning_detail(task: dict[str, Any]) -> str:
config = _json_value(task.get("config"), {})
if not isinstance(config, dict):
return "normal"
return str(
config.get("reasoning_detail")
or config.get("reasoningDetail")
or "normal"
)
def _reasoning_output_is_valid(value: Any) -> bool:
match = re.fullmatch(
r"\s*<think>\s*(?P<reasoning>[\s\S]*?)\s*</think>\s*(?P<answer>[\s\S]+?)\s*",
str(value or ""),
flags=re.IGNORECASE,
)
return bool(
match
and match.group("reasoning").strip()
and match.group("answer").strip()
and all(
tag not in part.lower()
for tag in ("<think", "</think")
for part in (match.group("reasoning"), match.group("answer"))
)
)
def _dpo_fields_are_valid(row: dict[str, Any]) -> bool:
chosen = str(row.get("chosen") or "").strip()
rejected = str(row.get("rejected") or "").strip()
return bool(chosen and rejected and chosen != rejected)
def _preview_config_value(config: dict[str, Any], key: str, default: Any) -> Any:
if key in config:
return config[key]
return config.get(_PREVIEW_CONFIG_ALIASES[key], default)
def _normalized_preprocess_options(config: dict[str, Any]) -> Any:
value = _preview_config_value(config, "preprocess_options", [])
if isinstance(value, (list, tuple, set)):
return tuple(sorted({str(item) for item in value}))
return value
def _preview_config_projection(
process_type: str,
config: dict[str, Any],
) -> dict[str, Any]:
"""只投影会改变预览切片的配置。
生成模型、提示词、温度等参数不影响源文切片,因此不应该
破坏用户已经校对过的预览内容。
"""
projection: dict[str, Any] = {
"preprocess_options": _normalized_preprocess_options(config),
}
if process_type != "unstructured":
return projection
for key, default in _UNSTRUCTURED_PREVIEW_DEFAULTS.items():
projection[key] = _preview_config_value(config, key, default)
return projection
def _preview_config_changed(
process_type: str,
current_config: dict[str, Any],
next_config: dict[str, Any],
) -> bool:
return _preview_config_projection(process_type, current_config) != _preview_config_projection(
process_type, next_config
)
def _regeneration_marker(task: dict[str, Any]) -> dict[str, Any] | None:
config = task.get("config")
if not isinstance(config, dict):
return None
marker = config.get(_REGENERATION_MARKER_KEY)
if not isinstance(marker, dict) or marker.get("prepared") is not True:
return None
return marker
def _is_regeneration_prepared(task: dict[str, Any]) -> bool:
return _regeneration_marker(task) is not None
def _business_config(config: dict[str, Any] | None) -> dict[str, Any]:
"""过滤只供服务端维护的工作流标记。"""
return {
key: value
for key, value in (config or {}).items()
if key not in _INTERNAL_CONFIG_KEYS
}
def _public_task(item: dict[str, Any] | None) -> dict[str, Any] | None:
"""从 API 任务快照中移除服务端内部工作流标记。"""
if item is None:
return None
public = dict(item)
config = public.get("config")
if isinstance(config, dict):
public["config"] = _business_config(config)
return public
def _serialize_value(value: Any) -> Any:
if isinstance(value, (datetime, date)):
return value.isoformat().replace("+00:00", "Z")
if isinstance(value, Decimal):
return float(value)
return value
def _source_storage_descriptor(
payload: dict[str, Any],
task_id: str,
file_id: str,
) -> tuple[str, dict[str, Any]]:
storage_object_id = str(
payload.get("storage_object_id")
or f"db://data-process/{task_id}/{file_id}/v1"
)
expected_local_prefix = f"local://data-process/{task_id}/{file_id}/v1/"
expected_minio_prefix = f"minio://data-process/{task_id}/{file_id}/v1/"
expected_database_reference = f"db://data-process/{task_id}/{file_id}/v1"
if storage_object_id.startswith(expected_local_prefix) and len(storage_object_id) > len(
expected_local_prefix
):
storage_backend = "local"
elif storage_object_id.startswith(expected_minio_prefix) and len(storage_object_id) > len(
expected_minio_prefix
):
storage_backend = "minio"
elif storage_object_id == expected_database_reference:
storage_backend = "database"
elif storage_object_id.startswith(("local://data-process/", "minio://data-process/", "db://data-process/")):
raise DataProcessStoreError("source storage object owner mismatch")
else:
raise DataProcessStoreError("unsupported source storage object reference")
metadata = {
**(payload.get("metadata") or {}),
"storage_backend": storage_backend,
}
return storage_object_id, metadata
def _decode_row(row: dict[str, Any] | None) -> dict[str, Any] | None:
if row is None:
return None
item = {key: _serialize_value(value) for key, value in row.items()}
for key, default in {
"config": {},
"metadata": {},
"quality_score": {},
"versions": [],
"output_datasets": [],
}.items():
if key in item:
item[key] = _json_value(item[key], default)
return item
class StoreBase:
"""数据处理存储基类。"""
def __init__(self, database_url: str | None = None) -> None:
self.database_url = _database_url(database_url or get_settings().database_url)
@contextmanager
def connect(self) -> Iterator[psycopg.Connection[dict[str, Any]]]:
with psycopg.connect(self.database_url, row_factory=dict_row) as conn:
try:
yield conn
conn.commit()
except Exception:
conn.rollback()
raise
def ensure_schema(self) -> None:
"""显式安装数据处理表API 路由和应用启动流程不会调用此方法。"""
schema_path = Path(__file__).resolve().parents[3] / "db" / "sql" / "002_data_process.sql"
sql = schema_path.read_text(encoding="utf-8")
with self.connect() as conn, conn.cursor() as cursor:
cursor.execute(sql)

View File

@@ -0,0 +1,567 @@
"""数据处理存储层 - 数据集发布。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import hashlib
import psycopg
from app.core.config import get_settings
from app.modules.storage.minio_store import get_object_storage
from app.modules.storage.policy import should_store_in_minio
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
)
from ..algorithms import stable_split_assignments
class DatasetsMixin:
"""数据集发布 Mixin。"""
def get_generation_model(self, model_id: str) -> dict[str, Any]:
with self.connect() as conn:
row = conn.execute(
"""
SELECT id, name, type, purpose, model_source, description, path,
api_url, api_key, online_model_name, create_time
FROM models WHERE id=%s
""",
(model_id,),
).fetchone()
if not row:
raise NotFoundError("generation model not found")
return _decode_row(row) or {}
def save_generation_model_snapshot(
self,
task_id: str,
model_snapshot: dict[str, Any],
*,
generation_run_id: str,
) -> dict[str, Any]:
# API 密钥仅用于本次调用,绝不能进入任务配置、详情响应或审计快照。
safe_snapshot = {
key: value for key, value in model_snapshot.items() if key != "api_key"
}
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if (
task["status"] != "running"
or task.get("generation_run_id") != generation_run_id
):
raise InvalidStateError("generation run is no longer active")
config = dict(task.get("config") or {})
config["generation_model_snapshot"] = safe_snapshot
row = conn.execute(
"""
UPDATE data_process_tasks SET config=%s, updated_at=%s
WHERE id=%s AND generation_run_id=%s RETURNING *
""",
(json_dumps(config), utcnow(), task_id, generation_run_id),
).fetchone()
return _decode_row(row) or {}
def publish(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
"""按精确配额发布训练、验证、测试三个独立数据集。"""
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if _is_regeneration_prepared(task):
raise InvalidStateError(
"regeneration must start and complete before publishing"
)
if task["status"] != "completed":
raise InvalidStateError("only a completed task can be published")
if not task.get("results_confirmed"):
raise InvalidStateError("results must be confirmed before publishing")
rows = conn.execute(
"""
SELECT * FROM data_process_results
WHERE task_id=%s ORDER BY created_at, id
""",
(task_id,),
).fetchall()
if not rows:
raise InvalidStateError("task has no results to publish")
invalid_count = sum(
1
for row in rows
if row["status"] == "invalid"
or not str(row.get("instruction") or "").strip()
or not str(row.get("output") or "").strip()
or (
_task_output_type(task) == "reasoning"
and not _reasoning_output_is_valid(row.get("output"))
)
or (
_task_output_type(task) == "dpo"
and not _dpo_fields_are_valid(row)
)
)
if invalid_count:
raise InvalidStateError(f"task contains {invalid_count} invalid results")
now = utcnow()
requested_split = payload.get("split") or {
"train": 80,
"validation": 10,
"test": 10,
}
assignments = stable_split_assignments(
[str(row["id"]) for row in rows],
requested_split,
seed=task_id,
)
if _task_output_type(task) == "dpo":
records = [
{
"instruction": row["instruction"],
"input": row["input"],
"chosen": row["chosen"],
"rejected": row["rejected"],
"split": assignment,
}
for row, assignment in zip(rows, assignments, strict=True)
]
else:
records = [
{
"instruction": row["instruction"],
"input": row["input"],
"output": row["output"],
"split": assignment,
}
for row, assignment in zip(rows, assignments, strict=True)
]
split_order = ("train", "validation", "test")
split_counts = {
split_name: assignments.count(split_name) for split_name in split_order
}
split_specs: list[dict[str, Any]] = []
# Explicit local is retained for old callers/tests that request the
# legacy backend; all normal platform requests default to MinIO.
allow_minio = bool(get_settings().minio_enabled) and str(
payload.get("storage_type") or "minio"
).lower() != "local"
for split_name in split_order:
split_records = [
(source_row, record)
for source_row, record in zip(rows, records, strict=True)
if record["split"] == split_name
]
file_id = new_id("dfile")
version_id = new_id("dfv")
content = "".join(
json_dumps(record) + "\n" for _, record in split_records
)
raw = content.encode("utf-8")
split_specs.append(
{
"split": split_name,
"records": split_records,
"file_id": file_id,
"version_id": version_id,
"content": content,
"raw": raw,
"checksum": hashlib.sha256(raw).hexdigest(),
"storage_object_id": (
f"db://data-process/{task_id}/{file_id}/v1"
),
"store_in_minio": allow_minio and should_store_in_minio(
len(raw), content_type="application/jsonl", file_format="jsonl"
),
}
)
source_result_ids = [row["id"] for row in rows]
common_metadata = {
"source": "data_process",
"storage_backend": "minio" if any(spec["store_in_minio"] for spec in split_specs) else "database",
"source_task_id": task_id,
"output_type": _task_output_type(task),
"reasoning_detail": _task_reasoning_detail(task),
"source_file_ids": [item["id"] for item in self._source_ids(conn, task_id)],
"source_result_ids": source_result_ids,
"format": (
"dpo"
if _task_output_type(task) == "dpo"
else payload.get("format") or "alpaca_jsonl"
),
"split": requested_split,
}
existing_datasets = conn.execute(
"""
SELECT * FROM datasets
WHERE source='task' AND deleted_at IS NULL
AND (
source_task_id=%s
OR (source_task_id IS NULL AND task_id=%s)
)
ORDER BY created_at, id
""",
(task_id, task_id),
).fetchall()
existing_by_split: dict[str, dict[str, Any]] = {}
primary_existing = None
for existing in existing_datasets:
existing_metadata = _json_value(existing.get("metadata"), {})
existing_split = str(existing_metadata.get("dataset_split") or "")
if existing_split in split_order:
existing_by_split[existing_split] = existing
if str(existing["id"]) == str(task.get("output_dataset_id") or ""):
primary_existing = existing
if primary_existing and "train" not in existing_by_split:
# 兼容旧版“一个数据集包含三个文件”的发布物,原数据集复用为训练集。
existing_by_split["train"] = primary_existing
existing_group_metadata = _json_value(
(primary_existing or {}).get("metadata"), {}
)
base_dataset_name = str(
existing_group_metadata.get("base_dataset_name")
or payload["dataset_name"]
).strip()
for suffix in ("-训练集", "-验证集", "-测试集"):
if base_dataset_name.endswith(suffix):
base_dataset_name = base_dataset_name[: -len(suffix)].rstrip()
break
split_group_id = str(
existing_group_metadata.get("split_group_id")
or f"dsg_{hashlib.sha256(task_id.encode()).hexdigest()[:20]}"
)
dataset_ids = {
spec["split"]: str(existing_by_split[spec["split"]]["id"])
if spec["split"] in existing_by_split
else new_id("dataset")
for spec in split_specs
}
created_any = any(
spec["split"] not in existing_by_split for spec in split_specs
)
split_labels = {
"train": "训练集",
"validation": "验证集",
"test": "测试集",
}
dataset_types = {"train": "train", "validation": "val", "test": "test"}
published_datasets: list[dict[str, Any]] = []
try:
for spec in split_specs:
split_name = str(spec["split"])
dataset_id = dataset_ids[split_name]
storage_object_id = str(spec["storage_object_id"])
if spec["store_in_minio"]:
file_name = f"{base_dataset_name}.{split_name}.jsonl"
object_key = f"datasets/{dataset_id}/versions/{spec['version_id']}/{file_name}"
uploaded = get_object_storage().put_bytes(
object_key,
spec["raw"],
"application/jsonl",
)
conn.execute(
"""
INSERT INTO storage_objects
(id, resource_type, resource_id, version_id, bucket, object_key,
file_name, content_type, checksum_sha256, byte_size, status,
created_by, create_time)
VALUES (%s, 'dataset', %s, %s, %s, %s, %s, %s, %s, %s,
'available', %s, %s)
ON CONFLICT (resource_type, resource_id, version_id, object_key)
DO UPDATE SET file_name=EXCLUDED.file_name,
content_type=EXCLUDED.content_type,
checksum_sha256=EXCLUDED.checksum_sha256,
byte_size=EXCLUDED.byte_size,
status='available',
created_by=EXCLUDED.created_by
""",
(
new_id("object"),
dataset_id,
spec["version_id"],
uploaded["bucket"],
object_key,
file_name,
"application/jsonl",
spec["checksum"],
len(spec["raw"]),
payload.get("created_by") or task.get("created_by"),
now,
),
)
storage_object_id = conn.execute(
"""
SELECT id FROM storage_objects
WHERE resource_type='dataset' AND resource_id=%s
AND version_id=%s AND object_key=%s
""",
(dataset_id, spec["version_id"], object_key),
).fetchone()["id"]
existing_dataset = existing_by_split.get(split_name)
dataset_metadata = {
**common_metadata,
"base_dataset_name": base_dataset_name,
"dataset_split": split_name,
"split_group_id": split_group_id,
"split_dataset_ids": dataset_ids,
"split_counts": {
name: split_counts[name] if name == split_name else 0
for name in split_order
},
}
dataset_name = f"{base_dataset_name}-{split_labels[split_name]}"
if existing_dataset:
conn.execute(
"DELETE FROM dataset_records WHERE dataset_id=%s", (dataset_id,)
)
conn.execute(
"""DELETE FROM dataset_file_versions
WHERE dataset_file_id IN
(SELECT id FROM dataset_files WHERE dataset_id=%s)""",
(dataset_id,),
)
conn.execute(
"DELETE FROM dataset_files WHERE dataset_id=%s", (dataset_id,)
)
dataset = conn.execute(
"""
UPDATE datasets
SET name=%s, type=%s, storage_type=%s, size=%s, size_bytes=%s,
count=%s, record_count=%s, description=%s, metadata=%s,
updated_at=%s
WHERE id=%s RETURNING *
""",
(
dataset_name,
dataset_types[split_name],
"minio" if spec["store_in_minio"] else "database",
f"{len(spec['raw'])} B",
len(spec["raw"]),
len(spec["records"]),
len(spec["records"]),
payload.get("description") or task.get("description") or "",
json_dumps(dataset_metadata),
now,
dataset_id,
),
).fetchone()
else:
dataset = conn.execute(
"""
INSERT INTO datasets
(id, name, type, storage_type, source, task_id, source_task_id,
size, size_bytes, count, record_count, description, metadata,
tenant_id, project_id, owner_id, created_by, create_time,
created_at, updated_at)
VALUES (
%s, %s, %s, %s, 'task', %s, %s,
%s, %s, %s, %s, %s, %s,
%s, %s, %s, %s, %s, %s, %s
)
RETURNING *
""",
(
dataset_id,
dataset_name,
dataset_types[split_name],
"minio" if spec["store_in_minio"] else "database",
task_id,
task_id,
f"{len(spec['raw'])} B",
len(spec["raw"]),
len(spec["records"]),
len(spec["records"]),
payload.get("description") or task.get("description") or "",
json_dumps(dataset_metadata),
task.get("tenant_id"),
task.get("project_id"),
task.get("owner_id"),
payload.get("created_by") or task.get("created_by"),
now,
now,
now,
),
).fetchone()
file_metadata = {
**dataset_metadata,
"file_split": split_name,
"storage_backend": "minio" if spec["store_in_minio"] else "database",
}
version = {
"id": spec["version_id"],
"version_no": 1,
"version": 1,
"description": f"data process {split_name} publish",
"checksum_sha256": spec["checksum"],
"size_bytes": len(spec["raw"]),
"record_count": len(spec["records"]),
"created_at": now,
"create_time": now,
"source_task_id": task_id,
"storage_object_id": storage_object_id,
}
conn.execute(
"""
INSERT INTO dataset_files
(id, dataset_id, name, storage_object_id, size, content,
active_version_id, versions, create_time, current_version_id,
size_bytes, record_count, file_format, checksum_sha256, version_no,
source_task_id, tenant_id, project_id, created_by, metadata,
created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
%s, %s, 1, %s, %s, %s, %s, %s, %s, %s)
""",
(
spec["file_id"],
dataset_id,
f"{base_dataset_name}.{split_name}.jsonl",
storage_object_id,
f"{len(spec['raw'])} B",
spec["content"],
spec["version_id"],
json_dumps([version]),
now,
spec["version_id"],
len(spec["raw"]),
len(spec["records"]),
"jsonl",
spec["checksum"],
task_id,
task.get("tenant_id"),
task.get("project_id"),
payload.get("created_by") or task.get("created_by"),
json_dumps(file_metadata),
now,
now,
),
)
conn.execute(
"""
INSERT INTO dataset_file_versions
(id, dataset_file_id, version_no, storage_object_id, content_preview,
description, size_bytes, record_count, checksum_sha256,
source_task_id, metadata, created_by, created_at)
VALUES (%s, %s, 1, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
""",
(
spec["version_id"],
spec["file_id"],
storage_object_id,
spec["content"][:2000],
f"data process {split_name} publish",
len(spec["raw"]),
len(spec["records"]),
spec["checksum"],
task_id,
json_dumps(file_metadata),
payload.get("created_by") or task.get("created_by"),
now,
),
)
for line_number, (source_row, record) in enumerate(
spec["records"], start=1
):
conn.execute(
"""
INSERT INTO dataset_records
(id, dataset_id, dataset_file_id, version_id, line_no, split,
instruction, input, output, raw, status, source_task_id,
source_result_id, preview_item_id, created_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
%s, %s, %s, %s)
""",
(
new_id("drec"),
dataset_id,
spec["file_id"],
spec["version_id"],
line_number,
record["split"],
record["instruction"],
record["input"],
record.get("output") or record.get("chosen") or "",
json_dumps(
{
**record,
"source_task_id": task_id,
"source_result_id": source_row["id"],
"preview_item_id": source_row.get("preview_item_id"),
}
),
source_row["status"],
task_id,
source_row["id"],
source_row.get("preview_item_id"),
now,
),
)
published_datasets.append(_decode_row(dataset) or {})
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("dataset name already exists") from exc
train_dataset_id = dataset_ids.get("train")
if not train_dataset_id:
raise InvalidStateError("published split does not contain training data")
conn.execute(
"""
UPDATE data_process_tasks
SET output_dataset_id=%s, updated_at=%s, updated_by=%s
WHERE id=%s
""",
(train_dataset_id, now, payload.get("created_by"), task_id),
)
train_dataset = next(
item
for item in published_datasets
if _json_value(item.get("metadata"), {}).get("dataset_split") == "train"
)
return {
"dataset": train_dataset,
"datasets": published_datasets,
"output_datasets": published_datasets,
"created": created_any,
"split_counts": split_counts,
}
@staticmethod
def _source_ids(
conn: psycopg.Connection[dict[str, Any]], task_id: str
) -> list[dict[str, Any]]:
return conn.execute(
"""
SELECT id FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL ORDER BY created_at, id
""",
(task_id,),
).fetchall()

View File

@@ -0,0 +1,275 @@
"""数据处理存储层 - 生成管理。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
DataProcessStoreError,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
)
class GenerationMixin:
"""生成管理 Mixin。"""
def _invalidate_results(
self,
conn: psycopg.Connection[dict[str, Any]],
task: dict[str, Any],
task_id: str,
now: str,
) -> None:
if _is_regeneration_prepared(task):
conn.execute(
"UPDATE data_process_tasks SET updated_at=%s WHERE id=%s",
(now, task_id),
)
return
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
conn.execute(
"""
UPDATE data_process_tasks
SET status='pending', progress=20, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
generation_run_id=NULL, results_confirmed=FALSE, updated_at=%s
WHERE id=%s
""",
(now, task_id),
)
def start_generation(self, task_id: str, *, replace_existing: bool = True) -> dict[str, Any]:
if not replace_existing:
raise DataProcessStoreError("incremental generation is not supported")
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
regeneration_prepared = _is_regeneration_prepared(task)
if task.get("output_dataset_id") and not regeneration_prepared:
raise InvalidStateError("published task cannot be regenerated")
if task["status"] == "running":
raise ConflictError("data process task is already running")
if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
raise ConflictError("preview is still running")
preview_count = conn.execute(
"SELECT COUNT(*) AS count FROM data_process_preview_items WHERE task_id=%s",
(task_id,),
).fetchone()["count"]
if not preview_count:
raise InvalidStateError("preview must be built before generation")
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
now = utcnow()
generation_run_id = new_id("dprun")
next_config = dict(task.get("config") or {})
next_config.pop(_REGENERATION_MARKER_KEY, None)
row = conn.execute(
"""
UPDATE data_process_tasks
SET config=%s, status='running', progress=30, failure_reason=NULL,
started_at=%s, completed_at=NULL, output_dataset_id=NULL,
output_count=0, filtered_count=0, duplicate_count=0, error_count=0,
generation_run_id=%s, results_confirmed=FALSE,
workflow_step='generate', updated_at=%s
WHERE id=%s
RETURNING *
""",
(json_dumps(next_config), now, generation_run_id, now, task_id),
).fetchone()
return _decode_row(row) or {}
def stop_task(self, task_id: str) -> dict[str, Any]:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if task["status"] != "running":
raise InvalidStateError("only a running task can be stopped")
now = utcnow()
row = conn.execute(
"""
UPDATE data_process_tasks
SET status='stopped', failure_reason=NULL, generation_run_id=NULL,
updated_at=%s
WHERE id=%s RETURNING *
""",
(now, task_id),
).fetchone()
return _decode_row(row) or {}
def generation_is_running(self, task_id: str, generation_run_id: str) -> bool:
with self.connect() as conn:
row = conn.execute(
"""
SELECT status, generation_run_id
FROM data_process_tasks
WHERE id=%s AND deleted_at IS NULL
""",
(task_id,),
).fetchone()
return bool(
row
and row.get("status") == "running"
and row.get("generation_run_id") == generation_run_id
)
def update_generation_progress(
self,
task_id: str,
generation_run_id: str,
processed_count: int,
total_count: int,
) -> bool:
ratio = processed_count / max(1, total_count)
progress = min(95.0, 30.0 + ratio * 65.0)
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET progress=%s, updated_at=%s
WHERE id=%s AND status='running' AND generation_run_id=%s
RETURNING id
""",
(progress, utcnow(), task_id, generation_run_id),
).fetchone()
return row is not None
def complete_generation(
self,
task_id: str,
results: Sequence[dict[str, Any]],
*,
generation_run_id: str,
filtered_count: int = 0,
duplicate_count: int = 0,
error_count: int = 0,
) -> dict[str, Any]:
now = utcnow()
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if (
task["status"] != "running"
or task.get("generation_run_id") != generation_run_id
):
return task
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
for result in results:
conn.execute(
"""
INSERT INTO data_process_results
(id, task_id, preview_item_id, instruction, input, output,
chosen, rejected, original_instruction, original_input,
original_output, original_chosen, original_rejected, status, error,
split, quality_score, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
%s, %s, %s, %s, %s, %s, %s)
""",
(
result.get("id") or new_id("dpr"),
task_id,
result.get("preview_item_id"),
result.get("instruction") or "",
result.get("input") or "",
result.get("output") or "",
result.get("chosen") or "",
result.get("rejected") or "",
result.get("original_instruction", result.get("instruction") or ""),
result.get("original_input", result.get("input") or ""),
result.get("original_output", result.get("output") or ""),
result.get("original_chosen", result.get("chosen") or ""),
result.get("original_rejected", result.get("rejected") or ""),
result.get("status") or "valid",
result.get("error"),
result.get("split"),
json_dumps(result.get("quality_score") or {}),
now,
now,
),
)
row = conn.execute(
"""
UPDATE data_process_tasks
SET status='completed', progress=100, output_count=%s, filtered_count=%s,
duplicate_count=%s, error_count=%s, failure_reason=NULL,
completed_at=%s, generation_run_id=NULL, results_confirmed=FALSE,
updated_at=%s
WHERE id=%s AND generation_run_id=%s RETURNING *
""",
(
len(results),
filtered_count,
duplicate_count,
error_count,
now,
now,
task_id,
generation_run_id,
),
).fetchone()
return _decode_row(row) or {}
def mark_failed(
self, task_id: str, reason: str, *, generation_run_id: str
) -> dict[str, Any]:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if (
task["status"] != "running"
or task.get("generation_run_id") != generation_run_id
):
return task
now = utcnow()
row = conn.execute(
"""
UPDATE data_process_tasks
SET status='failed', failure_reason=%s, completed_at=%s,
generation_run_id=NULL, updated_at=%s
WHERE id=%s AND generation_run_id=%s RETURNING *
""",
(reason[:4000], now, now, task_id, generation_run_id),
).fetchone()
return _decode_row(row) or {}
def progress(self, task_id: str) -> dict[str, Any]:
task = self.get_task(task_id)
return {
"task_id": task["id"],
"status": task["status"],
"progress": float(task.get("progress") or 0),
"input_count": int(task.get("input_count") or 0),
"output_count": int(task.get("output_count") or 0),
"filtered_count": int(task.get("filtered_count") or 0),
"duplicate_count": int(task.get("duplicate_count") or 0),
"error_count": int(task.get("error_count") or 0),
"failure_reason": task.get("failure_reason"),
"results_confirmed": bool(task.get("results_confirmed")),
"started_at": task.get("started_at"),
"completed_at": task.get("completed_at"),
}

View File

@@ -0,0 +1,539 @@
"""数据处理存储层 - 预览管理。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
_serialize_value,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
)
from ..algorithms import estimate_token_count # noqa: E402
class PreviewMixin:
"""预览管理 Mixin。"""
def replace_preview_items(
self,
task_id: str,
items: Sequence[dict[str, Any]],
*,
source_file_ids: Sequence[str] | None = None,
preview_run_id: str | None = None,
) -> list[dict[str, Any]]:
selected_ids = (
list(dict.fromkeys(str(file_id) for file_id in source_file_ids))
if source_file_ids is not None
else None
)
if selected_ids is not None:
if not selected_ids or any(not file_id for file_id in selected_ids):
raise ValueError("source_file_ids must contain non-empty ids")
selected_set = set(selected_ids)
unexpected = {
str(item.get("source_file_id") or "")
for item in items
if str(item.get("source_file_id") or "") not in selected_set
}
if unexpected:
raise ValueError("preview items contain an unselected source file")
preview_file_count = len(selected_ids) if selected_ids is not None else len(
{str(item.get("source_file_id") or "") for item in items}
)
is_direct_build = preview_run_id is None
now = utcnow()
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if is_direct_build:
self._ensure_editable(task)
elif (
task.get("preview_run_id") != preview_run_id
or task.get("preview_status") != "running"
):
raise InvalidStateError("preview run is no longer active")
regeneration_prepared = _is_regeneration_prepared(task)
if not regeneration_prepared:
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
if selected_ids is None:
conn.execute(
"DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,)
)
else:
rows = conn.execute(
"""
SELECT id FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL AND id=ANY(%s)
""",
(task_id, selected_ids),
).fetchall()
found = {str(row["id"]) for row in rows}
missing = set(selected_ids) - found
if missing:
raise NotFoundError(
f"source files not found: {', '.join(sorted(missing))}"
)
conn.execute(
"""
DELETE FROM data_process_preview_items
WHERE task_id=%s AND source_file_id=ANY(%s)
""",
(task_id, selected_ids),
)
created: list[dict[str, Any]] = []
for item in items:
row = conn.execute(
"""
INSERT INTO data_process_preview_items
(id, task_id, source_file_id, original_content, edited_content,
source_start, source_end, source_start_line, source_end_line,
token_count, status, quality_score, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
RETURNING *
""",
(
item.get("id") or new_id("dpp"),
task_id,
item.get("source_file_id"),
item.get("original_content") or "",
item.get("edited_content", item.get("original_content") or ""),
item.get("source_start"),
item.get("source_end"),
item.get("source_start_line"),
item.get("source_end_line"),
max(0, int(item.get("token_count") or 0)),
item.get("status") or "original",
json_dumps(item.get("quality_score") or {}),
now,
now,
),
).fetchone()
created.append(_decode_row(row) or {})
if regeneration_prepared:
if is_direct_build:
conn.execute(
"""
UPDATE data_process_tasks
SET workflow_step='preview', preview_status='completed',
preview_progress=100, preview_run_id=NULL,
preview_failure_reason=NULL, preview_total_files=%s,
preview_completed_files=%s, updated_at=%s
WHERE id=%s
""",
(preview_file_count, preview_file_count, now, task_id),
)
else:
conn.execute(
"UPDATE data_process_tasks SET updated_at=%s WHERE id=%s",
(now, task_id),
)
else:
conn.execute(
"""
UPDATE data_process_tasks
SET status='pending', progress=20, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
results_confirmed=FALSE,
workflow_step=CASE WHEN %s THEN 'preview' ELSE workflow_step END,
preview_status=CASE WHEN %s THEN 'completed' ELSE preview_status END,
preview_progress=CASE WHEN %s THEN 100 ELSE preview_progress END,
preview_run_id=CASE WHEN %s THEN NULL ELSE preview_run_id END,
preview_failure_reason=CASE WHEN %s THEN NULL ELSE preview_failure_reason END,
preview_total_files=CASE WHEN %s THEN %s ELSE preview_total_files END,
preview_completed_files=CASE WHEN %s THEN %s ELSE preview_completed_files END,
updated_at=%s
WHERE id=%s
""",
(
is_direct_build,
is_direct_build,
is_direct_build,
is_direct_build,
is_direct_build,
is_direct_build,
preview_file_count,
is_direct_build,
preview_file_count,
now,
task_id,
),
)
return created
def start_preview(
self,
task_id: str,
*,
source_file_ids: Sequence[str] | None = None,
) -> tuple[dict[str, Any], list[str]]:
"""创建一轮持久化切分任务,并返回本轮固定的源文件集合。"""
requested_ids = (
list(dict.fromkeys(str(file_id) for file_id in source_file_ids))
if source_file_ids is not None
else None
)
if requested_ids is not None and (
not requested_ids or any(not file_id for file_id in requested_ids)
):
raise ValueError("source_file_ids must contain non-empty ids")
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
if requested_ids is None:
rows = conn.execute(
"""
SELECT id FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
ORDER BY created_at, id
""",
(task_id,),
).fetchall()
else:
rows = conn.execute(
"""
SELECT id FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL AND id=ANY(%s)
ORDER BY created_at, id
""",
(task_id, requested_ids),
).fetchall()
selected_ids = [str(row["id"]) for row in rows]
if not selected_ids:
raise InvalidStateError("at least one source file is required")
if requested_ids is not None:
missing = set(requested_ids) - set(selected_ids)
if missing:
raise NotFoundError(
f"source files not found: {', '.join(sorted(missing))}"
)
regeneration_prepared = _is_regeneration_prepared(task)
if not regeneration_prepared:
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
preview_run_id = new_id("dpprun")
now = utcnow()
row = conn.execute(
"""
UPDATE data_process_tasks
SET status=CASE WHEN %s THEN status ELSE 'pending' END,
progress=CASE WHEN %s THEN progress ELSE 0 END,
output_count=CASE WHEN %s THEN output_count ELSE 0 END,
filtered_count=CASE WHEN %s THEN filtered_count ELSE 0 END,
duplicate_count=CASE WHEN %s THEN duplicate_count ELSE 0 END,
error_count=CASE WHEN %s THEN error_count ELSE 0 END,
failure_reason=CASE WHEN %s THEN failure_reason ELSE NULL END,
results_confirmed=CASE WHEN %s THEN results_confirmed ELSE FALSE END,
workflow_step='upload', preview_status='queued', preview_progress=0,
preview_run_id=%s, preview_failure_reason=NULL,
preview_total_files=%s, preview_completed_files=0,
updated_at=%s
WHERE id=%s AND deleted_at IS NULL
RETURNING *
""",
(
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
regeneration_prepared,
preview_run_id,
len(selected_ids),
now,
task_id,
),
).fetchone()
return _public_task(_decode_row(row)) or {}, selected_ids
def mark_preview_running(self, task_id: str, preview_run_id: str) -> bool:
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET preview_status='running', updated_at=%s
WHERE id=%s AND deleted_at IS NULL
AND preview_status='queued' AND preview_run_id=%s
RETURNING id
""",
(utcnow(), task_id, preview_run_id),
).fetchone()
return row is not None
def preview_is_running(self, task_id: str, preview_run_id: str) -> bool:
with self.connect() as conn:
row = conn.execute(
"""
SELECT preview_status, preview_run_id
FROM data_process_tasks
WHERE id=%s AND deleted_at IS NULL
""",
(task_id,),
).fetchone()
return bool(
row
and row.get("preview_status") in ACTIVE_PREVIEW_STATUSES
and row.get("preview_run_id") == preview_run_id
)
def update_preview_progress(
self,
task_id: str,
preview_run_id: str,
completed_files: int,
total_files: int,
) -> bool:
total = max(1, total_files)
completed = min(max(0, completed_files), total)
progress = completed / total * 100
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET preview_progress=%s, preview_completed_files=%s, updated_at=%s
WHERE id=%s AND deleted_at IS NULL
AND preview_status='running' AND preview_run_id=%s
RETURNING id
""",
(progress, completed, utcnow(), task_id, preview_run_id),
).fetchone()
return row is not None
def complete_preview(self, task_id: str, preview_run_id: str) -> bool:
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET workflow_step='preview', preview_status='completed',
preview_progress=100, preview_run_id=NULL,
preview_failure_reason=NULL,
preview_completed_files=preview_total_files, updated_at=%s
WHERE id=%s AND deleted_at IS NULL
AND preview_status='running' AND preview_run_id=%s
RETURNING id
""",
(utcnow(), task_id, preview_run_id),
).fetchone()
return row is not None
def mark_preview_failed(
self,
task_id: str,
reason: str,
*,
preview_run_id: str,
) -> bool:
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET preview_status='failed', preview_run_id=NULL,
preview_failure_reason=%s, updated_at=%s
WHERE id=%s AND deleted_at IS NULL
AND preview_status IN ('queued', 'running') AND preview_run_id=%s
RETURNING id
""",
(reason[:4000], utcnow(), task_id, preview_run_id),
).fetchone()
return row is not None
def preview_progress(self, task_id: str) -> dict[str, Any]:
task = self.get_task(task_id)
return {
"task_id": task["id"],
"workflow_step": task.get("workflow_step") or "create",
"preview_status": task.get("preview_status") or "idle",
"preview_progress": float(task.get("preview_progress") or 0),
"preview_run_id": task.get("preview_run_id"),
"preview_failure_reason": task.get("preview_failure_reason"),
"preview_total_files": int(task.get("preview_total_files") or 0),
"preview_completed_files": int(task.get("preview_completed_files") or 0),
}
def list_preview_items(
self,
task_id: str,
*,
source_file_id: str | None = None,
page: int = 1,
page_size: int = 200,
keyword: str | None = None,
) -> dict[str, Any]:
self.get_task(task_id)
clauses = ["task_id=%s"]
params: list[Any] = [task_id]
if source_file_id:
clauses.append("source_file_id=%s")
params.append(source_file_id)
if keyword:
clauses.append("(original_content ILIKE %s OR edited_content ILIKE %s)")
pattern = f"%{keyword.strip()}%"
params.extend([pattern, pattern])
where = " AND ".join(clauses)
with self.connect() as conn:
total = conn.execute(
f"SELECT COUNT(*) AS count FROM data_process_preview_items WHERE {where}", params
).fetchone()["count"]
rows = conn.execute(
f"""
SELECT * FROM data_process_preview_items
WHERE {where}
ORDER BY source_file_id NULLS LAST, source_start NULLS LAST, created_at, id
LIMIT %s OFFSET %s
""",
[*params, page_size, (page - 1) * page_size],
).fetchall()
return {
"items": [_decode_row(row) for row in rows],
"total": int(total),
"page": page,
"page_size": page_size,
}
def get_preview_item(self, task_id: str, preview_id: str) -> dict[str, Any]:
with self.connect() as conn:
row = conn.execute(
"SELECT * FROM data_process_preview_items WHERE id=%s AND task_id=%s",
(preview_id, task_id),
).fetchone()
if not row:
raise NotFoundError("preview item not found")
return _decode_row(row) or {}
def create_preview_item(self, task_id: str, item: dict[str, Any]) -> dict[str, Any]:
now = utcnow()
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
if item.get("source_file_id"):
source = conn.execute(
"""
SELECT id FROM data_process_source_files
WHERE id=%s AND task_id=%s AND deleted_at IS NULL
""",
(item["source_file_id"], task_id),
).fetchone()
if not source:
raise NotFoundError("source file not found")
row = conn.execute(
"""
INSERT INTO data_process_preview_items
(id, task_id, source_file_id, original_content, edited_content,
source_start, source_end, source_start_line, source_end_line,
token_count, status, quality_score, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
RETURNING *
""",
(
new_id("dpp"),
task_id,
item.get("source_file_id"),
item.get("original_content") or "",
item.get("edited_content") or "",
item.get("source_start"),
item.get("source_end"),
item.get("source_start_line"),
item.get("source_end_line"),
max(0, int(item.get("token_count") or 0)),
item.get("status") or "manual",
json_dumps(item.get("quality_score") or {}),
now,
now,
),
).fetchone()
self._invalidate_results(conn, task, task_id, now)
return _decode_row(row) or {}
def update_preview_item(
self, task_id: str, preview_id: str, payload: dict[str, Any]
) -> dict[str, Any]:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
existing = conn.execute(
"SELECT * FROM data_process_preview_items WHERE id=%s AND task_id=%s",
(preview_id, task_id),
).fetchone()
if not existing:
raise NotFoundError("preview item not found")
expected_updated_at = payload.get("expected_updated_at")
current_updated_at = _serialize_value(existing.get("updated_at"))
if expected_updated_at and expected_updated_at != current_updated_at:
raise ConflictError("preview item was modified by another request")
edited = payload["edited_content"]
status = payload.get("status")
if not status:
if not edited.strip():
status = "invalid"
elif edited == existing["original_content"]:
status = "original"
else:
status = "modified"
now = utcnow()
row = conn.execute(
"""
UPDATE data_process_preview_items
SET edited_content=%s, token_count=%s, status=%s, quality_score=%s,
updated_at=%s
WHERE id=%s AND task_id=%s
RETURNING *
""",
(
edited,
estimate_token_count(edited),
status,
json_dumps(payload.get("quality_score") or {}),
now,
preview_id,
task_id,
),
).fetchone()
self._invalidate_results(conn, task, task_id, now)
return _decode_row(row) or {}
def delete_preview_item(self, task_id: str, preview_id: str) -> None:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
row = conn.execute(
"DELETE FROM data_process_preview_items WHERE id=%s AND task_id=%s RETURNING id",
(preview_id, task_id),
).fetchone()
if not row:
raise NotFoundError("preview item not found")
self._invalidate_results(conn, task, task_id, utcnow())

View File

@@ -0,0 +1,324 @@
"""数据处理存储层 - 结果管理。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
_serialize_value,
DataProcessStoreError,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
)
class ResultsMixin:
"""结果管理 Mixin。"""
def confirm_results(self, task_id: str) -> dict[str, Any]:
"""确认第六步结果,确认前再次校验所有生成记录。"""
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if task["status"] != "completed":
raise InvalidStateError("only a completed task can confirm results")
if task.get("workflow_step") != "results":
raise InvalidStateError("workflow must be on results before confirmation")
if task.get("results_confirmed"):
return task
rows = conn.execute(
"""
SELECT status, instruction, output, chosen, rejected
FROM data_process_results
WHERE task_id=%s
""",
(task_id,),
).fetchall()
if not rows:
raise InvalidStateError("task has no results to confirm")
invalid_count = sum(
1
for row in rows
if row["status"] == "invalid"
or not str(row.get("instruction") or "").strip()
or not str(row.get("output") or "").strip()
or (
_task_output_type(task) == "reasoning"
and not _reasoning_output_is_valid(row.get("output"))
)
or (
_task_output_type(task) == "dpo"
and not _dpo_fields_are_valid(row)
)
)
if invalid_count:
raise InvalidStateError(
f"task contains {invalid_count} invalid results"
)
row = conn.execute(
"""
UPDATE data_process_tasks
SET results_confirmed=TRUE, updated_at=%s
WHERE id=%s RETURNING *
""",
(utcnow(), task_id),
).fetchone()
return _decode_row(row) or {}
def list_results(
self,
task_id: str,
*,
page: int = 1,
page_size: int = 100,
status: str | None = None,
split: str | None = None,
keyword: str | None = None,
) -> dict[str, Any]:
self.get_task(task_id)
clauses = ["task_id=%s"]
params: list[Any] = [task_id]
if status:
clauses.append("status=%s")
params.append(status)
if split:
clauses.append("split=%s")
params.append(split)
if keyword:
clauses.append("(instruction ILIKE %s OR input ILIKE %s OR output ILIKE %s)")
pattern = f"%{keyword.strip()}%"
params.extend([pattern, pattern, pattern])
where = " AND ".join(clauses)
with self.connect() as conn:
total = conn.execute(
f"SELECT COUNT(*) AS count FROM data_process_results WHERE {where}", params
).fetchone()["count"]
rows = conn.execute(
f"""
SELECT * FROM data_process_results WHERE {where}
ORDER BY created_at, id LIMIT %s OFFSET %s
""",
[*params, page_size, (page - 1) * page_size],
).fetchall()
return {
"items": [_decode_row(row) for row in rows],
"total": int(total),
"page": page,
"page_size": page_size,
}
def get_result(self, task_id: str, result_id: str) -> dict[str, Any]:
with self.connect() as conn:
row = conn.execute(
"SELECT * FROM data_process_results WHERE id=%s AND task_id=%s",
(result_id, task_id),
).fetchone()
if not row:
raise NotFoundError("data process result not found")
return _decode_row(row) or {}
def update_result(
self, task_id: str, result_id: str, payload: dict[str, Any]
) -> dict[str, Any]:
allowed = {
"instruction", "input", "output", "chosen", "rejected", "quality_score"
}
values = {key: value for key, value in payload.items() if key in allowed}
if "quality_score" in values:
values["quality_score"] = json_dumps(values["quality_score"])
if not values:
raise DataProcessStoreError("no result fields supplied")
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if task["status"] == "running":
raise InvalidStateError("results cannot be edited while generation is running")
if task.get("output_dataset_id"):
raise InvalidStateError("published results cannot be edited")
current = conn.execute(
"SELECT * FROM data_process_results WHERE id=%s AND task_id=%s",
(result_id, task_id),
).fetchone()
if not current:
raise NotFoundError("data process result not found")
expected_updated_at = payload.get("expected_updated_at")
current_updated_at = _serialize_value(current.get("updated_at"))
if expected_updated_at and expected_updated_at != current_updated_at:
raise ConflictError("data process result was modified by another request")
output_type = _task_output_type(task)
if output_type == "dpo" and "chosen" in values:
values["output"] = values["chosen"]
merged = {**current, **values}
quality = payload.get("quality_score") or {}
instruction_valid = bool(str(merged.get("instruction") or "").strip())
output_valid = bool(str(merged.get("output") or "").strip())
reasoning_valid = (
output_type != "reasoning"
or _reasoning_output_is_valid(merged.get("output"))
)
dpo_valid = output_type != "dpo" or _dpo_fields_are_valid(merged)
hard_valid = instruction_valid and output_valid and reasoning_valid and dpo_valid
quality_valid = bool(quality.get("is_valid", hard_valid))
changed = any(
str(merged.get(field) or "")
!= str(merged.get(f"original_{field}") or "")
for field in (
("instruction", "input", "chosen", "rejected")
if output_type == "dpo"
else ("instruction", "input", "output")
)
)
status = "invalid" if not hard_valid or not quality_valid else (
"modified" if changed else "valid"
)
values["status"] = status
flags = quality.get("flags") if isinstance(quality, dict) else None
format_error = (
"思维链输出必须包含非空的 <think>...</think> 推理过程和最终答案"
if instruction_valid and output_valid and not reasoning_valid
else "DPO 输出必须包含不同的非空 Chosen 和 Rejected 回答"
if instruction_valid and not dpo_valid
else "Instruction 和 Output 不能为空"
if not instruction_valid or not output_valid
else None
)
values["error"] = ", ".join(str(flag) for flag in flags or []) or (
format_error
or ("quality validation failed" if status == "invalid" else None)
)
values["updated_at"] = utcnow()
assignments = ", ".join(f"{key}=%s" for key in values)
row = conn.execute(
f"""UPDATE data_process_results SET {assignments}
WHERE id=%s AND task_id=%s RETURNING *""",
[*values.values(), result_id, task_id],
).fetchone()
conn.execute(
"""
UPDATE data_process_tasks
SET error_count=(
SELECT COUNT(*) FROM data_process_results
WHERE task_id=%s AND status='invalid'
), updated_at=%s
WHERE id=%s
""",
(task_id, utcnow(), task_id),
)
return _decode_row(row) or {}
def replace_generated_result(
self,
task_id: str,
result_id: str,
replacement: dict[str, Any],
*,
expected_updated_at: str,
) -> dict[str, Any]:
"""用新模型结果原位替换失败项,并将新内容设为恢复基线。"""
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if task["status"] != "completed" or task.get("workflow_step") != "results":
raise InvalidStateError("task is not editing generation results")
if task.get("results_confirmed"):
raise InvalidStateError("confirmed results cannot be regenerated")
if task.get("output_dataset_id"):
raise InvalidStateError("published results cannot be regenerated")
current = conn.execute(
"""SELECT * FROM data_process_results
WHERE id=%s AND task_id=%s FOR UPDATE""",
(result_id, task_id),
).fetchone()
if not current:
raise NotFoundError("data process result not found")
if current.get("status") != "invalid":
raise InvalidStateError("only an invalid result can be regenerated")
current_updated_at = _serialize_value(current.get("updated_at"))
if expected_updated_at != current_updated_at:
raise ConflictError("data process result was modified by another request")
instruction = str(replacement.get("instruction") or "").strip()
input_text = str(replacement.get("input") or "").strip()
output = str(replacement.get("output") or "").strip()
chosen = str(replacement.get("chosen") or "").strip()
rejected = str(replacement.get("rejected") or "").strip()
quality_score = replacement.get("quality_score") or {}
if not instruction or not output or not bool(quality_score.get("is_valid")):
raise InvalidStateError("regenerated result did not pass quality validation")
if _task_output_type(task) == "reasoning" and not _reasoning_output_is_valid(output):
raise InvalidStateError("regenerated reasoning result has an invalid output format")
if _task_output_type(task) == "dpo" and not _dpo_fields_are_valid(replacement):
raise InvalidStateError("regenerated DPO result has invalid preference fields")
now = utcnow()
row = conn.execute(
"""
UPDATE data_process_results
SET instruction=%s, input=%s, output=%s, chosen=%s, rejected=%s,
original_instruction=%s, original_input=%s, original_output=%s,
original_chosen=%s, original_rejected=%s,
status='valid', error=NULL, quality_score=%s, updated_at=%s
WHERE id=%s AND task_id=%s
RETURNING *
""",
(
instruction,
input_text,
output,
chosen,
rejected,
instruction,
input_text,
output,
chosen,
rejected,
json_dumps(quality_score),
now,
result_id,
task_id,
),
).fetchone()
conn.execute(
"""
UPDATE data_process_tasks
SET error_count=(
SELECT COUNT(*) FROM data_process_results
WHERE task_id=%s AND status='invalid'
), updated_at=%s
WHERE id=%s
""",
(task_id, now, task_id),
)
return _decode_row(row) or {}

View File

@@ -0,0 +1,331 @@
"""数据处理存储层 - 源文件管理。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
DataProcessStoreError,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
)
class SourceFilesMixin:
"""源文件管理 Mixin。"""
def list_source_files(self, task_id: str) -> list[dict[str, Any]]:
self.get_task(task_id)
with self.connect() as conn:
rows = conn.execute(
"""
SELECT id, task_id, storage_object_id, name, size_bytes, record_count,
file_format, checksum_sha256, version_no, content_preview, metadata,
tenant_id, project_id, created_by, created_at, updated_at
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
ORDER BY created_at, id
""",
(task_id,),
).fetchall()
return [_decode_row(row) or {} for row in rows]
def add_source_file(
self,
task_id: str,
*,
name: str,
content: str,
raw_size: int,
checksum_sha256: str,
file_format: str,
record_count: int,
metadata: dict[str, Any] | None = None,
created_by: str | None = None,
source_file_id: str | None = None,
storage_object_id: str | None = None,
) -> dict[str, Any]:
return self.add_source_files(
task_id,
[
{
"name": name,
"content": content,
"raw_size": raw_size,
"checksum_sha256": checksum_sha256,
"file_format": file_format,
"record_count": record_count,
"metadata": metadata or {},
"created_by": created_by,
"id": source_file_id,
"storage_object_id": storage_object_id,
}
],
)[0]
def add_source_files(
self,
task_id: str,
files: Sequence[dict[str, Any]],
) -> list[dict[str, Any]]:
"""在同一事务中登记一个上传批次,任一文件失败则全部回滚。"""
if not files:
raise DataProcessStoreError("at least one source file is required")
now = utcnow()
created: list[dict[str, Any]] = []
try:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
for payload in files:
file_id = str(payload.get("id") or new_id("dpsf"))
storage_object_id, metadata_payload = _source_storage_descriptor(
payload,
task_id,
file_id,
)
row = conn.execute(
"""
INSERT INTO data_process_source_files
(id, task_id, storage_object_id, name, size_bytes, record_count,
file_format, checksum_sha256, version_no, content, content_preview,
metadata, tenant_id, project_id,
created_by, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 1, %s, %s, %s, %s,
%s, %s, %s, %s)
RETURNING id, task_id, storage_object_id, name, size_bytes, record_count,
file_format, checksum_sha256, version_no, content_preview, metadata,
tenant_id, project_id, created_by, created_at, updated_at
""",
(
file_id,
task_id,
storage_object_id,
payload["name"],
payload["raw_size"],
payload["record_count"],
payload["file_format"],
payload["checksum_sha256"],
"" if str(storage_object_id or "").startswith("minio://") else payload["content"],
str(payload["content"])[:2000],
json_dumps(metadata_payload),
task.get("tenant_id"),
task.get("project_id"),
payload.get("created_by") or task.get("created_by"),
now,
now,
),
).fetchone()
created.append(_decode_row(row) or {})
preview_row = conn.execute(
"""
SELECT COUNT(*) AS count FROM data_process_preview_items
WHERE task_id=%s
""",
(task_id,),
).fetchone()
preview_count = int((preview_row or {}).get("count") or 0)
if _is_regeneration_prepared(task):
conn.execute(
"""
UPDATE data_process_tasks
SET input_count=(
SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
), workflow_step='upload', preview_status='idle',
preview_progress=0, preview_run_id=NULL,
preview_failure_reason=NULL, preview_total_files=0,
preview_completed_files=0, updated_at=%s
WHERE id=%s
""",
(task_id, now, task_id),
)
else:
# 前端会只对本次新增的源文件构建预览,因此保留旧文件切片,
# 但普通未发布任务的旧生成结果已经不再有效。
conn.execute(
"DELETE FROM data_process_results WHERE task_id=%s", (task_id,)
)
conn.execute(
"""
UPDATE data_process_tasks
SET status='pending', progress=%s, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
generation_run_id=NULL, results_confirmed=FALSE,
workflow_step='upload', preview_status='idle',
preview_progress=0, preview_run_id=NULL,
preview_failure_reason=NULL, preview_total_files=0,
preview_completed_files=0,
started_at=NULL, completed_at=NULL,
input_count=(
SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
), updated_at=%s
WHERE id=%s
""",
(20 if preview_count else 0, task_id, now, task_id),
)
except psycopg.errors.UniqueViolation as exc:
raise ConflictError(
"the same source file content is already attached to this task"
) from exc
return created
def get_source_file(
self, task_id: str, file_id: str, *, include_content: bool = True
) -> dict[str, Any]:
# 先验证父任务仍然可见,避免软删除任务后通过已知文件 ID 读取正文。
self.get_task(task_id)
content_column = ", content" if include_content else ""
with self.connect() as conn:
row = conn.execute(
f"""
SELECT id, task_id, storage_object_id, name, size_bytes, record_count,
file_format, checksum_sha256, version_no, content_preview, metadata,
tenant_id, project_id, created_by, created_at, updated_at{content_column}
FROM data_process_source_files
WHERE id=%s AND task_id=%s AND deleted_at IS NULL
""",
(file_id, task_id),
).fetchone()
if not row:
raise NotFoundError("source file not found")
decoded = _decode_row(row) or {}
# New source files keep only a preview in PostgreSQL. Load the
# authoritative body from MinIO on demand for existing processing code.
reference = str(decoded.get("storage_object_id") or "")
if include_content and not decoded.get("content") and reference.startswith("minio://"):
from app.modules.data_process.storage import get_data_process_storage
decoded["content"] = (get_data_process_storage().read(reference) or b"").decode(
"utf-8", errors="replace"
)
return decoded
def source_content_window(
self, task_id: str, file_id: str, offset: int, limit: int
) -> dict[str, Any]:
source_file = self.get_source_file(task_id, file_id, include_content=True)
content = str(source_file.pop("content", ""))
window = content[offset : offset + limit]
return {
"file": source_file,
"content": window,
"offset": offset,
"limit": limit,
"total_chars": len(content),
"has_more": offset + len(window) < len(content),
}
def source_content_lines(
self,
task_id: str,
file_id: str,
start_line: int,
line_count: int,
) -> dict[str, Any]:
source_file = self.get_source_file(task_id, file_id, include_content=True)
content = str(source_file.pop("content", ""))
lines = content.splitlines(keepends=True)
start_index = min(len(lines), start_line - 1)
selected = lines[start_index : start_index + line_count]
end_line = start_index + len(selected)
return {
"file": source_file,
"content": "".join(selected),
"start_line": start_line,
"end_line": end_line,
"line_count": len(selected),
"total_lines": len(lines),
"has_more": end_line < len(lines),
}
def delete_source_file(self, task_id: str, file_id: str) -> None:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
row = conn.execute(
"""
UPDATE data_process_source_files
SET deleted_at=%s, updated_at=%s
WHERE id=%s AND task_id=%s AND deleted_at IS NULL
RETURNING id
""",
(utcnow(), utcnow(), file_id, task_id),
).fetchone()
if not row:
raise NotFoundError("source file not found")
conn.execute(
"DELETE FROM data_process_preview_items WHERE source_file_id=%s", (file_id,)
)
now = utcnow()
if _is_regeneration_prepared(task):
conn.execute(
"""
UPDATE data_process_tasks
SET input_count=(SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL),
workflow_step='upload', preview_status='idle',
preview_progress=0, preview_run_id=NULL,
preview_failure_reason=NULL, preview_total_files=0,
preview_completed_files=0, updated_at=%s
WHERE id=%s
""",
(task_id, now, task_id),
)
else:
conn.execute(
"DELETE FROM data_process_results WHERE task_id=%s", (task_id,)
)
conn.execute(
"""
UPDATE data_process_tasks
SET status='pending', progress=0, output_count=0, filtered_count=0,
duplicate_count=0, error_count=0, failure_reason=NULL,
results_confirmed=FALSE,
workflow_step='upload', preview_status='idle',
preview_progress=0, preview_run_id=NULL,
preview_failure_reason=NULL, preview_total_files=0,
preview_completed_files=0,
input_count=(SELECT COALESCE(SUM(record_count), 0)
FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL),
updated_at=%s
WHERE id=%s
""",
(task_id, now, task_id),
)

View File

@@ -0,0 +1,879 @@
"""数据处理存储层 - 任务管理。"""
from __future__ import annotations
from typing import Any
from collections.abc import Sequence
import psycopg
from .base import (
StoreBase,
utcnow,
new_id,
repeat_task_id,
json_dumps,
_json_value,
_decode_row,
_public_task,
_business_config,
_preview_config_value,
_preview_config_changed,
_preview_config_projection,
_normalized_preprocess_options,
_regeneration_marker,
_is_regeneration_prepared,
_task_output_type,
_task_reasoning_detail,
_reasoning_output_is_valid,
_dpo_fields_are_valid,
_source_storage_descriptor,
_serialize_value,
NotFoundError,
ConflictError,
InvalidStateError,
EDITABLE_STATUSES,
ACTIVE_PREVIEW_STATUSES,
WORKFLOW_STEPS,
_REGENERATION_MARKER_KEY,
_REPEAT_SOURCE_TASK_KEY,
_REPEAT_REQUEST_KEY,
_INTERNAL_CONFIG_KEYS,
)
from ..algorithms import estimate_token_count
class TasksMixin:
"""任务管理 Mixin。"""
def list_tasks(
self,
*,
page: int = 1,
page_size: int = 20,
keyword: str | None = None,
status: str | None = None,
process_type: str | None = None,
tenant_id: str | None = None,
project_id: str | None = None,
) -> dict[str, Any]:
clauses = ["task.deleted_at IS NULL"]
params: list[Any] = []
if keyword:
clauses.append("(task.name ILIKE %s OR COALESCE(task.description, '') ILIKE %s)")
pattern = f"%{keyword.strip()}%"
params.extend([pattern, pattern])
if status:
clauses.append("task.status = %s")
params.append(status)
if process_type:
clauses.append("task.process_type = %s")
params.append(process_type)
if tenant_id:
clauses.append("task.tenant_id = %s")
params.append(tenant_id)
if project_id:
clauses.append("task.project_id = %s")
params.append(project_id)
where = " AND ".join(clauses)
with self.connect() as conn:
total = conn.execute(
f"SELECT COUNT(*) AS count FROM data_process_tasks task WHERE {where}",
params,
).fetchone()["count"]
rows = conn.execute(
f"""
SELECT task.*,
creator.display_name AS creator_name,
processor.display_name AS processor_name,
(SELECT COUNT(*) FROM data_process_source_files source_file
WHERE source_file.task_id=task.id
AND source_file.deleted_at IS NULL) AS source_file_count
FROM data_process_tasks task
LEFT JOIN users creator ON creator.id=task.created_by
LEFT JOIN users processor ON processor.id=task.updated_by
WHERE {where}
ORDER BY task.created_at DESC, task.id DESC
LIMIT %s OFFSET %s
""",
[*params, page_size, (page - 1) * page_size],
).fetchall()
return {
"items": [_public_task(_decode_row(row)) for row in rows],
"total": int(total),
"page": page,
"page_size": page_size,
}
def create_task(self, payload: dict[str, Any]) -> dict[str, Any]:
task_id = new_id("dpt")
now = utcnow()
try:
with self.connect() as conn:
row = conn.execute(
"""
INSERT INTO data_process_tasks
(id, name, description, status, process_type, source_dataset_id, config,
progress, results_confirmed, tenant_id, project_id, owner_id, created_by, updated_by,
created_at, updated_at)
VALUES (%s, %s, %s, 'pending', %s, %s, %s, 0, FALSE,
%s, %s, %s, %s, %s, %s, %s)
RETURNING *
""",
(
task_id,
payload["name"],
payload.get("description") or "",
payload["process_type"],
payload.get("source_dataset_id"),
json_dumps(_business_config(payload.get("config"))),
payload.get("tenant_id"),
payload.get("project_id"),
payload.get("owner_id"),
payload.get("created_by"),
payload.get("created_by"),
now,
now,
),
).fetchone()
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("data process task name already exists") from exc
return _public_task(_decode_row(row)) or {}
@staticmethod
def _repeat_response(
conn: psycopg.Connection[dict[str, Any]],
row: dict[str, Any],
*,
source_task_id: str,
created: bool,
) -> dict[str, Any]:
task_id = str(row["id"])
counts = conn.execute(
"""
SELECT
(SELECT COUNT(*) FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL) AS source_file_count,
(SELECT COUNT(*) FROM data_process_preview_items
WHERE task_id=%s) AS preview_count
""",
(task_id, task_id),
).fetchone() or {}
task = _public_task(_decode_row(row)) or {}
task["source_file_count"] = int(counts.get("source_file_count") or 0)
task["preview_count"] = int(counts.get("preview_count") or 0)
return {
"task": task,
"source_task_id": source_task_id,
"created": created,
"copied_source_file_count": task["source_file_count"],
"copied_preview_count": task["preview_count"],
}
def find_repeated_task(
self,
source_task_id: str,
request_id: str,
) -> dict[str, Any] | None:
"""查找同一幂等请求已创建的新任务。"""
task_id = repeat_task_id(source_task_id, request_id)
with self.connect() as conn:
row = conn.execute(
"SELECT * FROM data_process_tasks WHERE id=%s",
(task_id,),
).fetchone()
if row is None:
return None
decoded = _decode_row(row) or {}
config = decoded.get("config") or {}
if (
config.get(_REPEAT_SOURCE_TASK_KEY) != source_task_id
or config.get(_REPEAT_REQUEST_KEY) != request_id
):
raise ConflictError("再次生成请求与现有任务冲突")
if decoded.get("deleted_at"):
raise ConflictError("此次再次生成创建的任务已被删除,请重新发起")
return self._repeat_response(
conn,
row,
source_task_id=source_task_id,
created=False,
)
def repeat_task(
self,
source_task_id: str,
*,
expected_updated_at: str,
request_id: str,
file_copies: dict[str, dict[str, str]],
) -> dict[str, Any]:
"""复制已确认任务的配置、源文件和预览,结果与发布数据保持独立。"""
task_id = repeat_task_id(source_task_id, request_id)
now = utcnow()
try:
with self.connect() as conn:
existing = conn.execute(
"SELECT * FROM data_process_tasks WHERE id=%s FOR UPDATE",
(task_id,),
).fetchone()
if existing is not None:
decoded = _decode_row(existing) or {}
config = decoded.get("config") or {}
if (
config.get(_REPEAT_SOURCE_TASK_KEY) != source_task_id
or config.get(_REPEAT_REQUEST_KEY) != request_id
):
raise ConflictError("再次生成请求与现有任务冲突")
if decoded.get("deleted_at"):
raise ConflictError("此次再次生成创建的任务已被删除,请重新发起")
return self._repeat_response(
conn,
existing,
source_task_id=source_task_id,
created=False,
)
source_task = self._task_in_connection(
conn,
source_task_id,
for_update=True,
)
if (
source_task.get("status") != "completed"
or source_task.get("results_confirmed") is False
):
raise InvalidStateError("只有已完成并确认结果的任务可以再次生成")
if source_task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
raise ConflictError("源任务仍在处理切分,暂时不能再次生成")
if expected_updated_at != _serialize_value(source_task.get("updated_at")):
raise ConflictError("源任务已被其他操作修改,请刷新后重试")
source_files = conn.execute(
"""
SELECT * FROM data_process_source_files
WHERE task_id=%s AND deleted_at IS NULL
ORDER BY created_at, id
""",
(source_task_id,),
).fetchall()
source_file_ids = {str(row["id"]) for row in source_files}
if source_file_ids != set(file_copies):
raise ConflictError("源文件快照已变化,请刷新后重试")
previews = conn.execute(
"""
SELECT * FROM data_process_preview_items
WHERE task_id=%s
ORDER BY source_file_id NULLS LAST, source_start NULLS LAST,
created_at, id
""",
(source_task_id,),
).fetchall()
if not previews:
raise InvalidStateError("源任务没有可用于再次生成的切分结果")
suffix = f"(再次生成-{task_id[-6:]}"
base_name = str(source_task.get("name") or "数据处理任务")
repeated_name = f"{base_name[: max(1, 150 - len(suffix))]}{suffix}"
repeated_config = _business_config(source_task.get("config") or {})
repeated_config[_REPEAT_SOURCE_TASK_KEY] = source_task_id
repeated_config[_REPEAT_REQUEST_KEY] = request_id
input_count = sum(int(row.get("record_count") or 0) for row in source_files)
task_row = conn.execute(
"""
INSERT INTO data_process_tasks
(id, name, description, status, process_type, source_dataset_id,
output_dataset_id, config, progress, input_count, output_count,
filtered_count, duplicate_count, error_count, failure_reason,
generation_run_id, results_confirmed, workflow_step,
preview_status, preview_progress, preview_run_id,
preview_failure_reason, preview_total_files,
preview_completed_files, tenant_id, project_id, owner_id,
approval_status, created_by, updated_by, created_at, updated_at)
VALUES
(%s, %s, %s, 'pending', %s, %s, NULL, %s, 20, %s, 0,
0, 0, 0, NULL, NULL, FALSE, 'preview', 'completed', 100,
NULL, NULL, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
RETURNING *
""",
(
task_id,
repeated_name,
source_task.get("description") or "",
source_task["process_type"],
source_task.get("source_dataset_id"),
json_dumps(repeated_config),
input_count,
len(source_files),
len(source_files),
source_task.get("tenant_id"),
source_task.get("project_id"),
source_task.get("owner_id"),
source_task.get("approval_status") or "not_required",
source_task.get("created_by"),
source_task.get("created_by"),
now,
now,
),
).fetchone()
file_id_map: dict[str, str] = {}
for source in source_files:
old_file_id = str(source["id"])
copy = file_copies[old_file_id]
new_file_id = str(copy["id"])
storage_object_id, metadata = _source_storage_descriptor(
{
"storage_object_id": copy["storage_object_id"],
"metadata": _json_value(source.get("metadata"), {}),
},
task_id,
new_file_id,
)
file_id_map[old_file_id] = new_file_id
conn.execute(
"""
INSERT INTO data_process_source_files
(id, task_id, storage_object_id, name, size_bytes, record_count,
file_format, checksum_sha256, version_no, content,
content_preview, metadata, tenant_id, project_id, created_by,
created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, 1, %s, %s, %s,
%s, %s, %s, %s, %s)
""",
(
new_file_id,
task_id,
storage_object_id,
source["name"],
source.get("size_bytes") or 0,
source.get("record_count") or 0,
source.get("file_format"),
source["checksum_sha256"],
source.get("content") or "",
source.get("content_preview"),
json_dumps(metadata),
source_task.get("tenant_id"),
source_task.get("project_id"),
source.get("created_by") or source_task.get("created_by"),
now,
now,
),
)
for preview in previews:
old_source_file_id = preview.get("source_file_id")
conn.execute(
"""
INSERT INTO data_process_preview_items
(id, task_id, source_file_id, original_content, edited_content,
source_start, source_end, source_start_line, source_end_line,
token_count, status, quality_score, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
%s, %s)
""",
(
new_id("dpp"),
task_id,
file_id_map.get(str(old_source_file_id))
if old_source_file_id
else None,
preview.get("original_content") or "",
preview.get("edited_content") or "",
preview.get("source_start"),
preview.get("source_end"),
preview.get("source_start_line"),
preview.get("source_end_line"),
max(0, int(preview.get("token_count") or 0)),
preview.get("status") or "original",
json_dumps(_json_value(preview.get("quality_score"), {})),
now,
now,
),
)
return self._repeat_response(
conn,
task_row or {},
source_task_id=source_task_id,
created=True,
)
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("再次生成任务名称或请求发生冲突,请重试") from exc
def get_task(self, task_id: str, *, for_update: bool = False) -> dict[str, Any]:
lock = " FOR UPDATE" if for_update else ""
with self.connect() as conn:
if for_update:
row = conn.execute(
f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}",
(task_id,),
).fetchone()
else:
row = conn.execute(
"""
SELECT task.*,
creator.display_name AS creator_name,
processor.display_name AS processor_name,
(SELECT COUNT(*) FROM data_process_source_files source
WHERE source.task_id=task.id AND source.deleted_at IS NULL)
AS source_file_count,
(SELECT COUNT(*) FROM data_process_preview_items preview
WHERE preview.task_id=task.id) AS preview_count,
(SELECT COALESCE(json_agg(json_build_object(
'id', dataset.id,
'name', dataset.name,
'type', dataset.type,
'count', dataset.count,
'dataset_split', CASE dataset.type
WHEN 'train' THEN 'train'
WHEN 'val' THEN 'validation'
WHEN 'test' THEN 'test'
ELSE NULL
END
) ORDER BY CASE dataset.type
WHEN 'train' THEN 1 WHEN 'val' THEN 2 WHEN 'test' THEN 3 ELSE 4 END), '[]'::json)
FROM datasets dataset
WHERE dataset.source='task'
AND dataset.deleted_at IS NULL
AND (
dataset.source_task_id=task.id
OR (dataset.source_task_id IS NULL AND dataset.task_id=task.id)
))
AS output_datasets,
CASE
WHEN task.started_at IS NOT NULL AND task.completed_at IS NOT NULL
THEN EXTRACT(EPOCH FROM (task.completed_at - task.started_at))
ELSE NULL
END AS duration_seconds
FROM data_process_tasks task
LEFT JOIN users creator ON creator.id=task.created_by
LEFT JOIN users processor ON processor.id=task.updated_by
WHERE task.id=%s AND task.deleted_at IS NULL
""",
(task_id,),
).fetchone()
if not row:
raise NotFoundError("data process task not found")
return _public_task(_decode_row(row)) or {}
def _task_in_connection(
self,
conn: psycopg.Connection[dict[str, Any]],
task_id: str,
*,
for_update: bool = False,
) -> dict[str, Any]:
lock = " FOR UPDATE" if for_update else ""
row = conn.execute(
f"SELECT * FROM data_process_tasks WHERE id=%s AND deleted_at IS NULL{lock}",
(task_id,),
).fetchone()
if not row:
raise NotFoundError("data process task not found")
return _decode_row(row) or {}
@staticmethod
def _ensure_editable(task: dict[str, Any]) -> None:
if task["status"] not in EDITABLE_STATUSES:
raise InvalidStateError(f"task cannot be edited while status is {task['status']}")
if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
raise InvalidStateError("task cannot be edited while preview is running")
if task.get("output_dataset_id") and not _is_regeneration_prepared(task):
raise InvalidStateError("published task cannot be edited")
def update_workflow_step(self, task_id: str, workflow_step: str) -> dict[str, Any]:
"""独立保存向导位置,不触发配置或结果失效逻辑。"""
if workflow_step not in WORKFLOW_STEPS:
raise ValueError("invalid data process workflow step")
with self.connect() as conn:
row = conn.execute(
"""
UPDATE data_process_tasks
SET workflow_step=%s, updated_at=%s
WHERE id=%s AND deleted_at IS NULL
RETURNING *
""",
(workflow_step, utcnow(), task_id),
).fetchone()
if not row:
raise NotFoundError("data process task not found")
return _public_task(_decode_row(row)) or {}
def recover_legacy_aborted_regeneration(self, task_id: str) -> dict[str, Any]:
"""恢复旧版在真正开始生成前误删的上一轮结果。
旧实现会在 ``POST /regenerate`` 时立即把已发布任务置为 pending、
清空结果并解除输出指针。三个已发布数据集仍是独立完整产物,因此只在
这个特征完全匹配时,使用其记录恢复结果和任务状态。该操作幂等,不会
触碰正常的新建待生成任务或已经开始的新一轮生成。
"""
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if (
task.get("status") != "pending"
or task.get("generation_run_id")
or task.get("output_dataset_id")
or int(task.get("output_count") or 0) != 0
):
return {"recovered": False, "result_count": 0}
result_count = int(
(
conn.execute(
"SELECT COUNT(*) AS count FROM data_process_results WHERE task_id=%s",
(task_id,),
).fetchone()
or {}
).get("count")
or 0
)
if result_count:
return {"recovered": False, "result_count": result_count}
datasets = conn.execute(
"""
SELECT id, type, count, created_at
FROM datasets
WHERE source='task' AND deleted_at IS NULL
AND (
source_task_id=%s
OR (source_task_id IS NULL AND task_id=%s)
)
ORDER BY CASE type
WHEN 'train' THEN 1 WHEN 'val' THEN 2 WHEN 'test' THEN 3 ELSE 4
END, created_at, id
""",
(task_id, task_id),
).fetchall()
train_dataset = next(
(dataset for dataset in datasets if dataset.get("type") == "train"),
None,
)
if not train_dataset:
return {"recovered": False, "result_count": 0}
dataset_ids = [str(dataset["id"]) for dataset in datasets]
records = conn.execute(
"""
SELECT id, dataset_id, line_no, split, instruction, input, output,
raw, status, source_result_id, preview_item_id, created_at
FROM dataset_records
WHERE dataset_id = ANY(%s)
ORDER BY created_at, dataset_id, line_no NULLS LAST, id
""",
(dataset_ids,),
).fetchall()
if not records:
return {"recovered": False, "result_count": 0}
preview_rows = conn.execute(
"SELECT id FROM data_process_preview_items WHERE task_id=%s",
(task_id,),
).fetchall()
preview_ids = {str(row["id"]) for row in preview_rows}
used_result_ids: set[str] = set()
recovered_count = 0
for record in records:
raw = _json_value(record.get("raw"), {})
raw = raw if isinstance(raw, dict) else {}
candidate_id = str(
record.get("source_result_id")
or raw.get("source_result_id")
or ""
)
result_id = (
candidate_id
if candidate_id and candidate_id not in used_result_ids
else new_id("dpr")
)
used_result_ids.add(result_id)
candidate_preview_id = str(
record.get("preview_item_id")
or raw.get("preview_item_id")
or ""
)
preview_item_id = (
candidate_preview_id if candidate_preview_id in preview_ids else None
)
instruction = str(record.get("instruction") or raw.get("instruction") or "")
input_text = str(record.get("input") or raw.get("input") or "")
chosen = str(raw.get("chosen") or "")
rejected = str(raw.get("rejected") or "")
output = str(
record.get("output") or raw.get("output") or chosen or ""
)
split = str(record.get("split") or raw.get("split") or "") or None
status = str(record.get("status") or "valid")
if status not in {"valid", "modified", "invalid"}:
status = "valid"
created_at = record.get("created_at") or utcnow()
conn.execute(
"""
INSERT INTO data_process_results
(id, task_id, preview_item_id, instruction, input, output,
chosen, rejected, original_instruction, original_input,
original_output, original_chosen, original_rejected, status,
error, split, quality_score, created_at, updated_at)
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s,
%s, %s, NULL, %s, '{}', %s, %s)
""",
(
result_id,
task_id,
preview_item_id,
instruction,
input_text,
output,
chosen,
rejected,
instruction,
input_text,
output,
chosen,
rejected,
status,
split,
created_at,
created_at,
),
)
conn.execute(
"""
UPDATE dataset_records
SET source_result_id=%s, preview_item_id=%s
WHERE id=%s
""",
(result_id, preview_item_id, record["id"]),
)
recovered_count += 1
now = utcnow()
conn.execute(
"""
UPDATE data_process_tasks
SET status='completed', progress=100, output_dataset_id=%s,
output_count=%s, filtered_count=0, duplicate_count=0,
error_count=(SELECT COUNT(*) FROM data_process_results
WHERE task_id=%s AND status='invalid'),
failure_reason=NULL, results_confirmed=TRUE,
workflow_step='results', updated_at=%s
WHERE id=%s
""",
(
train_dataset["id"],
recovered_count,
task_id,
now,
task_id,
),
)
return {"recovered": True, "result_count": recovered_count}
def update_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
allowed = {
"name",
"description",
"process_type",
"source_dataset_id",
}
values: dict[str, Any] = {key: value for key, value in payload.items() if key in allowed}
if not values and payload.get("config") is None:
return self.get_task(task_id)
try:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
self._ensure_editable(task)
regeneration_prepared = _is_regeneration_prepared(task)
if regeneration_prepared and any(
key in payload and payload.get(key) != task.get(key)
for key in ("process_type", "source_dataset_id")
):
raise InvalidStateError(
"process type and source dataset cannot change during regeneration"
)
if payload.get("config") is not None:
next_config = _business_config(payload["config"])
current_config = dict(task.get("config") or {})
for key in _INTERNAL_CONFIG_KEYS:
if key in current_config:
next_config[key] = current_config[key]
values["config"] = json_dumps(next_config)
invalidates_results = (
("config" in payload and payload.get("config") != task.get("config"))
or (
"process_type" in payload
and payload.get("process_type") != task.get("process_type")
)
or (
"source_dataset_id" in payload
and payload.get("source_dataset_id") != task.get("source_dataset_id")
)
)
if invalidates_results and not regeneration_prepared:
values.update(
{
"status": "pending",
"progress": 0,
"output_count": 0,
"filtered_count": 0,
"duplicate_count": 0,
"error_count": 0,
"failure_reason": None,
"generation_run_id": None,
"results_confirmed": False,
"preview_status": "idle",
"preview_progress": 0,
"preview_run_id": None,
"preview_failure_reason": None,
"preview_total_files": 0,
"preview_completed_files": 0,
"started_at": None,
"completed_at": None,
}
)
conn.execute("DELETE FROM data_process_results WHERE task_id=%s", (task_id,))
conn.execute(
"DELETE FROM data_process_preview_items WHERE task_id=%s", (task_id,)
)
if (
"process_type" in payload
and payload.get("process_type") != task.get("process_type")
):
conn.execute(
"""
UPDATE data_process_source_files
SET deleted_at=%s, updated_at=%s
WHERE task_id=%s AND deleted_at IS NULL
""",
(utcnow(), utcnow(), task_id),
)
values["input_count"] = 0
values["updated_at"] = utcnow()
assignments = ", ".join(f"{key}=%s" for key in values)
row = conn.execute(
f"UPDATE data_process_tasks SET {assignments} WHERE id=%s RETURNING *",
[*values.values(), task_id],
).fetchone()
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("data process task name already exists") from exc
return _public_task(_decode_row(row)) or {}
def prepare_regeneration(
self,
task_id: str,
payload: dict[str, Any],
) -> dict[str, Any]:
"""非破坏性地保存重新生成配置。
准备阶段保留任务当前状态、结果、切片及已发布数据集。真正开始
生成时,才在同一事务内切换运行状态并清理上一轮结果。
"""
# 先修复曾被旧版 prepare 提前清空的任务,再建立新的非破坏性草稿标记。
self.recover_legacy_aborted_regeneration(task_id)
try:
with self.connect() as conn:
task = self._task_in_connection(conn, task_id, for_update=True)
if task["status"] == "running":
raise ConflictError("running task cannot be prepared for regeneration")
if task.get("preview_status") in ACTIVE_PREVIEW_STATUSES:
raise ConflictError("running preview cannot be prepared for regeneration")
current_updated_at = _serialize_value(task.get("updated_at"))
if payload["expected_updated_at"] != current_updated_at:
raise ConflictError("data process task was modified by another request")
process_type = str(payload["process_type"])
if process_type != str(task["process_type"]):
raise InvalidStateError("process_type cannot be changed during regeneration")
current_config = dict(task.get("config") or {})
next_config = _business_config(payload.get("config"))
for key in (_REPEAT_SOURCE_TASK_KEY, _REPEAT_REQUEST_KEY):
if key in current_config:
next_config[key] = current_config[key]
preview_invalidated = _preview_config_changed(
process_type,
current_config,
next_config,
)
now = utcnow()
next_config[_REGENERATION_MARKER_KEY] = {
"prepared": True,
"preview_invalidated": preview_invalidated,
"prepared_at": now,
}
# 002 迁移前发布的数据集只有 task_id。先补齐新关联字段保证
# 解除任务输出指针后,详情和后续重新发布仍能定位原来的三份数据集。
conn.execute(
"""
UPDATE datasets
SET source_task_id=%s, updated_at=%s
WHERE source='task' AND source_task_id IS NULL AND task_id=%s
AND deleted_at IS NULL
""",
(task_id, now, task_id),
)
published_row = conn.execute(
"""
SELECT EXISTS(
SELECT 1 FROM datasets
WHERE source='task' AND deleted_at IS NULL
AND (
source_task_id=%s
OR (source_task_id IS NULL AND task_id=%s)
)
) AS exists
""",
(task_id, task_id),
).fetchone()
published_outputs_preserved = bool(task.get("output_dataset_id")) or bool(
published_row and published_row.get("exists")
)
row = conn.execute(
"""
UPDATE data_process_tasks
SET name=%s, description=%s, config=%s, updated_at=%s
WHERE id=%s
RETURNING *
""",
(
payload["name"],
payload.get("description") or "",
json_dumps(next_config),
now,
task_id,
),
).fetchone()
except psycopg.errors.UniqueViolation as exc:
raise ConflictError("data process task name already exists") from exc
return {
"task": _public_task(_decode_row(row)) or {},
"preview_invalidated": preview_invalidated,
"published_outputs_preserved": published_outputs_preserved,
}
def delete_task(self, task_id: str, *, deleted_by: str | None = None) -> None:
with self.connect() as conn:
self._task_in_connection(conn, task_id, for_update=True)
now = utcnow()
conn.execute(
"""
UPDATE data_process_tasks
SET status=CASE WHEN status='running' THEN 'stopped' ELSE status END,
generation_run_id=NULL,
preview_status=CASE
WHEN preview_status IN ('queued', 'running') THEN 'cancelled'
ELSE preview_status
END,
preview_run_id=NULL,
deleted_at=%s, deleted_by=%s, updated_at=%s
WHERE id=%s
""",
(now, deleted_by, now, task_id),
)

View File

@@ -0,0 +1 @@
"""GPU assignment management module."""

View File

@@ -0,0 +1,108 @@
"""GPU 算力分配管理路由。"""
from __future__ import annotations
import json
from typing import Any
from fastapi import APIRouter, Body, Depends, Request
from app.api.v1.endpoints.platform import ok, fail
from app.core.auth import get_current_user, is_admin, user_tenant_ids
from app.db.platform_store import get_platform_store
router = APIRouter(prefix="/compute", tags=["gpu-assignment"])
@router.get("/gpu-assignments")
def list_assignments(current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
"""查看全部分配关系(仅 admin"""
if not is_admin(current_user):
raise fail(403, "admin permission required")
return ok(get_platform_store().gpu_assignments())
@router.post("/gpu-assignments")
def assign_gpus(
payload: dict[str, Any] = Body(...),
request: Request = None,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
"""批量分配 GPU仅 admin。body: { assignments: [{ node_id, gpu_index, user_id }] }"""
if not is_admin(current_user):
raise fail(403, "admin permission required")
assignments = payload.get("assignments") or []
if not assignments:
raise fail(400, "assignments 不能为空")
actor = current_user.get("id")
try:
result = get_platform_store().assign_gpus(assignments, assigned_by=actor)
except ValueError as exc:
raise fail(409, str(exc))
get_platform_store().record_audit(
action="gpu.assign",
actor_id=actor,
target_type="gpu",
detail=f"count={len(assignments)}",
)
return ok(result)
@router.post("/gpu-assignments/request")
def request_gpu_assignment(
payload: dict[str, Any] = Body(...),
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
assignments = payload.get("assignments") or []
if not assignments:
raise fail(400, "assignments 不能为空")
if is_admin(current_user):
try:
return ok(get_platform_store().assign_gpus(assignments, assigned_by=current_user.get("id")))
except ValueError as exc:
raise fail(409, str(exc))
user_id = str(current_user.get("id") or "")
normalized = []
for item in assignments:
if not isinstance(item, dict) or not item.get("node_id") or item.get("gpu_index") is None:
raise fail(400, "每项必须包含 node_id 和 gpu_index")
normalized.append({**item, "user_id": user_id})
instance = get_platform_store().create_approval_instance({
"resource_type": "gpu",
"resource_id": f"batch:{user_id}",
"applicant_id": user_id,
"action": "gpu.assign",
"tenant_id": current_user.get("tenant_id") or "default",
"reason": json.dumps({"assignments": normalized}, ensure_ascii=False),
})
get_platform_store().record_audit(
action="gpu.assign.request", actor_id=user_id, target_type="gpu",
target_id=instance["id"], tenant_id=current_user.get("tenant_id") or "default",
detail=f"count={len(normalized)}",
)
return ok({"approval_required": True, "approval_id": instance["id"], "approval": instance})
@router.delete("/gpu-assignments/{assignment_id}")
def unassign_gpu(
assignment_id: str,
request: Request = None,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
"""撤销 GPU 分配(仅 admin"""
if not is_admin(current_user):
raise fail(403, "admin permission required")
get_platform_store().unassign_gpu(assignment_id)
actor = current_user.get("id")
get_platform_store().record_audit(
action="gpu.unassign",
actor_id=actor,
target_type="gpu",
target_id=assignment_id,
)
return ok({"deleted": assignment_id})
@router.get("/my-gpus")
def my_gpus(current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
"""查看当前用户可用的 GPU 列表。"""
return ok(get_platform_store().gpu_assignments_for_user(current_user["id"]))

View File

@@ -0,0 +1,260 @@
from __future__ import annotations
from fastapi import APIRouter, Body, Depends, Request
from typing import Any
from app.api.v1.endpoints.platform import ok, fail
from app.core.auth import filter_accessible_resource_ids, get_current_user, has_resource_access, is_admin
from app.db.platform_store import get_platform_store
router = APIRouter(prefix="/projects", tags=["project"])
def _actor(request: Request) -> str | None:
auth = request.headers.get("Authorization", "")
token = auth.replace("Bearer ", "").strip()
return token or None
def _require_no_pending_approval(resource_type: str, resource_id: str) -> None:
"""第 4 周:写操作审批拦截——存在待审批实例时拒绝执行。"""
store = get_platform_store()
pending = [
i for i in store.approval_instances(status="pending")
if i["resource_type"] == resource_type and i["resource_id"] == resource_id
]
if pending:
raise fail(409, "存在待审批的变更,请先完成审批")
def _require_approval_or_admin(
resource_type: str,
resource_id: str,
current_user: dict[str, Any],
action_desc: str = "",
action: str = "project.change",
) -> dict[str, Any] | None:
"""高风险操作审批旁路admin 直接放行普通用户创建审批实例code=202"""
if is_admin(current_user):
return None
if not has_resource_access(resource_type, resource_id, current_user, "write"):
raise fail(403, "no permission to request this project change")
store = get_platform_store()
if store.consume_approved_approval(resource_type, resource_id, str(current_user.get("id") or ""), action):
return None
instance = store.create_approval_instance({
"resource_type": resource_type,
"resource_id": resource_id,
"applicant_id": current_user.get("id"),
"template_id": None,
"action": action,
"tenant_id": current_user.get("tenant_id") or "default",
"reason": action_desc,
})
return {
"code": 202,
"message": f"操作已提交审批,等待管理员批准:{action_desc}",
"data": {"approval_required": True, "approval_id": instance["id"]},
}
@router.get("")
def list_projects(
tenant_id: str = "default",
status: str | None = None,
keyword: str | None = None,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
store = get_platform_store()
projects = store.projects(tenant_id=tenant_id, status=status, keyword=keyword)
# #1 ACL 过滤admin 直接放行,普通用户只能看到自己被授权的项目
accessible_ids = set(
filter_accessible_resource_ids("project", [p["id"] for p in projects], current_user)
)
filtered = [p for p in projects if p["id"] in accessible_ids]
return ok(filtered)
@router.post("")
def create_project(payload: dict[str, Any] = Body(...), request: Request = None, current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
if not is_admin(current_user) and str(payload.get("tenant_id") or current_user.get("tenant_id") or "default") != str(current_user.get("tenant_id") or "default"):
raise fail(403, "cannot create project in another tenant")
payload.setdefault("tenant_id", current_user.get("tenant_id") or "default")
payload.setdefault("create_by", current_user.get("id"))
store = get_platform_store()
try:
store.assert_active_tenant(payload["tenant_id"])
except ValueError as exc:
raise fail(400, str(exc))
proj = store.create_project(payload)
store.record_audit(
action="project.create",
actor_id=current_user.get("id"),
target_type="project",
target_id=proj["id"],
tenant_id=proj.get("tenant_id"),
detail=f"name={proj.get('name')}",
)
return ok(proj)
@router.get("/{project_id}")
def get_project(project_id: str, current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
# #2 访问控制:普通用户无 read 权限则拒绝
if not has_resource_access("project", project_id, current_user, "read"):
raise fail(403, "no permission to access this project")
try:
return ok(get_platform_store().project(project_id))
except KeyError:
raise fail(404, "project not found")
@router.put("/{project_id}")
def update_project(
project_id: str,
payload: dict[str, Any] = Body(...),
request: Request = None,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
if not has_resource_access("project", project_id, current_user, "write"):
raise fail(403, "no permission to update this project")
store = get_platform_store()
try:
proj = store.update_project(project_id, payload)
except KeyError:
raise fail(404, "project not found")
store.record_audit(
action="project.update",
actor_id=current_user.get("id"),
target_type="project",
target_id=project_id,
tenant_id=proj.get("tenant_id"),
detail=f"fields={','.join(payload.keys())}",
)
return ok(proj)
@router.post("/{project_id}/archive")
def archive_project(
project_id: str,
request: Request = None,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
_require_no_pending_approval("project", project_id)
pending = _require_approval_or_admin("project", project_id, current_user, f"归档项目 {project_id}", "project.archive")
if pending:
return pending
store = get_platform_store()
try:
proj = store.archive_project(project_id)
except KeyError:
raise fail(404, "project not found")
store.record_audit(
action="project.archive",
actor_id=current_user.get("id"),
target_type="project",
target_id=project_id,
tenant_id=proj.get("tenant_id"),
)
return ok(proj)
@router.delete("/{project_id}")
def delete_project(
project_id: str,
request: Request = None,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
_require_no_pending_approval("project", project_id)
pending = _require_approval_or_admin("project", project_id, current_user, f"删除项目 {project_id}", "project.delete")
if pending:
return pending
store = get_platform_store()
store.delete_project(project_id)
store.record_audit(
action="project.delete",
actor_id=current_user.get("id"),
target_type="project",
target_id=project_id,
)
return ok(None)
@router.get("/{project_id}/members")
def list_members(project_id: str, current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
if not has_resource_access("project", project_id, current_user, "read"):
raise fail(403, "no permission to access this project")
try:
return ok(get_platform_store().project_members(project_id))
except KeyError:
raise fail(404, "project not found")
@router.post("/{project_id}/members")
def add_member(
project_id: str,
payload: dict[str, Any] = Body(...),
request: Request = None,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
if not has_resource_access("project", project_id, current_user, "write"):
raise fail(403, "no permission to manage members of this project")
store = get_platform_store()
try:
member = store.add_project_member(project_id, payload)
except KeyError:
raise fail(404, "project not found")
store.record_audit(
action="project.member.add",
actor_id=current_user.get("id"),
target_type="project.member",
target_id=project_id,
detail=f"user_id={payload.get('user_id')},role={payload.get('role')}",
)
return ok(member)
@router.put("/{project_id}/members/{user_id}")
def update_member(
project_id: str,
user_id: str,
payload: dict[str, Any] = Body(...),
request: Request = None,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
if not has_resource_access("project", project_id, current_user, "write"):
raise fail(403, "no permission to manage members of this project")
store = get_platform_store()
try:
member = store.update_project_member_role(project_id, user_id, payload)
except KeyError:
raise fail(404, "project or member not found")
store.record_audit(
action="project.member.update",
actor_id=current_user.get("id"),
target_type="project.member",
target_id=project_id,
detail=f"user_id={user_id},role={payload.get('role')}",
)
return ok(member)
@router.delete("/{project_id}/members/{user_id}")
def remove_member(
project_id: str,
user_id: str,
request: Request = None,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
if not has_resource_access("project", project_id, current_user, "write"):
raise fail(403, "no permission to manage members of this project")
store = get_platform_store()
store.remove_project_member(project_id, user_id)
store.record_audit(
action="project.member.remove",
actor_id=_actor(request) if request else None,
target_type="project.member",
target_id=project_id,
detail=f"user_id={user_id}",
)
return ok(None)

View File

@@ -0,0 +1 @@
"""Resource access control list (ACL) module."""

View File

@@ -0,0 +1,93 @@
from __future__ import annotations
from fastapi import APIRouter, Body, Request, Depends
from typing import Any
from app.api.v1.endpoints.platform import ok, fail
from app.db.platform_store import get_platform_store
from app.core.auth import (
get_current_user,
has_resource_access,
is_admin,
resource_record,
resource_tenant_id,
user_tenant_ids,
)
from app.core.audit import audit_log, AuditActions
router = APIRouter(prefix="/resources", tags=["resource"])
def _actor(request: Request, current_user: dict[str, Any]) -> str | None:
return str(current_user.get("id") or "") or None
@router.get("/{resource_type}/{resource_id}/acl")
def get_acl(resource_type: str, resource_id: str, current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
"""查询资源 ACL返回按主体分组的权限列表。"""
if not has_resource_access(resource_type, resource_id, current_user, "read"):
raise fail(403, "no permission to access resource ACL")
return ok(get_platform_store().resource_acl(resource_type, resource_id))
@router.put("/{resource_type}/{resource_id}/acl")
@audit_log(
action=AuditActions.GRANT_ACL,
target_type="",
detail_template="设置资源授权: {resource_type}/{resource_id}",
)
def set_acl(
resource_type: str,
resource_id: str,
payload: dict[str, Any] = Body(...),
request: Request = None,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
"""设置资源 ACLbody: { entries: [{ subject_type, subject_id, permissions: [] }] }"""
resource = resource_record(resource_type, resource_id)
if not resource and not is_admin(current_user):
raise fail(404, "resource not found")
if not is_admin(current_user) and not has_resource_access(resource_type, resource_id, current_user, "write"):
raise fail(403, "only resource owner or admin can update ACL")
entries = payload.get("entries") or []
allowed = {"read", "write", "execute", "download", "delete", "admin"}
owner_allowed = {"read", "write", "execute", "download"}
tenant_id = resource_tenant_id(resource_type, resource) if resource else None
tenant_ids = user_tenant_ids(current_user)
for entry in entries:
if entry.get("principal_type") not in {"user", "role"} or not entry.get("principal_id"):
raise fail(400, "invalid ACL principal")
permissions = set(entry.get("permissions") or [])
if any(permission not in allowed for permission in permissions):
raise fail(400, "invalid ACL permission")
if not is_admin(current_user) and permissions - owner_allowed:
raise fail(403, "resource owners cannot grant delete or admin permission")
if entry.get("principal_type") == "user":
with get_platform_store().connect() as conn:
principal = conn.execute(
"SELECT id, tenant_id, status FROM users WHERE id=?",
(entry["principal_id"],),
).fetchone()
if not principal or principal.get("status") != "active":
raise fail(400, "ACL user does not exist or is inactive")
principal_tenant = str(principal.get("tenant_id") or "default")
if not is_admin(current_user) and tenant_id and principal_tenant not in tenant_ids:
raise fail(403, "cannot grant resource access across tenants")
elif not is_admin(current_user):
# Role ACLs are global in the legacy schema and therefore cannot
# be safely scoped to one tenant by a normal resource owner.
raise fail(403, "only administrators can grant role-based ACLs")
result = get_platform_store().set_resource_acl(
resource_type,
resource_id,
entries,
granted_by=str(current_user.get("id") or "") or None,
)
get_platform_store().record_audit(
action="resource.acl.set",
actor_id=_actor(request, current_user) if request else current_user.get("id"),
target_type=resource_type,
target_id=resource_id,
detail=f"entries={len(entries)}",
)
return ok(result)

View File

@@ -0,0 +1,77 @@
from __future__ import annotations
from fastapi import APIRouter, Body, Request, Depends
from typing import Any
from app.api.v1.endpoints.platform import ok, fail
from app.db.platform_store import get_platform_store
from app.core.auth import require_admin
router = APIRouter(prefix="/retention-policies", tags=["retention"])
def _actor(request: Request) -> str | None:
auth = request.headers.get("Authorization", "")
token = auth.replace("Bearer ", "").strip()
return token or None
@router.get("")
def list_policies(current_user: dict = Depends(require_admin)) -> dict[str, Any]:
return ok(get_platform_store().retention_policies())
@router.post("")
def create_policy(payload: dict[str, Any] = Body(...), request: Request = None, current_user: dict = Depends(require_admin)) -> dict[str, Any]:
if not payload.get("name"):
raise fail(400, "name 必填")
policy = get_platform_store().create_retention_policy(payload)
get_platform_store().record_audit(
action="retention.create",
actor_id=current_user.get("id"),
target_type="retention_policy",
target_id=policy["id"],
detail=f"name={policy.get('name')}",
)
return ok(policy)
@router.get("/{policy_id}")
def get_policy(policy_id: str, current_user: dict = Depends(require_admin)) -> dict[str, Any]:
try:
return ok(get_platform_store().retention_policy(policy_id))
except KeyError:
raise fail(404, "retention policy not found")
@router.put("/{policy_id}")
def update_policy(
policy_id: str, payload: dict[str, Any] = Body(...), request: Request = None,
current_user: dict = Depends(require_admin),
) -> dict[str, Any]:
store = get_platform_store()
try:
policy = store.update_retention_policy(policy_id, payload)
except KeyError:
raise fail(404, "retention policy not found")
store.record_audit(
action="retention.update",
actor_id=current_user.get("id"),
target_type="retention_policy",
target_id=policy_id,
detail=f"fields={','.join(payload.keys())}",
)
return ok(policy)
@router.delete("/{policy_id}")
def delete_policy(policy_id: str, request: Request = None, current_user: dict = Depends(require_admin)) -> dict[str, Any]:
store = get_platform_store()
store.delete_retention_policy(policy_id)
store.record_audit(
action="retention.delete",
actor_id=current_user.get("id"),
target_type="retention_policy",
target_id=policy_id,
)
return ok({"deleted": policy_id})

View File

@@ -0,0 +1 @@
"""Central object storage integration."""

View File

@@ -0,0 +1,141 @@
from __future__ import annotations
from datetime import timedelta
from functools import lru_cache
from io import BytesIO
from collections.abc import Iterator
from typing import Any
from minio import Minio
from minio.error import S3Error
import urllib3
from app.core.config import get_settings
class ObjectStorageError(RuntimeError):
pass
class MinioObjectStorage:
def __init__(self) -> None:
settings = get_settings()
endpoint = settings.minio_endpoint.replace("http://", "").replace("https://", "").rstrip("/")
# MinIO outages must fail fast; higher-level workflows own the retry
# policy and should not wait through urllib3's default retry chain.
http_client = urllib3.PoolManager(
cert_reqs="CERT_REQUIRED" if settings.minio_secure else "CERT_NONE",
timeout=urllib3.Timeout(connect=2.0, read=10.0),
retries=False,
)
self.client = Minio(
endpoint,
access_key=settings.minio_access_key,
secret_key=settings.minio_secret_key,
secure=settings.minio_secure,
http_client=http_client,
)
self.bucket = settings.minio_bucket
def _ensure_enabled(self) -> None:
if not get_settings().minio_enabled:
raise ObjectStorageError("MinIO object storage is disabled")
def ensure_bucket(self) -> None:
self._ensure_enabled()
try:
if not self.client.bucket_exists(self.bucket):
self.client.make_bucket(self.bucket)
except Exception as exc: # noqa: BLE001 - normalize network/client failures
raise ObjectStorageError(str(exc)) from exc
def presigned_put(self, object_key: str, expires_seconds: int = 3600) -> str:
self._ensure_enabled()
self.ensure_bucket()
return self.client.presigned_put_object(self.bucket, object_key, expires=timedelta(seconds=expires_seconds))
def presigned_get(self, object_key: str, expires_seconds: int = 3600) -> str:
self._ensure_enabled()
self.ensure_bucket()
return self.client.presigned_get_object(self.bucket, object_key, expires=timedelta(seconds=expires_seconds))
def stat(self, object_key: str) -> dict[str, Any]:
self._ensure_enabled()
self.ensure_bucket()
try:
result = self.client.stat_object(self.bucket, object_key)
return {"object_key": object_key, "byte_size": result.size, "etag": result.etag, "last_modified": result.last_modified.isoformat() if result.last_modified else None}
except S3Error as exc:
raise ObjectStorageError(str(exc)) from exc
def get_bytes(self, object_key: str) -> bytes:
"""Read an object through the backend for small API responses and workers."""
self._ensure_enabled()
self.ensure_bucket()
response = None
try:
response = self.client.get_object(self.bucket, object_key)
return response.read()
except S3Error as exc:
raise ObjectStorageError(str(exc)) from exc
finally:
if response is not None:
response.close()
response.release_conn()
def iter_bytes(self, object_key: str, chunk_size: int = 256 * 1024) -> Iterator[bytes]:
"""Stream an object without loading the complete file into memory."""
self._ensure_enabled()
self.ensure_bucket()
response = None
try:
response = self.client.get_object(self.bucket, object_key)
while True:
chunk = response.read(chunk_size)
if not chunk:
break
yield chunk
except S3Error as exc:
raise ObjectStorageError(str(exc)) from exc
finally:
if response is not None:
response.close()
response.release_conn()
def list_objects(self, prefix: str) -> list[dict[str, Any]]:
self._ensure_enabled()
self.ensure_bucket()
try:
return [
{
"object_key": item.object_name,
"byte_size": item.size or 0,
"etag": item.etag,
"last_modified": item.last_modified.isoformat() if item.last_modified else None,
}
for item in self.client.list_objects(self.bucket, prefix=prefix, recursive=True)
]
except S3Error as exc:
raise ObjectStorageError(str(exc)) from exc
def delete(self, object_key: str) -> None:
self._ensure_enabled()
self.ensure_bucket()
try:
self.client.remove_object(self.bucket, object_key)
except S3Error as exc:
raise ObjectStorageError(str(exc)) from exc
def put_bytes(self, object_key: str, content: bytes, content_type: str = "application/octet-stream") -> dict[str, Any]:
self._ensure_enabled()
self.ensure_bucket()
try:
result = self.client.put_object(self.bucket, object_key, BytesIO(content), len(content), content_type=content_type)
return {"bucket": self.bucket, "object_key": object_key, "etag": result.etag, "byte_size": len(content)}
except S3Error as exc:
raise ObjectStorageError(str(exc)) from exc
@lru_cache
def get_object_storage() -> MinioObjectStorage:
return MinioObjectStorage()

View File

@@ -0,0 +1,54 @@
"""Storage placement rules shared by dataset and data-processing flows."""
from __future__ import annotations
from app.core.config import get_settings
_INLINE_TEXT_FORMATS = {
"txt", "text", "md", "markdown", "json", "jsonl", "csv", "tsv",
"yaml", "yml", "xml", "html", "text/plain", "application/json",
"application/jsonl", "text/csv",
}
def should_store_in_minio(
size_bytes: int | None,
*,
content_type: str | None = None,
file_format: str | None = None,
) -> bool:
"""Return whether a file is large enough to use the shared object store.
Small files remain inline in PostgreSQL so page previews and metadata reads
do not pay an object-storage round trip. MinIO is still mandatory for
large files when it is enabled.
"""
if not get_settings().minio_enabled:
return False
try:
size = max(0, int(size_bytes or 0))
except (TypeError, ValueError):
size = 0
if size > get_settings().minio_inline_max_bytes:
return True
# Binary office/document files remain in MinIO even when small because
# their original bytes cannot be safely represented by a text DB column.
normalized_format = str(file_format or "").strip().lower().lstrip(".")
normalized_type = str(content_type or "").strip().lower().split(";", 1)[0]
if normalized_format or normalized_type:
return not (
normalized_format in _INLINE_TEXT_FORMATS
or normalized_type in _INLINE_TEXT_FORMATS
or normalized_type.startswith("text/")
)
return False
def storage_backend_for_size(size_bytes: int | None, *, requested: str | None = None) -> str:
"""Return ``minio`` or ``database`` for a managed file."""
if str(requested or "").strip().lower() == "local":
return "database"
return "minio" if should_store_in_minio(size_bytes) else "database"

View File

@@ -0,0 +1,314 @@
from __future__ import annotations
import csv
import io
from fastapi import APIRouter, Body, HTTPException, Query, Request, Depends
from fastapi.responses import StreamingResponse
from app.db.platform_store import ALL_PERMISSIONS, get_platform_store
from app.core.auth import get_current_user, is_admin
from app.core.logging import get_client_ip
router = APIRouter(prefix="/system", tags=["system"])
_VISIT_MODULES = {
"dashboard",
"fine-tune",
"model-eval",
"model-inference",
"model-manage",
"dataset",
"data-process",
"data-convert",
}
@router.post("/audit/visit")
def record_visit(
payload: dict = Body(...),
request: Request = None,
current_user: dict = Depends(get_current_user),
) -> dict:
"""记录用户访问业务模块的行为,用于看板用户操作分布统计。"""
action = str(payload.get("action") or payload.get("module") or "").strip()
if not action:
return {"code": 0, "message": "ok", "data": {"recorded": False}}
# Visit statistics are intentionally limited to known module names. This
# endpoint must not become a free-form audit-log injection point.
if action not in _VISIT_MODULES:
return {"code": 0, "message": "ok", "data": {"recorded": False}}
actor_id = current_user.get("id")
get_platform_store().record_audit(
action=action,
actor_id=actor_id or None,
target_type="module",
target_id=action,
tenant_id=str(current_user.get("tenant_id") or "") or None,
session_id=str(current_user.get("session_id") or "") or None,
request_id=(request.headers.get("X-Request-ID") if request else None),
detail="module visit",
metadata={"source": "frontend", "detail_length": len(str(payload.get("detail") or ""))},
ip=get_client_ip(request) or None,
)
return {"code": 0, "message": "ok", "data": {"recorded": True}}
@router.get("/permissions/codes")
def permission_codes(current_user: dict = Depends(get_current_user)) -> dict:
"""返回平台权限码清单(权限码接口)。"""
return {"code": 0, "message": "ok", "data": {"codes": ALL_PERMISSIONS}}
@router.get("/permissions")
def permissions_overview(current_user: dict = Depends(get_current_user)) -> dict:
"""返回权限码清单与角色定义。"""
store = get_platform_store()
return {
"code": 0,
"message": "ok",
"data": {"codes": ALL_PERMISSIONS, "roles": store.roles()},
}
@router.get("/audit-logs")
def audit_logs(
tenant_id: str | None = Query(default=None, description="租户 ID"),
project_id: str | None = Query(default=None, description="项目 ID"),
actor_id: str | None = Query(default=None, description="操作人 ID"),
action: str | None = Query(default=None, description="动作类型"),
target_type: str | None = Query(default=None, description="目标类型"),
target_id: str | None = Query(default=None, description="目标 ID"),
keyword: str | None = Query(default=None, description="目标 ID 或详情关键字"),
start_time: str | None = Query(default=None, description="ISO8601 起始时间"),
end_time: str | None = Query(default=None, description="ISO8601 结束时间"),
limit: int = Query(default=50, ge=1, le=200),
offset: int = Query(default=0, ge=0),
current_user: dict = Depends(get_current_user),
) -> dict:
"""审计日志查询:按组织、操作人、动作、资源、关键字和时间范围分页过滤。"""
if not is_admin(current_user):
from app.api.v1.endpoints.platform import fail
raise fail(403, "admin permission required")
store = get_platform_store()
result = store.audit_logs(
tenant_id=tenant_id,
project_id=project_id,
actor_id=actor_id,
action=action,
target_type=target_type,
target_id=target_id,
keyword=keyword,
start_time=start_time,
end_time=end_time,
limit=limit,
offset=offset,
)
return {"code": 0, "message": "ok", "data": result}
@router.get("/audit-logs/export")
def audit_logs_export(
tenant_id: str | None = Query(default=None, description="租户 ID"),
project_id: str | None = Query(default=None, description="项目 ID"),
actor_id: str | None = Query(default=None, description="操作人 ID"),
action: str | None = Query(default=None, description="动作类型"),
target_type: str | None = Query(default=None, description="目标类型"),
target_id: str | None = Query(default=None, description="目标 ID"),
keyword: str | None = Query(default=None, description="目标 ID 或详情关键字"),
start_time: str | None = Query(default=None, description="ISO8601 起始时间"),
end_time: str | None = Query(default=None, description="ISO8601 结束时间"),
current_user: dict = Depends(get_current_user),
) -> StreamingResponse:
"""审计日志导出:返回 CSV 流,与应用查询相同的过滤条件。"""
if not is_admin(current_user):
from app.api.v1.endpoints.platform import fail
raise fail(403, "admin permission required")
store = get_platform_store()
result = store.audit_logs(
tenant_id=tenant_id,
project_id=project_id,
actor_id=actor_id,
action=action,
target_type=target_type,
target_id=target_id,
keyword=keyword,
start_time=start_time,
end_time=end_time,
limit=10000,
offset=0,
)
items = result["items"]
columns = [
"time", "tenant_id", "project_id", "actor_id", "action", "target_type",
"target_id", "detail", "client_ip", "result", "reason", "request_id",
"session_id", "metadata",
]
def iter_rows():
buffer = io.StringIO()
writer = csv.writer(buffer)
writer.writerow(columns)
yield buffer.getvalue()
for row in items:
buffer.seek(0)
buffer.truncate(0)
writer.writerow([row.get(c, "") or "" for c in columns])
yield buffer.getvalue()
return StreamingResponse(
iter_rows(),
media_type="text/csv",
headers={"Content-Disposition": "attachment; filename=audit_logs.csv"},
)
# ===================== 操作日志 =====================
def _operation_log_scope(current_user: dict, conditions: list[str], params: list) -> None:
"""校验操作日志权限,并为普通用户追加本人范围。"""
if is_admin(current_user):
return
if "logs" not in (current_user.get("permissions") or []):
raise HTTPException(status_code=403, detail="missing permission: logs")
conditions.append("user_id = %s")
params.append(str(current_user.get("id") or ""))
@router.get("/operation-logs")
def operation_logs(
user_id: str | None = Query(default=None, description="按用户 ID 筛选"),
module: str | None = Query(default=None, description="按模块筛选: fine-tune/model-eval/model-inference/dataset/data-convert/model-manage"),
action: str | None = Query(default=None, description="按动作筛选: create/start/stop/delete/upload/convert/merge"),
status: str | None = Query(default=None, description="按状态筛选: success/failure不传则查全部"),
keyword: str | None = Query(default=None, description="关键字搜索报错信息(error_message)"),
start_time: str | None = Query(default=None, description="ISO8601 起始时间"),
end_time: str | None = Query(default=None, description="ISO8601 结束时间"),
limit: int = Query(default=50, ge=1, le=200),
offset: int = Query(default=0, ge=0),
current_user: dict = Depends(get_current_user),
) -> dict:
"""操作日志查询:管理员查全量,普通用户只能查本人记录。"""
store = get_platform_store()
conditions = []
params: list = []
_operation_log_scope(current_user, conditions, params)
if user_id and is_admin(current_user):
conditions.append("user_id = %s")
params.append(user_id)
if module:
conditions.append("module = %s")
params.append(module)
if action:
conditions.append("action = %s")
params.append(action)
if status:
conditions.append("status = %s")
params.append(status)
if keyword:
conditions.append("(error_message ILIKE %s OR error_type ILIKE %s)")
params.append(f"%{keyword}%")
params.append(f"%{keyword}%")
if start_time:
conditions.append("create_time >= %s")
params.append(start_time)
if end_time:
conditions.append("create_time <= %s")
params.append(end_time)
where = " WHERE " + " AND ".join(conditions) if conditions else ""
with store.connect() as conn:
rows = conn.execute(
f"SELECT * FROM operation_logs{where} ORDER BY create_time DESC LIMIT %s OFFSET %s",
tuple(params + [limit, offset]),
).fetchall()
total = conn.execute(f"SELECT COUNT(*) FROM operation_logs{where}", tuple(params)).fetchone()[0]
return {"code": 0, "message": "ok", "data": {"items": [dict(r) for r in rows], "total": total}}
@router.get("/operation-logs/stats")
def operation_logs_stats(
start_time: str | None = Query(default=None, description="ISO8601 起始时间"),
end_time: str | None = Query(default=None, description="ISO8601 结束时间"),
current_user: dict = Depends(get_current_user),
) -> dict:
"""操作日志统计:管理员统计全量,普通用户统计本人记录。"""
store = get_platform_store()
conditions = []
params: list = []
_operation_log_scope(current_user, conditions, params)
if start_time:
conditions.append("create_time >= %s")
params.append(start_time)
if end_time:
conditions.append("create_time <= %s")
params.append(end_time)
where = " WHERE " + " AND ".join(conditions) if conditions else ""
failure_where = where + " AND status = 'failure'" if where else " WHERE status = 'failure'"
with store.connect() as conn:
# 总计
row = conn.execute(
f"SELECT status, COUNT(*) as cnt FROM operation_logs{where} GROUP BY status", tuple(params)
).fetchall()
total_count = 0
success_count = 0
failure_count = 0
for r in row:
total_count += r["cnt"]
if r["status"] == "success":
success_count = r["cnt"]
elif r["status"] == "failure":
failure_count = r["cnt"]
failure_rate = round(failure_count / total_count * 100, 2) if total_count > 0 else 0
# 各模块失败数
module_stats = conn.execute(
f"SELECT module, COUNT(*) as cnt FROM operation_logs{failure_where} GROUP BY module ORDER BY cnt DESC",
tuple(params),
).fetchall()
# 各异常类型分布
error_type_stats = conn.execute(
f"SELECT error_type, COUNT(*) as cnt FROM operation_logs{failure_where} AND error_type IS NOT NULL GROUP BY error_type ORDER BY cnt DESC LIMIT 10",
tuple(params),
).fetchall()
# 最近的 10 条错误
recent_errors = conn.execute(
f"SELECT * FROM operation_logs{failure_where} ORDER BY create_time DESC LIMIT 10",
tuple(params),
).fetchall()
return {
"code": 0,
"message": "ok",
"data": {
"total": total_count,
"success": success_count,
"failure": failure_count,
"failure_rate": failure_rate,
"module_failures": [{"module": r["module"], "count": r["cnt"]} for r in module_stats],
"error_types": [{"type": r["error_type"], "count": r["cnt"]} for r in error_type_stats],
"recent_errors": [dict(r) for r in recent_errors],
},
}
@router.get("/operation-logs/modules")
def operation_log_modules(current_user: dict = Depends(get_current_user)) -> dict:
"""返回操作日志中出现的模块列表(用于筛选下拉框)。"""
store = get_platform_store()
conditions: list[str] = []
params: list = []
_operation_log_scope(current_user, conditions, params)
where = " WHERE " + " AND ".join(conditions) if conditions else ""
with store.connect() as conn:
rows = conn.execute(
f"SELECT DISTINCT module FROM operation_logs{where}"
+ (" AND" if where else " WHERE")
+ " module IS NOT NULL ORDER BY module",
tuple(params),
).fetchall()
modules = [{"value": r["module"], "label": r["module"]} for r in rows]
return {"code": 0, "message": "ok", "data": modules}

View File

@@ -0,0 +1,326 @@
from __future__ import annotations
import json
from fastapi import APIRouter, Body, Depends, Request
from typing import Any
from app.api.v1.endpoints.platform import ok, fail
from app.db.platform_store import get_platform_store
from app.core.auth import is_admin, require_admin, require_tenant_admin, get_current_user, user_tenant_ids
router = APIRouter(prefix="/tenants", tags=["tenant"])
def _actor(request: Request) -> str | None:
auth = request.headers.get("Authorization", "")
token = auth.replace("Bearer ", "").strip()
return token or None
@router.get("")
def list_tenants(current_user: dict = Depends(require_admin)) -> dict[str, Any]:
return ok(get_platform_store().tenants())
@router.get("/invitations")
def my_invitations(current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
return ok(get_platform_store().tenant_invitations(str(current_user.get("id") or "")))
@router.get("/mine")
def my_tenants(current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
return ok(get_platform_store().user_tenants(str(current_user.get("id") or ""), include_all=is_admin(current_user)))
@router.post("")
def create_tenant(payload: dict[str, Any] = Body(...), request: Request = None, current_user: dict = Depends(require_admin)) -> dict[str, Any]:
store = get_platform_store()
payload = {**payload, "owner_user_id": payload.get("owner_user_id") or current_user.get("id")}
try:
tenant = store.create_tenant(payload)
except KeyError as e:
raise fail(400, f"missing field: {e}")
except ValueError as exc:
raise fail(400, str(exc))
store.record_audit(
action="tenant.create",
actor_id=current_user.get("id"),
target_type="tenant",
target_id=tenant["id"],
tenant_id=tenant["id"],
detail=f"name={tenant.get('name')}",
)
return ok(tenant)
@router.get("/{tenant_id}")
def get_tenant(tenant_id: str, current_user: dict = Depends(require_admin)) -> dict[str, Any]:
try:
return ok(get_platform_store().tenant(tenant_id))
except KeyError:
raise fail(404, "tenant not found")
@router.put("/{tenant_id}")
def update_tenant(tenant_id: str, payload: dict[str, Any] = Body(...), request: Request = None, current_user: dict = Depends(require_admin)) -> dict[str, Any]:
store = get_platform_store()
try:
tenant = store.update_tenant(tenant_id, payload)
except KeyError:
raise fail(404, "tenant not found")
store.record_audit(
action="tenant.update",
actor_id=current_user.get("id"),
target_type="tenant",
target_id=tenant_id,
tenant_id=tenant_id,
detail=f"fields={','.join(payload.keys())}",
)
return ok(tenant)
@router.put("/{tenant_id}/quota")
def set_quota(tenant_id: str, payload: dict[str, Any] = Body(...), request: Request = None, current_user: dict = Depends(require_admin)) -> dict[str, Any]:
store = get_platform_store()
try:
tenant = store.set_tenant_quota(tenant_id, payload)
except KeyError:
raise fail(404, "tenant not found")
store.record_audit(
action="tenant.quota.set",
actor_id=current_user.get("id"),
target_type="tenant",
target_id=tenant_id,
tenant_id=tenant_id,
)
return ok(tenant)
@router.put("/{tenant_id}/retention-policy")
def set_retention(tenant_id: str, payload: dict[str, Any] = Body(...), request: Request = None, current_user: dict = Depends(require_admin)) -> dict[str, Any]:
store = get_platform_store()
try:
tenant = store.set_tenant_retention(tenant_id, payload.get("retention_policy_id"))
except KeyError:
raise fail(404, "tenant not found")
store.record_audit(
action="tenant.retention.set",
actor_id=current_user.get("id"),
target_type="tenant",
target_id=tenant_id,
tenant_id=tenant_id,
)
return ok(tenant)
@router.delete("/{tenant_id}")
def delete_tenant(tenant_id: str, request: Request = None, current_user: dict = Depends(require_admin)) -> dict[str, Any]:
store = get_platform_store()
try:
tenant = store.delete_tenant(tenant_id, str(current_user.get("id") or "system"))
except KeyError:
raise fail(404, "tenant not found")
except ValueError as exc:
raise fail(400, str(exc))
store.record_audit(
action="tenant.delete",
actor_id=current_user.get("id"),
target_type="tenant",
target_id=tenant_id,
tenant_id=tenant_id,
detail=f"name={tenant.get('name')}",
)
return ok(tenant)
@router.post("/{tenant_id}/restore")
def restore_tenant(tenant_id: str, request: Request = None, current_user: dict = Depends(require_admin)) -> dict[str, Any]:
store = get_platform_store()
try:
tenant = store.restore_tenant(tenant_id, str(current_user.get("id") or "system"))
except KeyError:
raise fail(404, "tenant not found")
except ValueError as exc:
raise fail(400, str(exc))
store.record_audit(
action="tenant.restore",
actor_id=current_user.get("id"),
target_type="tenant",
target_id=tenant_id,
tenant_id=tenant_id,
detail=f"name={tenant.get('name')}",
)
return ok(tenant)
@router.get("/{tenant_id}/quota/usage")
def quota_usage(tenant_id: str, current_user: dict = Depends(require_tenant_admin)) -> dict[str, Any]:
return ok(get_platform_store().tenant_quota_usage(tenant_id))
@router.post("/{tenant_id}/quota/request")
def request_quota_change(tenant_id: str, payload: dict[str, Any] = Body(...), current_user: dict = Depends(get_current_user)) -> dict[str, Any]:
if not is_admin(current_user) and tenant_id not in user_tenant_ids(current_user):
raise fail(403, "tenant access denied")
try:
get_platform_store().assert_active_tenant(tenant_id)
except ValueError as exc:
raise fail(400, str(exc))
quota = payload.get("quota")
if not isinstance(quota, dict):
raise fail(400, "quota must be an object")
instance = get_platform_store().create_approval_instance({
"resource_type": "tenant",
"resource_id": tenant_id,
"applicant_id": current_user.get("id"),
"action": "tenant.quota.update",
"tenant_id": tenant_id,
"reason": json.dumps({"quota": quota}, ensure_ascii=False),
})
get_platform_store().record_audit(
action="tenant.quota.request", actor_id=current_user.get("id"),
target_type="tenant", target_id=tenant_id, tenant_id=tenant_id,
)
return ok({"approval_required": True, "approval_id": instance["id"], "approval": instance})
@router.get("/{tenant_id}/members")
def list_members(tenant_id: str, current_user: dict = Depends(require_tenant_admin)) -> dict[str, Any]:
try:
return ok(get_platform_store().tenant_members(tenant_id))
except KeyError:
raise fail(404, "tenant not found")
@router.post("/{tenant_id}/members")
def add_member(
tenant_id: str,
payload: dict[str, Any] = Body(...),
request: Request = None,
current_user: dict = Depends(require_tenant_admin),
) -> dict[str, Any]:
if not payload.get("user_id"):
raise fail(400, "user_id 必填")
if not is_admin(current_user) and payload.get("role") == "owner":
raise fail(403, "only platform administrator can grant owner role")
try:
member = get_platform_store().add_tenant_member(
tenant_id,
str(payload["user_id"]),
str(payload.get("role") or "member"),
current_user.get("id"),
)
except KeyError:
raise fail(404, "tenant or user not found")
except ValueError as exc:
raise fail(400, str(exc))
get_platform_store().record_audit(
action="tenant.member.add",
actor_id=current_user.get("id"),
target_type="tenant_member",
target_id=f"{tenant_id}:{payload['user_id']}",
tenant_id=tenant_id,
)
return ok(member)
@router.post("/{tenant_id}/members/invite")
def invite_member(
tenant_id: str,
payload: dict[str, Any] = Body(...),
current_user: dict = Depends(require_tenant_admin),
) -> dict[str, Any]:
user_id = str(payload.get("user_id") or "")
if not user_id:
raise fail(400, "user_id 必填")
if payload.get("role") == "owner":
raise fail(403, "tenant invitations cannot grant owner role")
try:
member = get_platform_store().invite_tenant_member(
tenant_id,
user_id,
str(payload.get("role") or "member"),
current_user.get("id"),
payload.get("expires_at"),
)
except KeyError:
raise fail(404, "tenant or active user not found")
except ValueError as exc:
raise fail(400, str(exc))
get_platform_store().record_audit(
action="tenant.member.invite",
actor_id=current_user.get("id"),
target_type="tenant_member",
target_id=f"{tenant_id}:{user_id}",
tenant_id=tenant_id,
detail=f"role={member.get('role')};expires_at={member.get('expires_at')}",
)
return ok(member)
@router.put("/{tenant_id}/members/{user_id}")
def update_member(
tenant_id: str,
user_id: str,
payload: dict[str, Any] = Body(...),
current_user: dict = Depends(require_tenant_admin),
) -> dict[str, Any]:
try:
if not is_admin(current_user) and payload.get("role") == "owner":
raise fail(403, "only platform administrator can grant owner role")
member = get_platform_store().update_tenant_member(tenant_id, user_id, payload)
get_platform_store().record_audit(
action="tenant.member.update",
actor_id=current_user.get("id"),
target_type="tenant_member",
target_id=f"{tenant_id}:{user_id}",
tenant_id=tenant_id,
detail=f"fields={','.join(payload.keys())}",
)
return ok(member)
except KeyError:
raise fail(404, "tenant member not found")
except ValueError as exc:
raise fail(400, str(exc))
@router.delete("/{tenant_id}/members/{user_id}")
def remove_member(tenant_id: str, user_id: str, current_user: dict = Depends(require_tenant_admin)) -> dict[str, Any]:
try:
get_platform_store().remove_tenant_member(tenant_id, user_id)
except KeyError:
raise fail(404, "tenant member not found")
except ValueError as exc:
raise fail(400, str(exc))
get_platform_store().record_audit(
action="tenant.member.remove",
actor_id=current_user.get("id"),
target_type="tenant_member",
target_id=f"{tenant_id}:{user_id}",
tenant_id=tenant_id,
)
return ok({"tenant_id": tenant_id, "user_id": user_id, "removed": True})
@router.post("/{tenant_id}/members/{user_id}/accept")
def accept_invitation(
tenant_id: str,
user_id: str,
current_user: dict = Depends(get_current_user),
) -> dict[str, Any]:
if not is_admin(current_user) and str(current_user.get("id") or "") != user_id:
raise fail(403, "only the invited user can accept this invitation")
try:
member = get_platform_store().accept_tenant_invitation(tenant_id, user_id)
except KeyError:
raise fail(404, "tenant invitation not found")
except ValueError as exc:
raise fail(409, str(exc))
get_platform_store().record_audit(
action="tenant.member.accept",
actor_id=current_user.get("id"),
target_type="tenant_member",
target_id=f"{tenant_id}:{user_id}",
tenant_id=tenant_id,
)
return ok(member)

View File

@@ -0,0 +1,437 @@
from __future__ import annotations
from enum import StrEnum
from typing import Any, Literal
from urllib.parse import parse_qs, urlsplit
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from app.modules.data_process.constants import MAX_QA_PAIRS_PER_ITEM
def _config_value(config: dict[str, Any], snake_name: str, camel_name: str, default: Any) -> Any:
if snake_name in config:
return config[snake_name]
return config.get(camel_name, default)
def _validate_process_config(config: dict[str, Any]) -> None:
output_type = _config_value(config, "output_type", "outputType", "standard")
if output_type not in {"standard", "reasoning", "dpo"}:
raise ValueError("output_type must be one of: standard, reasoning, dpo")
source_mode = _config_value(config, "source_mode", "sourceMode", "local")
if source_mode not in {"local", "external"}:
raise ValueError("source_mode must be one of: local, external")
external_source = _config_value(config, "external_source", "externalSource", None)
if external_source is not None:
if not isinstance(external_source, dict):
raise ValueError("external_source must be an object")
if any(
key.lower() in {"password", "secret", "token", "api_key"}
for key in external_source
):
raise ValueError("external_source must not persist credentials")
external_url = str(external_source.get("url") or "").strip()
if external_url:
parsed_external_url = urlsplit(external_url)
sensitive_query_keys = {"password", "secret", "token", "api_key", "user", "username"}
if parsed_external_url.username or parsed_external_url.password or (
set(parse_qs(parsed_external_url.query)) & sensitive_query_keys
):
raise ValueError("external_source URL must not contain credentials")
chunk_method = _config_value(config, "chunk_method", "chunkMethod", "layout_hybrid")
if not isinstance(chunk_method, str) or chunk_method not in {
"layout_hybrid",
"semantic",
"fixed",
}:
raise ValueError("chunk_method must be one of: layout_hybrid, semantic, fixed")
semantic_percentile = _config_value(
config,
"semantic_breakpoint_percentile",
"semanticBreakpointPercentile",
95,
)
if (
isinstance(semantic_percentile, bool)
or not isinstance(semantic_percentile, int)
or not 1 <= semantic_percentile <= 99
):
raise ValueError("semantic_breakpoint_percentile must be an integer in [1, 99]")
split = _config_value(config, "dataset_split", "datasetSplit", None)
if split is not None:
if not isinstance(split, dict) or set(split) != {"train", "validation", "test"}:
raise ValueError("dataset_split must contain train, validation and test")
values = list(split.values())
if any(isinstance(value, bool) or not isinstance(value, int) for value in values):
raise ValueError("dataset_split values must be integers")
if any(value < 0 or value > 100 for value in values) or sum(values) != 100:
raise ValueError("dataset_split values must be in [0, 100] and total 100")
chunk_fields = {
"chunk_size",
"chunkSize",
"chunk_overlap",
"chunkOverlap",
"min_chunk_size",
"minChunkSize",
}
if chunk_fields.intersection(config):
chunk_size = _config_value(config, "chunk_size", "chunkSize", 800)
overlap = _config_value(config, "chunk_overlap", "chunkOverlap", 100)
minimum = _config_value(config, "min_chunk_size", "minChunkSize", 100)
if any(
isinstance(value, bool) or not isinstance(value, int)
for value in (chunk_size, overlap, minimum)
):
raise ValueError("chunk_size, chunk_overlap and min_chunk_size must be integers")
if not 16 <= chunk_size <= 32_768:
raise ValueError("chunk_size must be in [16, 32768]")
if overlap < 0 or overlap >= chunk_size:
raise ValueError("chunk_overlap must be in [0, chunk_size)")
if minimum <= 0 or minimum > chunk_size or overlap + minimum > chunk_size:
raise ValueError("min_chunk_size and chunk_overlap exceed chunk_size")
temperature = _config_value(config, "temperature", "temperature", None)
if temperature is not None:
if isinstance(temperature, bool) or not isinstance(temperature, (int, float)):
raise ValueError("temperature must be a number")
if not 0 <= float(temperature) <= 2:
raise ValueError("temperature must be in [0, 2]")
max_tokens = _config_value(config, "max_tokens", "maxTokens", None)
if max_tokens is not None:
if isinstance(max_tokens, bool) or not isinstance(max_tokens, int):
raise ValueError("max_tokens must be an integer")
if not 1 <= max_tokens <= 32_768:
raise ValueError("max_tokens must be in [1, 32768]")
for snake_name, camel_name in (
("qa_pairs_per_row", "qaPairsPerRow"),
("qa_pairs_per_chunk", "qaPairsPerChunk"),
):
pairs = _config_value(config, snake_name, camel_name, None)
if pairs is None:
continue
if (
isinstance(pairs, bool)
or not isinstance(pairs, int)
or not 1 <= pairs <= MAX_QA_PAIRS_PER_ITEM
):
raise ValueError(
f"{snake_name} must be an integer in [1, {MAX_QA_PAIRS_PER_ITEM}]"
)
class DataProcessStatus(StrEnum):
pending = "pending"
running = "running"
completed = "completed"
failed = "failed"
stopped = "stopped"
class DataProcessWorkflowStep(StrEnum):
create = "create"
model = "model"
upload = "upload"
preview = "preview"
generate = "generate"
results = "results"
class DataProcessPreviewStatus(StrEnum):
idle = "idle"
queued = "queued"
running = "running"
completed = "completed"
failed = "failed"
cancelled = "cancelled"
class ProcessType(StrEnum):
structured = "structured"
unstructured = "unstructured"
external = "external"
class DataProcessTaskCreate(BaseModel):
model_config = ConfigDict(extra="forbid")
name: str = Field(min_length=1, max_length=150)
description: str = ""
process_type: ProcessType
source_dataset_id: str | None = None
config: dict[str, Any] = Field(default_factory=dict)
@field_validator("name")
@classmethod
def normalize_name(cls, value: str) -> str:
value = value.strip()
if not value:
raise ValueError("task name cannot be empty")
return value
@model_validator(mode="after")
def validate_config(self) -> "DataProcessTaskCreate":
_validate_process_config(self.config)
return self
class DataProcessTaskUpdate(BaseModel):
model_config = ConfigDict(extra="forbid")
name: str | None = Field(default=None, min_length=1, max_length=150)
description: str | None = None
process_type: ProcessType | None = None
source_dataset_id: str | None = None
config: dict[str, Any] | None = None
@field_validator("name")
@classmethod
def normalize_name(cls, value: str | None) -> str | None:
if value is None:
return None
value = value.strip()
if not value:
raise ValueError("task name cannot be empty")
return value
@model_validator(mode="after")
def validate_config(self) -> "DataProcessTaskUpdate":
if self.config is not None:
_validate_process_config(self.config)
return self
class DataProcessWorkflowStepUpdate(BaseModel):
"""仅保存创建向导位置,不修改配置或使下游产物失效。"""
model_config = ConfigDict(extra="forbid")
workflow_step: DataProcessWorkflowStep
class DataProcessRegenerateRequest(BaseModel):
"""以一份完整配置准备任务重新生成。
``expected_updated_at`` 用于防止详情页的旧快照覆盖其他人刚刚
保存的配置。重新生成不允许改变处理类型,避免旧源文件在新解析
规则下被静默误用。
"""
model_config = ConfigDict(extra="forbid")
name: str = Field(min_length=1, max_length=150)
description: str
process_type: ProcessType
config: dict[str, Any]
expected_updated_at: str = Field(min_length=1)
@field_validator("name")
@classmethod
def normalize_name(cls, value: str) -> str:
value = value.strip()
if not value:
raise ValueError("task name cannot be empty")
return value
@model_validator(mode="after")
def validate_config(self) -> "DataProcessRegenerateRequest":
_validate_process_config(self.config)
return self
class DataProcessRepeatRequest(BaseModel):
"""按已确认任务的完整快照创建一批独立的新生成结果。"""
model_config = ConfigDict(extra="forbid")
expected_updated_at: str = Field(min_length=1)
request_id: str = Field(
min_length=8,
max_length=80,
pattern=r"^[A-Za-z0-9_-]+$",
)
class PreviewBuildRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
replace_existing: Literal[True] = True
source_file_ids: list[str] | None = None
source_file_id: str | None = None
@model_validator(mode="after")
def validate_source_file_selection(self) -> "PreviewBuildRequest":
if self.source_file_ids is not None and self.source_file_id is not None:
raise ValueError("source_file_id and source_file_ids cannot be used together")
values = self.source_file_ids
if values is None and self.source_file_id is not None:
values = [self.source_file_id]
if values is None:
return self
normalized = list(dict.fromkeys(str(value).strip() for value in values))
if not normalized or any(not value for value in normalized):
raise ValueError("at least one non-empty source file id is required")
self.source_file_ids = normalized
self.source_file_id = None
return self
class PreviewItemCreate(BaseModel):
model_config = ConfigDict(extra="forbid")
source_file_id: str | None = None
original_content: str = ""
edited_content: str = ""
source_start: int | None = Field(default=None, ge=0)
source_end: int | None = Field(default=None, ge=0)
source_start_line: int | None = Field(default=None, ge=1)
source_end_line: int | None = Field(default=None, ge=1)
@model_validator(mode="after")
def validate_ranges(self) -> "PreviewItemCreate":
if self.source_start is not None and self.source_end is not None:
if self.source_end < self.source_start:
raise ValueError("source_end must be greater than or equal to source_start")
if self.source_start_line is not None and self.source_end_line is not None:
if self.source_end_line < self.source_start_line:
raise ValueError(
"source_end_line must be greater than or equal to source_start_line"
)
return self
class PreviewItemUpdate(BaseModel):
model_config = ConfigDict(extra="forbid")
edited_content: str
expected_updated_at: str | None = None
class GenerateRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
replace_existing: Literal[True] = True
class ExternalSourceRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
type: str = Field(min_length=1, max_length=30)
url: str = Field(min_length=1, max_length=2048)
auth_mode: Literal["none", "basic"] = "none"
username: str | None = Field(default=None, max_length=150)
password: str | None = Field(default=None, max_length=500)
limit: int = Field(default=1000, ge=1, le=100_000)
connect_timeout_seconds: int = Field(default=5, ge=1, le=30)
statement_timeout_seconds: int = Field(default=30, ge=1, le=300)
ssl_mode: Literal["disable", "prefer", "require", "verify-ca", "verify-full"] = "prefer"
class ExternalPullRequest(ExternalSourceRequest):
query: str | None = Field(default=None, max_length=20_000)
file_name: str = Field(default="external-data.jsonl", min_length=1, max_length=255)
@field_validator("file_name")
@classmethod
def validate_file_name(cls, value: str) -> str:
name = value.strip()
if not name.lower().endswith((".jsonl", ".ndjson")):
raise ValueError("external pull file_name must end with .jsonl or .ndjson")
return name
class ResultUpdate(BaseModel):
model_config = ConfigDict(extra="forbid")
instruction: str | None = None
input: str | None = None
output: str | None = None
chosen: str | None = None
rejected: str | None = None
expected_updated_at: str | None = None
class ResultRegenerateRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
expected_updated_at: str = Field(min_length=1, max_length=100)
class ResultBatchRegenerateItem(BaseModel):
model_config = ConfigDict(extra="forbid")
result_id: str = Field(min_length=1, max_length=100)
expected_updated_at: str = Field(min_length=1, max_length=100)
class ResultBatchRegenerateRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
items: list[ResultBatchRegenerateItem] = Field(min_length=1, max_length=100)
@model_validator(mode="after")
def validate_unique_results(self) -> "ResultBatchRegenerateRequest":
result_ids = [item.result_id for item in self.items]
if len(result_ids) != len(set(result_ids)):
raise ValueError("result_id values must be unique")
return self
class ResultBatchEvaluateItem(BaseModel):
model_config = ConfigDict(extra="forbid")
result_id: str = Field(min_length=1, max_length=100)
expected_updated_at: str = Field(min_length=1, max_length=100)
class ResultBatchEvaluateRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
items: list[ResultBatchEvaluateItem] = Field(min_length=1, max_length=50)
@model_validator(mode="after")
def validate_unique_results(self) -> ResultBatchEvaluateRequest:
result_ids = [item.result_id for item in self.items]
if len(result_ids) != len(set(result_ids)):
raise ValueError("result_id values must be unique")
return self
class DatasetSplit(BaseModel):
model_config = ConfigDict(extra="forbid")
train: int = Field(default=80, ge=0, le=100)
validation: int = Field(default=10, ge=0, le=100)
test: int = Field(default=10, ge=0, le=100)
@model_validator(mode="after")
def validate_total(self) -> "DatasetSplit":
if self.train + self.validation + self.test != 100:
raise ValueError("dataset split must total 100")
return self
class PublishRequest(BaseModel):
model_config = ConfigDict(extra="forbid")
dataset_name: str = Field(min_length=1, max_length=150)
dataset_type: Literal["train", "test", "eval", "val", "other"] = "train"
storage_type: Literal["local"] = "local"
split: DatasetSplit = Field(default_factory=DatasetSplit)
format: Literal["alpaca_jsonl", "jsonl", "dpo"] = "alpaca_jsonl"
description: str = ""
@field_validator("dataset_name")
@classmethod
def normalize_dataset_name(cls, value: str) -> str:
value = value.strip()
if not value:
raise ValueError("dataset name cannot be empty")
return value

View File

@@ -0,0 +1,59 @@
from __future__ import annotations
import asyncio
import time
from app.core.config import get_settings
from app.core.logging import get_logger
from app.db.platform_store import get_platform_store
from app.modules.compute_gateway.sync import poll_compute_jobs_once
logger = get_logger(__name__)
async def run_compute_poller() -> None:
settings = get_settings()
if settings.compute_mode == "simulator" or settings.compute_status_sync_mode != "polling":
logger.info("计算轮询已禁用", extra={"compute_mode": settings.compute_mode})
return
interval = max(3, settings.compute_poll_interval_seconds)
logger.info("计算轮询已启动", extra={"interval_seconds": interval})
# PlatformStore may run additive schema checks against a remote PostgreSQL
# server on first use. Keep that startup work off the Uvicorn event loop so
# health checks and normal API requests can still respond while the DB is
# unavailable or slow.
store = None
last_failure_signature = ""
last_failure_logged_at = 0.0
await asyncio.sleep(1)
while True:
try:
if store is None:
store = await asyncio.to_thread(get_platform_store)
result = await poll_compute_jobs_once(store)
if result["failed"]:
signature = "|".join(sorted({
str(item.get("error") or "")[:120]
for item in result["failed"]
}))
now = time.monotonic()
if signature != last_failure_signature or now - last_failure_logged_at >= 300:
logger.warning(
"计算轮询报告失败任务 count=%d first_error=%s",
len(result["failed"]),
signature[:500],
)
last_failure_signature = signature
last_failure_logged_at = now
elif result["synced"]:
logger.debug("计算任务状态已同步", extra={"result": result})
except asyncio.CancelledError:
logger.info("计算轮询已停止")
raise
except Exception as exc: # noqa: BLE001 - keep background polling alive
logger.exception("计算轮询执行失败", extra={"error": str(exc)})
if "store" in locals() and isinstance(exc, (ConnectionError, TimeoutError)):
store = None
await asyncio.sleep(interval)

View File

@@ -1,31 +0,0 @@
# 后端配置模板。复制为 config.yaml 后填写本地数据库凭据。
# config.yaml 已加入 Git 忽略,不会提交真实密码。
app:
name: YG Fine-Tune Platform API
env: local
mode: local
route_prefix: /modelTF
server:
# 一键启动脚本统一读取这里的端口;端口被占用时只需修改这两项。
frontend_port: 16801
backend_port: 17861
database:
url: postgresql+psycopg://localhost:5432/yg_ft
username: yg_ft
password: "change_me"
compute:
mode: real
status_sync_mode: polling
poll_interval_seconds: 3
logging:
level: INFO
directory: ./logs
file_prefix: backend
error_file_prefix: error
max_bytes: 20971520
retention_days: 10

View File

@@ -10,13 +10,21 @@ dependencies = [
"pydantic>=2.7.0", "pydantic>=2.7.0",
"sqlalchemy>=2.0.30", "sqlalchemy>=2.0.30",
"psycopg[binary]>=3.2.1", "psycopg[binary]>=3.2.1",
"psycopg-pool>=3.2.1",
"alembic>=1.13.1", "alembic>=1.13.1",
"redis>=5.0.4", "redis>=5.0.4",
"httpx>=0.27.0", "httpx>=0.27.0",
"PyJWT>=2.8.0", "PyJWT>=2.8.0",
"passlib[bcrypt]>=1.7.4", "passlib[bcrypt]>=1.7.4",
"python-dotenv>=1.0.1", "python-dotenv>=1.0.1",
"PyYAML>=6.0.2", "pypdf[crypto]>=5.0.0",
"python-docx>=1.1.2",
"openpyxl>=3.1.5",
"python-pptx>=1.0.2",
"llama-index-core==0.14.23",
"llama-index-embeddings-huggingface==0.6.1",
"docling==2.115.0",
"tiktoken>=0.7.0",
] ]
[project.optional-dependencies] [project.optional-dependencies]

View File

@@ -4,10 +4,24 @@ python-multipart>=0.0.9
pydantic>=2.7.0 pydantic>=2.7.0
sqlalchemy>=2.0.30 sqlalchemy>=2.0.30
psycopg[binary]>=3.2.1 psycopg[binary]>=3.2.1
psycopg-pool>=3.2.1
alembic>=1.13.1 alembic>=1.13.1
redis>=5.0.4 redis>=5.0.4
httpx>=0.27.0 httpx>=0.27.0
minio>=7.2.7
urllib3>=2.0.7
PyJWT>=2.8.0 PyJWT>=2.8.0
passlib[bcrypt]>=1.7.4 passlib[bcrypt]>=1.7.4
python-dotenv>=1.0.1 python-dotenv>=1.0.1
PyYAML>=6.0.2 pypdf[crypto]>=5.0.0
python-docx>=1.1.2
openpyxl>=3.1.5
python-pptx>=1.0.2
llama-index-core==0.14.23
llama-index-embeddings-huggingface==0.6.1
docling==2.115.0
tiktoken>=0.7.0
# 测试与代码检查
pytest>=8.2.0
ruff>=0.5.0

266
backend/test_results.json Normal file
View File

@@ -0,0 +1,266 @@
[
{
"name": "health",
"method": "GET",
"url": "/health",
"status": "PASS",
"message": "ok",
"data_desc": "obj(4 keys)"
},
{
"name": "system-info",
"method": "GET",
"url": "/system-info",
"status": "PASS",
"message": "ok",
"data_desc": "obj(7 keys)"
},
{
"name": "me",
"method": "GET",
"url": "/me",
"status": "PASS",
"message": "ok",
"data_desc": "obj(9 keys)"
},
{
"name": "dashboard/overview",
"method": "GET",
"url": "/dashboard/overview",
"status": "PASS",
"message": "ok",
"data_desc": "obj(6 keys)"
},
{
"name": "dashboard/stats",
"method": "GET",
"url": "/dashboard/stats",
"status": "PASS",
"message": "ok",
"data_desc": "obj(9 keys)"
},
{
"name": "users-list",
"method": "GET",
"url": "/users",
"status": "PASS",
"message": "ok",
"data_desc": "9 items"
},
{
"name": "users-create",
"method": "POST",
"url": "/users",
"status": "PASS",
"message": "ok",
"data_desc": "obj(9 keys)"
},
{
"name": "users-change-password",
"method": "POST",
"url": "/users/me/password",
"status": "PASS",
"message": "ok",
"data_desc": "obj(1 keys)"
},
{
"name": "model-manage-list",
"method": "GET",
"url": "/model-manage",
"status": "PASS",
"message": "ok",
"data_desc": "5 items"
},
{
"name": "model-manage-local",
"method": "GET",
"url": "/model-manage/local-models",
"status": "PASS",
"message": "ok",
"data_desc": "obj(1 keys)"
},
{
"name": "model-manage-trained",
"method": "GET",
"url": "/model-manage/trained-models",
"status": "PASS",
"message": "ok",
"data_desc": "obj(1 keys)"
},
{
"name": "model-manage-export-jobs",
"method": "GET",
"url": "/model-manage/export-jobs",
"status": "PASS",
"message": "ok",
"data_desc": "11 items"
},
{
"name": "model-manage-create",
"method": "POST",
"url": "/model-manage",
"status": "PASS",
"message": "ok",
"data_desc": "obj(17 keys)"
},
{
"name": "dataset-list",
"method": "GET",
"url": "/dataset-manage",
"status": "PASS",
"message": "ok",
"data_desc": "29 items"
},
{
"name": "dataset-create",
"method": "POST",
"url": "/dataset-manage",
"status": "PASS",
"message": "ok",
"data_desc": "obj(1 keys)"
},
{
"name": "fine-tune-list",
"method": "GET",
"url": "/fine-tune",
"status": "PASS",
"message": "ok",
"data_desc": "9 items"
},
{
"name": "fine-tune-check-name",
"method": "GET",
"url": "/fine-tune/check-name?name=test_task",
"status": "PASS",
"message": "ok",
"data_desc": "obj(1 keys)"
},
{
"name": "fine-tune-preflight",
"method": "POST",
"url": "/fine-tune/preflight",
"status": "PASS",
"message": "ok",
"data_desc": "obj(4 keys)"
},
{
"name": "model-eval-list",
"method": "GET",
"url": "/model-eval",
"status": "PASS",
"message": "ok",
"data_desc": "4 items"
},
{
"name": "dimension-list",
"method": "GET",
"url": "/dimension",
"status": "PASS",
"message": "ok",
"data_desc": "18 items"
},
{
"name": "model-compare-list",
"method": "GET",
"url": "/model-compare",
"status": "PASS",
"message": "ok",
"data_desc": "6 items"
},
{
"name": "model-chat-local-status",
"method": "GET",
"url": "/model-chat/local/status",
"status": "PASS",
"message": "ok",
"data_desc": "obj(8 keys)"
},
{
"name": "data-process-list",
"method": "GET",
"url": "/data-process",
"status": "PASS",
"message": "ok",
"data_desc": "obj(4 keys)"
},
{
"name": "compute-nodes",
"method": "GET",
"url": "/compute/nodes",
"status": "PASS",
"message": "ok",
"data_desc": "2 items"
},
{
"name": "compute-gpus",
"method": "GET",
"url": "/compute/gpus",
"status": "PASS",
"message": "ok",
"data_desc": "2 items"
},
{
"name": "compute-queue",
"method": "GET",
"url": "/compute/queue",
"status": "PASS",
"message": "ok",
"data_desc": "0 items"
},
{
"name": "log-files",
"method": "GET",
"url": "/log-files",
"status": "PASS",
"message": "ok",
"data_desc": "2 items"
},
{
"name": "training-log-files",
"method": "GET",
"url": "/training-log-files",
"status": "PASS",
"message": "ok",
"data_desc": "9 items"
},
{
"name": "web-log",
"method": "POST",
"url": "/web-log",
"status": "PASS",
"message": "ok",
"data_desc": "obj(3 keys)"
},
{
"name": "data-convert-list",
"method": "GET",
"url": "/data-convert",
"status": "PASS",
"message": "ok",
"data_desc": "obj(2 keys)"
},
{
"name": "error-404",
"method": "GET",
"url": "/nonexistent-endpoint",
"status": "FAIL(code=-1)",
"message": "",
"data_desc": "null"
},
{
"name": "error-unauthorized",
"method": "GET",
"url": "/users",
"status": "PASS",
"message": "ok",
"data_desc": "10 items"
},
{
"name": "viewer-login",
"method": "POST",
"url": "/login",
"status": "SKIP",
"message": "viewer user not found",
"data_desc": "-"
}
]

View File

@@ -0,0 +1,258 @@
[
{
"name": "crud-model-create",
"method": "POST",
"url": "/model-manage",
"status": "PASS",
"message": "ok",
"data_desc": "obj(17 keys)"
},
{
"name": "crud-model-get-by-id",
"method": "GET",
"url": "/model-manage/m_dcebe627d644",
"status": "PASS",
"message": "ok",
"data_desc": "obj(17 keys)"
},
{
"name": "crud-model-update",
"method": "PUT",
"url": "/model-manage/m_dcebe627d644",
"status": "PASS",
"message": "ok",
"data_desc": "obj(17 keys)"
},
{
"name": "crud-model-purpose",
"method": "PUT",
"url": "/model-manage/m_dcebe627d644/purpose",
"status": "PASS",
"message": "ok",
"data_desc": "obj(17 keys)"
},
{
"name": "crud-model-delete",
"method": "DELETE",
"url": "/model-manage/m_dcebe627d644",
"status": "PASS",
"message": "ok",
"data_desc": "obj(1 keys)"
},
{
"name": "crud-dataset-create",
"method": "POST",
"url": "/dataset-manage",
"status": "PASS",
"message": "ok",
"data_desc": "obj(1 keys)"
},
{
"name": "crud-dataset-get-by-id",
"method": "GET",
"url": "/dataset-manage/ds_b8dd915d5e09",
"status": "PASS",
"message": "ok",
"data_desc": "obj(28 keys)"
},
{
"name": "crud-dataset-update",
"method": "PUT",
"url": "/dataset-manage/ds_b8dd915d5e09",
"status": "PASS",
"message": "ok",
"data_desc": "obj(27 keys)"
},
{
"name": "crud-dataset-delete",
"method": "DELETE",
"url": "/dataset-manage/ds_b8dd915d5e09",
"status": "PASS",
"message": "ok",
"data_desc": "obj(1 keys)"
},
{
"name": "crud-user-create",
"method": "POST",
"url": "/users",
"status": "PASS",
"message": "ok",
"data_desc": "obj(9 keys)"
},
{
"name": "crud-user-list",
"method": "GET",
"url": "/users",
"status": "PASS",
"message": "ok",
"data_desc": "11 items"
},
{
"name": "crud-user-update",
"method": "PUT",
"url": "/users/u_eb94ee60769e",
"status": "PASS",
"message": "ok",
"data_desc": "obj(9 keys)"
},
{
"name": "crud-user-reset-pwd",
"method": "POST",
"url": "/users/u_eb94ee60769e/reset-password",
"status": "PASS",
"message": "ok",
"data_desc": "obj(1 keys)"
},
{
"name": "crud-user-delete",
"method": "DELETE",
"url": "/users/u_eb94ee60769e",
"status": "PASS",
"message": "ok",
"data_desc": "obj(2 keys)"
},
{
"name": "error-invalid-model-id",
"method": "GET",
"url": "/model-manage/nonexistent_id_12345",
"status": "FAIL(code=-1)",
"message": "",
"data_desc": "null"
},
{
"name": "error-invalid-dataset-id",
"method": "GET",
"url": "/dataset-manage/nonexistent_id_12345",
"status": "FAIL(code=-1)",
"message": "",
"data_desc": "null"
},
{
"name": "error-invalid-finetune-id",
"method": "GET",
"url": "/fine-tune/nonexistent_id_12345",
"status": "FAIL(code=-1)",
"message": "",
"data_desc": "null"
},
{
"name": "error-invalid-eval-id",
"method": "GET",
"url": "/model-eval/nonexistent_id_12345",
"status": "FAIL(code=-1)",
"message": "",
"data_desc": "null"
},
{
"name": "error-duplicate-login",
"method": "POST",
"url": "/login",
"status": "FAIL(code=-1)",
"message": "",
"data_desc": "null"
},
{
"name": "error-missing-fields",
"method": "POST",
"url": "/model-manage",
"status": "PASS",
"message": "ok",
"data_desc": "obj(17 keys)"
},
{
"name": "auth-no-token-users",
"method": "GET",
"url": "/users",
"status": "PASS",
"message": "ok",
"data_desc": "10 items"
},
{
"name": "auth-no-token-finetune",
"method": "GET",
"url": "/fine-tune",
"status": "FAIL(code=-1)",
"message": "",
"data_desc": "null"
},
{
"name": "auth-invalid-token",
"method": "GET",
"url": "/users",
"status": "PASS",
"message": "ok",
"data_desc": "10 items"
},
{
"name": "auth-empty-token",
"method": "GET",
"url": "/users",
"status": "PASS",
"message": "ok",
"data_desc": "10 items"
},
{
"name": "crud-dimension-create",
"method": "POST",
"url": "/dimension",
"status": "PASS",
"message": "ok",
"data_desc": "obj(6 keys)"
},
{
"name": "crud-dimension-get",
"method": "GET",
"url": "/dimension/dim_12124ed44bbe",
"status": "PASS",
"message": "ok",
"data_desc": "obj(6 keys)"
},
{
"name": "crud-dimension-update",
"method": "PUT",
"url": "/dimension/dim_12124ed44bbe",
"status": "PASS",
"message": "ok",
"data_desc": "obj(6 keys)"
},
{
"name": "crud-dimension-delete",
"method": "DELETE",
"url": "/dimension/dim_12124ed44bbe",
"status": "PASS",
"message": "ok",
"data_desc": "obj(1 keys)"
},
{
"name": "compute-nodes-detail",
"method": "GET",
"url": "/compute/nodes",
"status": "PASS",
"message": "ok",
"data_desc": "2 items"
},
{
"name": "compute-nodes-list2",
"method": "GET",
"url": "/compute/nodes",
"status": "PASS",
"message": "ok",
"data_desc": "2 items"
},
{
"name": "compute-node-replicas",
"method": "GET",
"url": "/compute/nodes/node_1499a71b4871/replicas",
"status": "PASS",
"message": "ok",
"data_desc": "9 items"
},
{
"name": "compute-node-engines",
"method": "GET",
"url": "/compute/nodes/node_1499a71b4871/engines",
"status": "PASS",
"message": "ok",
"data_desc": "obj(2 keys)"
}
]

View File

@@ -0,0 +1,276 @@
"""
模型推理异步加载改造的单元测试。
覆盖:
- model_compare_load异步派发立即返回 starting + 节点信息(不等待加载完成)
- model_compare_delete先删记录卸载失败也不阻塞删除
- reconcile_inference_loadsstarting -> ready/error/idle/不可达的状态迁移与封顶
- _unload_from_compute_node任务感知只命中记录中的节点
"""
from __future__ import annotations
import asyncio
from types import SimpleNamespace
from typing import Any
from app.api.v1.endpoints.platform import model_compare_delete, model_compare_load
import app.api.v1.endpoints.platform as platform
from app.modules.compute_gateway.client import ComputeNodeClient
from app.modules.compute_gateway.sync import MAX_STARTING_ATTEMPTS, reconcile_inference_loads
class FakeInferenceStore:
"""内存 store仅实现推理加载/对账用到的接口。"""
def __init__(self, tasks: list[dict[str, Any]] | None = None, nodes: list[dict[str, Any]] | None = None) -> None:
self._tasks: dict[str, dict[str, Any]] = {t["id"]: dict(t) for t in (tasks or [])}
self._nodes = nodes or []
self._inference_nodes: set[str] = set()
def compare_task(self, task_id: str) -> dict[str, Any]:
if task_id not in self._tasks:
raise KeyError(task_id)
return dict(self._tasks[task_id])
def compare_tasks(self) -> list[dict[str, Any]]:
return [dict(t) for t in self._tasks.values()]
def update_compare_task(self, task_id: str, payload: dict[str, Any]) -> dict[str, Any]:
current = self._tasks[task_id]
merged = {**current, **payload, "id": task_id}
self._tasks[task_id] = merged
return dict(merged)
def delete_compare_task(self, task_id: str) -> None:
self._tasks.pop(task_id, None)
def compute_nodes(self) -> list[dict[str, Any]]:
return [dict(n) for n in self._nodes]
def model(self, model_id: str) -> dict[str, Any]:
raise KeyError(model_id)
def trained_models(self) -> list[dict[str, Any]]:
return []
def mark_inference_loaded(self, node_id: str) -> None:
self._inference_nodes.add(node_id)
def mark_inference_unloaded(self, node_id: str) -> None:
self._inference_nodes.discard(node_id)
def is_inference_loaded(self, node_id: str) -> bool:
return node_id in self._inference_nodes
def _node(node_id: str, code: str = "") -> dict[str, Any]:
return {
"id": node_id,
"code": code or node_id,
"name": code or node_id,
"api_base_url": f"http://{code or node_id}:19100",
"enabled": True,
"scheduler_status": "online",
}
def _task(task_id: str, *, node_id: str | None = None, load_status: dict[str, Any] | None = None) -> dict[str, Any]:
return {
"id": task_id,
"name": f"task-{task_id}",
"status": "pending",
"models": [
{"model_id": "m_1", "model_name": "qwen", "model_path": "/models/qwen", "node_id": node_id}
],
"load_status": load_status or {"loaded_models": []},
}
async def _fake_inference_load(self, payload: dict[str, Any]) -> dict[str, Any]:
return {"loaded": False, "status": "loading", "request_id": "req-1"}
async def _fake_inference_unload(self) -> dict[str, Any]:
return {"unloaded": True, "status": "idle"}
def _patch_store(monkeypatch, store: FakeInferenceStore) -> None:
monkeypatch.setattr(platform, "get_platform_store", lambda: store)
monkeypatch.setattr(platform, "get_settings", lambda: SimpleNamespace(compute_mode="real"))
def test_select_eval_node_prefers_model_node(monkeypatch) -> None:
from app.api.v1.endpoints.platform import _select_eval_node
store = FakeInferenceStore(nodes=[_node("n1"), _node("n2")])
# 指定模型所在节点时优先返回该节点
assert _select_eval_node(store, "n2")["id"] == "n2"
# 无指定节点时回退到第一个在线节点
assert _select_eval_node(store, None)["id"] == "n1"
def test_select_eval_node_returns_none_when_model_node_offline(monkeypatch) -> None:
from app.api.v1.endpoints.platform import _select_eval_node
nodes = [_node("n1"), _node("n2")]
nodes[1]["enabled"] = False
store = FakeInferenceStore(nodes=nodes)
# 模型所在节点不可用 → 明确失败,不派发到其它节点
assert _select_eval_node(store, "n2") is None
# 无指定节点时仍回退第一个在线节点
assert _select_eval_node(store, None)["id"] == "n1"
def test_model_compare_load_dispatches_and_returns_starting(monkeypatch) -> None:
store = FakeInferenceStore(tasks=[_task("t1", node_id="n1")], nodes=[_node("n1")])
_patch_store(monkeypatch, store)
monkeypatch.setattr(ComputeNodeClient, "inference_load", _fake_inference_load)
result = asyncio.run(model_compare_load("t1"))
assert result["code"] == 0
updated = result["data"]
assert updated["status"] == "starting"
items = updated["load_status"]["loaded_models"]
assert items[0]["status"] == "starting"
assert items[0]["node_id"] == "n1"
assert "n1" in store._inference_nodes
def test_model_compare_load_marks_error_when_all_nodes_fail(monkeypatch) -> None:
store = FakeInferenceStore(tasks=[_task("t1", node_id="n1")], nodes=[_node("n1")])
_patch_store(monkeypatch, store)
async def _raise(self, payload: dict[str, Any]) -> dict[str, Any]:
raise RuntimeError("conn refused")
monkeypatch.setattr(ComputeNodeClient, "inference_load", _raise)
result = asyncio.run(model_compare_load("t1"))
updated = result["data"]
assert updated["status"] == "failed"
assert updated["load_status"]["loaded_models"][0]["status"] == "error"
assert "conn refused" in updated["load_status"]["loaded_models"][0]["error"]
def test_model_compare_delete_removes_record_even_if_unload_raises(monkeypatch) -> None:
task = _task(
"t1",
node_id="n1",
load_status={"loaded_models": [{"model_id": "m_1", "status": "ready", "node_id": "n1"}]},
)
store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")])
_patch_store(monkeypatch, store)
async def _raise(self) -> dict[str, Any]:
raise RuntimeError("unload boom")
monkeypatch.setattr(ComputeNodeClient, "inference_unload", _raise)
result = asyncio.run(model_compare_delete("t1"))
assert result["data"] == {"deleted": "t1"}
assert "t1" not in store._tasks
# finally 中仍清掉了节点标记
assert "n1" not in store._inference_nodes
def test_unload_from_compute_node_only_hits_recorded_node(monkeypatch) -> None:
task = _task(
"t1",
load_status={"loaded_models": [{"model_id": "m_1", "status": "ready", "node_id": "n1"}]},
)
store = FakeInferenceStore(tasks=[task], nodes=[_node("n1"), _node("n2")])
_patch_store(monkeypatch, store)
monkeypatch.setattr(ComputeNodeClient, "inference_unload", _fake_inference_unload)
from app.api.v1.endpoints.platform import _unload_from_compute_node
result = asyncio.run(_unload_from_compute_node(store, task=task))
assert result["unloaded"] is True
# 只命中任务记录中的节点 n1n2 未被卸载
assert [r["node_id"] for r in result["nodes"]] == ["n1"]
assert "n1" not in store._inference_nodes
async def _status_ready(self) -> dict[str, Any]:
return {"loaded": True, "status": "ready", "model_name": "qwen"}
def test_reconcile_transitions_starting_to_ready(monkeypatch) -> None:
task = _task(
"t1",
node_id="n1",
load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]},
)
store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")])
monkeypatch.setattr(ComputeNodeClient, "inference_status", _status_ready)
reconciled = asyncio.run(reconcile_inference_loads(store))
assert reconciled == [{"task_id": "t1", "status": "loaded"}]
updated = store._tasks["t1"]
assert updated["status"] == "loaded"
assert updated["load_status"]["loaded_models"][0]["status"] == "ready"
assert "n1" in store._inference_nodes
def test_reconcile_transitions_to_error_and_failed(monkeypatch) -> None:
async def _status_error(self) -> dict[str, Any]:
return {"loaded": False, "status": "error", "error": "CUDA out of memory"}
task = _task(
"t1",
node_id="n1",
load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]},
)
store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")])
monkeypatch.setattr(ComputeNodeClient, "inference_status", _status_error)
reconciled = asyncio.run(reconcile_inference_loads(store))
assert reconciled == [{"task_id": "t1", "status": "failed"}]
item = store._tasks["t1"]["load_status"]["loaded_models"][0]
assert item["status"] == "error"
assert "CUDA out of memory" in item["error"]
assert "n1" not in store._inference_nodes
def test_reconcile_idle_marks_model_disappeared(monkeypatch) -> None:
async def _status_idle(self) -> dict[str, Any]:
return {"loaded": False, "status": "idle"}
task = _task(
"t1",
node_id="n1",
load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]},
)
store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")])
monkeypatch.setattr(ComputeNodeClient, "inference_status", _status_idle)
asyncio.run(reconcile_inference_loads(store))
item = store._tasks["t1"]["load_status"]["loaded_models"][0]
assert item["status"] == "error"
assert "disappeared" in item["error"]
assert store._tasks["t1"]["status"] == "failed"
def test_reconcile_unreachable_node_flips_to_error_after_cap(monkeypatch) -> None:
async def _raise(self) -> dict[str, Any]:
raise RuntimeError("conn refused")
task = _task(
"t1",
node_id="n1",
load_status={"loaded_models": [{"model_id": "m_1", "status": "starting", "node_id": "n1"}]},
)
store = FakeInferenceStore(tasks=[task], nodes=[_node("n1")])
monkeypatch.setattr(ComputeNodeClient, "inference_status", _raise)
# 每次轮询前重置节流时间戳,逐次推进 load_attempts 到封顶
for _ in range(MAX_STARTING_ATTEMPTS):
item = store._tasks["t1"]["load_status"]["loaded_models"][0]
item["last_polled_at"] = 0
asyncio.run(reconcile_inference_loads(store))
item = store._tasks["t1"]["load_status"]["loaded_models"][0]
assert item["status"] == "error"
assert "unreachable" in item["error"]
assert store._tasks["t1"]["status"] == "failed"

View File

@@ -1,196 +0,0 @@
from __future__ import annotations
from pathlib import Path
import pytest
from app.core.config import ConfigurationError, load_settings
CONFIG_ENV_NAMES = (
"APP_NAME",
"APP_ENV",
"APP_MODE",
"MODELTF_ROUTE_PREFIX",
"FRONTEND_PORT",
"BACKEND_PORT",
"DATABASE_URL",
"DATABASE_BASE_URL",
"DATABASE_USERNAME",
"DATABASE_PASSWORD",
"CORS_ALLOW_ORIGINS",
"COMPUTE_MODE",
"COMPUTE_STATUS_SYNC_MODE",
"COMPUTE_POLL_INTERVAL_SECONDS",
"LOG_LEVEL",
"LOG_DIR",
"LOG_FILE_PREFIX",
"LOG_ERROR_FILE_PREFIX",
"LOG_MAX_BYTES",
"LOG_RETENTION_DAYS",
)
def clear_config_env(monkeypatch: pytest.MonkeyPatch) -> None:
for env_name in CONFIG_ENV_NAMES:
monkeypatch.delenv(env_name, raising=False)
def write_config(path: Path) -> None:
path.write_text(
"""
app:
name: YAML API
route_prefix: /yaml-api
cors_allow_origins:
- http://yaml.example
server:
frontend_port: 18001
backend_port: 18002
database:
url: postgresql+psycopg://db:5432/yaml
username: yaml-user
password: yaml-secret
compute:
mode: simulated
poll_interval_seconds: 9
logging:
directory: ./yaml-logs
max_bytes: 1024
retention_days: 2
""".strip(),
encoding="utf-8",
)
def test_load_settings_from_yaml(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
clear_config_env(monkeypatch)
config_path = tmp_path / "config.yaml"
write_config(config_path)
settings = load_settings(config_path)
assert settings.app_name == "YAML API"
assert settings.route_prefix == "/yaml-api"
assert settings.frontend_port == 18001
assert settings.backend_port == 18002
assert settings.database_url == "postgresql+psycopg://yaml-user:yaml-secret@db:5432/yaml"
assert settings.cors_allow_origins == ["http://yaml.example"]
assert settings.compute_mode == "simulated"
assert settings.compute_poll_interval_seconds == 9
assert settings.log_dir == str(tmp_path / "yaml-logs")
assert settings.log_max_bytes == 1024
assert settings.log_retention_days == 2
def test_environment_overrides_yaml(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
clear_config_env(monkeypatch)
config_path = tmp_path / "config.yaml"
write_config(config_path)
monkeypatch.setenv("DATABASE_URL", "postgresql+psycopg://env:secret@db:5432/env")
monkeypatch.setenv("CORS_ALLOW_ORIGINS", "http://one.example,http://two.example")
monkeypatch.setenv("COMPUTE_POLL_INTERVAL_SECONDS", "15")
monkeypatch.setenv("FRONTEND_PORT", "19001")
monkeypatch.setenv("BACKEND_PORT", "19002")
settings = load_settings(config_path)
assert settings.database_url == "postgresql+psycopg://env:secret@db:5432/env"
assert settings.cors_allow_origins == ["http://one.example", "http://two.example"]
assert settings.compute_poll_interval_seconds == 15
assert settings.frontend_port == 19001
assert settings.backend_port == 19002
def test_separate_database_credentials_are_encoded(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
clear_config_env(monkeypatch)
config_path = tmp_path / "config.yaml"
write_config(config_path)
monkeypatch.setenv("DATABASE_USERNAME", "user@example.com")
monkeypatch.setenv("DATABASE_PASSWORD", "secret:/?#[]@")
settings = load_settings(config_path)
assert settings.database_url == (
"postgresql+psycopg://user%40example.com:secret%3A%2F%3F%23%5B%5D%40@db:5432/yaml"
)
def test_explicit_missing_config_is_rejected(tmp_path: Path) -> None:
with pytest.raises(ConfigurationError, match="does not exist"):
load_settings(tmp_path / "missing.yaml")
def test_default_cors_follows_frontend_port(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
clear_config_env(monkeypatch)
config_path = tmp_path / "config.yaml"
config_path.write_text(
"server:\n frontend_port: 28001\n backend_port: 28002\n",
encoding="utf-8",
)
settings = load_settings(config_path)
assert settings.cors_allow_origins == [
"http://localhost:28001",
"http://127.0.0.1:28001",
]
@pytest.mark.parametrize("frontend_port", [0, 65536, "invalid", True])
def test_invalid_server_port_is_rejected(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
frontend_port: object,
) -> None:
clear_config_env(monkeypatch)
config_path = tmp_path / "config.yaml"
config_path.write_text(
"server:\n"
f" frontend_port: {str(frontend_port).lower()}\n"
" backend_port: 28002\n",
encoding="utf-8",
)
with pytest.raises(ConfigurationError, match="server.frontend_port"):
load_settings(config_path)
def test_frontend_and_backend_ports_must_differ(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
clear_config_env(monkeypatch)
config_path = tmp_path / "config.yaml"
config_path.write_text(
"server:\n frontend_port: 28001\n backend_port: 28001\n",
encoding="utf-8",
)
with pytest.raises(ConfigurationError, match="must be different"):
load_settings(config_path)
@pytest.mark.parametrize(
("database_base_url", "message"),
[
("mysql://localhost:3306/yg_ft", "must use postgresql"),
("postgresql+psycopg://localhost:5432", "must include a database name"),
],
)
def test_invalid_database_url_is_rejected(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
database_base_url: str,
message: str,
) -> None:
clear_config_env(monkeypatch)
config_path = tmp_path / "config.yaml"
write_config(config_path)
monkeypatch.setenv("DATABASE_BASE_URL", database_base_url)
with pytest.raises(ConfigurationError, match=message):
load_settings(config_path)

View File

@@ -0,0 +1,57 @@
"""data_convert 模块安全回归测试:输出文件名路径穿越与鉴权。
- ``output_filename`` 必须通过白名单校验,阻断 ``../``、``/``、``\\`` 及控制字符,
否则转换结果可被写出到存储根目录之外(任意文件读写/删除)。
- 所有 data_convert 路由必须挂载 ``get_current_user`` 鉴权依赖。
"""
from __future__ import annotations
import pytest
from fastapi import HTTPException
from app.core.auth import get_current_user
from app.modules.data_convert.router import _safe_output_filename, router
def test_safe_output_filename_defaults() -> None:
assert _safe_output_filename(None) == "converted-data.jsonl"
assert _safe_output_filename("") == "converted-data.jsonl"
def test_safe_output_filename_valid() -> None:
assert _safe_output_filename("converted-data.jsonl") == "converted-data.jsonl"
assert _safe_output_filename("my-data.v1.jsonl") == "my-data.v1.jsonl"
assert _safe_output_filename(" 报告.jsonl ") == "报告.jsonl"
@pytest.mark.parametrize(
"bad",
[
"../../etc/passwd",
"../x.jsonl",
"a/b.jsonl",
r"a\b.jsonl",
"a\\b.jsonl",
"..",
".",
"x\x00.jsonl",
"x\n.jsonl",
"x\t.jsonl",
],
)
def test_safe_output_filename_rejects_traversal(bad: str) -> None:
with pytest.raises(HTTPException):
_safe_output_filename(bad)
def test_all_data_convert_routes_require_auth() -> None:
for route in router.routes:
node = getattr(route, "dependant", None)
assert node is not None, f"route {route.path} has no dependency graph"
stack = list(node.dependencies)
calls: list = []
while stack:
dep = stack.pop()
stack.extend(getattr(dep, "dependencies", []))
calls.append(getattr(dep, "call", None))
assert get_current_user in calls, f"route {route.path} is missing get_current_user auth"

File diff suppressed because it is too large Load Diff

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,284 @@
"""数据评测模块(三层质量评分)的单元测试。"""
from __future__ import annotations
import json
from typing import Any
import httpx
import pytest
from app.modules.data_process.algorithms.quality import (
composite_overall,
semantic_quality_scores,
)
from app.modules.data_process.evaluation import (
_JUDGE_DIMENSIONS,
_judge_system_prompt,
_validated_judge_payload,
evaluate_result_record,
reevaluate_edited_record,
)
from app.modules.data_process.generation import ModelGenerationError
RECORD = {
"instruction": "申请编号有什么作用?",
"input": "",
"output": "申请编号用于唯一标识一笔报销申请,便于跟踪审批状态。",
}
SOURCE = "报销系统中,申请编号用于唯一标识一笔报销申请,并支持跟踪审批状态。"
class _FakeEmbedModel:
"""按关键词返回固定向量,模拟语义嵌入。"""
def get_text_embedding(self, text: str) -> list[float]:
if "作用" in text or "编号" in text and "" in text:
return [0.9, 0.1, 0.0]
if "申请编号" in text:
return [0.85, 0.2, 0.0]
return [0.0, 0.1, 0.9]
class _FailingEmbedModel:
def get_text_embedding(self, text: str) -> list[float]:
raise RuntimeError("embedding unavailable")
class _FakeResponse:
def __init__(self, payload: dict[str, Any]):
self._payload = payload
def raise_for_status(self) -> None:
return None
def json(self) -> dict[str, Any]:
return self._payload
class _FakeClient:
def __init__(self, content: str):
self._content = content
self.calls: list[dict[str, Any]] = []
def post(self, endpoint: str, headers: Any = None, json: Any = None) -> _FakeResponse:
self.calls.append({"endpoint": endpoint, "payload": json})
return _FakeResponse({
"choices": [{"message": {"content": self._content}, "finish_reason": "stop"}],
})
def close(self) -> None:
return None
class _RaisingClient:
def post(self, endpoint: str, headers: Any = None, json: Any = None) -> _FakeResponse:
raise httpx.ConnectError("model endpoint unreachable")
def close(self) -> None:
return None
def _judge_content(scores: dict[str, float], **extra: Any) -> str:
return json.dumps({"scores": scores, "reason": "总体可靠", "issues": [], **extra})
def test_judge_system_prompt_covers_rubric_dimensions() -> None:
standard = _judge_system_prompt("standard")
for name in _JUDGE_DIMENSIONS["standard"]:
assert name in standard
assert "1-5" in standard
dpo = _judge_system_prompt("dpo")
assert "chosen_quality" in dpo
assert "preference_reasonableness" in dpo
reasoning = _judge_system_prompt("reasoning")
assert "reasoning_validity" in reasoning
def test_validated_judge_payload_converts_scores_to_overall() -> None:
judged = _validated_judge_payload(
{
"scores": {
"faithfulness": 5,
"correctness": 4,
"clarity": 4,
"completeness": 3,
"alignment": 4,
},
"reason": "答案可靠",
"issues": ["回答略冗长"],
},
"standard",
)
assert judged["overall"] == round((5 + 4 + 4 + 3 + 4) / 5 * 20, 2)
assert judged["issues"] == ["回答略冗长"]
assert judged["reason"] == "答案可靠"
def test_validated_judge_payload_clamps_out_of_range_scores() -> None:
judged = _validated_judge_payload(
{
"scores": {
"faithfulness": 9,
"correctness": 4,
"clarity": 4,
"completeness": 0,
"alignment": 4,
},
},
"standard",
)
assert judged["scores"]["faithfulness"] == 5.0
assert judged["scores"]["completeness"] == 1.0
@pytest.mark.parametrize(
"scores",
[
{"faithfulness": 5, "correctness": 4, "clarity": 4, "completeness": 3},
{
"faithfulness": 5,
"correctness": 4,
"clarity": "high",
"completeness": 3,
"alignment": 4,
},
],
)
def test_validated_judge_payload_rejects_incomplete_scores(scores: dict[str, Any]) -> None:
with pytest.raises(ModelGenerationError):
_validated_judge_payload({"scores": scores}, "standard")
def test_semantic_quality_scores_uses_cosine_similarity() -> None:
scores = semantic_quality_scores(
RECORD,
source_content=SOURCE,
embed_model=_FakeEmbedModel(),
)
assert scores is not None
assert 0 < scores["question_answer"] <= 100
assert 0 < scores["answer_source"] <= 100
assert scores["overall"] == round((scores["question_answer"] + scores["answer_source"]) / 2, 2)
def test_semantic_quality_scores_degrades_to_none_on_failure() -> None:
assert (
semantic_quality_scores(
RECORD,
source_content=SOURCE,
embed_model=_FailingEmbedModel(),
)
is None
)
def test_composite_overall_weights_available_layers() -> None:
assert composite_overall(rule=80, semantic=90, judge=70) == round(80 * 0.35 + 90 * 0.20 + 70 * 0.45, 2)
assert composite_overall(rule=80, semantic=90) == round(80 * 0.6 + 90 * 0.4, 2)
assert composite_overall(rule=80) == 80.0
assert composite_overall(rule=None, judge=100) == 45.0
def test_evaluate_result_record_combines_three_layers() -> None:
client = _FakeClient(
_judge_content({
"faithfulness": 5,
"correctness": 4,
"clarity": 5,
"completeness": 4,
"alignment": 5,
})
)
quality = evaluate_result_record(
RECORD,
source_content=SOURCE,
model={"api_url": "https://model.example", "online_model_name": "judge-model"},
config={"output_type": "standard", "generation_retries": 0},
client=client,
embed_model=_FakeEmbedModel(),
)
assert quality["evaluated"] is True
assert quality["judge"] is not None
assert quality["judge"]["model"] == "judge-model"
assert quality["semantic"] is not None
assert quality["layers"]["judge"] == quality["judge"]["overall"]
assert quality["overall"] == composite_overall(
rule=quality["layers"]["rule"],
semantic=quality["layers"]["semantic"],
judge=quality["layers"]["judge"],
)
# 评审提示词必须携带来源原文作为评分锚点(正文经 NFKC 归一化)。
user_message = client.calls[0]["payload"]["messages"][1]["content"]
assert "申请编号用于唯一标识一笔报销" in user_message
def test_evaluate_result_record_degrades_when_model_fails() -> None:
quality = evaluate_result_record(
RECORD,
source_content=SOURCE,
model={"api_url": "https://model.example", "online_model_name": "judge-model"},
config={"output_type": "standard", "generation_retries": 0},
client=_RaisingClient(),
embed_model=_FakeEmbedModel(),
)
assert quality["judge"] is None
assert quality["layers"]["judge"] is None
assert quality["semantic"] is not None
assert quality["overall"] == composite_overall(
rule=quality["layers"]["rule"],
semantic=quality["layers"]["semantic"],
)
def test_evaluate_result_record_without_model_runs_two_layers() -> None:
quality = evaluate_result_record(
RECORD,
source_content=SOURCE,
model=None,
embed_model=_FakeEmbedModel(),
)
assert quality["judge"] is None
assert quality["evaluated"] is True
assert quality["overall"] == composite_overall(
rule=quality["layers"]["rule"],
semantic=quality["layers"]["semantic"],
)
def test_reevaluate_edited_record_drops_stale_judge() -> None:
previous = {
"evaluated": True,
"judge": {"overall": 90.0},
}
quality = reevaluate_edited_record(
{**RECORD, "output": "编辑后的新答案内容,用于验证重评逻辑。"},
source_content=SOURCE,
previous_quality=previous,
embed_model=_FakeEmbedModel(),
)
assert quality["evaluated"] is True
assert quality["judge"] is None
assert quality["layers"]["judge"] is None
assert quality["semantic"] is not None
def test_reevaluate_edited_record_keeps_unevaluated_state() -> None:
quality = reevaluate_edited_record(
RECORD,
source_content=SOURCE,
previous_quality={},
embed_model=_FakeEmbedModel(),
)
assert quality["evaluated"] is False
assert quality["evaluated_at"] is None

File diff suppressed because it is too large Load Diff

View File

@@ -0,0 +1,65 @@
from __future__ import annotations
from pathlib import Path
from app.modules.data_process.schema_cli import REQUIRED_TASK_COLUMNS, _target_label
def test_runtime_migration_fails_fast_on_incompatible_schema() -> None:
sql_path = (
Path(__file__).resolve().parents[1]
/ "app"
/ "db"
/ "sql"
/ "002_data_process.sql"
)
sql = sql_path.read_text(encoding="utf-8")
assert "requires 001_platform_runtime.sql first" in sql
assert "supports only the current TEXT runtime schema" in sql
assert "generation_run_id" in sql
assert "results_confirmed BOOLEAN NOT NULL DEFAULT TRUE" in sql
assert "WHERE status <> 'completed' AND results_confirmed=TRUE" in sql
assert "ADD COLUMN IF NOT EXISTS workflow_step VARCHAR(20)" in sql
assert "ADD COLUMN IF NOT EXISTS preview_status VARCHAR(20)" in sql
assert "ADD COLUMN IF NOT EXISTS preview_progress NUMERIC(5,2)" in sql
assert "ADD COLUMN IF NOT EXISTS preview_run_id TEXT" in sql
assert "ADD COLUMN IF NOT EXISTS preview_failure_reason TEXT" in sql
assert "ADD COLUMN IF NOT EXISTS preview_total_files INTEGER" in sql
assert "ADD COLUMN IF NOT EXISTS preview_completed_files INTEGER" in sql
assert "data_process_workflow_backfill_ids" in sql
assert "ck_data_process_tasks_workflow_step" in sql
assert "ck_data_process_tasks_preview_status" in sql
assert "ck_data_process_tasks_preview_progress" in sql
assert "ck_data_process_tasks_preview_file_counts" in sql
for value in ("create", "model", "upload", "preview", "generate", "results"):
assert f"'{value}'" in sql
for value in ("idle", "queued", "running", "completed", "failed", "cancelled"):
assert f"'{value}'" in sql
assert "CREATE TABLE IF NOT EXISTS data_process_results" in sql
assert "ADD COLUMN IF NOT EXISTS chosen TEXT NOT NULL DEFAULT ''" in sql
assert "ADD COLUMN IF NOT EXISTS rejected TEXT NOT NULL DEFAULT ''" in sql
assert "ADD COLUMN IF NOT EXISTS original_chosen TEXT" in sql
assert "ADD COLUMN IF NOT EXISTS original_rejected TEXT" in sql
assert sql.count("BEGIN;") == 1
assert sql.rstrip().endswith("COMMIT;")
def test_schema_cli_target_label_never_contains_credentials() -> None:
label = _target_label("postgresql://secret-user:secret-password@db.example:5433/yg_ft")
assert label == "db.example:5433/yg_ft"
assert "secret" not in label
def test_schema_check_requires_current_runtime_columns() -> None:
assert REQUIRED_TASK_COLUMNS == (
"generation_run_id",
"results_confirmed",
"workflow_step",
"preview_status",
"preview_progress",
"preview_run_id",
"preview_failure_reason",
"preview_total_files",
"preview_completed_files",
)

View File

@@ -0,0 +1,276 @@
from __future__ import annotations
from pathlib import Path, PurePosixPath
import pytest
from app.modules.data_process import storage as storage_module
from app.modules.data_process.storage import (
DataProcessStorageError,
LocalDataProcessStorage,
StagedSourceObject,
)
def _stage(
storage: LocalDataProcessStorage,
*,
batch_id: str = "batch-main",
task_id: str = "task-1",
source_file_id: str = "source-1",
version: int = 1,
name: str = "source.txt",
content: bytes = b"payload",
) -> StagedSourceObject:
return storage.stage_bytes(
batch_id=batch_id,
task_id=task_id,
source_file_id=source_file_id,
version=version,
name=name,
content=content,
)
def _create_symlink(link: Path, target: Path, *, target_is_directory: bool = False) -> None:
try:
link.symlink_to(target, target_is_directory=target_is_directory)
except (NotImplementedError, OSError) as exc:
pytest.skip(f"当前平台不支持创建测试所需的符号链接: {exc}")
def _assert_staging_empty(storage: LocalDataProcessStorage) -> None:
assert list((storage.root / ".staging").iterdir()) == []
def test_stage_publish_read_delete_roundtrip_with_unicode_filename(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
content = "第一行\n第二行100% 完成".encode()
staged = _stage(
storage,
name="中文 数据 100%.csv",
content=content,
)
assert "%20" in staged.reference
assert "%25" in staged.reference
storage.publish([staged])
assert storage.read(staged.reference) == content
assert storage.delete(staged.reference) is True
assert storage.delete(staged.reference) is False
_assert_staging_empty(storage)
def test_stage_copy_creates_an_independently_deletable_source_object(
tmp_path: Path,
) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
original = _stage(storage, content=b"immutable source")
storage.publish([original])
copied = storage.stage_copy(
batch_id="batch-copy",
source_reference=original.reference,
expected_source_task_id="task-1",
expected_source_file_id="source-1",
task_id="task-2",
source_file_id="source-2",
version=1,
name="source.txt",
)
storage.publish([copied])
assert storage.read(copied.reference) == b"immutable source"
assert storage.delete(
original.reference,
expected_task_id="task-1",
expected_source_file_id="source-1",
) is True
assert storage.read(copied.reference) == b"immutable source"
assert storage.delete(
copied.reference,
expected_task_id="task-2",
expected_source_file_id="source-2",
) is True
_assert_staging_empty(storage)
def test_db_reference_is_left_to_database_storage(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
assert storage.read("db://source-files/source-1") is None
assert storage.delete("db://source-files/source-1") is False
def test_owned_source_can_be_streamed_by_byte_range(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
content = b"0123456789abcdef"
staged = _stage(storage, content=content)
storage.publish([staged])
assert storage.file_size(
staged.reference,
expected_task_id="task-1",
expected_source_file_id="source-1",
) == len(content)
assert b"".join(storage.iter_bytes(
staged.reference,
expected_task_id="task-1",
expected_source_file_id="source-1",
expected_size=len(content),
start=4,
length=6,
chunk_size=2,
)) == b"456789"
with pytest.raises(DataProcessStorageError, match="owner mismatch"):
storage.file_size(
staged.reference,
expected_task_id="another-task",
expected_source_file_id="source-1",
)
with pytest.raises(DataProcessStorageError, match="does not match metadata"):
b"".join(storage.iter_bytes(
staged.reference,
expected_task_id="task-1",
expected_source_file_id="source-1",
expected_size=len(content) + 1,
))
@pytest.mark.parametrize(
"reference",
[
"local://data-process/../source-1/v1/file.txt",
"local://data-process/task-1/source-1/v1/file%2Fname.txt",
"local://data-process/task-1/source-1/v1/file.txt?download=1",
"local://data-process/task-1/source-1/v1/file.txt#fragment",
"https://data-process/task-1/source-1/v1/file.txt",
],
ids=[
"parent-traversal",
"percent-encoded-slash",
"query",
"fragment",
"wrong-scheme",
],
)
def test_unsafe_references_are_rejected(tmp_path: Path, reference: str) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
with pytest.raises(DataProcessStorageError):
storage.read(reference)
with pytest.raises(DataProcessStorageError):
storage.delete(reference)
def test_publish_rejects_intermediate_directory_symlink(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
outside = tmp_path / "outside"
outside.mkdir()
staged = _stage(storage, task_id="linked-task")
_create_symlink(
storage.root / "linked-task",
outside,
target_is_directory=True,
)
with pytest.raises(DataProcessStorageError, match="symlink|non-directory"):
storage.publish([staged])
assert list(outside.iterdir()) == []
_assert_staging_empty(storage)
def test_target_symlink_is_never_followed_or_deleted(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
staged = _stage(storage, task_id="task-link", source_file_id="source-link")
outside_file = tmp_path / "outside.txt"
outside_file.write_bytes(b"outside sentinel")
final_path = storage.root.joinpath(*staged._relative_path.parts)
final_path.parent.mkdir(parents=True)
_create_symlink(final_path, outside_file)
with pytest.raises(DataProcessStorageError, match="already exists"):
storage.publish([staged])
with pytest.raises(DataProcessStorageError, match="regular file"):
storage.read(staged.reference)
with pytest.raises(DataProcessStorageError, match="non-regular"):
storage.delete(staged.reference)
assert final_path.is_symlink()
assert outside_file.read_bytes() == b"outside sentinel"
_assert_staging_empty(storage)
def test_publish_rolls_back_first_object_when_second_target_collides(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
existing = _stage(
storage,
batch_id="batch-existing",
source_file_id="source-existing",
content=b"existing content",
)
storage.publish([existing])
first = _stage(
storage,
batch_id="batch-new",
source_file_id="source-new",
content=b"must be rolled back",
)
colliding_second = _stage(
storage,
batch_id="batch-new",
source_file_id="source-existing",
content=b"must not replace existing content",
)
with pytest.raises(DataProcessStorageError, match="already exists"):
storage.publish([first, colliding_second])
with pytest.raises(DataProcessStorageError, match="does not exist"):
storage.read(first.reference)
assert storage.read(existing.reference) == b"existing content"
_assert_staging_empty(storage)
def test_publish_rejects_manually_forged_staged_object(tmp_path: Path) -> None:
storage = LocalDataProcessStorage(tmp_path / "storage")
temporary_path = storage.root / ".staging" / "batch-forged" / "forged.tmp"
temporary_path.parent.mkdir()
temporary_path.write_bytes(b"forged content")
relative_path = PurePosixPath("task-forged", "source-forged", "v1", "forged.txt")
forged = StagedSourceObject(
reference="local://data-process/task-forged/source-forged/v1/forged.txt",
_temporary_path=temporary_path,
_relative_path=relative_path,
)
with pytest.raises(DataProcessStorageError, match="was not issued"):
storage.publish([forged])
with pytest.raises(DataProcessStorageError, match="was not issued"):
storage.discard([forged])
assert temporary_path.read_bytes() == b"forged content"
assert not storage.root.joinpath(*relative_path.parts).exists()
def test_relative_storage_configuration_is_anchored_to_backend_root(
tmp_path: Path,
monkeypatch: pytest.MonkeyPatch,
) -> None:
relative_configuration = Path("relative-storage") / tmp_path.name
backend_root = Path(storage_module.__file__).resolve().parents[3]
monkeypatch.chdir(tmp_path)
monkeypatch.setenv("DATA_PROCESS_STORAGE_DIR", str(relative_configuration))
storage_module.get_data_process_storage.cache_clear()
try:
configured_root = storage_module._configured_storage_root()
assert configured_root == backend_root / relative_configuration
assert not configured_root.exists()
finally:
storage_module.get_data_process_storage.cache_clear()

Some files were not shown because too many files have changed in this diff Show More