From 4114a168b47ac54af2fc64cc683459efc9b7cbf8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E8=88=AA=E5=AE=87?= <3364451258@qq.com> Date: Sun, 5 Jul 2026 01:36:40 +0800 Subject: [PATCH] refactor: FastAPI Depends injection replacing globals, real health check, logging config --- src/server/app.py | 149 ++++++++++++++++----------------------------- src/server/deps.py | 61 +++++++++++++++++++ 2 files changed, 112 insertions(+), 98 deletions(-) create mode 100644 src/server/deps.py diff --git a/src/server/app.py b/src/server/app.py index 0d24080..ad9f083 100644 --- a/src/server/app.py +++ b/src/server/app.py @@ -7,18 +7,32 @@ from fastapi import FastAPI, HTTPException, Depends, Request from fastapi.responses import RedirectResponse from pydantic import BaseModel, Field -from src.core.config import load_config, EmbedConfig, DEFAULT_CONFIG_PATH -from src.core.db import VectorDB -from src.core.embedder import create_embedder -from src.core.ingest import DocumentIngestor -from src.core.search import Searcher from src.server.auth import verify_api_key, rate_limiter +from src.server.deps import get_state, AppState +logging.basicConfig( + level=logging.INFO, + format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", +) logger = logging.getLogger("md-vector-db") +# -- 请求模型 -- +class IngestRequest(BaseModel): + file_path: str | None = None + content: str | None = None + file_name: str | None = Field( + default=None, max_length=255, pattern=r"^[^\\/:*?\"<>|]+\.md$" + ) + + +class SearchRequest(BaseModel): + query: str = Field(..., min_length=1, max_length=2000) + top_k: int = Field(default=10, ge=1, le=100) + + +# -- 路径安全检查 -- def _is_safe_path(path_str: str) -> bool: - """检查路径是否安全: 仅允许相对路径且不含 .. 穿越.""" normalized = os.path.normpath(path_str) if os.path.isabs(normalized): return False @@ -27,76 +41,8 @@ def _is_safe_path(path_str: str) -> bool: return True -# -- 请求模型 -- -class IngestRequest(BaseModel): - file_path: str | None = None - content: str | None = None - file_name: str | None = Field(default=None, max_length=255, pattern=r"^[^\\/:*?\"<>|]+\.md$") - - -class SearchRequest(BaseModel): - query: str = Field(..., min_length=1, max_length=2000) - top_k: int = Field(default=10, ge=1, le=100) - - -# -- 懒加载单例 -- -_db: VectorDB | None = None -_embedder = None -_searcher: Searcher | None = None -_ingestor: DocumentIngestor | None = None - - -def _get_db(): - global _db - if _db is None: - data_dir = os.getenv("MD_VECTOR_DB_DATA_DIR", "./data") - _db = VectorDB(persist_dir=data_dir) - return _db - - -def _get_embedder(): - global _embedder - if _embedder is None: - # 先尝试从配置文件加载, 失败则用默认 local - try: - cfg = load_config(DEFAULT_CONFIG_PATH) - embed_cfg = cfg.embed - except Exception: - embed_cfg = EmbedConfig(mode="local") - _embedder = create_embedder(embed_cfg) - return _embedder - - -def _init_services(): - global _searcher, _ingestor - collection = os.getenv("MD_VECTOR_DB_COLLECTION", "markdown_docs") - try: - cfg = load_config(DEFAULT_CONFIG_PATH) - collection = cfg.chroma.collection_name - except Exception: - pass - _searcher = Searcher(_get_db(), _get_embedder(), collection) - _ingestor = DocumentIngestor(_get_db(), _get_embedder(), collection) - - -def _get_searcher(): - if _searcher is None: - _init_services() - return _searcher - - -def _get_ingestor(): - if _ingestor is None: - _init_services() - return _ingestor - - # -- App -- -app = FastAPI( - title="md-vector-db", - description="Markdown 文档向量数据库 API", - version="0.1.0", -) +app = FastAPI(title="md-vector-db", version="0.1.0") @app.middleware("http") @@ -108,43 +54,44 @@ async def rate_limit_middleware(request: Request, call_next): @app.get("/") def root(): - """根路径重定向到 API 文档.""" - from fastapi.responses import RedirectResponse return RedirectResponse(url="/docs") @app.get("/api/v1/health") -def health(): - return {"status": "ok"} +def health(state: AppState = Depends(get_state)): + return state.is_healthy() @app.get("/api/v1/collections") -def list_collections(): - searcher = _get_searcher() - info = searcher.get_collection_info() - sources = searcher.list_sources() +def list_collections(state: AppState = Depends(get_state)): + info = state.searcher.get_collection_info() + sources = state.searcher.list_sources() return {"collections": [info], "sources": sources} @app.post("/api/v1/ingest") -def ingest_document(req: IngestRequest, _: bool = Depends(verify_api_key)): - ingestor = _get_ingestor() +def ingest_document( + req: IngestRequest, + state: AppState = Depends(get_state), + _: bool = Depends(verify_api_key), +): try: if req.file_path: if not _is_safe_path(req.file_path): raise HTTPException(status_code=400, detail="不允许的路径") - path = Path(req.file_path) + path = Path(req.file_path).resolve() + cwd = Path.cwd().resolve() + if not str(path).startswith(str(cwd)): + raise HTTPException(status_code=400, detail="不允许访问当前目录外的路径") if not path.exists(): - raise HTTPException(status_code=404, detail=f"文件不存在: {req.file_path}") - count = ingestor.ingest_file(str(path)) + raise HTTPException(status_code=404, detail=f"文件不存在: {path.name}") + count = state.ingestor.ingest_file(str(path)) file_name = path.name elif req.content: file_name = req.file_name or "untitled.md" - count = ingestor.ingest_content(req.content, file_name) + count = state.ingestor.ingest_content(req.content, file_name) else: - raise HTTPException( - status_code=400, detail="需要提供 file_path 或 content" - ) + raise HTTPException(status_code=400, detail="需要提供 file_path 或 content") return {"status": "ok", "chunks": count, "file": file_name} except HTTPException: raise @@ -154,16 +101,22 @@ def ingest_document(req: IngestRequest, _: bool = Depends(verify_api_key)): @app.post("/api/v1/search") -def search_documents(req: SearchRequest, _: bool = Depends(verify_api_key)): - searcher = _get_searcher() - results = searcher.search(req.query, top_k=req.top_k) +def search_documents( + req: SearchRequest, + state: AppState = Depends(get_state), + _: bool = Depends(verify_api_key), +): + results = state.searcher.search(req.query, top_k=req.top_k) return {"results": results} @app.delete("/api/v1/documents/{file_name}") -def delete_document(file_name: str, _: bool = Depends(verify_api_key)): - searcher = _get_searcher() - deleted = searcher.delete_by_source(file_name) +def delete_document( + file_name: str, + state: AppState = Depends(get_state), + _: bool = Depends(verify_api_key), +): + deleted = state.searcher.delete_by_source(file_name) if not deleted: raise HTTPException(status_code=404, detail=f"文档不存在: {file_name}") return {"status": "ok", "file": file_name} diff --git a/src/server/deps.py b/src/server/deps.py new file mode 100644 index 0000000..0e3746c --- /dev/null +++ b/src/server/deps.py @@ -0,0 +1,61 @@ +"""FastAPI 依赖注入 — 集中管理应用状态, 替代模块级全局变量.""" +import os +import logging + +from src.core.config import load_config +from src.core.db import VectorDB +from src.core.embedder import create_embedder +from src.core.ingest import DocumentIngestor +from src.core.search import Searcher + +logger = logging.getLogger(__name__) + + +class AppState: + """应用级共享状态.""" + + def __init__(self): + config_path = os.environ.get("MD_VECTOR_CONFIG", "config.yaml") + self.config = load_config(config_path) + + data_dir = os.environ.get( + "MD_VECTOR_DB_DATA_DIR", self.config.chroma.persist_dir + ) + self.db = VectorDB(persist_dir=data_dir) + + self.embedder = create_embedder(self.config.embed) + + collection = os.environ.get( + "MD_VECTOR_DB_COLLECTION", self.config.chroma.collection_name + ) + self.searcher = Searcher(self.db, self.embedder, collection) + self.ingestor = DocumentIngestor(self.db, self.embedder, collection) + + def is_healthy(self) -> dict: + """真实健康检查: 验证 ChromaDB 和 Embedder 是否可用.""" + status = {"status": "ok", "checks": {}} + try: + count = self.searcher.collection.count() + status["checks"]["chromadb"] = {"status": "ok", "count": count} + except Exception as e: + status["checks"]["chromadb"] = {"status": "error", "detail": str(e)} + status["status"] = "degraded" + try: + dim = self.embedder.dimension + status["checks"]["embedder"] = {"status": "ok", "dimension": dim} + except Exception as e: + status["checks"]["embedder"] = {"status": "error", "detail": str(e)} + status["status"] = "degraded" + return status + + +_state: AppState | None = None + + +def get_state() -> AppState: + """获取应用状态单例 (懒初始化).""" + global _state + if _state is None: + logger.info("初始化应用状态...") + _state = AppState() + return _state