agent-management/backend/embedding-bridge/server.py

#!/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")