#GPU 模型部署：vLLM + BGE 自定义服务版

"""BGE-M3 向量化 HTTP 服务。

本文件提供一个最小可部署的 FastAPI 服务，用于把已经下载好的本地
BGE-M3 模型发布为 HTTP 接口。
"""

import logging
import os
from typing import Dict, List, Optional

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
from pymilvus.model.hybrid import BGEM3EmbeddingFunction

load_dotenv()

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

# 全局模型对象：服务启动后只加载一次，避免每次请求重复加载模型。
bge_m3_ef: Optional[BGEM3EmbeddingFunction] = None


class HybridEmbeddingRequest(BaseModel):
    """混合向量生成请求体。

    字段保持和 knowledge.utils.bge_m3_embedding_util.generate_hybrid_embeddings
    的 embedding_documents 参数一致，传入需要生成向量的文本列表。
    """

    embedding_documents: List[str] = Field(..., description="需要生成向量的文本列表")


class HybridEmbeddingResponse(BaseModel):
    """混合向量生成响应体。

    dense 为稠密向量列表，sparse 为稀疏向量列表。
    sparse 中每个字典的 key 是 token_id，value 是对应权重。
    """

    dense: List[List[float]]
    sparse: List[Dict[int, float]]


def get_bge_m3_embedding_model() -> Optional[BGEM3EmbeddingFunction]:
    """获取 BGE-M3 嵌入模型实例。

    模型会被缓存在全局变量中，后续请求复用同一个模型对象。
    默认模型路径为 /usr-data/models/BAAI/bge-m3，可通过 BGE_M3_PATH 覆盖。
    默认使用 cuda 和 fp16，适合云 GPU 服务器；如需调整可设置环境变量：
    BGE_DEVICE=cpu 或 BGE_FP16=False。

    :return: BGEM3EmbeddingFunction 实例；加载失败时返回 None。
    """

    global bge_m3_ef

    if bge_m3_ef is not None:
        return bge_m3_ef

    model_name = os.getenv("BGE_M3_PATH", "/usr-data/models/BAAI/bge-m3")
    device = os.getenv("BGE_DEVICE", "cuda")
    use_fp16_str = os.getenv("BGE_FP16", "True")
    use_fp16 = use_fp16_str.lower() in ("true", "1", "yes")

    try:
        logger.info("开始加载 BGE-M3 模型，model_name=%s, device=%s, use_fp16=%s", model_name, device, use_fp16)
        bge_m3_ef = BGEM3EmbeddingFunction(
            model_name=model_name,
            device=device,
            use_fp16=use_fp16,
        )
        logger.info("BGE-M3 模型加载完成")
        return bge_m3_ef
    except Exception as e:
        logger.exception("BGE-M3 模型加载失败: %s", e)
        return None


def generate_hybrid_embeddings(
    embedding_model: BGEM3EmbeddingFunction,
    embedding_documents: List[str],
) -> Optional[Dict[str, List]]:
    """为文本列表生成 BGE-M3 混合向量。

    本函数参考 knowledge.utils.bge_m3_embedding_util.generate_hybrid_embeddings
    的处理逻辑，返回结构保持为：
    {
        "dense": [[...], ...],
        "sparse": [{token_id: weight, ...}, ...]
    }

    :param embedding_model: 已加载的 BGE-M3 嵌入模型。
    :param embedding_documents: 需要生成嵌入向量的文本列表。
    :return: 包含 dense 和 sparse 的字典；生成失败时返回 None。
    """

    try:
        embedding_result = embedding_model.encode_documents(embedding_documents)
        csr_array = embedding_result["sparse"]

        processed_sparse_result = []
        for index in range(len(embedding_documents)):
            # csr_array 使用 CSR 稀疏矩阵格式存储，这里按行取出每条文本的稀疏向量。
            start_ind_ptr = csr_array.indptr[index]
            end_ind_ptr = csr_array.indptr[index + 1]

            token_ids = csr_array.indices[start_ind_ptr:end_ind_ptr].tolist()
            weights = csr_array.data[start_ind_ptr:end_ind_ptr].tolist()
            sparse_vector = dict(zip(token_ids, weights))

            processed_sparse_result.append(sparse_vector)

        return {
            "dense": [dense_vector.tolist() for dense_vector in embedding_result["dense"]],
            "sparse": processed_sparse_result,
        }
    except Exception as e:
        logger.exception("BGE-M3 向量生成失败: %s", e)
        return None


