123 lines
3.6 KiB
Python
123 lines
3.6 KiB
Python
"""FastAPI 服务层."""
|
|
import os
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
from fastapi import FastAPI, HTTPException, Depends, Request
|
|
from fastapi.responses import RedirectResponse
|
|
from pydantic import BaseModel, Field
|
|
|
|
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
|
|
)
|
|
|
|
|
|
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
|
|
if ".." in normalized.split(os.sep):
|
|
return False
|
|
return True
|
|
|
|
|
|
# -- App --
|
|
app = FastAPI(title="md-vector-db", version="0.1.0")
|
|
|
|
|
|
@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)):
|
|
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,
|
|
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).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"文件不存在: {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 = state.ingestor.ingest_content(req.content, file_name)
|
|
else:
|
|
raise HTTPException(status_code=400, detail="需要提供 file_path 或 content")
|
|
return {"status": "ok", "chunks": count, "file": file_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),
|
|
):
|
|
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,
|
|
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}
|