tools-api/core/registrar.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
import asyncio
import threading
from contextlib import asynccontextmanager

from fastapi import Depends, FastAPI, APIRouter
from fastapi_limiter import FastAPILimiter
from fastapi_pagination import add_pagination
from starlette.middleware.authentication import AuthenticationMiddleware

from app.router import route, call_center_route
from batch_task.batch_task import start_batch_task, stop_batch_task, execute_task, \
    periodically_execute
from batch_task.update_llm_intent import update_llm_intent
from common.exception.exception_handler import register_exception
from common.log import set_customize_logfile, setup_logging, log
from core.conf import settings
from core.path_conf import STATIC_DIR
from database.db_redis import redis_client
from middleware.jwt_call_center_auth_middleware import JwtCallCenterAuthMiddleware

from middleware.opera_log_middleware import OperaLogMiddleware
from utils.demo_site import demo_site
from utils.health_check import ensure_unique_route_names, http_limit_callback
from utils.openapi import simplify_operation_ids
from utils.serializers import MsgSpecJSONResponse

@asynccontextmanager
async def register_init(app: FastAPI):
    """
    启动初始化

    :return:
    """
    # 创建数据库表
    # await create_table()
    # 连接 redis
    await redis_client.open()
    # 初始化 limiter
    await FastAPILimiter.init(redis_client, prefix=settings.LIMITER_REDIS_PREFIX, http_callback=http_limit_callback)
    yield

    # 关闭 redis 连接
    await redis_client.close()
    # 关闭 limiter
    await FastAPILimiter.close()


def register_app():
    # FastAPI
    app = FastAPI(
        title=settings.TITLE,
        version=settings.VERSION,
        description=settings.DESCRIPTION,
        docs_url=settings.DOCS_URL,
        redoc_url=settings.REDOCS_URL,
        openapi_url=settings.OPENAPI_URL,
        default_response_class=MsgSpecJSONResponse,
        lifespan=register_init,
    )

    admin_app = FastAPI(
        title=settings.TITLE,
        version=settings.VERSION,
        description=settings.DESCRIPTION,
        default_response_class=MsgSpecJSONResponse,
    )

    call_center_app = FastAPI(
        title="Call Center",  # 子应用的标题
        version="0.1.0",  # 子应用的版本
        default_response_class=MsgSpecJSONResponse,
    )

    # 中间件
    register_call_center_middleware(call_center_app)

    # 路由
    register_router(call_center_app, call_center_route)

    # 分页
    register_page(call_center_app)

    # 全局异常处理
    register_exception(call_center_app)

    app.mount(f"/call_center", call_center_app)

    # 日志
    register_logger()

    # 静态文件
    register_static_file(admin_app)

    # 中间件
    register_middleware(admin_app)

    # 路由
    register_router(admin_app, route)

    # 分页
    register_page(admin_app)

    # 全局异常处理
    register_exception(admin_app)

    app.mount(f"{settings.API_V1_STR}/gpt", admin_app)

    return app


def register_logger() -> None:
    """
    系统日志

    :return:
    """
    setup_logging()
    set_customize_logfile()


def register_static_file(app: FastAPI):
    """
    静态文件交互开发模式, 生产使用 nginx 静态资源服务

    :param app:
    :return:
    """
    if settings.STATIC_FILES:
        import os

        from fastapi.staticfiles import StaticFiles

        if not os.path.exists(STATIC_DIR):
            os.mkdir(STATIC_DIR)
        app.mount('/static', StaticFiles(directory=STATIC_DIR), name='static')


def register_middleware(app: FastAPI):
    """
    中间件，执行顺序从下往上

    :param app:
    :return:
    """
    # Opera log
    # app.add_middleware(OperaLogMiddleware)
    # JWT auth, required
    # app.add_middleware(
    #     AuthenticationMiddleware, backend=JwtAuthMiddleware(), on_error=JwtAuthMiddleware.auth_exception_handler
    # )
    # Access log
    if settings.MIDDLEWARE_ACCESS:
        from middleware.access_middleware import AccessMiddleware

        app.add_middleware(AccessMiddleware)
    # CORS: Always at the end
    if settings.MIDDLEWARE_CORS:
        from fastapi.middleware.cors import CORSMiddleware

        app.add_middleware(
            CORSMiddleware,
            allow_origins=['*'],
            allow_credentials=True,
            allow_methods=['*'],
            allow_headers=['*'],
        )

def register_call_center_middleware(app: FastAPI):
    """
    中间件，执行顺序从下往上

    :param app:
    :return:
    """
    # Opera log (required)
    app.add_middleware(OperaLogMiddleware)
    # JWT auth (required)
    app.add_middleware(
        AuthenticationMiddleware, backend=JwtCallCenterAuthMiddleware(), on_error=JwtCallCenterAuthMiddleware.auth_exception_handler
    )
    # Access log
    if settings.MIDDLEWARE_ACCESS:
        from middleware.access_middleware import AccessMiddleware

        app.add_middleware(AccessMiddleware)
    # CORS: Always at the end
    if settings.MIDDLEWARE_CORS:
        from fastapi.middleware.cors import CORSMiddleware

        app.add_middleware(
            CORSMiddleware,
            allow_origins=['*'],
            allow_credentials=True,
            allow_methods=['*'],
            allow_headers=['*'],
        )

def register_router(app: FastAPI, route: APIRouter):
    """
    路由

    :param app: FastAPI
    :param route: APIRouter
    :return:
    """
    dependencies = [Depends(demo_site)] if settings.DEMO_MODE else None

    # API
    app.include_router(route, dependencies=dependencies)

    # Extra
    ensure_unique_route_names(app)
    simplify_operation_ids(app)


def register_page(app: FastAPI):
    """
    分页查询

    :param app:
    :return:
    """
    add_pagination(app)