def create_app() -> FastAPI:
    """创建 FastAPI 应用实例。"""

    app = FastAPI(title="BGE-M3 Embedding Service", description="BGE-M3 混合向量生成服务")

    @app.post("/v1/embeddings", response_model=HybridEmbeddingResponse)
    def generate_hybrid_embeddings_api(request: HybridEmbeddingRequest):
        """生成 BGE-M3 混合向量。

        请求示例：
        {
            "embedding_documents": ["我是中国人", "我喜欢自由"]
        }

        响应结构和 generate_hybrid_embeddings 函数保持一致，包含 dense 与 sparse。
        """

        if not request.embedding_documents:
            raise HTTPException(status_code=400, detail="embedding_documents 不能为空")

        embedding_model = get_bge_m3_embedding_model()
        if embedding_model is None:
            raise HTTPException(status_code=500, detail="BGE-M3 模型加载失败")

        embedding_result = generate_hybrid_embeddings(
            embedding_model=embedding_model,
            embedding_documents=request.embedding_documents,
        )
        if embedding_result is None:
            raise HTTPException(status_code=500, detail="BGE-M3 向量生成失败")

        return embedding_result

    return app


app = create_app()


if __name__ == "__main__":
    uvicorn.run(app=app, host="0.0.0.0", port=6006)

"""Reranker HTTP 服务。

本文件提供一个最小可部署的 FastAPI 服务，用于把本地已经下载好的
BGE Reranker 模型发布为 HTTP 接口。

接口路径：POST /v1/rerank
默认端口：6008
"""

import logging
import os
from typing import List, Optional, Tuple

import uvicorn
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException
from FlagEmbedding import FlagReranker
from pydantic import BaseModel, Field
from transformers import XLMRobertaTokenizer

load_dotenv()

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 全局模型实例：服务进程内只加载一次，避免每次请求重复加载 GPU 模型。
_reranker_model: Optional[FlagReranker] = None


class RerankRequest(BaseModel):
    """Rerank 请求体。

    sentence_pairs 与 FlagEmbedding 的 reranker.compute_score 入参保持一致：
    每个元素都是一个二元组，格式为 (query, document_content)。

    请求示例：
    {
        "sentence_pairs": [
            ["用户问题", "候选文档内容1"],
            ["用户问题", "候选文档内容2"]
        ]
    }
    """

    sentence_pairs: List[Tuple[str, str]] = Field(
        ...,
        description="待打分的 query-document 成对列表，格式为 [[query, document], ...]",
    )


class RerankResponse(BaseModel):
    """Rerank 响应体。

    scores 与 reranker.compute_score 的返回结果保持同序：
    scores[i] 对应 sentence_pairs[i] 的相关性分数。
    """

    scores: List[float] = Field(..., description="每组 query-document pair 对应的 rerank 分数")


def _patch_xlm_roberta_tokenizer() -> None:
    """为 transformers 5.x 兼容 FlagEmbedding 的旧 tokenizer 调用。

    当前业务中的 knowledge.utils.bge_rerank_util.py 已包含相同兼容逻辑。
    FlagEmbedding 的部分版本仍会调用 XLMRobertaTokenizer.prepare_for_model，
    但该方法在 transformers 5.x 中可能不存在，因此这里按现有业务代码补齐。
    """

    if hasattr(XLMRobertaTokenizer, "prepare_for_model"):
        return

    def prepare_for_model(self, ids, pair_ids=None, **kwargs):
        truncation = kwargs.get("truncation")
        max_length = kwargs.get("max_length")
        pair_ids = pair_ids or []

        if max_length and truncation in (True, "only_second"):
            special_tokens_count = 4 if pair_ids else 2
            max_pair_length = max_length - len(ids) - special_tokens_count
            if max_pair_length < len(pair_ids):
                pair_ids = pair_ids[: max(0, max_pair_length)]

        cls_token_id = getattr(self, "cls_token_id", None) or 0
        sep_token_id = getattr(self, "sep_token_id", None) or 2

        if pair_ids:
            input_ids = [cls_token_id] + ids + [sep_token_id, sep_token_id] + pair_ids + [sep_token_id]
        else:
            input_ids = [cls_token_id] + ids + [sep_token_id]
        return {"input_ids": input_ids}

    XLMRobertaTokenizer.prepare_for_model = prepare_for_model


