优化推送任务

main.py

@@ -2,22 +2,69 @@ from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
 from contextlib import asynccontextmanager
 from sqlalchemy import text
+import redis.asyncio as redis
+import threading
+import subprocess
+import sys
+import os
 
 from app.config import settings
 from app.database import engine
 from app.routes import router
 from app.logger import get_logger, LoggerManager
+from app.celery_app import celery_app
 
 
 # 初始化日志系统
 LoggerManager.setup_logging()
 logger = get_logger("main")
 
+# Celery Worker 进程管理
+celery_worker_process = None
+
+
+def start_celery_worker():
+    """启动 Celery Worker"""
+    global celery_worker_process
+    try:
+        logger.info("🌿 启动Celery Worker...")
+        
+        # 启动Celery worker
+        celery_app.start([
+            'worker',
+            '--loglevel=info',
+            '--concurrency=4',
+            '--prefetch-multiplier=1',
+            '--max-tasks-per-child=1000',
+            '--time-limit=300',  # 5分钟任务超时
+            '--soft-time-limit=240',  # 4分钟软超时
+        ])
+    except Exception as e:
+        logger.error(f"❌ Celery Worker 启动失败: {e}")
+
+
+def start_celery_beat():
+    """启动 Celery Beat 调度器"""
+    try:
+        logger.info("📅 启动Celery Beat调度器...")
+        
+        # 启动Celery beat
+        celery_app.start([
+            'beat',
+            '--loglevel=info',
+            '--schedule=/tmp/celerybeat-schedule',
+        ])
+    except Exception as e:
+        logger.error(f"❌ Celery Beat 启动失败: {e}")
+
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # 启动时初始化
     logger.info("🚀 应用启动中...")
+    
+    # Redis 连接对象
+    redis_client = None
 
     try:
         # 验证数据库连接
@@ -35,6 +82,35 @@ async def lifespan(app: FastAPI):
             logger.error(f"❌ 数据库连接验证失败: {db_error}")
             raise
 
+        # 验证 Redis 连接
+        logger.info("🔴 验证 Redis 连接...")
+        try:
+            redis_client = redis.from_url(settings.celery_broker_url)
+            # 执行 ping 命令验证连接
+            await redis_client.ping()
+            logger.info("✅ Redis 连接验证成功")
+            
+            # 存储到应用状态中供其他组件使用
+            app.state.redis_client = redis_client
+            
+        except Exception as redis_error:
+            logger.error(f"❌ Redis 连接验证失败: {redis_error}")
+            raise
+
+        # 检查启动模式
+        if len(sys.argv) > 1:
+            mode = sys.argv[1].replace("--mode=", "")
+            if mode == "all":
+                # 启动 Celery Worker (在后台线程中)
+                logger.info("🌿 启动Celery Worker...")
+                celery_worker_thread = threading.Thread(target=start_celery_worker, daemon=True)
+                celery_worker_thread.start()
+                
+                # 启动 Celery Beat 调度器 (在后台线程中)
+                logger.info("📅 启动Celery Beat调度器...")
+                celery_beat_thread = threading.Thread(target=start_celery_beat, daemon=True)
+                celery_beat_thread.start()
+
         logger.info(f"🎉 {settings.app_name} 启动完成!")
         yield
 
@@ -45,6 +121,15 @@ async def lifespan(app: FastAPI):
     finally:
         # 关闭时清理
         logger.info("🛑 应用关闭中...")
+        
+        # 关闭 Redis 连接
+        if redis_client:
+            try:
+                await redis_client.close()
+                logger.info("🔴 Redis 连接已关闭")
+            except Exception as e:
+                logger.warning(f"⚠️ 关闭 Redis 连接时出现警告: {e}")
+        
         logger.info("👋 应用已关闭")
 
 
@@ -101,10 +186,46 @@ async def health_check():
         raise HTTPException(status_code=404, detail="Not Found")
 
     logger.debug("💓 健康检查接口被访问")
-    return {"status": "healthy"}
+    
+    health_status = {"status": "healthy"}
+    
+    # 检查 Redis 连接状态
+    try:
+        redis_client = getattr(app.state, 'redis_client', None)
+        if redis_client:
+            await redis_client.ping()
+            health_status["redis"] = "connected"
+        else:
+            health_status["redis"] = "disconnected"
+    except Exception as e:
+        health_status["redis"] = f"error: {str(e)}"
+        health_status["status"] = "degraded"
+    
+    return health_status
 
 
 if __name__ == "__main__":
     import uvicorn
-
-    uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=settings.debug)
+    import argparse
+    
+    parser = argparse.ArgumentParser(description="AI Talk Callback API")
+    parser.add_argument("--mode", choices=["api", "worker", "beat", "all"], 
+                       default="api", help="启动模式: api(仅API), worker(仅Celery Worker), beat(仅Celery Beat), all(全部)")
+    args = parser.parse_args()
+    
+    if args.mode == "api":
+        # 仅启动 FastAPI 应用
+        logger.info("🚀 启动FastAPI应用...")
+        uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=settings.debug)
+    elif args.mode == "worker":
+        # 仅启动 Celery Worker
+        logger.info("🌿 启动Celery Worker...")
+        start_celery_worker()
+    elif args.mode == "beat":
+        # 仅启动 Celery Beat
+        logger.info("📅 启动Celery Beat调度器...")
+        start_celery_beat()
+    elif args.mode == "all":
+        # 启动 FastAPI + Celery Worker + Celery Beat
+        logger.info("🚀 启动完整服务栈(FastAPI + Celery Worker + Celery Beat)...")
+        uvicorn.run("main:app", host="0.0.0.0", port=8000, reload=settings.debug)