diff --git a/src/pipeline/config.py b/src/pipeline/config.py index 646ef3a..038f615 100644 --- a/src/pipeline/config.py +++ b/src/pipeline/config.py @@ -1,6 +1,6 @@ from pydantic_settings import BaseSettings from pydantic import ConfigDict -from pydantic import MySQLDsn +from pydantic import MySQLDsn, RedisDsn class Settings(BaseSettings): @@ -25,7 +25,11 @@ class Settings(BaseSettings): es_port: int es_user: str es_password: str + mysql_dsn: MySQLDsn + celery_backend: RedisDsn + celery_broker: RedisDsn + diff --git a/src/pipeline/models/__init__.py b/src/pipeline/models/__init__.py index c5c5309..0736274 100644 --- a/src/pipeline/models/__init__.py +++ b/src/pipeline/models/__init__.py @@ -113,7 +113,7 @@ class File(BaseModel): size = fields.IntField(default=0, description="文件大小(字节)") hash_md5 = fields.CharField(max_length=32, description="MD5", db_index=True) # 业务属性 - status = fields.IntField(default=1, description="状态:1正常 2禁用 3删除") + status = fields.IntField(default=1, description="状态:1待处理/刚上传完成 2处理中 3已完成 4失败") biz_type = fields.CharField(max_length=50, description="业务类型", null=True) class Meta: diff --git a/src/pipeline/services/rag.py b/src/pipeline/services/rag.py index 30b03ae..6ca2d3d 100644 --- a/src/pipeline/services/rag.py +++ b/src/pipeline/services/rag.py @@ -60,4 +60,7 @@ class RagService: res = await TMessage.select({"chat_id": t_chat.id}, order_by=["id"]) return res + async def run_task_embedding(self): + pass + rag_service = RagService() diff --git a/src/pipeline/tasks/celery_app.py b/src/pipeline/tasks/celery_app.py new file mode 100644 index 0000000..3280f4f --- /dev/null +++ b/src/pipeline/tasks/celery_app.py @@ -0,0 +1,23 @@ +from src.pipeline.config import config +from celery import Celery + +celery_app = Celery( + "ap_pipeline", + broker=str(config.celery_broker), + backend=str(config.celery_backend), +) + +celery_app.conf.update( + task_serializer="json", + result_serializer="json", + accept_content=["json"], + # 高性能关键参数 + worker_prefetch_multiplier=1, # 防止某个 worker 吃太多任务 + task_acks_late=True, # 执行完才 ack,防止任务丢失 + task_reject_on_worker_lost=True, + broker_transport_options={ + "visibility_timeout": 3600, # 超时任务回队列 + }, + result_expires=3600, + timezone="Asia/Shanghai", +) diff --git a/src/pipeline/tasks/rag.py b/src/pipeline/tasks/rag.py index e69de29..07313cc 100644 --- a/src/pipeline/tasks/rag.py +++ b/src/pipeline/tasks/rag.py @@ -0,0 +1,41 @@ +import asyncio +from src.pipeline.tasks.celery_app import celery_app +from src.pipeline.models import File as TFile + + +@celery_app.task(bind=True, autoretry_for=(Exception,), retry_kwargs={"max_retries": 1, "countdown": 5}) +def embedding_file(self, file_uuid: str): + """ + 文件 embedding 任务 + + 状态流转: + - 1: 待处理/刚上传完成 -> 2: 处理中 + - 2: 处理中 -> 3: 已完成 (成功) + - 2: 处理中 -> 4: 失败 (失败) + """ + async def _process(): + # 查找文件记录 + file = await TFile.filter(uuid=file_uuid).first() + if not file: + raise ValueError(f"TFile not found: {file_uuid}") + + # 更新状态为处理中 + file.status = 2 + await file.save() + + try: + # TODO: 这里执行实际的 embedding 操作 + await asyncio.sleep(2) # 模拟耗时任务 + + # 更新状态为已完成 + file.status = 3 + await file.save() + return {"file_uuid": str(file_uuid), "status": "success"} + except Exception as e: + # 更新状态为失败 + file.status = 4 + await file.save() + raise e + + loop = asyncio.get_event_loop() + return loop.run_until_complete(_process())