def get_reranker_model() -> Optional[FlagReranker]:
    """获取 Reranker 模型实例。

    模型使用单例方式缓存到全局变量中。第一次请求时加载模型，后续请求直接复用。

    环境变量：
    - BGE_RERANKER_LARGE：模型路径，默认 /usr-data/models/BAAI/bge-reranker-large
    - BGE_RERANKER_DEVICE：运行设备，默认 cuda
    - BGE_RERANKER_FP16：是否使用 fp16，默认 True

    :return: FlagReranker 实例；加载失败时返回 None。
    """

    global _reranker_model

    if _reranker_model is not None:
        return _reranker_model

    _patch_xlm_roberta_tokenizer()

    model_path = os.getenv("BGE_RERANKER_LARGE", "/usr-data/models/BAAI/bge-reranker-large")
    device = os.getenv("BGE_RERANKER_DEVICE", "cuda")
    use_fp16 = os.getenv("BGE_RERANKER_FP16", "True").lower() in ("true", "1", "yes")

    try:
        logger.info(
            "开始加载 Reranker 模型：model_path=%s, device=%s, use_fp16=%s",
            model_path,
            device,
            use_fp16,
        )
        _reranker_model = FlagReranker(
            model_name_or_path=model_path,
            device=device,
            use_fp16=use_fp16,
        )
        logger.info("Reranker 模型加载完成")
        return _reranker_model
    except Exception as exc:
        logger.exception("Reranker 模型加载失败：%s", exc)
        return None


def compute_rerank_scores(reranker: FlagReranker, sentence_pairs: List[Tuple[str, str]]) -> List[float]:
    """计算 query-document pair 的 rerank 分数。

    本函数直接对齐当前业务代码中的调用方式：
    reranker.compute_score(sentence_pairs=query_doc_content_pairs)

    :param reranker: 已加载的 FlagReranker 模型实例。
    :param sentence_pairs: 待打分的 (query, document_content) 列表。
    :return: 与 sentence_pairs 顺序一致的分数列表。
    """

    scores = reranker.compute_score(sentence_pairs=sentence_pairs)

    # compute_score 在单条输入或不同版本依赖下，可能返回单个数值、list 或 numpy array。
    if isinstance(scores, (int, float)):
        return [float(scores)]
    if hasattr(scores, "tolist"):
        scores = scores.tolist()

    return [float(score) for score in scores]


def create_app() -> FastAPI:
    """创建 FastAPI 应用实例。

    这里只注册 /v1/rerank 一个接口，保持部署实现最小化。
    """

    app = FastAPI(title="BGE Reranker Service", description="BGE Reranker 重排序打分服务")

    @app.post("/v1/rerank", response_model=RerankResponse)
    def rerank_api(request: RerankRequest) -> RerankResponse:
        """Rerank 打分接口。

        入参结构参考 reranker.compute_score：
        {
            "sentence_pairs": [["query", "document"], ...]
        }

        出参结构：
        {
            "scores": [0.1, 0.2, ...]
        }
        """

        if not request.sentence_pairs:
            raise HTTPException(status_code=400, detail="sentence_pairs 不能为空")

        reranker = get_reranker_model()
        if reranker is None:
            raise HTTPException(status_code=500, detail="Reranker 模型加载失败")

        try:
            scores = compute_rerank_scores(reranker=reranker, sentence_pairs=request.sentence_pairs)
            return RerankResponse(scores=scores)
        except Exception as exc:
            logger.exception("Rerank 打分失败：%s", exc)
            raise HTTPException(status_code=500, detail="Rerank 打分失败") from exc

    return app


app = create_app()


if __name__ == "__main__":
    uvicorn.run(app=app, host="0.0.0.0", port=6008)