feat: add FastAPI HTTP API layer

- Implement FastAPI server with health, collections, ingest, search, and delete endpoints
- Add Pydantic request models for ingest and search
- Add lazy singleton initialization for DB, embedder, and services
- Support environment variable override for data dir and collection name
- Add integration tests with TestClient and temp directory
- Set TRANSFORMERS_OFFLINE/HF_HUB_OFFLINE for offline model loading in tests
This commit is contained in:
2026-07-05 01:09:07 +08:00
parent 7309b0a437
commit d021390fd3
2 changed files with 229 additions and 0 deletions
+137
View File
@@ -0,0 +1,137 @@
"""FastAPI 服务层."""
import os
from pathlib import Path
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from src.core.config import load_config, EmbedConfig
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
# -- 请求模型 --
class IngestRequest(BaseModel):
file_path: str | None = None
content: str | None = None
file_name: str | None = None
class SearchRequest(BaseModel):
query: str
top_k: int = 10
# -- 懒加载单例 --
_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:
# 先尝试从 config.yaml 加载, 失败则用默认 local
try:
cfg = load_config("config.yaml")
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("config.yaml")
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.get("/api/v1/health")
def health():
return {"status": "ok"}
@app.get("/api/v1/collections")
def list_collections():
searcher = _get_searcher()
info = searcher.get_collection_info()
sources = searcher.list_sources()
return {"collections": [info], "sources": sources}
@app.post("/api/v1/ingest")
def ingest_document(req: IngestRequest):
ingestor = _get_ingestor()
try:
if req.file_path:
path = Path(req.file_path)
if not path.exists():
raise HTTPException(status_code=404, detail=f"文件不存在: {req.file_path}")
count = 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)
else:
raise HTTPException(
status_code=400, detail="需要提供 file_path 或 content"
)
return {"status": "ok", "chunks": count, "file": file_name}
except HTTPException:
raise
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/v1/search")
def search_documents(req: SearchRequest):
searcher = _get_searcher()
results = 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):
searcher = _get_searcher()
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}