QMDSearch/app/main.py

from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from pathlib import Path

import structlog
from fastapi import FastAPI, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse
from redis import asyncio as redis_async

from app.api.response import ApiError, error
from app.api.v1.auth import router as auth_router
from app.api.v1.document import router as document_router
from app.api.v1.knowledge import router as knowledge_router
from app.api.v1.search import router as search_router
from app.api.v1.settings import router as settings_router
from app.config import settings
from app.core.users import UserStore, bootstrap_admin
from app.services.qdrant import QdrantService

logger = structlog.get_logger()


@asynccontextmanager
async def lifespan(_: FastAPI) -> AsyncIterator[None]:
    """应用生命周期：启动时初始化 Qdrant 集合与默认管理员

    初始化失败仅记录日志、不阻止启动（本地开发可能无 Qdrant/Redis）。
    """
    try:
        await QdrantService().ensure_collections()
        logger.info("Qdrant 集合初始化完成")
    except Exception:
        logger.error("Qdrant 集合初始化失败，跳过初始化继续启动")
    try:
        # Redis 客户端创建失败（如 URL 非法）时传 None，UserStore 降级为内存模式
        redis_client = redis_async.from_url(settings.redis_url, decode_responses=True)
    except Exception:
        redis_client = None
    try:
        # 空库引导默认管理员；明文密码由 bootstrap_admin 内部 warning 打印一次
        await bootstrap_admin(UserStore(redis_client), logger)
    except Exception:
        logger.warning("默认管理员初始化失败，跳过", exc_info=True)
    try:
        Path(settings.upload_dir).mkdir(parents=True, exist_ok=True)
        logger.info("上传目录已就绪", upload_dir=settings.upload_dir)
    except Exception:
        logger.warning("上传目录初始化失败，文件落盘将按需创建", exc_info=True)
    yield


app = FastAPI(
    title=settings.app_name,
    description="AI Agent 分层信息检索服务",
    version="0.1.0",
    lifespan=lifespan,
)

app.include_router(auth_router)
app.include_router(search_router)
app.include_router(document_router)
app.include_router(knowledge_router)
app.include_router(settings_router)


@app.exception_handler(ApiError)
async def api_error_handler(request: Request, exc: ApiError) -> JSONResponse:
    """业务异常 → 统一错误响应（code 取自异常）"""
    return JSONResponse(content=error(exc.code, exc.message))


@app.exception_handler(RequestValidationError)
async def validation_error_handler(
    request: Request, exc: RequestValidationError
) -> JSONResponse:
    """请求参数校验失败 → code 1001"""
    errors = exc.errors()
    message = f"请求参数校验失败: {errors[0]['msg']}" if errors else "请求参数校验失败"
    return JSONResponse(content=error(1001, message))


@app.exception_handler(Exception)
async def unhandled_error_handler(request: Request, exc: Exception) -> JSONResponse:
    """未捕获异常 → code 2000"""
    logger.exception("未捕获异常", path=request.url.path)
    return JSONResponse(content=error(2000, "服务器内部错误"))


@app.get("/api/v1/health")
async def health() -> dict[str, str]:
    """健康检查"""
    return {"status": "ok"}


_STATIC_DIR = Path(__file__).resolve().parent / "static"
_ADMIN_HTML = _STATIC_DIR / "admin.html"
_ADMIN_SPA_DIR = _STATIC_DIR / "admin"


def _spa_index() -> Path | None:
    """SPA 入口文件路径；不存在返回 None（回退旧 admin.html）"""
    index = _ADMIN_SPA_DIR / "index.html"
    return index if index.is_file() else None


@app.get("/admin", include_in_schema=False)
async def admin_page():
    """管理后台：优先服务 Vue SPA，回退到旧单文件 admin.html"""
    spa = _spa_index()
    if spa is not None:
        return FileResponse(spa, media_type="text/html")
    return FileResponse(_ADMIN_HTML, media_type="text/html")


@app.get("/admin/", include_in_schema=False)
async def admin_page_trailing_slash():
    """带斜杠的 /admin/ 重定向到 /admin"""
    return RedirectResponse(url="/admin", status_code=307)


@app.get("/admin/{rest:path}", include_in_schema=False)
async def admin_spa(rest: str):
    """SPA 静态资源与客户端路由兜底

    - 真实文件（assets/*.js, *.css 等）→ 直接返回
    - 其余路径（/overview, /documents 等客户端路由）→ 返回 index.html
    - SPA 目录不存在时 → 404（旧 admin.html 无子路由需求）
    """
    spa_dir = _ADMIN_SPA_DIR
    if not spa_dir.is_dir():
        return JSONResponse(
            status_code=404, content={"code": 1002, "message": "Not found"}
        )

    # 安全校验：防止路径越界
    file_path = (spa_dir / rest).resolve()
    try:
        file_path.relative_to(spa_dir.resolve())
    except ValueError:
        return FileResponse(spa_dir / "index.html", media_type="text/html")

    if file_path.is_file():
        return FileResponse(file_path)
    # SPA 客户端路由兜底
    return FileResponse(spa_dir / "index.html", media_type="text/html")


# AI Agent Skill 压缩包（用于 Agent 上传文档与检索），随 ./app 卷挂载，restart 即生效
_AGENT_SKILL_ZIP = _STATIC_DIR / "agent-skill" / "QMDSearch-Agent-Skill.zip"


@app.get("/agent-skill", include_in_schema=False)
async def agent_skill_download() -> FileResponse | JSONResponse:
    """下载 AI Agent Skill 压缩包（qmdsearch-agent skill，含 SKILL.md 与示例）"""
    if not _AGENT_SKILL_ZIP.is_file():
        return JSONResponse(
            status_code=404, content={"code": 1002, "message": "Skill 包不存在"}
        )
    return FileResponse(
        _AGENT_SKILL_ZIP,
        media_type="application/zip",
        filename="QMDSearch-Agent-Skill.zip",
    )