deepdata: core product code (backend, frontend, runtime, deepflow module)

This commit is contained in:
figmar
2026-08-30 19:23:38 +08:00
commit bd3db30be9
1405 changed files with 271727 additions and 0 deletions
@@ -0,0 +1,109 @@
"""切片处理器模块 - 负责文本切片和持久化操作"""
import uuid
from typing import Any
from sqlalchemy.ext.asyncio import AsyncSession
from app.db.models.data_synthesis import (
DataSynthInstance,
DataSynthesisFileInstance,
DataSynthesisChunkInstance,
)
from app.db.session import logger
class ChunkProcessor:
"""切片处理器"""
def __init__(self, db: AsyncSession):
self.db = db
async def persist_chunks(
self,
synthesis_task: DataSynthInstance,
file_task: DataSynthesisFileInstance,
file_id: str,
chunks: list[Any],
) -> int:
"""将切片结果保存到数据库。
Args:
synthesis_task: 合成任务实例
file_task: 文件任务实例
file_id: 原始文件ID
chunks: 切片后的文档列表
Returns:
保存的切片数量
"""
for idx, doc in enumerate(chunks, start=1):
base_metadata = dict(getattr(doc, "metadata", {}) or {})
base_metadata.update({
"task_id": str(synthesis_task.id),
"file_id": file_id
})
chunk_record = DataSynthesisChunkInstance(
id=str(uuid.uuid4()),
synthesis_file_instance_id=file_task.id,
chunk_index=idx,
chunk_content=doc.page_content,
chunk_metadata=base_metadata,
)
self.db.add(chunk_record)
file_task.total_chunks = len(chunks)
file_task.status = "processing"
await self.db.commit()
await self.db.refresh(file_task)
logger.info(f"Persisted {len(chunks)} chunks for file_task={file_task.id}")
return len(chunks)
async def load_chunk_batch(
self,
file_task_id: str,
start_index: int,
end_index: int,
) -> list[DataSynthesisChunkInstance]:
"""按索引范围加载指定文件任务下的一批切片记录。
Args:
file_task_id: 文件任务ID
start_index: 起始索引(包含)
end_index: 结束索引(包含)
Returns:
切片实例列表
"""
from sqlalchemy import select
result = await self.db.execute(
select(DataSynthesisChunkInstance)
.where(
DataSynthesisChunkInstance.synthesis_file_instance_id == file_task_id,
DataSynthesisChunkInstance.chunk_index >= start_index,
DataSynthesisChunkInstance.chunk_index <= end_index,
)
.order_by(DataSynthesisChunkInstance.chunk_index.asc())
)
return list(result.scalars().all())
async def count_chunks_for_file(self, synth_file_instance_id: str) -> int:
"""统计指定任务与文件下的切片总数。
Args:
synth_file_instance_id: 文件实例ID
Returns:
切片总数
"""
from sqlalchemy import func, select
result = await self.db.execute(
select(func.count(DataSynthesisChunkInstance.id)).where(
DataSynthesisChunkInstance.synthesis_file_instance_id == synth_file_instance_id
)
)
return int(result.scalar() or 0)
@@ -0,0 +1,67 @@
"""
切片查询服务 - 负责切片相关的查询操作
"""
from typing import List, Tuple
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from app.db.models.data_synthesis import DataSynthesisChunkInstance
from app.module.generation.schema.generation import (
DataSynthesisChunkItem,
PagedDataSynthesisChunkResponse,
)
class ChunkQueryService:
"""切片查询服务 - 负责切片相关的查询操作"""
def __init__(self, db: AsyncSession):
self.db = db
async def get_by_id(self, chunk_id: str) -> DataSynthesisChunkInstance | None:
"""
根据 ID 获取切片
Args:
chunk_id: 切片 ID
Returns:
切片实例,不存在则返回 None
"""
return await self.db.get(DataSynthesisChunkInstance, chunk_id)
async def list_by_file(
self, file_id: str, page: int, page_size: int
) -> Tuple[List[DataSynthesisChunkItem], int]:
"""
分页查询文件任务下的切片列表
Args:
file_id: 文件任务 ID
page: 页码
page_size: 每页数量
Returns:
(切片列表, 总数)
"""
base_query = select(DataSynthesisChunkInstance).where(
DataSynthesisChunkInstance.synthesis_file_instance_id == file_id
)
count_result = await self.db.execute(
select(func.count()).select_from(base_query.subquery())
)
total = count_result.scalar_one() or 0
offset = max(page, 1) - 1
limit = max(page_size, 1)
result = await self.db.execute(
base_query.order_by(DataSynthesisChunkInstance.chunk_index.asc())
.offset(offset * limit)
.limit(limit)
)
rows = result.scalars().all()
items = [DataSynthesisChunkItem.model_validate(row) for row in rows]
return items, total
@@ -0,0 +1,274 @@
"""
数据导出服务 - 负责将合成数据导出为各种格式
"""
import datetime
import json
import os
from typing import Iterable, List, Optional, Sequence
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.logging import get_logger
from app.db.models.data_synthesis import (
DataSynthInstance,
DataSynthesisFileInstance,
SynthesisData,
)
from app.db.models.dataset_management import Dataset, DatasetFiles
from app.module.generation.schema.generation import (
ExportSynthesisDataResponse,
)
logger = get_logger(__name__)
class SynthesisExportError(Exception):
"""导出失败时抛出的异常"""
class SynthesisDatasetExporter:
"""
数据导出器 - 将合成数据导出为各种格式
导出规则:
- 维度:原始文件 (DatasetFiles)
- 每个原始文件生成一个 JSONL 文件
- JSONL 文件名称与原始文件名称一致
- 支持格式转换:alpaca、sharegpt、raw
"""
SUPPORTED_FORMATS = ["alpaca", "sharegpt", "raw"]
DEFAULT_FORMAT = "alpaca"
def __init__(self, db: AsyncSession, format: str = "alpaca", output_path: Optional[str] = None):
self._db = db
self._format = format if format in self.SUPPORTED_FORMATS else self.DEFAULT_FORMAT
self._output_path = output_path
async def export_task_to_dataset(
self,
task_id: str,
dataset_id: str,
) -> Dataset:
"""
将任务的全部合成数据导出到已有数据集
Args:
task_id: 任务 ID
dataset_id: 目标数据集 ID
Returns:
更新后的数据集
"""
task = await self._db.get(DataSynthInstance, task_id)
if not task:
raise SynthesisExportError(f"Synthesis task {task_id} not found")
dataset = await self._db.get(Dataset, dataset_id)
if not dataset:
raise SynthesisExportError(f"Dataset {dataset_id} not found")
file_instances = await self._load_file_instances(task_id)
if not file_instances:
raise SynthesisExportError("No synthesis file instances found for task")
base_path = self._ensure_dataset_path(dataset)
created_files: list[DatasetFiles] = []
total_size = 0
for file_instance in file_instances:
records = await self._load_synthesis_data_for_file(file_instance.id)
if not records:
continue
original_name = file_instance.file_name or "unknown"
base_name, _ = os.path.splitext(original_name)
archived_file_name = f"{base_name}.jsonl"
file_path = os.path.join(base_path, archived_file_name)
os.makedirs(os.path.dirname(file_path), exist_ok=True)
self._write_jsonl(file_path, records, self._format)
try:
file_size = os.path.getsize(file_path)
except OSError:
file_size = 0
df = DatasetFiles(
dataset_id=dataset.id,
file_name=archived_file_name,
file_path=file_path,
file_type="jsonl",
file_size=file_size,
last_access_time=datetime.datetime.now(),
)
self._db.add(df)
created_files.append(df)
total_size += file_size
if created_files:
dataset.file_count = (dataset.file_count or 0) + len(created_files)
dataset.size_bytes = (dataset.size_bytes or 0) + total_size
dataset.status = "ACTIVE"
await self._db.commit()
logger.info(
"Exported synthesis task %s to dataset %s with %d files (total %d bytes)",
task_id,
dataset.id,
len(created_files),
total_size,
)
return dataset
async def export_data(
self,
task_id: str,
file_instance_ids: Optional[List[str]] = None,
) -> ExportSynthesisDataResponse:
"""
将合成数据导出为指定格式的 JSONL 文件
Args:
task_id: 任务 ID
file_instance_ids: 文件实例 ID 列表,为空则导出全部
Returns:
导出结果
"""
if file_instance_ids:
file_instances = await self._load_specific_file_instances(file_instance_ids)
else:
file_instances = await self._load_file_instances(task_id)
if not file_instances:
raise SynthesisExportError("No file instances to export")
output_dir = self._output_path or self._get_default_output_dir()
os.makedirs(output_dir, exist_ok=True)
file_paths: List[str] = []
total_records = 0
for file_instance in file_instances:
records = await self._load_synthesis_data_for_file(file_instance.id)
if not records:
continue
original_name = file_instance.file_name or "unknown"
base_name, _ = os.path.splitext(original_name)
output_file_name = f"{base_name}.jsonl"
file_path = os.path.join(output_dir, output_file_name)
self._write_jsonl(file_path, records, self._format)
file_paths.append(file_path)
total_records += len(records)
return ExportSynthesisDataResponse(
file_paths=file_paths,
total_records=total_records,
format=self._format,
)
async def _load_file_instances(self, task_id: str) -> Sequence[DataSynthesisFileInstance]:
"""加载任务下的所有文件实例"""
result = await self._db.execute(
select(DataSynthesisFileInstance).where(
DataSynthesisFileInstance.synthesis_instance_id == task_id
)
)
return result.scalars().all()
async def _load_specific_file_instances(
self, file_ids: List[str]
) -> Sequence[DataSynthesisFileInstance]:
"""加载指定的文件实例"""
result = await self._db.execute(
select(DataSynthesisFileInstance).where(
DataSynthesisFileInstance.id.in_(file_ids)
)
)
return result.scalars().all()
async def _load_synthesis_data_for_file(
self, file_instance_id: str
) -> List[dict]:
"""加载单个文件实例的所有合成数据"""
result = await self._db.execute(
select(SynthesisData).where(
SynthesisData.synthesis_file_instance_id == file_instance_id
)
)
rows: Sequence[SynthesisData] = result.scalars().all()
records: List[dict] = []
for row in rows:
payload = row.data or {}
records.append(payload)
return records
def _write_jsonl(self, path: str, records: Iterable[dict], format: Optional[str] = None) -> None:
"""写入 JSONL 文件"""
fmt = format or self._format
os.makedirs(os.path.dirname(path), exist_ok=True)
with open(path, "w", encoding="utf-8") as f:
for record in records:
formatted = self._format_record(record, fmt)
f.write(json.dumps(formatted, ensure_ascii=False))
f.write("\n")
def _format_record(self, record: dict, format: str) -> dict:
"""根据格式转换记录"""
if format == "alpaca":
return self._format_as_alpaca(record)
elif format == "sharegpt":
return self._format_as_sharegpt(record)
else:
return self._format_as_raw(record)
def _format_as_alpaca(self, record: dict) -> dict:
"""转换为 Alpaca 格式"""
return {
"instruction": record.get("instruction", ""),
"output": record.get("output", ""),
}
def _format_as_sharegpt(self, record: dict) -> dict:
"""转换为 ShareGPT 格式"""
instruction = record.get("instruction", "")
output = record.get("output", "")
input_text = record.get("input", "")
conversations = [
{"from": "human", "value": instruction},
{"from": "gpt", "value": output},
]
result = {"conversations": conversations}
if input_text:
result["context"] = input_text
return result
def _format_as_raw(self, record: dict) -> dict:
"""原始格式"""
return record
def _get_default_output_dir(self) -> str:
"""获取默认输出目录"""
import tempfile
return tempfile.gettempdir()
@staticmethod
def _ensure_dataset_path(dataset: Dataset) -> str:
"""确保数据集路径存在"""
if not dataset.path:
raise SynthesisExportError("Dataset path is empty")
os.makedirs(dataset.path, exist_ok=True)
return dataset.path
@@ -0,0 +1,65 @@
"""
文件查询服务 - 负责文件任务相关的查询操作
"""
from typing import List, Tuple
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from app.db.models.data_synthesis import DataSynthesisFileInstance
from app.module.generation.schema.generation import (
DataSynthesisFileTaskItem,
PagedDataSynthesisFileTaskResponse,
)
class FileQueryService:
"""文件查询服务 - 负责文件任务相关的查询操作"""
def __init__(self, db: AsyncSession):
self.db = db
async def get_by_id(self, file_id: str) -> DataSynthesisFileInstance | None:
"""
根据 ID 获取文件任务
Args:
file_id: 文件任务 ID
Returns:
文件任务实例,不存在则返回 None
"""
return await self.db.get(DataSynthesisFileInstance, file_id)
async def list_by_task(
self, task_id: str, page: int, page_size: int
) -> Tuple[List[DataSynthesisFileTaskItem], int]:
"""
分页查询任务下的文件任务列表
Args:
task_id: 任务 ID
page: 页码
page_size: 每页数量
Returns:
(文件任务列表, 总数)
"""
base_query = select(DataSynthesisFileInstance).where(
DataSynthesisFileInstance.synthesis_instance_id == task_id
)
count_result = await self.db.execute(
select(func.count()).select_from(base_query.subquery())
)
total = count_result.scalar_one() or 0
offset = max(page, 1) - 1
limit = max(page_size, 1)
result = await self.db.execute(
base_query.offset(offset * limit).limit(limit)
)
rows = result.scalars().all()
items = [DataSynthesisFileTaskItem.model_validate(row) for row in rows]
return items, total
@@ -0,0 +1,608 @@
"""数据合成服务 - 核心业务逻辑层"""
import asyncio
from contextlib import asynccontextmanager
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker
from app.db.models.base_entity import LineageNode, LineageEdge
from app.db.models.data_synthesis import (
DataSynthInstance,
DataSynthesisFileInstance,
DataSynthesisChunkInstance,
SynthesisData,
)
from app.db.models.dataset_management import DatasetFiles, Dataset
from app.db.session import logger, AsyncSessionLocal
from app.module.generation.schema.generation import Config, SyntheConfig
from app.module.generation.service.chunk_processor import ChunkProcessor
from app.module.generation.service.qa_generator import QAGenerator
from app.module.generation.service.qa_generator import _filter_docs_by_size
from app.module.generation.service.task_executor import run_in_thread
from app.module.shared.common.document_loaders import load_documents
from app.module.shared.common.text_split import DocumentSplitter
from app.module.shared.llm import LLMFactory
from app.module.system.service.common_service import get_model_by_id
from app.module.shared.common.lineage import LineageService
from app.module.shared.schema import NodeType, EdgeType
class GenerationService:
"""数据合成服务 - 使用模块化架构"""
def __init__(self, db: AsyncSession):
self.db = db
self.chunk_processor = ChunkProcessor(db)
self.qa_generator = QAGenerator(db)
@asynccontextmanager
async def _session_scope(self):
"""
创建独立的数据库会话上下文管理器
用于长时间运行的任务,避免长时间占用连接
"""
session = AsyncSessionLocal()
try:
yield session
await session.commit()
except Exception:
await session.rollback()
raise
finally:
await session.close()
def _load_and_split_sync(
self, file_path: str, chunk_size: int, chunk_overlap: int
):
"""
同步方法:加载并切分文件
在线程池中执行,不阻塞事件循环
"""
try:
docs = load_documents(file_path)
split_docs = DocumentSplitter.auto_split(docs, chunk_size, chunk_overlap)
return _filter_docs_by_size(split_docs, chunk_size)
except Exception as e:
logger.error(f"Error loading or splitting file {file_path}: {e}")
raise
async def process_task(self, task_id: str):
"""
处理数据合成任务入口 - 使用独立会话和线程池
Args:
task_id: 合成任务ID
"""
# 使用独立的数据库会话,不与API请求共享
async with self._session_scope() as session:
synth_task: DataSynthInstance | None = await session.get(
DataSynthInstance, task_id
)
if not synth_task:
logger.error(f"Synthesis task {task_id} not found, abort processing")
return
logger.info(f"Start processing synthesis task {task_id}")
# 更新任务状态为处理中
synth_task.status = "processing"
await session.commit()
# 从 synth_config 中读取 max_qa_pairs
max_qa_pairs = self._parse_max_qa_pairs(synth_task)
# 获取任务关联的文件ID列表
file_ids = await self._get_file_ids_for_task(session, task_id)
if not file_ids:
logger.warning(
f"No files associated with task {task_id}, abort processing"
)
await self._mark_task_failed(session, task_id, "no_files_associated")
return
# 逐个文件处理
processed_count = 0
total_task_chunks = 0
processed_task_chunks = 0
for file_id in file_ids:
try:
file_task = await self._get_or_create_file_instance(
session, str(synth_task.id), file_id
)
success = await self._process_single_file(
session, synth_task, file_id, max_qa_pairs=max_qa_pairs
)
# 累加任务级别的切片统计
if file_task:
total_task_chunks += file_task.total_chunks or 0
processed_task_chunks += file_task.processed_chunks or 0
except Exception as e:
logger.exception(
f"Unexpected error when processing file {file_id} for task {task_id}: {e}"
)
await self._mark_file_failed(
session, str(synth_task.id), file_id, str(e)
)
success = False
if success:
processed_count += 1
synth_task.processed_files = processed_count
# 更新任务级别的切片统计
synth_task.total_chunks = total_task_chunks
synth_task.processed_chunks = processed_task_chunks
await session.commit()
# 更新最终任务状态
if processed_count == len(file_ids):
synth_task.status = "completed"
else:
synth_task.status = "partially_completed"
# 确保任务完成时切片统计准确
if synth_task.status in ("completed", "partially_completed"):
synth_task.total_chunks = total_task_chunks
synth_task.processed_chunks = processed_task_chunks
await session.commit()
logger.info(f"Finished processing synthesis task {synth_task.id}")
def _parse_max_qa_pairs(self, synth_task: DataSynthInstance) -> int | None:
"""解析最大QA对数量配置"""
try:
cfg = Config(**(synth_task.synth_config or {}))
max_qa_pairs = (
cfg.max_qa_pairs
if (cfg and cfg.max_qa_pairs and cfg.max_qa_pairs > 0)
else None
)
except Exception:
max_qa_pairs = None
return max_qa_pairs
async def _process_single_file(
self,
session: AsyncSession,
synth_task: DataSynthInstance,
file_id: str,
max_qa_pairs: int | None = None,
) -> bool:
"""
处理单个源文件 - 使用线程池执行阻塞操作
Args:
session: 数据库会话
synth_task: 合成任务实例
file_id: 源文件ID
max_qa_pairs: 最大QA对数量限制
Returns:
处理是否成功
"""
# 解析文件路径
file_path = await self._resolve_file_path(session, file_id)
if not file_path:
logger.warning(f"File path not found for file_id={file_id}, skip")
await self._mark_file_failed(
session, str(synth_task.id), file_id, "file_path_not_found"
)
return False
logger.info(f"Processing file_id={file_id}, path={file_path}")
# 解析配置
try:
config = Config(**(synth_task.synth_config or {}))
except Exception as e:
logger.error(f"Invalid synth_config for task={synth_task.id}: {e}")
await self._mark_file_failed(
session, str(synth_task.id), file_id, "invalid_synth_config"
)
return False
# 1. 加载并切片 - 在线程池中执行阻塞操作
try:
chunks = await run_in_thread(
self._load_and_split_sync,
file_path,
config.text_split_config.chunk_size,
config.text_split_config.chunk_overlap,
)
except Exception as e:
logger.error(f"Failed to load and split file {file_id}: {e}")
await self._mark_file_failed(
session, str(synth_task.id), file_id, f"load_split_failed: {e}"
)
return False
if not chunks:
logger.warning(f"No chunks generated for file_id={file_id}")
await self._mark_file_failed(
session, str(synth_task.id), file_id, "no_chunks_generated"
)
return False
logger.info(f"File {file_id} split into {len(chunks)} chunks")
# 2. 获取文件实例并持久化切片
file_task = await self._get_or_create_file_instance(
session,
synthesis_task_id=str(synth_task.id),
source_file_id=file_id,
)
if not file_task:
logger.error(
f"DataSynthesisFileInstance not found for task={synth_task.id}, file_id={file_id}"
)
await self._mark_file_failed(
session, str(synth_task.id), file_id, "file_instance_not_found"
)
return False
# 使用 ChunkProcessor 持久化切片
chunk_processor = ChunkProcessor(session)
await chunk_processor.persist_chunks(synth_task, file_task, file_id, chunks)
total_chunks = len(chunks)
del chunks # 释放内存
# 3. 验证问答配置
question_cfg: SyntheConfig | None = config.question_synth_config
answer_cfg: SyntheConfig | None = config.answer_synth_config
if not question_cfg or not answer_cfg:
logger.error(
f"Question/Answer synth config missing for task={synth_task.id}, file={file_id}"
)
await self._mark_file_failed(
session, str(synth_task.id), file_id, "qa_config_missing"
)
return False
logger.info(
f"Start QA generation for task={synth_task.id}, file={file_id}, total_chunks={total_chunks}"
)
# 4. 创建模型客户端
question_model = await get_model_by_id(session, question_cfg.model_id)
answer_model = await get_model_by_id(session, answer_cfg.model_id)
question_chat = LLMFactory.create_chat(
question_model.model_name, question_model.base_url, question_model.api_key
)
answer_chat = LLMFactory.create_chat(
answer_model.model_name, answer_model.base_url, answer_model.api_key
)
# 5. 分批次处理切片
await self._process_chunks_in_batches(
session=session,
file_task=file_task,
total_chunks=total_chunks,
question_cfg=question_cfg,
answer_cfg=answer_cfg,
question_chat=question_chat,
answer_chat=answer_chat,
synth_task_id=str(synth_task.id),
max_qa_pairs=max_qa_pairs,
)
# 6. 标记文件任务完成
file_task.status = "completed"
await session.commit()
await session.refresh(file_task)
return True
async def _process_chunks_in_batches(
self,
session: AsyncSession,
file_task: DataSynthesisFileInstance,
total_chunks: int,
question_cfg: SyntheConfig,
answer_cfg: SyntheConfig,
question_chat,
answer_chat,
synth_task_id: str,
max_qa_pairs: int | None,
) -> None:
"""分批次处理切片
Args:
session: 数据库会话
file_task: 文件任务实例
total_chunks: 总切片数
question_cfg: 问题生成配置
answer_cfg: 答案生成配置
question_chat: 问题生成模型
answer_chat: 答案生成模型
synth_task_id: 合成任务ID
max_qa_pairs: 最大QA对数量限制
"""
batch_size = 50 # 减小批次大小,降低内存占用
current_index = 1
chunk_processor = ChunkProcessor(session)
qa_generator = QAGenerator(session)
while current_index <= total_chunks:
end_index = min(current_index + batch_size - 1, total_chunks)
# 使用 ChunkProcessor 加载切片批次
chunk_batch = await chunk_processor.load_chunk_batch(
file_task_id=file_task.id,
start_index=current_index,
end_index=end_index,
)
if not chunk_batch:
logger.warning(
f"Empty chunk batch loaded for file={file_task.id}, range=[{current_index}, {end_index}]"
)
current_index = end_index + 1
continue
# 顺序处理每个切片,避免并发过高
for chunk in chunk_batch:
try:
success = await self._process_single_chunk(
session=session,
file_task=file_task,
chunk=chunk,
question_cfg=question_cfg,
answer_cfg=answer_cfg,
question_chat=question_chat,
answer_chat=answer_chat,
synth_task_id=synth_task_id,
max_qa_pairs=max_qa_pairs,
qa_generator=qa_generator,
)
if success:
# 更新已处理的切片计数
file_task.processed_chunks = (file_task.processed_chunks or 0) + 1
# 实时刷新进度到数据库
await session.commit()
await session.refresh(file_task)
logger.info(f"File {file_task.id} progress: {file_task.processed_chunks}/{file_task.total_chunks} chunks processed")
except Exception as e:
logger.error(f"Error processing chunk {chunk.id}: {e}")
# 即使处理失败,也要更新进度
file_task.processed_chunks = (file_task.processed_chunks or 0) + 1
await session.commit()
await session.refresh(file_task)
# 短暂释放控制权,让出事件循环,处理其他请求
await asyncio.sleep(0)
current_index = end_index + 1
async def _process_single_chunk(
self,
session: AsyncSession,
file_task: DataSynthesisFileInstance,
chunk: DataSynthesisChunkInstance,
question_cfg: SyntheConfig,
answer_cfg: SyntheConfig,
question_chat,
answer_chat,
synth_task_id: str,
max_qa_pairs: int | None,
qa_generator: QAGenerator,
) -> bool:
"""处理单个切片
Args:
session: 数据库会话
file_task: 文件任务实例
chunk: 切片实例
question_cfg: 问题生成配置
answer_cfg: 答案生成配置
question_chat: 问题生成模型
answer_chat: 答案生成模型
synth_task_id: 合成任务ID
max_qa_pairs: 最大QA对数量限制
qa_generator: QA生成器
Returns:
处理是否成功
"""
# 检查QA对数量上限
if max_qa_pairs and max_qa_pairs > 0:
if await qa_generator.check_qa_limit(synth_task_id, max_qa_pairs):
logger.info(f"max_qa_pairs reached, skipping chunk {chunk.chunk_index}")
file_task.status = "completed"
if file_task.total_chunks is not None:
file_task.processed_chunks = file_task.total_chunks
await session.commit()
await session.refresh(file_task)
return False
# 使用 QAGenerator 处理问答生成
success_count = await qa_generator.process_chunk_qa(
file_task=file_task,
chunk=chunk,
question_cfg=question_cfg,
answer_cfg=answer_cfg,
question_chat=question_chat,
answer_chat=answer_chat,
)
return success_count > 0
async def _resolve_file_path(self, session: AsyncSession, file_id: str) -> str | None:
"""根据文件ID获取文件路径
Args:
session: 数据库会话
file_id: 文件ID
Returns:
文件路径
"""
result = await session.execute(
select(DatasetFiles).where(DatasetFiles.id == file_id)
)
file_obj = result.scalar_one_or_none()
if not file_obj:
return None
return file_obj.file_path
async def _get_or_create_file_instance(
self,
session: AsyncSession,
synthesis_task_id: str,
source_file_id: str,
) -> DataSynthesisFileInstance | None:
"""获取或创建文件任务实例
Args:
session: 数据库会话
synthesis_task_id: 合成任务ID
source_file_id: 源文件ID
Returns:
文件任务实例
"""
result = await session.execute(
select(DataSynthesisFileInstance).where(
DataSynthesisFileInstance.synthesis_instance_id == synthesis_task_id,
DataSynthesisFileInstance.source_file_id == source_file_id,
)
)
return result.scalar_one_or_none()
async def _mark_file_failed(
self,
session: AsyncSession,
synth_task_id: str,
file_id: str,
reason: str | None = None,
) -> None:
"""标记文件任务失败
Args:
session: 数据库会话
synth_task_id: 合成任务ID
file_id: 文件ID
reason: 失败原因
"""
try:
result = await session.execute(
select(DataSynthesisFileInstance).where(
DataSynthesisFileInstance.synthesis_instance_id == synth_task_id,
DataSynthesisFileInstance.source_file_id == file_id,
)
)
file_task = result.scalar_one_or_none()
if not file_task:
logger.warning(
f"Failed to mark file as failed: no DataSynthesisFileInstance found for task={synth_task_id}, file_id={file_id}, reason={reason}"
)
return
file_task.status = "failed"
await session.commit()
await session.refresh(file_task)
logger.info(
f"Marked file task as failed for task={synth_task_id}, file_id={file_id}, reason={reason}"
)
except Exception as e:
logger.exception(
f"Unexpected error when marking file failed for task={synth_task_id}, file_id={file_id}, original_reason={reason}, error={e}"
)
async def _get_file_ids_for_task(
self, session: AsyncSession, synth_task_id: str
):
"""获取任务关联的文件ID列表
Args:
session: 数据库会话
synth_task_id: 合成任务ID
Returns:
文件ID列表
"""
result = await session.execute(
select(DataSynthesisFileInstance.source_file_id).where(
DataSynthesisFileInstance.synthesis_instance_id == synth_task_id
)
)
return result.scalars().all()
async def _mark_task_failed(
self,
session: AsyncSession,
task_id: str,
reason: str | None = None,
) -> None:
"""标记任务失败
Args:
session: 数据库会话
task_id: 任务ID
reason: 失败原因
"""
try:
task = await session.get(DataSynthInstance, task_id)
if task:
task.status = "failed"
await session.commit()
logger.info(
f"Marked task {task_id} as failed, reason={reason}"
)
except Exception as e:
logger.exception(f"Error marking task {task_id} as failed: {e}")
async def add_synthesis_to_graph(
self, db: AsyncSession, task_id: str, dest_dataset_id: str
) -> None:
"""记录数据合成血缘关系
Args:
db: 数据库会话
task_id: 任务ID
dest_dataset_id: 目标数据集ID
"""
try:
task = await db.get(DataSynthInstance, task_id)
src_dataset_result = await db.execute(
select(DatasetFiles.dataset_id)
.join(
DataSynthesisFileInstance,
DatasetFiles.id == DataSynthesisFileInstance.source_file_id,
)
.where(DataSynthesisFileInstance.synthesis_instance_id == task_id)
.limit(1)
)
src_dataset_id = src_dataset_result.scalar_one_or_none()
src_dataset = await db.get(Dataset, src_dataset_id)
dst_dataset = await db.get(Dataset, dest_dataset_id)
if not task or not dst_dataset:
logger.warning("Missing task or destination dataset for lineage graph")
return
src_node = LineageNode(
id=src_dataset.id,
node_type=NodeType.DATASET.value,
name=src_dataset.name,
description=src_dataset.description,
)
dest_node = LineageNode(
id=dst_dataset.id,
node_type=NodeType.DATASET.value,
name=dst_dataset.name,
description=dst_dataset.description,
)
synthesis_edge = LineageEdge(
process_id=task_id,
name=task.name,
edge_type=EdgeType.DATA_SYNTHESIS.value,
description=task.description,
from_node_id=src_node.id,
to_node_id=dst_dataset.id,
)
lineage_service = LineageService(db=db)
await lineage_service.generate_graph(src_node, synthesis_edge, dest_node)
await db.commit()
logger.info(f"Added synthesis lineage: {src_node.name} -> {dest_dataset.name}")
except Exception as exc:
logger.error(f"Failed to add synthesis lineage: {exc}")
@@ -0,0 +1,138 @@
from app.module.generation.schema.generation import SynthesisType
QUESTION_GENERATOR_PROMPT=f"""# Role: 文本问题生成专家
## Profile:
- Description: 你是一名专业的文本分析与问题设计专家,能够从复杂文本中提炼关键信息并产出可用于模型微调的高质量问题集合。
- Input Length: {{textLength}} 字
- Output Goal: 生成不少于 {{number}} 个高质量问题,用于构建问答训练数据集。
## Skills:
1. 能够全面理解原文内容,识别核心概念、事实与逻辑结构。
2. 擅长设计具有明确答案指向性的问题,覆盖文本多个侧面。
3. 善于控制问题难度与类型,保证多样性与代表性。
4. 严格遵守格式规范,确保输出可直接用于程序化处理。
## Workflow:
1. **文本解析**:通读全文,分段识别关键实体、事件、数值与结论。
2. **问题设计**:基于信息密度和重要性选择最佳提问切入点。
3. **质量检查**:逐条校验问题,确保:
- 问题答案可在原文中直接找到依据。
- 问题之间主题不重复、角度不雷同。
- 语言表述准确、无歧义且符合常规问句形式。
## Constraints:
1. 所有问题必须严格依据原文内容,不得添加外部信息或假设情境。
2. 问题需覆盖文本的不同主题、层级或视角,避免集中于单一片段。
3. 禁止输出与材料元信息相关的问题(如作者、章节、目录等)。
4. 提问时请假设没有相应的文章可供参考,因此不要在问题中使用"这个"或"这些"等指示代词,也不得包含“报告/文章/文献/表格中提到”等表述。
5. 输出不少于 {{number}} 个问题,问题语言与原文主要语言保持一致。
## Output Format:
- 使用合法的 JSON 数组,仅包含字符串元素。
- 字段必须使用英文双引号。
- 严格遵循以下结构:
```
["问题1", "问题2", "..."]
```
## Output Example:
```
["人工智能伦理框架应包含哪些核心要素", "民法典对个人数据保护有哪些新规定?"]
```
## 参考原文:
{{text}}
"""
ANSWER_GENERATOR_PROMPT=f"""# Role: 微调数据生成专家
## Profile:
- Description: 你是一名微调数据生成专家,擅长基于给定内容生成准确对应的问题答案,确保答案的准确性、相关性和完整性,能够直接输出符合模型训练要求的结构化数据。
## Skills:
1. 严格基于给定内容生成答案,不添加任何外部信息
2. 答案需准确无误、逻辑通顺,与问题高度相关
3. 能够精准提取给定内容中的关键信息,并整合为自然流畅的完整答案
4. 输出结果必须符合指定的结构化格式要求
## Workflow:
1. 分析给定的参考内容,梳理核心信息和逻辑框架
2. 结合提出的具体问题,从参考内容中提取与之匹配的关键依据
3. 基于提取的依据,生成准确、详尽且符合逻辑的答案
4. 将依据内容和答案分别填入指定字段,形成结构化输出
5. 校验输出内容,确保格式正确、信息完整、无引用性表述
## Output Format:
输出格式为固定字典结构:
```json
{{
"input": "此处填入回答问题所依据的完整参考内容",
"output": "此处填入基于参考内容生成的准确答案"
}}
```
## Constrains:
1. `input`字段必须根据给定的参考内容填入回答问题的依据,不得更改原文含义
2. `output`字段的答案必须完全基于`input`中的内容,严禁编造、添加外部信息
3. 答案需充分详细,包含回答问题的所有必要信息,满足大模型微调训练的数据要求
4. 答案中不得出现「参考」「依据」「文献中提到」等任何引用性表述,仅呈现最终结论
5. 必须严格遵守指定的字典输出格式,不得额外添加其他内容
## Reference Content
------ 参考内容 Start ------
{{text}}
------ 参考内容 End ------
## Question
{{question}}
"""
COT_GENERATOR_PROMPT=f"""# Role: 微调数据生成专家
## Profile:
- Description: 你是一名微调数据生成专家,擅长基于给定参考内容,通过**思维链(COT)逐步推理**生成准确、完整的答案,输出符合大模型微调训练要求的结构化COT数据,还原从信息提取到结论推导的全思考路径。
## Skills:
1. 严格基于给定参考内容开展推理,不引入任何外部信息
2. 能够拆解问题逻辑,按步骤提取关键信息并推导,确保推理过程连贯、无跳跃
3. 生成的答案精准对应问题,逻辑通顺,与参考内容高度一致
4. 输出结果严格符合指定的结构化COT格式要求
## Workflow:
1. 分析给定参考内容,梳理核心信息、概念及逻辑关联
2. 结合具体问题,明确推理起点与目标,划定参考内容中的相关信息范围
3. 分步推导:提取关键信息→推导中间结论→结合更多信息完善逻辑→形成最终结论
4. 将完整推理过程、最终答案填入指定字段,生成结构化COT数据
5. 校验:确保推理每一步均对应参考内容,无编造信息,格式合规,无引用性表述
## Output Format:
输出固定JSON结构,包含思维链推理、最终答案两部分:
```json
{{
"chain_of_thought": "基于参考内容逐步推理的完整思维链,详述每一步提取的信息和推导的逻辑过程",
"output": "此处填入基于思维链推理得出的准确、详细的最终结论"
}}
```
## Constrains:
2. `chain_of_thought`字段需还原完整推理路径,每一步推导均需对应`Reference Content`中的具体内容,严禁主观臆断
3. `output`字段的答案必须完全来源于`Reference Content`和`chain_of_thought`的推导,不添加任何外部信息,满足大模型微调对数据质量的要求
4. 整个输出中不得出现「参考」「依据」「文献中提到」等引用性表述,仅呈现推理过程与结论
5. 必须严格遵守指定JSON格式,字段顺序固定,无额外解释或标记内容
## Reference Content
------ 参考内容 Start ------
{{text}}
------ 参考内容 End ------
## Question
{{question}}
"""
def get_prompt(synth_type: SynthesisType):
if synth_type == SynthesisType.QA:
return ANSWER_GENERATOR_PROMPT
elif synth_type == SynthesisType.COT:
return COT_GENERATOR_PROMPT
elif synth_type == SynthesisType.QUESTION:
return QUESTION_GENERATOR_PROMPT
else:
raise ValueError(f"Unsupported synthesis type: {synth_type}")
@@ -0,0 +1,349 @@
"""问答生成器模块 - 负责问题生成和答案生成"""
import asyncio
import json
import re
import uuid
from typing import Any
from langchain_core.language_models import BaseChatModel
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
from app.db.models.data_synthesis import (
DataSynthesisChunkInstance,
DataSynthesisFileInstance,
SynthesisData,
)
from app.db.session import logger
from app.module.generation.schema.generation import SyntheConfig
from app.module.generation.service.prompt import (
ANSWER_GENERATOR_PROMPT,
QUESTION_GENERATOR_PROMPT,
)
from app.module.shared.util.model_chat import extract_json_substring
from app.module.shared.llm import LLMFactory
# 图片URL提取正则
IMG_URL_PATTERN = re.compile(r"!\[\]\((.*?)\)")
def extract_img_urls(text: str) -> list[str]:
"""提取文本中的图片URL"""
return IMG_URL_PATTERN.findall(text)
def _filter_docs_by_size(docs: list[Any], chunk_size: int, ratio: float = 0.7) -> list[Any]:
"""过滤文档,保留长度大于 chunk_size * ratio 的文档
Args:
docs: 文档列表
chunk_size: 切片大小
ratio: 长度比例阈值
Returns:
过滤后的文档列表
"""
min_length = chunk_size * ratio
return [doc for doc in docs if len(doc.page_content) >= min_length]
class QuestionGenerator:
"""问题生成器"""
def __init__(self, db: AsyncSession, semaphore: asyncio.Semaphore | None = None):
self.db = db
self.semaphore = semaphore or asyncio.Semaphore(10)
async def generate_questions(
self,
chunk_text: str,
question_cfg: SyntheConfig,
question_chat: BaseChatModel,
) -> list[str]:
"""针对单个切片文本,调用模型生成问题列表
Args:
chunk_text: 切片文本内容
question_cfg: 问题生成配置
question_chat: 问题生成模型
Returns:
问题列表
"""
# 计算问题数量:每千tokens生成指定数量的问题
base_number = question_cfg.number or 5
calculated_number = max(int(len(chunk_text) / 1000 * base_number), 1)
# 获取提示模板
template = getattr(question_cfg, "prompt_template", None)
template = template if template and template.strip() else QUESTION_GENERATOR_PROMPT
# 构建提示
prompt = (
template
.replace("{text}", chunk_text)
.replace("{number}", str(calculated_number))
.replace("{textLength}", str(len(chunk_text)))
)
# 调用模型
async with self.semaphore:
loop = asyncio.get_running_loop()
raw_answer = await loop.run_in_executor(
None,
LLMFactory.invoke_sync,
question_chat,
prompt,
)
# 解析问题列表
return self._parse_questions(raw_answer)
def _parse_questions(self, raw_answer: str) -> list[str]:
"""从模型返回中解析问题列表
Args:
raw_answer: 模型原始返回
Returns:
问题列表
"""
if not raw_answer:
return []
cleaned = extract_json_substring(raw_answer)
try:
data = json.loads(cleaned)
except Exception as e:
logger.error(f"Failed to parse question list JSON: {e}")
return []
if isinstance(data, list):
return [str(q) for q in data if isinstance(q, str) and q.strip()]
# 容错:如果是单个字符串
if isinstance(data, str) and data.strip():
return [data.strip()]
return []
class AnswerGenerator:
"""答案生成器"""
def __init__(self, db: AsyncSession, semaphore: asyncio.Semaphore | None = None):
self.db = db
self.semaphore = semaphore or asyncio.Semaphore(10)
async def generate_answers_for_chunk(
self,
file_task: DataSynthesisFileInstance,
chunk: DataSynthesisChunkInstance,
questions: list[str],
answer_cfg: SyntheConfig,
answer_chat: BaseChatModel,
) -> int:
"""为一个切片的所有问题生成答案并写入数据库
Args:
file_task: 文件任务实例
chunk: 切片实例
questions: 问题列表
answer_cfg: 答案生成配置
answer_chat: 答案生成模型
Returns:
成功生成的QA对数量
"""
if not questions:
return 0
chunk_text = chunk.chunk_content or ""
template = getattr(answer_cfg, "prompt_template", None)
template = template if template and template.strip() else ANSWER_GENERATOR_PROMPT
success_count = 0
async def process_question(question: str) -> bool:
nonlocal success_count
try:
prompt = template.replace("{text}", chunk_text).replace("{question}", question)
async with self.semaphore:
loop = asyncio.get_running_loop()
answer = await loop.run_in_executor(
None,
LLMFactory.invoke_sync,
answer_chat,
prompt,
)
# 构建数据对象
data_obj = self._build_synthesis_data(chunk_text, question, answer)
# 保存记录
record = SynthesisData(
id=str(uuid.uuid4()),
data=data_obj,
synthesis_file_instance_id=file_task.id,
chunk_instance_id=chunk.id,
)
self.db.add(record)
success_count += 1
return True
except Exception as e:
logger.error(f"Failed to generate answer for question '{question[:50]}...': {e}")
return False
# 并行处理所有问题
tasks = [process_question(q) for q in questions]
await asyncio.gather(*tasks, return_exceptions=True)
if success_count > 0:
await self.db.commit()
return success_count
def _build_synthesis_data(
self,
chunk_text: str,
question: str,
answer: str,
) -> dict[str, Any]:
"""构建合成数据结构
Args:
chunk_text: 切片原文
question: 问题
answer: 答案
Returns:
结构化数据字典
"""
# 基础结构
base_obj: dict[str, Any] = {
"input": chunk_text,
"output": answer,
}
# 尝试解析模型返回的JSON
parsed_obj: dict[str, Any] | None = None
if isinstance(answer, str):
cleaned = extract_json_substring(answer)
try:
parsed = json.loads(cleaned)
if isinstance(parsed, dict):
parsed_obj = parsed
except Exception:
parsed_obj = None
# 合并解析结果
if parsed_obj is not None:
parsed_obj["instruction"] = question
data_obj = parsed_obj
else:
base_obj["instruction"] = question
data_obj = base_obj
# 提取图片URL
img_urls = extract_img_urls(chunk_text)
if img_urls:
data_obj["img_urls"] = img_urls
return data_obj
class QAGenerator:
"""问答生成器 - 整合问题和答案生成"""
def __init__(self, db: AsyncSession):
self.db = db
self.question_semaphore = asyncio.Semaphore(10)
self.answer_semaphore = asyncio.Semaphore(10)
self.question_generator = QuestionGenerator(db, self.question_semaphore)
self.answer_generator = AnswerGenerator(db, self.answer_semaphore)
async def process_chunk_qa(
self,
file_task: DataSynthesisFileInstance,
chunk: DataSynthesisChunkInstance,
question_cfg: SyntheConfig,
answer_cfg: SyntheConfig,
question_chat: Any,
answer_chat: Any,
) -> int:
"""处理单个切片的QA生成
Args:
file_task: 文件任务实例
chunk: 切片实例
question_cfg: 问题生成配置
answer_cfg: 答案生成配置
question_chat: 问题生成模型
answer_chat: 答案生成模型
Returns:
成功生成的QA对数量
"""
chunk_text = chunk.chunk_content or ""
if not chunk_text.strip():
logger.warning(f"Empty chunk text for chunk_index={chunk.chunk_index}")
return 0
# 1. 生成问题
try:
questions = await self.question_generator.generate_questions(
chunk_text=chunk_text,
question_cfg=question_cfg,
question_chat=question_chat,
)
except Exception as e:
logger.error(f"Generate questions failed: {e}")
questions = []
if not questions:
logger.info(f"No questions generated for chunk_index={chunk.chunk_index}")
return 0
# 2. 生成答案
success_count = await self.answer_generator.generate_answers_for_chunk(
file_task=file_task,
chunk=chunk,
questions=questions,
answer_cfg=answer_cfg,
answer_chat=answer_chat,
)
return success_count
async def check_qa_limit(
self,
synth_task_id: str,
max_qa_pairs: int | None,
) -> bool:
"""检查是否已达到QA对数量上限
Args:
synth_task_id: 合成任务ID
max_qa_pairs: 最大QA对数量限制
Returns:
是否已达到上限
"""
if max_qa_pairs is None or max_qa_pairs <= 0:
return False
result = await self.db.execute(
select(func.count(SynthesisData.id)).where(
SynthesisData.synthesis_file_instance_id.in_(
select(DataSynthesisFileInstance.id).where(
DataSynthesisFileInstance.synthesis_instance_id == synth_task_id
)
)
)
)
current_count = int(result.scalar() or 0)
return current_count >= max_qa_pairs
@@ -0,0 +1,98 @@
"""
合成数据查询服务 - 负责合成数据相关的查询操作
"""
from typing import List
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.db.models.data_synthesis import (
DataSynthesisFileInstance,
SynthesisData,
)
from app.module.generation.schema.generation import (
SynthesisDataItem,
SynthesisDataPatchItem,
)
class SynthesisDataQueryService:
"""合成数据查询服务 - 负责合成数据相关的查询操作"""
def __init__(self, db: AsyncSession):
self.db = db
async def list_by_chunk(self, chunk_id: str) -> List[SynthesisDataItem]:
"""
根据切片 ID 查询所有合成数据
Args:
chunk_id: 切片 ID
Returns:
合成数据列表
"""
result = await self.db.execute(
select(SynthesisData).where(SynthesisData.chunk_instance_id == chunk_id)
)
rows = result.scalars().all()
return [SynthesisDataItem.model_validate(row) for row in rows]
async def list_by_file(self, file_id: str) -> List[SynthesisDataPatchItem]:
"""
根据文件任务 ID 查询所有合成数据
Args:
file_id: 文件任务 ID
Returns:
合成数据列表
"""
result = await self.db.execute(
select(SynthesisData).where(
SynthesisData.synthesis_file_instance_id == file_id
)
)
rows = result.scalars().all()
return [SynthesisDataPatchItem.model_validate(row) for row in rows]
async def list_by_task(self, task_id: str) -> List[SynthesisDataPatchItem]:
"""
根据任务 ID 查询所有合成数据
Args:
task_id: 任务 ID
Returns:
合成数据列表
"""
result = await self.db.execute(
select(SynthesisData).where(
SynthesisData.synthesis_file_instance_id.in_(
select(DataSynthesisFileInstance.id).where(
DataSynthesisFileInstance.synthesis_instance_id == task_id
)
)
)
)
rows = result.scalars().all()
return [SynthesisDataPatchItem.model_validate(row) for row in rows]
async def update(self, data_id: str, data: dict) -> SynthesisDataItem | None:
"""
更新合成数据
Args:
data_id: 数据 ID
data: 新的数据内容
Returns:
更新后的数据,不存在则返回 None
"""
record = await self.db.get(SynthesisData, data_id)
if not record:
return None
record.data = data
await self.db.commit()
await self.db.refresh(record)
return SynthesisDataItem.model_validate(record)
@@ -0,0 +1,89 @@
"""
任务命令服务 - 负责任务的创建操作
"""
import uuid
from typing import List
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.exception import transaction
from app.db.models.data_synthesis import (
DataSynthInstance,
DataSynthesisFileInstance,
save_synthesis_task,
)
from app.db.models.dataset_management import DatasetFiles
from app.module.generation.schema.generation import CreateSynthesisTaskRequest
class TaskCreateService:
"""任务创建服务 - 负责任务的创建操作"""
def __init__(self, db: AsyncSession):
self.db = db
async def create(self, request: CreateSynthesisTaskRequest) -> DataSynthInstance:
"""
创建数据合成任务
Args:
request: 创建任务请求
Returns:
创建的任务实例
"""
dataset_files = await self._query_source_files(request.source_file_id)
request.source_file_id = [str(f.id) for f in dataset_files]
async with transaction(self.db):
task = await save_synthesis_task(self.db, request)
await self._create_file_instances(task.id, dataset_files)
return task
async def _query_source_files(self, file_ids: List[str]) -> List[DatasetFiles]:
"""
查询源文件列表
Args:
file_ids: 文件 ID 列表
Returns:
源文件列表
"""
if not file_ids:
return []
result = await self.db.execute(
select(DatasetFiles).where(DatasetFiles.id.in_(file_ids))
)
return list(result.scalars().all())
async def _create_file_instances(
self, task_id: str, dataset_files: List[DatasetFiles]
) -> None:
"""
创建文件任务实例
Args:
task_id: 任务 ID
dataset_files: 源文件列表
"""
if not dataset_files:
return
file_instances = [
DataSynthesisFileInstance(
id=str(uuid.uuid4()),
synthesis_instance_id=task_id,
file_name=f.file_name,
source_file_id=str(f.id),
status="pending",
total_chunks=0,
processed_chunks=0,
created_by="system",
updated_by="system",
)
for f in dataset_files
]
self.db.add_all(file_instances)
@@ -0,0 +1,219 @@
"""
任务删除服务 - 负责任务的删除操作
"""
from typing import List
from sqlalchemy import delete, select
from sqlalchemy.ext.asyncio import AsyncSession
from app.core.exception import BusinessError, ErrorCodes
from app.db.models.data_synthesis import (
DataSynthInstance,
DataSynthesisFileInstance,
DataSynthesisChunkInstance,
SynthesisData,
)
class TaskDeleteService:
"""任务删除服务 - 负责任务的删除操作"""
def __init__(self, db: AsyncSession):
self.db = db
async def delete_task(self, task_id: str) -> None:
"""
删除任务及其所有关联数据
Args:
task_id: 任务 ID
"""
task = await self._get_task(task_id)
if not task:
raise BusinessError(ErrorCodes.GENERATION_TASK_NOT_FOUND)
file_ids = await self._get_file_instance_ids(task_id)
if file_ids:
await self._delete_synthesis_data_by_files(file_ids)
await self._delete_chunks_by_files(file_ids)
await self._delete_file_instances(file_ids)
await self.db.delete(task)
await self.db.commit()
async def delete_file_task(self, task_id: str, file_id: str) -> None:
"""
删除文件任务并更新任务统计
Args:
task_id: 任务 ID
file_id: 文件任务 ID
"""
task = await self._get_task(task_id)
if not task:
raise BusinessError(ErrorCodes.GENERATION_TASK_NOT_FOUND)
file_task = await self._get_file_task(file_id)
if not file_task:
raise BusinessError(ErrorCodes.GENERATION_FILE_NOT_FOUND)
await self._delete_synthesis_data_by_files([file_id])
await self._delete_chunks_by_files([file_id])
await self.db.delete(file_task)
self._decrement_file_count(task)
await self.db.commit()
async def delete_chunk(self, chunk_id: str) -> None:
"""
删除切片及其关联的合成数据
Args:
chunk_id: 切片 ID
"""
chunk = await self._get_chunk(chunk_id)
if not chunk:
raise BusinessError(ErrorCodes.GENERATION_CHUNK_NOT_FOUND)
await self._delete_synthesis_data_by_chunk(chunk_id)
await self._delete_chunk(chunk_id)
await self.db.commit()
async def delete_synthesis_data_by_chunk(self, chunk_id: str) -> int:
"""
删除切片下的所有合成数据
Args:
chunk_id: 切片 ID
Returns:
删除的记录数
"""
result = await self._delete_synthesis_data_by_chunk(chunk_id)
deleted = int(getattr(result, "rowcount", 0) or 0)
await self.db.commit()
return deleted
async def batch_delete_synthesis_data(self, ids: List[str]) -> int:
"""
批量删除合成数据
Args:
ids: 合成数据 ID 列表
Returns:
删除的记录数
"""
if not ids:
return 0
result = await self.db.execute(
delete(SynthesisData).where(SynthesisData.id.in_(ids))
)
deleted = int(getattr(result, "rowcount", 0) or 0)
await self.db.commit()
return deleted
async def delete_chunks_batch(self, chunk_ids: List[str]) -> None:
"""
批量删除切片
Args:
chunk_ids: 切片 ID 列表
"""
for chunk_id in chunk_ids:
chunk = await self._get_chunk(chunk_id)
if chunk:
await self._delete_synthesis_data_by_chunk(chunk_id)
await self._delete_chunk(chunk_id)
await self.db.commit()
async def delete_chunks_by_file(self, file_id: str) -> None:
"""
按文件任务维度删除切片
Args:
file_id: 文件任务 ID
"""
await self._delete_synthesis_data_by_files([file_id])
await self._delete_chunks_by_files([file_id])
await self.db.commit()
async def delete_chunks_by_task(self, task_id: str) -> None:
"""
按任务维度删除切片
Args:
task_id: 任务 ID
"""
file_ids = await self._get_file_instance_ids(task_id)
if file_ids:
await self._delete_synthesis_data_by_files(file_ids)
await self._delete_chunks_by_files(file_ids)
await self.db.commit()
async def _get_task(self, task_id: str) -> DataSynthInstance | None:
"""获取任务"""
return await self.db.get(DataSynthInstance, task_id)
async def _get_file_task(self, file_id: str) -> DataSynthesisFileInstance | None:
"""获取文件任务"""
return await self.db.get(DataSynthesisFileInstance, file_id)
async def _get_chunk(self, chunk_id: str) -> DataSynthesisChunkInstance | None:
"""获取切片"""
return await self.db.get(DataSynthesisChunkInstance, chunk_id)
async def _get_file_instance_ids(self, task_id: str) -> List[str]:
"""获取任务下的所有文件任务 ID"""
result = await self.db.execute(
select(DataSynthesisFileInstance.id).where(
DataSynthesisFileInstance.synthesis_instance_id == task_id
)
)
return [row[0] for row in result.all()]
def _decrement_file_count(self, task: DataSynthInstance) -> None:
"""递减任务的文件数量"""
if task.total_files and task.total_files > 0:
task.total_files -= 1
if task.total_files < 0:
task.total_files = 0
async def _delete_synthesis_data_by_files(self, file_ids: List[str]) -> None:
"""批量删除文件任务下的合成数据"""
await self.db.execute(
delete(SynthesisData).where(
SynthesisData.synthesis_file_instance_id.in_(file_ids)
)
)
async def _delete_synthesis_data_by_chunk(self, chunk_id: str) -> None:
"""删除切片下的合成数据"""
await self.db.execute(
delete(SynthesisData).where(SynthesisData.chunk_instance_id == chunk_id)
)
async def _delete_chunks_by_files(self, file_ids: List[str]) -> None:
"""批量删除文件任务下的切片"""
await self.db.execute(
delete(DataSynthesisChunkInstance).where(
DataSynthesisChunkInstance.synthesis_file_instance_id.in_(file_ids)
)
)
async def _delete_chunk(self, chunk_id: str) -> None:
"""删除切片"""
await self.db.execute(
delete(DataSynthesisChunkInstance).where(
DataSynthesisChunkInstance.id == chunk_id
)
)
async def _delete_file_instances(self, file_ids: List[str]) -> None:
"""批量删除文件任务"""
await self.db.execute(
delete(DataSynthesisFileInstance).where(
DataSynthesisFileInstance.id.in_(file_ids)
)
)
@@ -0,0 +1,137 @@
"""
独立任务执行器 - 真正隔离后台任务与主事件循环
使用线程池执行CPU密集型/阻塞操作,避免阻塞主事件循环
"""
import asyncio
import concurrent.futures
from typing import Any, Callable
from functools import partial
from app.core.logging import get_logger
logger = get_logger(__name__)
# 全局线程池 - 用于执行阻塞操作
# 使用 ProcessPoolExecutor 处理CPU密集型任务,ThreadPoolExecutor 处理I/O密集型任务
_thread_pool: concurrent.futures.ThreadPoolExecutor | None = None
_process_pool: concurrent.futures.ProcessPoolExecutor | None = None
# 任务队列 - 限制并发任务数量
_task_semaphore: asyncio.Semaphore | None = None
def init_executor(
max_workers: int = 10,
max_concurrent_tasks: int = 5,
):
"""
初始化执行器
Args:
max_workers: 线程池最大工作线程数
max_concurrent_tasks: 最大并发任务数(同时执行的数据合成任务数)
"""
global _thread_pool, _process_pool, _task_semaphore
if _thread_pool is None:
_thread_pool = concurrent.futures.ThreadPoolExecutor(
max_workers=max_workers,
thread_name_prefix="generation_task_"
)
logger.info(f"Initialized thread pool with {max_workers} workers")
if _task_semaphore is None:
_task_semaphore = asyncio.Semaphore(max_concurrent_tasks)
logger.info(f"Initialized task semaphore with limit {max_concurrent_tasks}")
def get_thread_pool() -> concurrent.futures.ThreadPoolExecutor:
"""获取线程池,如果未初始化则自动初始化"""
global _thread_pool
if _thread_pool is None:
init_executor()
return _thread_pool
def get_task_semaphore() -> asyncio.Semaphore:
"""获取任务信号量,如果未初始化则自动初始化"""
global _task_semaphore
if _task_semaphore is None:
init_executor()
return _task_semaphore
async def run_in_thread[
T
](func: Callable[..., T], *args, **kwargs) -> T:
"""
在线程池中执行阻塞函数
Args:
func: 要执行的阻塞函数
*args: 位置参数
**kwargs: 关键字参数
Returns:
函数执行结果
"""
pool = get_thread_pool()
loop = asyncio.get_running_loop()
# 包装函数,添加超时和异常处理
def wrapper():
try:
return func(*args, **kwargs)
except Exception as e:
logger.error(f"Error in thread execution: {e}")
raise
try:
# 使用 asyncio.wait_for 添加超时控制
return await asyncio.wait_for(
loop.run_in_executor(pool, wrapper),
timeout=300 # 5分钟超时
)
except asyncio.TimeoutError:
logger.error(f"Thread execution timeout after 300s")
raise
async def execute_generation_task(task_id: str, process_task_func: Callable[..., Any]):
"""
执行数据合成任务,带并发控制和资源隔离
Args:
task_id: 任务ID
process_task_func: 异步任务处理函数,接收task_id作为参数
"""
semaphore = get_task_semaphore()
async with semaphore:
logger.info(f"Starting generation task {task_id} (semaphore acquired)")
try:
# 直接调用异步函数,而不是通过run_in_thread
# 因为process_task_func已经是异步函数
await process_task_func(task_id)
logger.info(f"Generation task {task_id} completed successfully")
except Exception as e:
logger.exception(f"Generation task {task_id} failed: {e}")
raise
def shutdown_executor():
"""关闭执行器,释放资源"""
global _thread_pool, _process_pool, _task_semaphore
if _thread_pool:
_thread_pool.shutdown(wait=True)
_thread_pool = None
logger.info("Thread pool shutdown")
if _process_pool:
_process_pool.shutdown(wait=True)
_process_pool = None
logger.info("Process pool shutdown")
_task_semaphore = None
@@ -0,0 +1,59 @@
from sqlalchemy import select, func
from sqlalchemy.ext.asyncio import AsyncSession
from app.db.models.data_synthesis import DataSynthInstance
from app.module.generation.schema.generation import (
DataSynthesisTaskItem,
PagedDataSynthesisTaskResponse,
)
class TaskQueryService:
"""任务查询服务 - 专注于任务列表和详情查询"""
def __init__(self, db: AsyncSession):
self.db = db
async def get_task(self, task_id: str) -> DataSynthInstance | None:
"""获取单个任务"""
return await self.db.get(DataSynthInstance, task_id)
async def list_tasks(
self,
page: int,
page_size: int,
synthesis_type: str | None = None,
status: str | None = None,
name: str | None = None,
) -> PagedDataSynthesisTaskResponse:
"""分页查询任务列表"""
query = select(DataSynthInstance)
if synthesis_type:
query = query.where(DataSynthInstance.synth_type == synthesis_type)
if status:
query = query.where(DataSynthInstance.status == status)
if name:
query = query.where(DataSynthInstance.name.like(f"%{name}%"))
query = query.order_by(DataSynthInstance.created_at.desc())
count_result = await self.db.execute(
select(func.count()).select_from(query.subquery())
)
total = count_result.scalar_one() or 0
offset = max(page, 1) - 1
limit = max(page_size, 1)
result = await self.db.execute(query.offset(offset * limit).limit(limit))
rows = result.scalars().all()
items = [DataSynthesisTaskItem.from_orm_with_config(row) for row in rows]
return PagedDataSynthesisTaskResponse(
content=items,
totalElements=total,
totalPages=(total + page_size - 1) // page_size,
page=page,
size=page_size,
)