From d021390fd39b9d53fa13f75d28630f3c20287403 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=88=98=E8=88=AA=E5=AE=87?= <3364451258@qq.com> Date: Sun, 5 Jul 2026 01:09:07 +0800 Subject: [PATCH] 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 --- src/server/app.py | 137 ++++++++++++++++++++++++++++++++++++++++++++++ tests/test_api.py | 92 +++++++++++++++++++++++++++++++ 2 files changed, 229 insertions(+) create mode 100644 src/server/app.py create mode 100644 tests/test_api.py diff --git a/src/server/app.py b/src/server/app.py new file mode 100644 index 0000000..338fcce --- /dev/null +++ b/src/server/app.py @@ -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} diff --git a/tests/test_api.py b/tests/test_api.py new file mode 100644 index 0000000..824a601 --- /dev/null +++ b/tests/test_api.py @@ -0,0 +1,92 @@ +"""API 端点集成测试.""" +import os +import tempfile + +import pytest +from fastapi.testclient import TestClient + + +@pytest.fixture +def client(): + """创建测试客户端（使用临时目录 + 本地嵌入模型）.""" + tmpdir = tempfile.mkdtemp() + # 设置环境变量使 app 使用测试配置 + os.environ["MD_VECTOR_DB_DATA_DIR"] = tmpdir + os.environ["MD_VECTOR_DB_COLLECTION"] = "test_api" + # 离线模式: 模型已缓存, 无需联网 + os.environ["TRANSFORMERS_OFFLINE"] = "1" + os.environ["HF_HUB_OFFLINE"] = "1" + + from src.server.app import app + + with TestClient(app) as c: + yield c + + +class TestHealthEndpoint: + """健康检查.""" + + def test_health_returns_ok(self, client): + response = client.get("/api/v1/health") + assert response.status_code == 200 + data = response.json() + assert data["status"] == "ok" + + +class TestCollectionEndpoint: + """Collection 端点.""" + + def test_list_collections(self, client): + response = client.get("/api/v1/collections") + assert response.status_code == 200 + data = response.json() + assert "collections" in data + + +class TestSearchEndpoint: + """搜索端点.""" + + def test_search_requires_query(self, client): + response = client.post("/api/v1/search", json={}) + assert response.status_code == 422 # FastAPI 自动校验 + + def test_search_empty_collection(self, client): + response = client.post( + "/api/v1/search", + json={"query": "test", "top_k": 5}, + ) + assert response.status_code == 200 + data = response.json() + assert "results" in data + assert data["results"] == [] + + +class TestIngestEndpoint: + """入库端点.""" + + def test_ingest_content(self, client): + response = client.post( + "/api/v1/ingest", + json={ + "content": "# Test\nHello world.", + "file_name": "test.md", + }, + ) + assert response.status_code == 200 + data = response.json() + assert data["status"] == "ok" + assert data["chunks"] > 0 + + def test_ingest_then_search(self, client): + """入库后能检索到.""" + client.post( + "/api/v1/ingest", + json={"content": "# 配置说明\nChromaDB 配置很简单。", "file_name": "config.md"}, + ) + response = client.post( + "/api/v1/search", + json={"query": "配置", "top_k": 3}, + ) + assert response.status_code == 200 + data = response.json() + assert len(data["results"]) > 0