feat: 任务处理

This commit is contained in:
李如威 2026-01-09 21:03:21 +08:00
parent bbf8ccfb4a
commit b1b236cb20
5 changed files with 73 additions and 2 deletions

View File

@ -1,6 +1,6 @@
from pydantic_settings import BaseSettings from pydantic_settings import BaseSettings
from pydantic import ConfigDict from pydantic import ConfigDict
from pydantic import MySQLDsn from pydantic import MySQLDsn, RedisDsn
class Settings(BaseSettings): class Settings(BaseSettings):
@ -25,7 +25,11 @@ class Settings(BaseSettings):
es_port: int es_port: int
es_user: str es_user: str
es_password: str es_password: str
mysql_dsn: MySQLDsn mysql_dsn: MySQLDsn
celery_backend: RedisDsn
celery_broker: RedisDsn

View File

@ -113,7 +113,7 @@ class File(BaseModel):
size = fields.IntField(default=0, description="文件大小(字节)") size = fields.IntField(default=0, description="文件大小(字节)")
hash_md5 = fields.CharField(max_length=32, description="MD5", db_index=True) 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) biz_type = fields.CharField(max_length=50, description="业务类型", null=True)
class Meta: class Meta:

View File

@ -60,4 +60,7 @@ class RagService:
res = await TMessage.select({"chat_id": t_chat.id}, order_by=["id"]) res = await TMessage.select({"chat_id": t_chat.id}, order_by=["id"])
return res return res
async def run_task_embedding(self):
pass
rag_service = RagService() rag_service = RagService()

View File

@ -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",
)

View File

@ -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())