ai-flow-python/app/main.py

import logging
from contextlib import asynccontextmanager

from fastapi import FastAPI

from app.config import settings
from app.controller.flow_controller import (
    init_controller,
    router,
    shutdown_controller,
)
from app.middleware.api_key import ApiKeyMiddleware
from app.repository.local_store import LocalRecordStore
from app.repository.redis_store import RedisRecordStore
from app.service.template.template_service import TemplateService

from app.util.tracking import TrackingFilter

_LOG_FORMAT = "%(asctime)s  %(levelname)-5s  [%(name)s]  %(tracking_id)s%(message)s"
_LOG_DATEFMT = "%Y-%m-%d %H:%M:%S"

logger = logging.getLogger(__name__)


def setup_logging() -> None:
    """配置根日志。使用 force=True 确保在 uvicorn dictConfig 之后仍能生效。"""
    level = getattr(logging, settings.log_level.upper(), logging.INFO)
    logging.basicConfig(
        level=level,
        format=_LOG_FORMAT,
        datefmt=_LOG_DATEFMT,
        force=True,
    )
    tracking_filter = TrackingFilter()
    for handler in logging.getLogger().handlers:
        handler.addFilter(tracking_filter)


@asynccontextmanager
async def lifespan(app: FastAPI):
    setup_logging()

    logger.info("========================================")
    logger.info("Starting %s ...", settings.app_name)
    logger.info("Log level     : %s", settings.log_level.upper())
    logger.info("Store backend : %s", settings.store_backend)
    logger.info("Executor pool : %d", settings.executor_pool_size)
    logger.info("========================================")

    if settings.store_backend == "redis":
        store = RedisRecordStore()
        logger.info("Redis record store initialized (host=%s, port=%d)",
                     settings.redis_host, settings.redis_port)
    else:
        store = LocalRecordStore()
        logger.info("Local record store initialized")

    template_service = TemplateService()
    init_controller(store, template_service)

    logger.info("%s is ready to accept requests", settings.app_name)
    yield

    shutdown_controller()
    logger.info("%s stopped.", settings.app_name)


app = FastAPI(title="AI Flow", lifespan=lifespan)
app.add_middleware(ApiKeyMiddleware)
app.include_router(router)