218 lines
8.0 KiB
Python
218 lines
8.0 KiB
Python
"""
|
||
FastAPI 应用入口
|
||
挂载 v1 API 路由、CORS 中间件、生命周期管理
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import logging
|
||
import threading
|
||
from contextlib import asynccontextmanager
|
||
from pathlib import Path
|
||
from typing import AsyncGenerator
|
||
|
||
from fastapi import FastAPI, Request
|
||
from fastapi.middleware.cors import CORSMiddleware
|
||
from fastapi.responses import FileResponse, JSONResponse, RedirectResponse, Response
|
||
|
||
from app.config import settings
|
||
from app.database import close_db, init_db
|
||
|
||
_logger = logging.getLogger(__name__)
|
||
|
||
|
||
class BotManager:
|
||
"""管理企业微信机器人的生命周期（在后台线程中运行）"""
|
||
|
||
def __init__(self):
|
||
self._thread: threading.Thread | None = None
|
||
self._running = False
|
||
|
||
def start(self):
|
||
"""在后台线程中启动机器人"""
|
||
if self._running:
|
||
_logger.info("企业微信机器人已在运行")
|
||
return
|
||
if not settings.WEWORK_BOT_ENABLED or not settings.WEWORK_BOT_ID or not settings.WEWORK_BOT_SECRET:
|
||
_logger.info("企业微信机器人未配置，跳过启动")
|
||
return
|
||
|
||
self._thread = threading.Thread(
|
||
target=self._run_bot, daemon=True, name="wework-bot"
|
||
)
|
||
self._thread.start()
|
||
_logger.info("企业微信机器人已在后台线程启动")
|
||
|
||
def _run_bot(self):
|
||
"""在线程中运行 bot（需要新建事件循环）"""
|
||
self._running = True
|
||
loop = asyncio.new_event_loop()
|
||
asyncio.set_event_loop(loop)
|
||
try:
|
||
loop.run_until_complete(self._run_bot_async())
|
||
except Exception as e:
|
||
_logger.error("企业微信机器人线程异常退出: %s", e, exc_info=True)
|
||
finally:
|
||
loop.close()
|
||
self._running = False
|
||
|
||
async def _run_bot_async(self):
|
||
"""异步运行 bot"""
|
||
from app.wework_bot import WeComBotService
|
||
service = WeComBotService()
|
||
await service.start_async()
|
||
|
||
async def restart(self) -> dict:
|
||
"""重启机器人（先停后启）"""
|
||
was_running = self._running
|
||
# 停止旧的（设置 enabled=False 让旧线程自然退出）
|
||
old_id = settings.WEWORK_BOT_ID
|
||
old_secret = settings.WEWORK_BOT_SECRET
|
||
old_enabled = settings.WEWORK_BOT_ENABLED
|
||
|
||
if not old_enabled or not old_id or not old_secret:
|
||
return {"success": False, "message": "请先配置机器人 ID 和 Secret 并启用"}
|
||
|
||
# 等待旧线程结束
|
||
if self._thread and self._thread.is_alive():
|
||
# SDK 的 run() 会阻塞，无法从外部优雅停止
|
||
# 通过设置 enabled=False 让下次检查时退出
|
||
_logger.info("等待旧机器人线程结束...")
|
||
# 无法强制停止，提示用户
|
||
if was_running:
|
||
return {"success": False, "message": "旧机器人正在运行，请稍后重试（SDK 不支持强制停止）"}
|
||
|
||
# 启动新的
|
||
self.start()
|
||
return {"success": True, "message": "机器人已启动"}
|
||
|
||
def get_status(self) -> dict:
|
||
"""获取机器人状态"""
|
||
configured = bool(settings.WEWORK_BOT_ID and settings.WEWORK_BOT_SECRET)
|
||
return {
|
||
"enabled": settings.WEWORK_BOT_ENABLED,
|
||
"configured": configured,
|
||
"running": self._running,
|
||
"bot_id": settings.WEWORK_BOT_ID,
|
||
}
|
||
|
||
|
||
# 全局 bot 管理器
|
||
bot_manager = BotManager()
|
||
|
||
|
||
@asynccontextmanager
|
||
async def lifespan(_app: FastAPI) -> AsyncGenerator[None, None]:
|
||
"""
|
||
应用生命周期管理：
|
||
- 启动时：初始化数据目录、验证数据库连接、预热嵌入服务
|
||
- 关闭时：释放数据库连接
|
||
"""
|
||
_logger = logging.getLogger(__name__)
|
||
|
||
# ── 启动 ──
|
||
settings.ensure_dirs()
|
||
await init_db()
|
||
|
||
# 启动企业微信机器人（后台线程）
|
||
bot_manager.start()
|
||
|
||
# 预热嵌入服务（懒加载，首次调用时初始化）
|
||
try:
|
||
from app.services.embedding_service import EmbeddingService
|
||
EmbeddingService.get_instance()
|
||
except Exception as exc:
|
||
_logger.warning("嵌入服务预热失败（将在首次使用时重试）: %s", exc)
|
||
|
||
yield
|
||
|
||
# ── 关闭 ──
|
||
await close_db()
|
||
|
||
|
||
# ──────────────────────────── 创建应用 ────────────────────────────
|
||
|
||
app = FastAPI(
|
||
title="HuiBrain - 中文直播教育知识库",
|
||
description="基于向量检索的中文直播教育知识库系统，支持直播文字稿导入、OCR 截图识别、语义搜索等功能。",
|
||
version="1.2.0",
|
||
lifespan=lifespan,
|
||
)
|
||
|
||
# ──────────────────────────── CORS 中间件 ────────────────────────────
|
||
|
||
app.add_middleware(
|
||
CORSMiddleware,
|
||
allow_origins=settings.CORS_ORIGINS,
|
||
allow_credentials=True,
|
||
allow_methods=["*"],
|
||
allow_headers=["*"],
|
||
)
|
||
|
||
|
||
# ──────────────────────────── 全局异常处理 ────────────────────────────
|
||
|
||
@app.exception_handler(Exception)
|
||
async def global_exception_handler(request: Request, exc: Exception) -> JSONResponse:
|
||
"""全局未捕获异常处理"""
|
||
logger = logging.getLogger(__name__)
|
||
logger.error("未处理异常: %s", exc, exc_info=True)
|
||
return JSONResponse(
|
||
status_code=500,
|
||
content={"detail": "服务器内部错误", "message": str(exc)},
|
||
)
|
||
|
||
|
||
# ──────────────────────────── 注册路由 ────────────────────────────
|
||
|
||
from app.api.v1.pages import router as pages_router
|
||
from app.api.v1.search import router as search_router
|
||
from app.api.v1.import_export import router as import_export_router
|
||
from app.api.v1.images import router as images_router
|
||
from app.api.v1.settings import router as settings_router
|
||
|
||
app.include_router(pages_router, prefix="/api/v1/pages", tags=["知识页面"])
|
||
app.include_router(search_router, prefix="/api/v1/search", tags=["语义搜索"])
|
||
app.include_router(import_export_router, prefix="/api/v1/import", tags=["导入导出"])
|
||
app.include_router(images_router, prefix="/api/v1/images", tags=["图片OCR"])
|
||
app.include_router(settings_router, prefix="/api/v1/settings", tags=["系统设置"])
|
||
|
||
|
||
# ──────────────────────────── 健康检查 ────────────────────────────
|
||
|
||
@app.get("/health", tags=["系统"])
|
||
async def health_check():
|
||
"""健康检查接口"""
|
||
return {"status": "ok", "version": "1.2.0"}
|
||
|
||
|
||
# ──────────────────────────── 图片文件访问 ────────────────────────────
|
||
|
||
@app.get("/data/images/{image_name}", tags=["前端"])
|
||
async def serve_image(image_name: str):
|
||
"""访问 images 目录下的图片文件（用于前端预览）"""
|
||
img_path = settings.images_dir / image_name
|
||
if not img_path.exists():
|
||
return JSONResponse(status_code=404, content={"detail": "图片不存在"})
|
||
return FileResponse(str(img_path))
|
||
|
||
|
||
# ──────────────────────────── 前端页面 ────────────────────────────
|
||
|
||
@app.get("/favicon.ico", include_in_schema=False)
|
||
async def favicon():
|
||
"""返回空响应，避免浏览器 404 报错"""
|
||
return Response(status_code=204)
|
||
|
||
|
||
@app.get("/", tags=["前端"])
|
||
async def index():
|
||
"""返回前端管理页面"""
|
||
static_dir = Path(__file__).parent.parent / "static"
|
||
index_path = static_dir / "index.html"
|
||
if index_path.exists():
|
||
return FileResponse(str(index_path))
|
||
return RedirectResponse(url="/docs")
|