main.py

import httpx
import requests
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, status
from fastapi.middleware.cors import CORSMiddleware

from logger import logger

from routes.swaggerui import setupSwaggerUI
from routes.file import router as file_router
from routes.search import router as search_router


@asynccontextmanager
async def lifespan(app: FastAPI):
    # init_embedding_db()
    logger.info("FastAPI application is starting up...")
    # 在应用对外可用之前，先在后台线程执行 PDF 加载与入库操作
    try:
        # await asyncio.get_running_loop().run_in_executor(None, load_all_pdfs)
        logger.info("load_pdf completed before FastAPI startup.")
    except Exception as e:
        logger.exception(f"Error while running load_all_pdfs before startup: {e}")
    yield
    # after the application stops
    logger.info("FastAPI application is shutting down.")
    pass


app = FastAPI(
    title="KB API",
    version="2.0.0",
    description="",
    lifespan=lifespan,
    docs_url=None,
    redoc_url=None,
)

app.add_middleware(
    CORSMiddleware,  # ty:ignore[invalid-argument-type]
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.exception_handler(requests.exceptions.Timeout)
@app.exception_handler(httpx.HTTPStatusError)
async def timeout_exception_handler(request, exc):
    logger.exception("Timeout error: {}", exc)
    raise HTTPException(
        status_code=status.HTTP_504_GATEWAY_TIMEOUT, detail="内部请求超时，请稍后重试"
    )


@app.exception_handler(httpx.HTTPStatusError)
async def http_status_error_handler(request, exc):
    logger.exception("HTTP status error: {}", exc)
    raise HTTPException(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        detail="内部请求发生错误，请稍后重试",
    )


setupSwaggerUI(app)

app.include_router(file_router, tags=["文件管理"])
app.include_router(search_router, tags=["知识检索"])

if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=8000)