"""FastAPI 服务层.""" import os import uuid import logging from pathlib import Path from contextlib import asynccontextmanager from fastapi import FastAPI, HTTPException, Depends, Request, Query from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import RedirectResponse from pydantic import BaseModel, Field, model_validator from src.core.security import is_path_within_workspace from src.server.auth import verify_api_key, rate_limiter from src.server.deps import get_state, AppState logger = logging.getLogger("md-vector-db") # -- 请求模型 -- class IngestRequest(BaseModel): file_path: str | None = None content: str | None = Field( default=None, max_length=500_000, description="Markdown 文本内容 (最多 500KB)", ) file_name: str | None = Field(default=None, min_length=1, max_length=255) collection: str | None = Field( default=None, max_length=128, pattern=r"^[a-zA-Z0-9_-]+$", description="目标 collection(仅允许字母数字下划线连字符)", ) @model_validator(mode="after") def _check_exclusive(self): """确保 file_path 和 content 至少提供一个.""" if not self.file_path and not self.content: raise ValueError("需要提供 file_path 或 content") return self class SearchRequest(BaseModel): query: str = Field(..., min_length=1, max_length=2000) top_k: int = Field(default=10, ge=1, le=100) collection: str | None = Field( default=None, max_length=128, pattern=r"^[a-zA-Z0-9_-]+$", description="检索的 collection(仅允许字母数字下划线连字符)", ) # -- App -- @asynccontextmanager async def lifespan(app: FastAPI): """应用启动/关闭时的日志和状态管理.""" logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", ) # 启动时检查 API Key 配置 if not os.environ.get("MD_VECTOR_API_KEY"): logger.warning("MD_VECTOR_API_KEY 未设置 — API 认证已禁用, 建议在生产环境设置密钥") yield app = FastAPI(title="md-vector-db", version="0.1.0", lifespan=lifespan) # CORS 中间件 app.add_middleware( CORSMiddleware, allow_origins=os.environ.get("CORS_ORIGINS", "http://localhost:3000").split(","), allow_credentials=False, allow_methods=["GET", "POST", "DELETE", "OPTIONS"], allow_headers=["Content-Type", "Authorization", "X-API-Key"], ) @app.middleware("http") async def security_headers_middleware(request: Request, call_next): """添加安全响应头.""" response = await call_next(request) response.headers["X-Content-Type-Options"] = "nosniff" response.headers["X-Frame-Options"] = "DENY" response.headers["X-XSS-Protection"] = "1; mode=block" response.headers["Referrer-Policy"] = "no-referrer" return response @app.middleware("http") async def rate_limit_middleware(request: Request, call_next): await rate_limiter(request) response = await call_next(request) return response @app.get("/") def root(): return RedirectResponse(url="/docs") @app.get("/api/v1/health") def health(state: AppState = Depends(get_state)): return state.is_healthy() @app.get("/api/v1/collections") def list_collections( state: AppState = Depends(get_state), _: bool = Depends(verify_api_key), ): return {"collections": state.list_collections_with_stats()} @app.post("/api/v1/ingest") def ingest_document( req: IngestRequest, state: AppState = Depends(get_state), _: bool = Depends(verify_api_key), ): ingestor = state.get_ingestor(req.collection) try: if req.file_path: if not is_path_within_workspace(req.file_path): raise HTTPException(status_code=400, detail="不允许的路径") path = Path(req.file_path).resolve() if not path.exists(): raise HTTPException(status_code=404, detail=f"文件不存在: {path.name}") count = ingestor.ingest_file(str(path)) file_name = path.name else: # content 模式 (file_path/content 互斥由 Pydantic 校验保证) file_name = req.file_name or f"untitled_{uuid.uuid4().hex[:8]}.md" count = ingestor.ingest_content(req.content, file_name) return {"status": "ok", "chunks": count, "file": file_name, "collection": ingestor.collection_name} except HTTPException: raise except Exception: logger.exception("入库失败") raise HTTPException(status_code=500, detail="服务器内部错误") @app.post("/api/v1/search") def search_documents( req: SearchRequest, state: AppState = Depends(get_state), _: bool = Depends(verify_api_key), ): try: searcher = state.get_searcher(req.collection) results = searcher.search(req.query, top_k=req.top_k) return {"results": results, "collection": searcher.collection_name} except HTTPException: raise except Exception: logger.exception("检索失败") raise HTTPException(status_code=500, detail="服务器内部错误") @app.delete("/api/v1/documents/{file_name}") def delete_document( file_name: str, state: AppState = Depends(get_state), _: bool = Depends(verify_api_key), collection: str | None = Query( default=None, max_length=128, pattern=r"^[a-zA-Z0-9_-]+$", ), ): if not file_name or len(file_name) > 512: raise HTTPException(status_code=400, detail="file_name 长度应在 1-512 之间") searcher = state.get_searcher(collection) deleted = searcher.delete_by_source(file_name) if not deleted: raise HTTPException(status_code=404, detail=f"文档不存在: {file_name}") return {"status": "ok", "file": file_name, "collection": searcher.collection_name}