from fastapi import FastAPI from fastapi.middleware.cors import CORSMiddleware from contextlib import asynccontextmanager from sqlalchemy import text import redis.asyncio as redis 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 LoggerManager, get_main_logger # 初始化日志系统 LoggerManager.setup_logging() logger = get_main_logger() def start_celery_worker(): """启动 Celery Worker""" try: logger.info("🌿 启动Celery Worker...") # 构建Worker启动命令 - 独立进程启动 worker_cmd = [ sys.executable, "-m", "celery", "-A", "app.celery_app", # 指定celery应用模块 "worker", '--loglevel=info', '--pool=solo', '--concurrency=1', '--time-limit=300', # 5分钟任务超时 '--soft-time-limit=240' # 4分钟软超时 ] subprocess.run(worker_cmd, check=True) except Exception as e: logger.error(f"❌ Celery Worker 启动失败: {e}") return None def start_celery_beat(): """启动 Celery Beat 调度器""" try: logger.info("📅 启动Celery Beat调度器...") # 使用跨平台的调度文件路径 import tempfile import os schedule_file = os.path.join(tempfile.gettempdir(), 'celerybeat-schedule') # 构建Beat启动命令 - 独立进程启动 beat_cmd = [ sys.executable, "-m", "celery", "-A", "app.celery_app", # 指定celery应用模块 "beat", '--loglevel=info', f'--schedule={schedule_file}' ] subprocess.run(beat_cmd, check=True) except Exception as e: logger.error(f"❌ Celery Beat 启动失败: {e}") def start_flower(): """启动 Flower 监控服务""" try: logger.info("📊 启动Flower监控服务...") # 构建Flower启动命令 - 独立进程启动 flower_cmd = [ sys.executable, "-m", "celery", "-A", "app.celery_app", # 指定celery应用模块 f"--broker={settings.celery_broker_url}", "flower", f"--port={settings.flower_port}" ] # 添加基础认证（如果配置了） if settings.flower_basic_auth: flower_cmd.append(f"--basic_auth={settings.flower_basic_auth}") # 添加URL前缀（如果配置了） if settings.flower_url_prefix: flower_cmd.append(f"--url_prefix={settings.flower_url_prefix}") subprocess.run(flower_cmd, check=True) except Exception as e: logger.error(f"❌ Flower 监控服务启动失败: {e}") @asynccontextmanager async def lifespan(app: FastAPI): # 启动时初始化 logger.info("🚀 应用启动中...") # Redis 连接对象 redis_client = None try: # 验证数据库连接 logger.info("📊 验证数据库连接...") try: async with engine.begin() as conn: # 执行简单的查询来验证连接 result = await conn.execute(text("SELECT 1")) connection_test = result.scalar() if connection_test == 1: logger.info("✅ 数据库连接验证成功") else: raise Exception("数据库连接测试失败") except Exception as db_error: 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 logger.info("✅ 服务初始化完成，FastAPI应用启动中...") logger.info(f"🎉 {settings.app_name} 启动完成!") yield except Exception as e: logger.error(f"❌ 应用启动失败: {e}") raise 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("👋 应用已关闭") app = FastAPI( title=settings.app_name, version="1.0.0", lifespan=lifespan, docs_url=( "/docs" if not settings.disable_docs and settings.environment != "production" else None ), redoc_url=( "/redoc" if not settings.disable_docs and settings.environment != "production" else None ), openapi_url="/openapi.json" if not settings.disable_docs else None, ) # 添加CORS中间件 app.add_middleware( CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # 注册路由 app.include_router(router) @app.get("/") async def root(): # 检查环境，生产环境下禁用根接口 if settings.environment == "production": logger.warning("🚫 生产环境下禁止访问根接口") from fastapi import HTTPException raise HTTPException(status_code=404, detail="Not Found") logger.info("📝 根接口被访问") return {"message": f"Welcome to {settings.app_name}"} @app.get("/health") async def health_check(): # 检查环境，生产环境下禁用健康检查接口 if settings.environment == "production": logger.warning("🚫 生产环境下禁止访问健康检查接口") from fastapi import HTTPException raise HTTPException(status_code=404, detail="Not Found") logger.debug("💓 健康检查接口被访问") return {"status": "healthy"} if __name__ == "__main__": import uvicorn import argparse parser = argparse.ArgumentParser(description="AI Talk Callback API") parser.add_argument("--mode", choices=["api", "worker", "beat", "flower"], help="启动模式: api(仅API), worker(仅Celery Worker), beat(仅Celery Beat), flower(仅Flower监控)") args = parser.parse_args() # 如果没有传递任何参数，输出完整的提示信息 if not args.mode: print("🚀 AI Talk Callback API 管理指南") print("=" * 50) print("\n📋 启动模式:") print(" api - 启动 FastAPI Web 应用服务 (端口: 8000)") print(" worker - 启动 Celery Worker 任务处理器") print(" beat - 启动 Celery Beat 定时任务调度器") print(" flower - 启动 Flower 监控服务") print("\n🔧 使用示例:") print(" python main.py --mode=api # 启动 Web API 服务") print(" python main.py --mode=worker # 启动任务处理器") print(" python main.py --mode=beat # 启动定时任务调度器") print(f" python main.py --mode=flower # 启动监控服务 (访问: {settings.flower_url})") print("\n🌐 服务地址:") print(" API 服务: http://localhost:8000") print(" API 文档: http://localhost:8000/docs") print(f" 任务监控界面: {settings.flower_url}") print("\n💡 提示:") print(" - 请确保 Redis 和 PostgreSQL 服务已启动") print(" - 生产环境请根据需要调整配置文件") print(" - 建议在多个终端中分别启动不同服务") sys.exit(0) # 处理启动服务请求 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 == "flower": # 仅启动 Flower 监控服务 logger.info("📊 启动Flower监控服务...") start_flower()