205 lines
7.0 KiB
Python
205 lines
7.0 KiB
Python
"""FastAPI 服务层."""
|
|
import logging
|
|
import os
|
|
import time
|
|
import uuid
|
|
from contextlib import asynccontextmanager
|
|
from pathlib import Path
|
|
|
|
from fastapi import Depends, FastAPI, HTTPException, Query, Request
|
|
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 rate_limiter, verify_api_key
|
|
from src.server.deps import AppState, get_state
|
|
|
|
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["Content-Security-Policy"] = "default-src 'self'"
|
|
response.headers["Referrer-Policy"] = "no-referrer"
|
|
return response
|
|
|
|
|
|
@app.middleware("http")
|
|
async def body_size_limit_middleware(request: Request, call_next):
|
|
"""限制请求体大小（防止内存耗尽攻击）."""
|
|
content_length = request.headers.get("content-length")
|
|
max_size = int(os.environ.get("MAX_REQUEST_BODY_SIZE", str(10 * 1024 * 1024))) # 默认 10MB
|
|
if content_length and int(content_length) > max_size:
|
|
raise HTTPException(status_code=413, detail="请求体过大")
|
|
return await call_next(request)
|
|
|
|
|
|
@app.middleware("http")
|
|
async def audit_log_middleware(request: Request, call_next):
|
|
"""记录所有 API 请求的审计日志."""
|
|
start = time.time()
|
|
response = await call_next(request)
|
|
duration_ms = (time.time() - start) * 1000
|
|
logger.info(
|
|
"audit: %s %s → %d (%.1fms) [%s]",
|
|
request.method, request.url.path,
|
|
response.status_code, duration_ms,
|
|
request.client.host if request.client else "unknown",
|
|
)
|
|
return response
|
|
|
|
|
|
@app.middleware("http")
|
|
async def rate_limit_middleware(request: Request, call_next):
|
|
# 健康检查和根路径不需要速率限制
|
|
if request.url.path in ("/api/v1/health", "/"):
|
|
return await call_next(request)
|
|
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),
|
|
):
|
|
try:
|
|
return {"collections": state.list_collections_with_stats()}
|
|
except Exception:
|
|
logger.exception("列出集合失败")
|
|
raise HTTPException(status_code=500, detail="服务器内部错误")
|
|
|
|
|
|
@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}
|