#!/usr/bin/env python3 """ Embedding Bridge - 本地 Embedding 与 Milvus 写入服务 供 Java Spring Boot 调用，完成： - 三级分块 - 本地 HuggingFace 稠密向量 - Milvus 2.5+ 原生 BM25 稀疏向量 - Leaf-only 向量化存储 启动方式： python server.py --port 18732 """ import argparse import faulthandler import logging import os import sys import threading import time from pathlib import Path from typing import List, Optional # 限制 PyTorch / OpenMP 线程数，避免在 Windows + uvicorn 多线程环境下出现 segfault os.environ.setdefault("OMP_NUM_THREADS", "1") os.environ.setdefault("MKL_NUM_THREADS", "1") os.environ.setdefault("OPENBLAS_NUM_THREADS", "1") os.environ.setdefault("NUMEXPR_NUM_THREADS", "1") os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE") faulthandler.enable() # --------------------------------------------------------------------------- # 父进程守护：Java 后端被 kill 后，子进程自动退出，释放端口与数据库连接 # --------------------------------------------------------------------------- def _start_parent_watcher(): """启动守护线程，当父进程退出时自杀。""" try: import psutil except ImportError: logging.warning("未安装 psutil，无法监听父进程状态；Java 退出后子进程可能残留") return try: parent = psutil.Process(os.getppid()) except Exception: return def _watch(): while True: time.sleep(2) try: if not parent.is_running() or parent.status() == psutil.STATUS_ZOMBIE: logging.info("父进程已退出，Embedding Bridge 自动终止") os._exit(0) except Exception: # 获取不到父进程信息时也退出，避免成为孤儿进程 logging.info("父进程状态不可获取，Embedding Bridge 自动终止") os._exit(0) watcher = threading.Thread(target=_watch, daemon=True, name="parent-watcher") watcher.start() _start_parent_watcher() # --------------------------------------------------------------------------- # 强制 stdout/stderr 使用 UTF-8 # --------------------------------------------------------------------------- for _stream in (sys.stdout, sys.stderr): try: _stream.reconfigure(encoding="utf-8", errors="replace") except Exception: pass # --------------------------------------------------------------------------- # 项目路径注入 # --------------------------------------------------------------------------- BRIDGE_DIR = Path(__file__).resolve().parent if str(BRIDGE_DIR) not in sys.path: sys.path.insert(0, str(BRIDGE_DIR)) # --------------------------------------------------------------------------- # 日志 # --------------------------------------------------------------------------- logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s", stream=sys.stderr, ) logger = logging.getLogger("embedding-bridge") # --------------------------------------------------------------------------- # FastAPI app # --------------------------------------------------------------------------- try: from fastapi import FastAPI, HTTPException from pydantic import BaseModel, Field except ImportError: logger.error("请先安装依赖: pip install -r requirements.txt") sys.exit(1) app = FastAPI(title="Embedding Bridge", version="1.0.0") # --------------------------------------------------------------------------- # 共享密钥（X-Bridge-Token 头校验） # --------------------------------------------------------------------------- _BRIDGE_AUTH_TOKEN = os.environ.get("EMBEDDING_BRIDGE_AUTH_TOKEN", "").strip() @app.middleware("http") async def _verify_token(request, call_next): """校验 X-Bridge-Token 头（/health 不校验）""" if _BRIDGE_AUTH_TOKEN and request.url.path != "/health": token = request.headers.get("X-Bridge-Token", "") if token != _BRIDGE_AUTH_TOKEN: from fastapi.responses import JSONResponse return JSONResponse(status_code=401, content={"detail": "invalid or missing X-Bridge-Token"}) return await call_next(request) # --------------------------------------------------------------------------- # 模型 # --------------------------------------------------------------------------- class IndexRequest(BaseModel): document_id: int text: str filename: str = "" file_type: str = "" file_path: str = "" page_number: int = 0 chunk_size: int = 800 chunk_overlap: int = 100 category_id: Optional[int] = None class VectorizeRequest(BaseModel): document_id: int filename: str = "" file_type: str = "" file_path: str = "" chunks: List[dict] class DeleteByDocumentRequest(BaseModel): document_id: int class DeleteByIdsRequest(BaseModel): vector_ids: List[str] class RetrieveRequest(BaseModel): query: str = Field(min_length=1) top_k: int = Field(default=5, ge=1, le=100) mode: str = "hybrid" filter_expr: str = "" class EmbedRequest(BaseModel): texts: List[str] = Field(min_length=1, max_length=500) # --------------------------------------------------------------------------- # 预导入含 C 扩展的依赖（必须在主线程完成，避免在 uvicorn 工作线程中首次加载触发 segfault） # --------------------------------------------------------------------------- import pyarrow # noqa: F401 import pandas # noqa: F401 from backend.indexing.milvus_writer import MilvusWriter from backend.indexing.text_splitter import HierarchicalTextSplitter # --------------------------------------------------------------------------- # 延迟初始化（避免启动时立即加载大模型） # --------------------------------------------------------------------------- _splitter: Optional[HierarchicalTextSplitter] = None _writer: Optional[MilvusWriter] = None _lock = threading.Lock() def _get_splitter(chunk_size: int, chunk_overlap: int): return HierarchicalTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) def _get_writer(): global _writer if _writer is None: _writer = MilvusWriter() return _writer # --------------------------------------------------------------------------- # 健康检查 # --------------------------------------------------------------------------- @app.get("/health") def health(): return {"status": "ok"} @app.post("/embed") def embed(req: EmbedRequest): try: from backend.indexing.embedding import embedding_service return {"vectors": embedding_service.get_embeddings(req.texts)} except Exception as exc: raise HTTPException(status_code=500, detail=f"embedding failed: {exc}") from exc # --------------------------------------------------------------------------- # 文档分块 + 向量化 # --------------------------------------------------------------------------- @app.post("/index") def index_document(req: IndexRequest): start = time.time() try: splitter = _get_splitter(req.chunk_size, req.chunk_overlap) chunks = splitter.split_text( text=req.text, document_id=req.document_id, filename=req.filename, file_type=req.file_type, file_path=req.file_path, page_number=req.page_number, ) if not chunks: return {"chunks": [], "collection": os.getenv("MILVUS_COLLECTION", "kb_documents"), "vector_count": 0} # 为所有 chunk 注入 document_id for c in chunks: c["document_id"] = req.document_id c["category_id"] = req.category_id writer = _get_writer() vector_ids = writer.write_chunks(chunks) leaf_index = 0 for c in chunks: if c.get("chunk_level") == 3: if leaf_index < len(vector_ids): c["vector_id"] = vector_ids[leaf_index] leaf_index += 1 # L1/L2 不设置 vector_id，避免 JSON 中出现 null cost = round((time.time() - start) * 1000) logger.info( "[index] document_id=%s chunks=%s leaf=%s cost=%sms", req.document_id, len(chunks), leaf_index, cost, ) return { "chunks": chunks, "collection": os.getenv("MILVUS_COLLECTION", "kb_documents"), "vector_count": leaf_index, } except Exception as e: logger.error("[index] document_id=%s 失败: %s", req.document_id, e, exc_info=True) raise HTTPException(status_code=500, detail=f"向量化失败: {e}") # --------------------------------------------------------------------------- # 批量 chunk 向量化（用于单个 chunk 重试） # --------------------------------------------------------------------------- @app.post("/vectorize") def vectorize_chunks(req: VectorizeRequest): try: for c in req.chunks: c["document_id"] = req.document_id c.setdefault("filename", req.filename) c.setdefault("file_type", req.file_type) c.setdefault("file_path", req.file_path) c.setdefault("page_number", 0) writer = _get_writer() vector_ids = writer.write_chunks(req.chunks) return {"vector_ids": vector_ids} except Exception as e: logger.error("[vectorize] document_id=%s 失败: %s", req.document_id, e, exc_info=True) raise HTTPException(status_code=500, detail=f"向量化失败: {e}") # --------------------------------------------------------------------------- # 按文档删除向量 # --------------------------------------------------------------------------- @app.post("/delete_by_document") def delete_by_document(req: DeleteByDocumentRequest): try: from backend.indexing.milvus_client import get_milvus_store store = get_milvus_store() count = store.delete_by_document(req.document_id) return {"deleted": count} except Exception as e: logger.error("[delete_by_document] document_id=%s 失败: %s", req.document_id, e, exc_info=True) raise HTTPException(status_code=500, detail=f"删除失败: {e}") # --------------------------------------------------------------------------- # 按 vector_id 删除向量 # --------------------------------------------------------------------------- @app.post("/delete_by_vector_ids") def delete_by_vector_ids(req: DeleteByIdsRequest): try: from backend.indexing.milvus_client import get_milvus_store store = get_milvus_store() count = store.delete_by_ids(req.vector_ids) return {"deleted": count} except Exception as e: logger.error("[delete_by_vector_ids] 失败: %s", e, exc_info=True) raise HTTPException(status_code=500, detail=f"删除失败: {e}") # --------------------------------------------------------------------------- # 主入口 # --------------------------------------------------------------------------- @app.post("/retrieve") def retrieve(req: RetrieveRequest): try: from backend.indexing.embedding import embedding_service from backend.indexing.milvus_client import get_milvus_store mode = req.mode.strip().lower() if mode not in ("dense", "hybrid"): raise HTTPException(status_code=400, detail="mode must be dense or hybrid") dense = embedding_service.get_embeddings([req.query])[0] store = get_milvus_store() if mode == "dense": hits = store.dense_retrieve(dense, req.top_k, req.filter_expr) else: hits = store.hybrid_retrieve(dense, req.query, req.top_k, filter_expr=req.filter_expr) return {"mode": mode, "top_k": req.top_k, "hits": hits} except HTTPException: raise except Exception as e: logger.error("[retrieve] failed: %s", e, exc_info=True) raise HTTPException(status_code=500, detail=f"retrieval failed: {e}") if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--port", type=int, default=int(os.getenv("EMBEDDING_BRIDGE_PORT", "18732"))) parser.add_argument("--host", type=str, default=os.getenv("EMBEDDING_BRIDGE_HOST", "127.0.0.1")) args = parser.parse_args() import uvicorn logger.info("Embedding Bridge 启动: %s:%s", args.host, args.port) uvicorn.run(app, host=args.host, port=args.port, log_level="info")