Compare commits
1 Commits
ft_wyt
...
1468834116
| Author | SHA1 | Date | |
|---|---|---|---|
| 1468834116 |
35
.github/workflows/ci.yml
vendored
35
.github/workflows/ci.yml
vendored
@@ -1,35 +0,0 @@
|
||||
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
|
||||
54
.gitignore
vendored
54
.gitignore
vendored
@@ -12,8 +12,6 @@ __pycache__/
|
||||
build/
|
||||
develop-eggs/
|
||||
dist/
|
||||
node_modules/
|
||||
*.tsbuildinfo
|
||||
downloads/
|
||||
eggs/
|
||||
.eggs/
|
||||
@@ -39,22 +37,13 @@ MANIFEST
|
||||
pip-log.txt
|
||||
pip-delete-this-directory.txt
|
||||
|
||||
# Runtime data and logs
|
||||
runtime/
|
||||
backend/runtime/
|
||||
backend/storage/
|
||||
logs/
|
||||
backend/logs/
|
||||
*.db
|
||||
*.sqlite
|
||||
*.sqlite3
|
||||
|
||||
# Unit test / coverage reports
|
||||
htmlcov/
|
||||
.tox/
|
||||
.nox/
|
||||
.coverage
|
||||
.coverage.*
|
||||
.cache
|
||||
nosetests.xml
|
||||
coverage.xml
|
||||
*.cover
|
||||
@@ -141,7 +130,6 @@ celerybeat.pid
|
||||
|
||||
# Environments
|
||||
.env
|
||||
!.env.example
|
||||
.venv
|
||||
env/
|
||||
venv/
|
||||
@@ -149,16 +137,6 @@ ENV/
|
||||
env.bak/
|
||||
venv.bak/
|
||||
|
||||
# Local backend config (含数据库账号密码等敏感信息,勿提交)
|
||||
backend/config.yaml
|
||||
|
||||
# Agent / IDE 工具产物,不应进版本库
|
||||
.codex-backups/
|
||||
.pnpm-store/
|
||||
.zcode/
|
||||
.claude/
|
||||
CLAUDE.md
|
||||
|
||||
# Spyder project settings
|
||||
.spyderproject
|
||||
.spyproject
|
||||
@@ -196,33 +174,3 @@ cython_debug/
|
||||
# PyPI configuration file
|
||||
.pypirc
|
||||
|
||||
docker/llamafactory-latest.tar.gz
|
||||
# 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
|
||||
|
||||
19
Dockerfile
Normal file
19
Dockerfile
Normal file
@@ -0,0 +1,19 @@
|
||||
# syntax=docker/dockerfile:1
|
||||
|
||||
FROM node:20-alpine AS build
|
||||
WORKDIR /app/frontend
|
||||
|
||||
COPY frontend/package*.json ./
|
||||
RUN npm ci
|
||||
|
||||
COPY frontend/ ./
|
||||
RUN npm run build
|
||||
|
||||
FROM nginx:1.27-alpine
|
||||
COPY docker/nginx.conf.template /etc/nginx/templates/default.conf.template
|
||||
COPY --from=build /app/frontend/dist /usr/share/nginx/html
|
||||
|
||||
EXPOSE 80
|
||||
|
||||
HEALTHCHECK --interval=30s --timeout=3s --start-period=10s --retries=3 \
|
||||
CMD wget -qO- http://127.0.0.1/ >/dev/null || exit 1
|
||||
132
README.md
132
README.md
@@ -1,59 +1,107 @@
|
||||
# YG_FT 模型微调平台
|
||||
|
||||
YG_FT 是面向多用户、多租户和多算力节点的模型训练与推理平台,提供数据集、模型、训练、权重合并、推理、评测、算力节点、项目隔离、权限和审计能力。
|
||||
YG_FT 是一个面向企业治理场景的完整模型微调平台,覆盖用户中心、多租户、项目隔离、数据集管理、模型管理、训练任务、评测、推理、审批流、审计留存、算力调度和训练引擎适配。当前前端已存在基础页面,后端与算力平台已按多人协作开发方式建立工程骨架。
|
||||
|
||||
## 架构
|
||||
## 总体架构
|
||||
|
||||
```text
|
||||
浏览器 -> Frontend Nginx:16801 -> Backend API:17861
|
||||
|-> PostgreSQL(元数据、权限、审计、任务状态)
|
||||
|-> Redis(缓存及任务辅助状态)
|
||||
|-> MinIO:19000(模型和数据唯一对象源)
|
||||
|-> Compute API:19100
|
||||
|-> Compute Agent/GPU/LLaMA-Factory
|
||||
`-> File Gateway:19101
|
||||
YG_FT/
|
||||
frontend/ # 前端应用,承载训练平台控制台页面
|
||||
backend/ # FastAPI 应用平台后端
|
||||
app/
|
||||
api/v1/ # 对前端暴露的 REST API
|
||||
core/ # 配置、日志、中间件、权限等基础能力
|
||||
db/ # 数据库连接、迁移、事务工具
|
||||
modules/ # 业务模块目录
|
||||
schemas/ # Pydantic 入参/出参模型
|
||||
services/ # 跨模块应用服务
|
||||
workers/ # 后台任务入口
|
||||
requirements.txt # 后端 Python 第三方依赖
|
||||
compute/ # 算力平台与训练框架适配层
|
||||
api/ # 内部 Compute API
|
||||
agent/ # 单机多 GPU 调度与进程管理
|
||||
engines/llama_factory/ # LLaMA-Factory 适配器
|
||||
file_gateway/ # 本地文件上传、下载、导入、产物管理
|
||||
docs/ # 需求、接口、数据库、开发计划和部署文档
|
||||
docker/ # Nginx 等容器化配置
|
||||
```
|
||||
|
||||
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/` | 架构、权限、部署和测试文档 |
|
||||
## 关键能力
|
||||
|
||||
## 端口
|
||||
- 多租户:租户级数据隔离、租户配置、租户成员和角色。
|
||||
- 权限控制:支持项目、模型、数据集级隔离,后续可扩展到字段级和操作级策略。
|
||||
- 审批流:覆盖数据集发布、模型发布、训练资源申请、推理服务上线等企业流程。
|
||||
- 审计留存:操作审计、安全审计、审批审计、任务审计,支持留存周期策略。
|
||||
- 训练任务:训练参数管理、单机多 GPU 调度、任务状态同步、训练日志、产物管理。
|
||||
- 引擎适配:默认 LLaMA-Factory,预留统一 Engine Adapter 接口接入其他微调框架。
|
||||
- 文件存储:当前使用本地磁盘,按租户/项目/数据集/任务分区。
|
||||
- 日志采集:后端 JSON Lines 日志,主日志和错误日志拆分,便于 ELK/日志平台采集。
|
||||
|
||||
| 服务 | 主机端口 | 容器端口 |
|
||||
| --- | ---: | ---: |
|
||||
| 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
|
||||
cd frontend && npm ci && npm run build && cd ..
|
||||
docker build -f docker/app/Dockerfile.backend -t yg-ft-backend-api:latest .
|
||||
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
|
||||
cd backend
|
||||
python -m venv .venv
|
||||
.venv\Scripts\activate
|
||||
pip install -r requirements.txt
|
||||
uvicorn app.main:app --reload
|
||||
```
|
||||
|
||||
## 权限与性能
|
||||
默认健康检查:
|
||||
|
||||
系统使用角色权限、资源 ACL、用户/项目/租户归属联合校验;删除为软删除,关键操作写入审计。模型合并前准备 base model 和 adapter,结果归档 MinIO;推理前按选择节点准备缓存。远程 PostgreSQL 延迟会影响全量列表和看板,页面慢时应检查浏览器 Network、Nginx、Backend 日志、连接池和节点可达性。
|
||||
```text
|
||||
GET /api/v1/health
|
||||
```
|
||||
|
||||
详细部署见 `docker/README.md`,测试见 `测试用例.md`,本次快照见 `docs/20260812/`。
|
||||
## 日志
|
||||
|
||||
后端日志模块位于 `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/backend-api-design.md`:FastAPI 接口分组、参数定义、权限说明。
|
||||
- `docs/postgres-schema.sql`:PostgreSQL 数据库脚本,包含权限、用户中心、多租户、审批、审计等模型。
|
||||
- `docs/system-development-plan.md`:多人协作开发计划,按前端、后端、DB、部署拆分。
|
||||
- `docs/backend-logging.md`:后端日志模块使用说明。
|
||||
- `docs/deployment-plan.md`:后期部署方案,覆盖单机算力服务器部署与应用/算力分离部署。
|
||||
|
||||
## 部署模式
|
||||
|
||||
平台支持两种主要部署模式:
|
||||
|
||||
1. 所有服务部署在算力服务器:适合 PoC、内网试点、小团队单机多 GPU 使用。
|
||||
2. 应用服务和算力/训练服务独立部署:适合企业生产环境,应用平台部署在业务服务区,算力平台和 LLaMA-Factory 部署在 GPU 服务器。
|
||||
|
||||
生产环境建议采用第二种模式。算力平台与训练框架应部署在 GPU 算力服务器上,应用平台不直接控制 GPU 进程,而是通过内部 Compute API 调度训练任务。
|
||||
|
||||
详细方案见 `docs/deployment-plan.md`。
|
||||
|
||||
## 后续开发原则
|
||||
|
||||
- 接口实现优先遵循 `docs/backend-api-design.md`。
|
||||
- 数据库实现优先遵循 `docs/postgres-schema.sql`,后续通过 Alembic 迁移管理变更。
|
||||
- 前端页面与后端接口、数据库表之间的映射以文档中的“对应页面/功能模块”为准。
|
||||
- 训练引擎适配必须通过 `compute/engines/` 下的标准接口,不在应用平台后端直接拼接训练命令。
|
||||
- 敏感信息不得写入日志,生产环境密钥通过环境变量或密钥管理系统注入。
|
||||
|
||||
9
backend/.env.example
Normal file
9
backend/.env.example
Normal file
@@ -0,0 +1,9 @@
|
||||
APP_NAME=YG Fine-Tune Platform API
|
||||
APP_ENV=local
|
||||
API_PREFIX=/api
|
||||
LOG_LEVEL=INFO
|
||||
LOG_DIR=./logs
|
||||
LOG_FILE_PREFIX=backend
|
||||
LOG_ERROR_FILE_PREFIX=error
|
||||
LOG_MAX_BYTES=20971520
|
||||
LOG_RETENTION_DAYS=10
|
||||
@@ -8,7 +8,7 @@
|
||||
backend/
|
||||
app/
|
||||
main.py # FastAPI 应用入口
|
||||
api/v1/ # 对前端暴露的 接口路由
|
||||
api/v1/ # 对前端暴露的 API 路由
|
||||
core/ # 配置、日志、中间件、权限等基础能力
|
||||
db/ # 数据库连接、迁移集成、事务工具
|
||||
modules/ # 业务模块
|
||||
@@ -48,7 +48,7 @@ uvicorn app.main:app --reload
|
||||
健康检查:
|
||||
|
||||
```text
|
||||
GET /modelTF/health
|
||||
GET /api/v1/health
|
||||
```
|
||||
|
||||
## 日志
|
||||
|
||||
@@ -1,10 +0,0 @@
|
||||
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
@@ -1,40 +1,12 @@
|
||||
from fastapi import APIRouter
|
||||
from fastapi import APIRouter
|
||||
|
||||
from app.core.logging import get_logger
|
||||
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()
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@router.get("/health")
|
||||
async def health_check() -> dict[str, object]:
|
||||
# Health endpoints are called frequently by Docker and the frontend.
|
||||
# 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},
|
||||
}
|
||||
|
||||
async def health_check() -> dict[str, str]:
|
||||
logger.info("health check requested")
|
||||
return {"status": "ok"}
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,26 +1,6 @@
|
||||
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.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.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(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"])
|
||||
|
||||
@@ -1,269 +0,0 @@
|
||||
"""
|
||||
审计日志装饰器模块
|
||||
|
||||
提供 @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"
|
||||
@@ -1,615 +0,0 @@
|
||||
"""鉴权依赖:从 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 覆盖 read;admin 覆盖一切。"""
|
||||
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
|
||||
@@ -1,46 +0,0 @@
|
||||
"""集中管理项目本地缓存目录(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)
|
||||
@@ -1,19 +1,6 @@
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass
|
||||
from functools import lru_cache
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
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
|
||||
|
||||
|
||||
def _int_env(name: str, default: int) -> int:
|
||||
@@ -23,92 +10,19 @@ def _int_env(name: str, default: int) -> int:
|
||||
return int(raw)
|
||||
|
||||
|
||||
def _list_env(name: str, default: list[str]) -> list[str]:
|
||||
raw = os.getenv(name)
|
||||
if raw is None or raw.strip() == "":
|
||||
return default
|
||||
return [item.strip() for item in raw.split(",") if item.strip()]
|
||||
|
||||
|
||||
def _bool_env(name: str, default: bool) -> bool:
|
||||
raw = os.getenv(name)
|
||||
if raw is None or raw.strip() == "":
|
||||
return default
|
||||
return raw.strip().lower() in {"1", "true", "yes", "on"}
|
||||
|
||||
|
||||
def docs_kwargs(enabled: bool) -> dict[str, Any]:
|
||||
"""Swagger UI / ReDoc / OpenAPI schema 路由开关。
|
||||
|
||||
关闭时 FastAPI 不注册 /docs、/redoc、/openapi.json,访问一律返回 404,
|
||||
避免未授权访问泄露 API 结构。
|
||||
"""
|
||||
if enabled:
|
||||
return {}
|
||||
return {"docs_url": None, "redoc_url": None, "openapi_url": None}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Settings:
|
||||
app_name: str = os.getenv("APP_NAME", "YG Zhilian API")
|
||||
app_name: str = os.getenv("APP_NAME", "YG Fine-Tune Platform API")
|
||||
app_env: str = os.getenv("APP_ENV", "local")
|
||||
route_prefix: str = os.getenv("MODELTF_ROUTE_PREFIX", "/modelTF")
|
||||
app_mode: str = os.getenv("APP_MODE", "local")
|
||||
database_url: str = os.getenv("DATABASE_URL", "postgresql+psycopg://yg_ft:change_me@localhost:15432/yg_ft")
|
||||
cors_allow_origins: list[str] = None # type: ignore[assignment]
|
||||
compute_mode: str = os.getenv("COMPUTE_MODE", "real")
|
||||
compute_status_sync_mode: str = os.getenv("COMPUTE_STATUS_SYNC_MODE", "polling")
|
||||
compute_poll_interval_seconds: int = _int_env("COMPUTE_POLL_INTERVAL_SECONDS", 3)
|
||||
compute_request_timeout_seconds: int = _int_env("COMPUTE_REQUEST_TIMEOUT_SECONDS", 5)
|
||||
minio_enabled: bool = _bool_env("MINIO_ENABLED", False)
|
||||
# MinIO is an independent service and may run on another host. The
|
||||
# endpoint must therefore be reachable from the Backend container.
|
||||
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")
|
||||
minio_bucket: str = os.getenv("MINIO_BUCKET", "yg-ft-resources")
|
||||
minio_secure: bool = _bool_env("MINIO_SECURE", False)
|
||||
# Small text/data files stay inline in PostgreSQL to avoid unnecessary
|
||||
# MinIO round trips. Larger files remain the shared canonical objects.
|
||||
minio_inline_max_bytes: int = _int_env("MINIO_INLINE_MAX_BYTES", 256 * 1024)
|
||||
minio_presign_max_bytes: int = _int_env("MINIO_PRESIGN_MAX_BYTES", 1024 * 1024 * 1024 * 1024)
|
||||
storage_wait_seconds: int = _int_env("STORAGE_WAIT_SECONDS", 300)
|
||||
storage_check_interval_seconds: int = _int_env("STORAGE_CHECK_INTERVAL_SECONDS", 10)
|
||||
compute_service_token: str = os.getenv("COMPUTE_SERVICE_TOKEN", "")
|
||||
api_prefix: str = os.getenv("API_PREFIX", "/api")
|
||||
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:
|
||||
object.__setattr__(
|
||||
self,
|
||||
"cors_allow_origins",
|
||||
_list_env(
|
||||
"CORS_ALLOW_ORIGINS",
|
||||
[
|
||||
"http://localhost:16801",
|
||||
"http://127.0.0.1:16801",
|
||||
"http://localhost:17861",
|
||||
"http://127.0.0.1:17861",
|
||||
],
|
||||
),
|
||||
)
|
||||
# Swagger UI / ReDoc / OpenAPI 文档路由开关:
|
||||
# 未显式配置 ENABLE_DOCS 时,仅本地/开发环境开放,生产环境默认关闭,
|
||||
# 避免未授权访问泄露 API 结构。从运行时环境读取 APP_ENV,而非类定义时
|
||||
# 缓存的默认值,保证生产默认关闭始终生效且便于测试。
|
||||
object.__setattr__(
|
||||
self,
|
||||
"enable_docs",
|
||||
_bool_env("ENABLE_DOCS", os.getenv("APP_ENV", "local") != "prod"),
|
||||
)
|
||||
|
||||
|
||||
@lru_cache
|
||||
def get_settings() -> Settings:
|
||||
return Settings()
|
||||
|
||||
|
||||
@@ -1,278 +1,58 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
import re
|
||||
import socket
|
||||
import sys
|
||||
import time
|
||||
from contextvars import ContextVar
|
||||
from datetime import date, datetime, timedelta
|
||||
import json
|
||||
import logging
|
||||
from logging import Handler, LogRecord
|
||||
from pathlib import Path
|
||||
from typing import Any, Callable
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
from fastapi import FastAPI, Request
|
||||
|
||||
from app.core.config import Settings, get_settings
|
||||
|
||||
# ==================== 链路追踪 ContextVar ====================
|
||||
|
||||
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="")
|
||||
|
||||
|
||||
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 到每条日志记录。"""
|
||||
|
||||
class RequestIdFilter(logging.Filter):
|
||||
def filter(self, record: LogRecord) -> bool:
|
||||
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")
|
||||
record.request_id = request_id_var.get()
|
||||
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):
|
||||
"""
|
||||
生产级 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",
|
||||
}
|
||||
"""Format one JSON object per line for ELK/Filebeat collection."""
|
||||
|
||||
def format(self, record: LogRecord) -> str:
|
||||
# 时间戳:ISO8601 带时区
|
||||
timestamp = datetime.fromtimestamp(record.created).astimezone().isoformat(
|
||||
timespec="milliseconds"
|
||||
)
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
"@timestamp": timestamp,
|
||||
"@timestamp": datetime.fromtimestamp(record.created).astimezone().isoformat(
|
||||
timespec="milliseconds"
|
||||
),
|
||||
"level": record.levelname,
|
||||
"logger": record.name,
|
||||
"traceId": getattr(record, "traceId", "-"),
|
||||
"message": record.getMessage(),
|
||||
"file": f"{Path(record.pathname).name}:{record.lineno}",
|
||||
"thread": record.threadName,
|
||||
"host": getattr(record, "host", ""),
|
||||
"app": getattr(record, "app", ""),
|
||||
"env": getattr(record, "env", ""),
|
||||
"module": record.module,
|
||||
"function": record.funcName,
|
||||
"file": record.pathname,
|
||||
"line": record.lineno,
|
||||
"process": record.process,
|
||||
"thread": record.thread,
|
||||
"thread_name": record.threadName,
|
||||
"request_id": getattr(record, "request_id", "-"),
|
||||
}
|
||||
|
||||
# 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:
|
||||
if 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=(",", ":"))
|
||||
|
||||
|
||||
# ==================== DateSizeRotatingFileHandler ====================
|
||||
|
||||
class DateSizeRotatingFileHandler(Handler):
|
||||
"""按日期+大小滚动的文件日志处理器。
|
||||
|
||||
- 按天创建文件,文件名包含日期
|
||||
- 单文件超过 max_bytes 时自动滚动(带序号后缀)
|
||||
- 自动清理超过 retention_days 的旧日志
|
||||
"""
|
||||
"""Rotate log files by date and size while keeping date in every file name."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -330,8 +110,10 @@ class DateSizeRotatingFileHandler(Handler):
|
||||
today = date.today()
|
||||
if not force and self._stream and self._current_date == today:
|
||||
return
|
||||
|
||||
if self._stream and not self._stream.closed:
|
||||
self._stream.close()
|
||||
|
||||
self._current_date = today
|
||||
self._current_path = self._dated_path(today)
|
||||
self._stream = self._current_path.open("a", encoding=self.encoding)
|
||||
@@ -346,9 +128,11 @@ class DateSizeRotatingFileHandler(Handler):
|
||||
def _rotate_by_size(self) -> None:
|
||||
if not self._current_path or not self._current_path.exists():
|
||||
return
|
||||
|
||||
if self._stream and not self._stream.closed:
|
||||
self._stream.close()
|
||||
self._stream = None
|
||||
|
||||
stem = self._current_path.stem
|
||||
suffix = self._current_path.suffix
|
||||
index = 1
|
||||
@@ -362,6 +146,7 @@ class DateSizeRotatingFileHandler(Handler):
|
||||
def _cleanup_expired_files(self) -> None:
|
||||
if self.retention_days <= 0:
|
||||
return
|
||||
|
||||
cutoff = date.today() - timedelta(days=self.retention_days - 1)
|
||||
pattern = re.compile(
|
||||
rf"^{re.escape(self.file_prefix)}-(\d{{4}}-\d{{2}}-\d{{2}})(?:\.\d+)?\.log$"
|
||||
@@ -375,308 +160,94 @@ class DateSizeRotatingFileHandler(Handler):
|
||||
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:
|
||||
"""
|
||||
生产级日志配置,符合方案文档 §二(分类分流)和 §六(性能安全)。
|
||||
|
||||
日志分类:
|
||||
- 业务日志 (app-biz): INFO+ 业务流程(保留 7 天)
|
||||
- 系统日志 (app-sys): 框架/中间件日志(保留 7 天)
|
||||
- 访问日志 (app-access): HTTP 请求日志(保留 15 天)
|
||||
- 错误日志 (app-error): ERROR 级别(保留 30 天)
|
||||
"""
|
||||
settings = settings or get_settings()
|
||||
|
||||
root_logger = logging.getLogger()
|
||||
root_logger.handlers.clear()
|
||||
root_logger.setLevel(settings.log_level.upper())
|
||||
|
||||
# ---- Formatter ----
|
||||
console_formatter = logging.Formatter(
|
||||
fmt=(
|
||||
"%(asctime)s | %(levelname)s | pid=%(process)d | %(threadName)s | "
|
||||
"traceId=%(traceId)s | %(name)s | %(pathname)s:%(lineno)d | %(message)s"
|
||||
"request_id=%(request_id)s | %(name)s | %(pathname)s:%(lineno)d | %(message)s"
|
||||
),
|
||||
datefmt="%Y-%m-%d %H:%M:%S",
|
||||
)
|
||||
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.setFormatter(console_formatter)
|
||||
console_handler.addFilter(trace_filter)
|
||||
console_handler.addFilter(request_filter)
|
||||
|
||||
# ---- 业务日志文件 Handler (app-biz) ----
|
||||
biz_file_handler = DateSizeRotatingFileHandler(
|
||||
file_handler = DateSizeRotatingFileHandler(
|
||||
log_dir=settings.log_dir,
|
||||
file_prefix="app-biz",
|
||||
file_prefix=settings.log_file_prefix,
|
||||
max_bytes=settings.log_max_bytes,
|
||||
retention_days=7,
|
||||
retention_days=settings.log_retention_days,
|
||||
)
|
||||
biz_file_handler.setFormatter(json_formatter)
|
||||
biz_file_handler.addFilter(trace_filter)
|
||||
file_handler.setFormatter(json_formatter)
|
||||
file_handler.addFilter(request_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(
|
||||
log_dir=settings.log_dir,
|
||||
file_prefix="app-error",
|
||||
file_prefix=settings.log_error_file_prefix,
|
||||
max_bytes=settings.log_max_bytes,
|
||||
retention_days=30,
|
||||
retention_days=settings.log_retention_days,
|
||||
)
|
||||
error_file_handler.setLevel(logging.ERROR)
|
||||
error_file_handler.setFormatter(json_formatter)
|
||||
error_file_handler.addFilter(trace_filter)
|
||||
error_file_handler.addFilter(request_filter)
|
||||
|
||||
# ---- 注册 Handler ----
|
||||
root_logger.addHandler(console_handler)
|
||||
root_logger.addHandler(biz_file_handler)
|
||||
root_logger.addHandler(access_file_handler)
|
||||
root_logger.addHandler(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"):
|
||||
lg = logging.getLogger(logger_name)
|
||||
lg.handlers.clear()
|
||||
lg.propagate = True
|
||||
logger = logging.getLogger(logger_name)
|
||||
logger.handlers.clear()
|
||||
logger.propagate = True
|
||||
|
||||
# 框架类日志归入系统日志,生产环境设为 WARN
|
||||
logging.getLogger("uvicorn.access").setLevel(logging.WARNING)
|
||||
logging.getLogger("psycopg.pool").setLevel(logging.ERROR)
|
||||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
|
||||
# ---- 兼容旧文件前缀(向后兼容)----
|
||||
# 如果配置了旧的 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 get_logger(name: str) -> logging.Logger:
|
||||
return logging.getLogger(name)
|
||||
|
||||
# 旧错误日志前缀兼容
|
||||
if settings.log_error_file_prefix and settings.log_error_file_prefix != "app-error":
|
||||
legacy_error_handler = DateSizeRotatingFileHandler(
|
||||
log_dir=settings.log_dir,
|
||||
file_prefix=settings.log_error_file_prefix,
|
||||
max_bytes=settings.log_max_bytes,
|
||||
retention_days=settings.log_retention_days,
|
||||
)
|
||||
legacy_error_handler.setLevel(logging.ERROR)
|
||||
legacy_error_handler.setFormatter(json_formatter)
|
||||
legacy_error_handler.addFilter(trace_filter)
|
||||
root_logger.addHandler(legacy_error_handler)
|
||||
|
||||
def set_request_id(request_id: str) -> None:
|
||||
request_id_var.set(request_id)
|
||||
|
||||
|
||||
def setup_request_logging(app: FastAPI) -> None:
|
||||
logger = get_logger("app.access")
|
||||
|
||||
@app.middleware("http")
|
||||
async def request_logging_middleware(request: Request, call_next): # type: ignore[no-untyped-def]
|
||||
request_id = request.headers.get("X-Request-ID") or str(uuid4())
|
||||
token = request_id_var.set(request_id)
|
||||
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)
|
||||
|
||||
@@ -1,503 +0,0 @@
|
||||
"""
|
||||
操作日志工具模块
|
||||
|
||||
提供 @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
@@ -1,40 +1,4 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Generator
|
||||
from contextlib import contextmanager
|
||||
|
||||
from sqlalchemy import create_engine
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
|
||||
|
||||
DATABASE_URL = os.getenv("DATABASE_URL", "postgresql+psycopg://yg_ft:change_me@localhost:15432/yg_ft")
|
||||
|
||||
engine = create_engine(
|
||||
DATABASE_URL,
|
||||
pool_pre_ping=True,
|
||||
future=True,
|
||||
)
|
||||
SessionLocal = sessionmaker(bind=engine, autoflush=False, autocommit=False, expire_on_commit=False, future=True)
|
||||
|
||||
|
||||
def get_db() -> Generator[Session, None, None]:
|
||||
db = SessionLocal()
|
||||
try:
|
||||
yield db
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def session_scope() -> Generator[Session, None, None]:
|
||||
db = SessionLocal()
|
||||
try:
|
||||
yield db
|
||||
db.commit()
|
||||
except Exception:
|
||||
db.rollback()
|
||||
raise
|
||||
finally:
|
||||
db.close()
|
||||
"""Database session factory placeholder.
|
||||
|
||||
Implement SQLAlchemy/SQLModel session management here when database development starts.
|
||||
"""
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,373 +0,0 @@
|
||||
CREATE TABLE IF NOT EXISTS users (
|
||||
id TEXT PRIMARY KEY,
|
||||
username TEXT NOT NULL UNIQUE,
|
||||
password_hash TEXT NOT NULL,
|
||||
display_name TEXT NOT NULL,
|
||||
role TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
permissions TEXT NOT NULL,
|
||||
create_time TEXT NOT NULL,
|
||||
last_login TEXT,
|
||||
protected INTEGER NOT NULL DEFAULT 0
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS models (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
type TEXT NOT NULL,
|
||||
purpose TEXT NOT NULL,
|
||||
model_source TEXT NOT NULL,
|
||||
description TEXT,
|
||||
path TEXT,
|
||||
api_url TEXT,
|
||||
api_key TEXT,
|
||||
online_model_name TEXT,
|
||||
create_time TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS trained_models (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
train_methods TEXT NOT NULL,
|
||||
base_model_path TEXT,
|
||||
create_time TEXT NOT NULL,
|
||||
merged INTEGER NOT NULL DEFAULT 0,
|
||||
merging INTEGER NOT NULL DEFAULT 0,
|
||||
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 (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
type TEXT NOT NULL,
|
||||
storage_type TEXT NOT NULL,
|
||||
source TEXT NOT NULL,
|
||||
task_id TEXT,
|
||||
size TEXT,
|
||||
count INTEGER NOT NULL DEFAULT 0,
|
||||
description TEXT,
|
||||
create_time TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS dataset_files (
|
||||
id TEXT PRIMARY KEY,
|
||||
dataset_id TEXT NOT NULL REFERENCES datasets(id) ON DELETE CASCADE,
|
||||
name TEXT NOT NULL,
|
||||
size TEXT,
|
||||
content TEXT NOT NULL,
|
||||
active_version_id TEXT NOT NULL,
|
||||
versions TEXT NOT NULL,
|
||||
create_time TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS compute_nodes (
|
||||
id TEXT PRIMARY KEY,
|
||||
code TEXT NOT NULL UNIQUE,
|
||||
name TEXT NOT NULL,
|
||||
api_base_url TEXT NOT NULL,
|
||||
file_gateway_url TEXT NOT NULL,
|
||||
enabled INTEGER NOT NULL DEFAULT 1,
|
||||
scheduler_status TEXT NOT NULL,
|
||||
scheduler_weight INTEGER NOT NULL DEFAULT 100,
|
||||
tags TEXT NOT NULL,
|
||||
gpu_count INTEGER NOT NULL DEFAULT 0,
|
||||
current_running_jobs INTEGER NOT NULL DEFAULT 0,
|
||||
max_parallel_jobs INTEGER NOT NULL DEFAULT 2,
|
||||
data_root TEXT NOT NULL,
|
||||
model_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,
|
||||
health_detail TEXT NOT NULL
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS gpus (
|
||||
id TEXT PRIMARY KEY,
|
||||
node_id TEXT NOT NULL REFERENCES compute_nodes(id) ON DELETE CASCADE,
|
||||
gpu_index INTEGER NOT NULL,
|
||||
uuid TEXT NOT NULL,
|
||||
name TEXT NOT NULL,
|
||||
memory_total_gb DOUBLE PRECISION NOT NULL,
|
||||
power_limit_w DOUBLE PRECISION NOT NULL,
|
||||
base_temperature INTEGER NOT NULL,
|
||||
last_seen_at TEXT
|
||||
);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS fine_tune_tasks (
|
||||
id TEXT PRIMARY KEY,
|
||||
name TEXT NOT NULL UNIQUE,
|
||||
payload TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
progress INTEGER NOT NULL DEFAULT 0,
|
||||
process_id INTEGER,
|
||||
create_time TEXT NOT NULL,
|
||||
start_time TEXT,
|
||||
completed_at TEXT,
|
||||
compute_node_id TEXT REFERENCES compute_nodes(id) ON DELETE SET NULL,
|
||||
gpus TEXT NOT NULL,
|
||||
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 (
|
||||
id TEXT PRIMARY KEY,
|
||||
node_id TEXT NOT NULL REFERENCES compute_nodes(id) ON DELETE CASCADE,
|
||||
resource_type TEXT NOT NULL,
|
||||
resource_id TEXT NOT NULL,
|
||||
local_path TEXT NOT NULL,
|
||||
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 TABLE IF NOT EXISTS resource_sync_jobs (
|
||||
id TEXT PRIMARY KEY,
|
||||
target_node_id TEXT NOT NULL,
|
||||
resources TEXT NOT NULL,
|
||||
status TEXT NOT NULL,
|
||||
progress INTEGER NOT NULL DEFAULT 0,
|
||||
create_time TEXT NOT NULL,
|
||||
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_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_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 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
|
||||
);
|
||||
@@ -1,352 +0,0 @@
|
||||
-- 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;
|
||||
@@ -1,111 +0,0 @@
|
||||
-- 平台治理:租户 / 审批 / 审计(字段以 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);
|
||||
@@ -1,20 +0,0 @@
|
||||
-- 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;
|
||||
@@ -1,3 +0,0 @@
|
||||
-- 租户配额与保留策略扩展(如后续治理表需补列,可在此追加)
|
||||
ALTER TABLE tenants ADD COLUMN IF NOT EXISTS gpu_quota TEXT;
|
||||
ALTER TABLE tenants ADD COLUMN IF NOT EXISTS storage_quota TEXT;
|
||||
@@ -1,22 +0,0 @@
|
||||
-- ============================================================
|
||||
-- 权限体系扩展: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;
|
||||
@@ -1,32 +0,0 @@
|
||||
-- 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;
|
||||
@@ -1,21 +0,0 @@
|
||||
-- 原子 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;
|
||||
@@ -1,33 +0,0 @@
|
||||
-- 平台可靠性闭环:导出权限、对象清理、缓存 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;
|
||||
@@ -1,53 +0,0 @@
|
||||
-- 权限 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';
|
||||
|
||||
@@ -1,116 +0,0 @@
|
||||
-- 权限 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 = '');
|
||||
@@ -1,21 +0,0 @@
|
||||
-- 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';
|
||||
@@ -1,20 +0,0 @@
|
||||
-- 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);
|
||||
@@ -1,63 +0,0 @@
|
||||
-- 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);
|
||||
@@ -1,47 +1,17 @@
|
||||
import asyncio
|
||||
from contextlib import suppress
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
|
||||
from app.core.cache_paths import setup_local_caches
|
||||
|
||||
# 在任何 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
|
||||
from app.api.v1.router import api_router
|
||||
from app.core.config import get_settings
|
||||
from app.core.logging import configure_logging, setup_request_logging
|
||||
|
||||
|
||||
def create_app() -> FastAPI:
|
||||
settings = get_settings()
|
||||
configure_logging(settings)
|
||||
|
||||
app = FastAPI(title=settings.app_name, **docs_kwargs(settings.enable_docs))
|
||||
app.add_middleware(
|
||||
CORSMiddleware,
|
||||
allow_origins=settings.cors_allow_origins,
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
)
|
||||
app = FastAPI(title=settings.app_name)
|
||||
setup_request_logging(app)
|
||||
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
|
||||
|
||||
app.include_router(api_router, prefix=settings.api_prefix)
|
||||
return app
|
||||
|
||||
|
||||
|
||||
@@ -1,396 +0,0 @@
|
||||
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)
|
||||
@@ -1,282 +0,0 @@
|
||||
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)
|
||||
@@ -1,454 +0,0 @@
|
||||
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}
|
||||
@@ -1,3 +0,0 @@
|
||||
from .router import router
|
||||
|
||||
__all__ = ["router"]
|
||||
@@ -1,570 +0,0 @@
|
||||
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
|
||||
@@ -1,46 +1 @@
|
||||
"""数据处理模块:从原始文件接入到生成标准训练数据的全流程。
|
||||
|
||||
整体分层
|
||||
--------
|
||||
- ``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"]
|
||||
"""Data processing module."""
|
||||
|
||||
@@ -1,151 +0,0 @@
|
||||
"""数据处理算法模块。
|
||||
|
||||
重构为多个按职责拆分的子模块:
|
||||
|
||||
- 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",
|
||||
]
|
||||
@@ -1,24 +0,0 @@
|
||||
"""数据处理算法 - 本地语义嵌入模型共享单例。"""
|
||||
|
||||
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,
|
||||
)
|
||||
@@ -1,119 +0,0 @@
|
||||
"""数据处理算法 - 格式检测。"""
|
||||
|
||||
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),
|
||||
)
|
||||
@@ -1,27 +0,0 @@
|
||||
"""文档解析器模块。"""
|
||||
|
||||
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',
|
||||
]
|
||||
@@ -1,89 +0,0 @@
|
||||
"""数据处理算法 - 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
|
||||
@@ -1,395 +0,0 @@
|
||||
"""数据处理算法 - 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)
|
||||
@@ -1,174 +0,0 @@
|
||||
"""基于 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)
|
||||
@@ -1,716 +0,0 @@
|
||||
"""数据处理算法 - 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
|
||||
@@ -1,299 +0,0 @@
|
||||
"""数据处理算法 - 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)
|
||||
@@ -1,428 +0,0 @@
|
||||
"""数据处理算法 - 质量评分和去重。"""
|
||||
|
||||
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)
|
||||
@@ -1,809 +0,0 @@
|
||||
"""数据处理算法 - 结构化数据处理。"""
|
||||
|
||||
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
|
||||
@@ -1,328 +0,0 @@
|
||||
"""数据处理算法 - 文本处理工具。"""
|
||||
|
||||
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_utils(normalize_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)
|
||||
@@ -1,81 +0,0 @@
|
||||
"""数据处理算法 - 数据集转换和分割。"""
|
||||
|
||||
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
|
||||
@@ -1,266 +0,0 @@
|
||||
"""数据处理模块使用的无副作用算法。
|
||||
|
||||
本模块不访问数据库、文件系统或网络,便于 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 解析时发现同一对象内的重复键。"""
|
||||
@@ -1,7 +0,0 @@
|
||||
"""数据处理模块的共享限制。"""
|
||||
|
||||
MAX_QA_PAIRS_PER_ITEM = 50
|
||||
MODEL_GENERATION_BATCH_SIZE = 10
|
||||
|
||||
|
||||
__all__ = ["MAX_QA_PAIRS_PER_ITEM", "MODEL_GENERATION_BATCH_SIZE"]
|
||||
@@ -1,158 +0,0 @@
|
||||
"""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)
|
||||
@@ -1,693 +0,0 @@
|
||||
"""基于 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,
|
||||
)
|
||||
@@ -1,313 +0,0 @@
|
||||
"""数据处理 - 生成结果的多层质量评测。
|
||||
|
||||
三层体系:规则层(确定性规则分)+ 语义层(本地嵌入向量)+ 评审层
|
||||
(复用生成模型按 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
|
||||
@@ -1,641 +0,0 @@
|
||||
"""数据处理任务的大模型生成适配器。"""
|
||||
|
||||
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"]
|
||||
@@ -1,319 +0,0 @@
|
||||
"""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",
|
||||
]
|
||||
@@ -1,76 +0,0 @@
|
||||
"""数据处理运行表的显式检查与安装命令。"""
|
||||
|
||||
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())
|
||||
@@ -1,669 +0,0 @@
|
||||
"""数据处理源文件的受控暂存与分层对象存储。"""
|
||||
|
||||
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",
|
||||
]
|
||||
@@ -1,69 +0,0 @@
|
||||
"""数据处理存储层。"""
|
||||
|
||||
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",
|
||||
]
|
||||
@@ -1,310 +0,0 @@
|
||||
"""数据处理存储层 - 基础设施。
|
||||
|
||||
包含:异常类、工具函数、常量定义、基类。
|
||||
"""
|
||||
|
||||
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)
|
||||
@@ -1,567 +0,0 @@
|
||||
"""数据处理存储层 - 数据集发布。"""
|
||||
|
||||
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()
|
||||
@@ -1,275 +0,0 @@
|
||||
"""数据处理存储层 - 生成管理。"""
|
||||
|
||||
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"),
|
||||
}
|
||||
@@ -1,539 +0,0 @@
|
||||
"""数据处理存储层 - 预览管理。"""
|
||||
|
||||
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())
|
||||
@@ -1,324 +0,0 @@
|
||||
"""数据处理存储层 - 结果管理。"""
|
||||
|
||||
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 {}
|
||||
@@ -1,331 +0,0 @@
|
||||
"""数据处理存储层 - 源文件管理。"""
|
||||
|
||||
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),
|
||||
)
|
||||
@@ -1,879 +0,0 @@
|
||||
"""数据处理存储层 - 任务管理。"""
|
||||
|
||||
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),
|
||||
)
|
||||
@@ -1 +0,0 @@
|
||||
"""GPU assignment management module."""
|
||||
@@ -1,108 +0,0 @@
|
||||
"""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"]))
|
||||
@@ -1,260 +0,0 @@
|
||||
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)
|
||||
@@ -1 +0,0 @@
|
||||
"""Resource access control list (ACL) module."""
|
||||
@@ -1,93 +0,0 @@
|
||||
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]:
|
||||
"""设置资源 ACL,body: { 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)
|
||||
@@ -1,77 +0,0 @@
|
||||
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})
|
||||
@@ -1 +0,0 @@
|
||||
"""Central object storage integration."""
|
||||
@@ -1,141 +0,0 @@
|
||||
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()
|
||||
@@ -1,54 +0,0 @@
|
||||
"""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"
|
||||
@@ -1,314 +0,0 @@
|
||||
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}
|
||||
@@ -1,326 +0,0 @@
|
||||
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)
|
||||
@@ -1,437 +0,0 @@
|
||||
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
|
||||
@@ -1,59 +0,0 @@
|
||||
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)
|
||||
@@ -2,29 +2,20 @@
|
||||
name = "yg-ft-backend"
|
||||
version = "0.1.0"
|
||||
description = "Backend service for the model fine-tuning platform"
|
||||
requires-python = ">=3.12"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"fastapi>=0.111.0",
|
||||
"uvicorn[standard]>=0.30.0",
|
||||
"python-multipart>=0.0.9",
|
||||
"pydantic>=2.7.0",
|
||||
"sqlalchemy>=2.0.30",
|
||||
"psycopg[binary]>=3.2.1",
|
||||
"psycopg-pool>=3.2.1",
|
||||
"asyncpg>=0.29.0",
|
||||
"alembic>=1.13.1",
|
||||
"redis>=5.0.4",
|
||||
"httpx>=0.27.0",
|
||||
"PyJWT>=2.8.0",
|
||||
"passlib[bcrypt]>=1.7.4",
|
||||
"python-dotenv>=1.0.1",
|
||||
"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]
|
||||
@@ -35,7 +26,7 @@ dev = [
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 100
|
||||
target-version = "py312"
|
||||
target-version = "py311"
|
||||
|
||||
[tool.pytest.ini_options]
|
||||
testpaths = ["tests"]
|
||||
|
||||
@@ -3,25 +3,10 @@ uvicorn[standard]>=0.30.0
|
||||
python-multipart>=0.0.9
|
||||
pydantic>=2.7.0
|
||||
sqlalchemy>=2.0.30
|
||||
psycopg[binary]>=3.2.1
|
||||
psycopg-pool>=3.2.1
|
||||
asyncpg>=0.29.0
|
||||
alembic>=1.13.1
|
||||
redis>=5.0.4
|
||||
httpx>=0.27.0
|
||||
minio>=7.2.7
|
||||
urllib3>=2.0.7
|
||||
PyJWT>=2.8.0
|
||||
passlib[bcrypt]>=1.7.4
|
||||
python-dotenv>=1.0.1
|
||||
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
|
||||
|
||||
@@ -1,266 +0,0 @@
|
||||
[
|
||||
{
|
||||
"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": "-"
|
||||
}
|
||||
]
|
||||
@@ -1,258 +0,0 @@
|
||||
[
|
||||
{
|
||||
"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)"
|
||||
}
|
||||
]
|
||||
@@ -1,276 +0,0 @@
|
||||
"""
|
||||
模型推理异步加载改造的单元测试。
|
||||
|
||||
覆盖:
|
||||
- model_compare_load:异步派发,立即返回 starting + 节点信息(不等待加载完成)
|
||||
- model_compare_delete:先删记录,卸载失败也不阻塞删除
|
||||
- reconcile_inference_loads:starting -> 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
|
||||
# 只命中任务记录中的节点 n1,n2 未被卸载
|
||||
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"
|
||||
@@ -1,57 +0,0 @@
|
||||
"""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
@@ -1,284 +0,0 @@
|
||||
"""数据评测模块(三层质量评分)的单元测试。"""
|
||||
|
||||
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
@@ -1,65 +0,0 @@
|
||||
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",
|
||||
)
|
||||
@@ -1,276 +0,0 @@
|
||||
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()
|
||||
File diff suppressed because it is too large
Load Diff
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user