md-vector-db/src/server/app.py

"""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}