enterprise-kb-rag/3.3 基于Mysql库的问答系统/integrated_qa_system/mysql_qa/retrieval/bm25_search.py

# 导入 BM25 算法
from rank_bm25 import BM25Okapi
# 导入数值计算库
import numpy as np
# 导入文本预处理
from mysql_qa.utils.preprocess import preprocess_text
# 导入日志
from base import logger


class BM25Search:
    def __init__(self, redis_client, mysql_client):
        # 初始化日志与数据客户端
        self.logger = logger
        self.redis_client = redis_client
        self.mysql_client = mysql_client
        self.bm25 = None
        self.questions = None
        self.original_questions = None
        self._load_data()

    def _load_data(self):
        # Redis 键的数据结构：str
        original_key = "qa_original_questions"
        tokenized_key = "qa_tokenized_questions"

        # Redis 中的值由 set_data() 使用 json.dumps() 序列化，因此实际存储的是 JSON 字符串：
        # qa_original_questions -> '["如何修改密码？", "如何申请退款？"]'
        # qa_tokenized_questions -> '[["如何", "修改密码", "？"], ["如何", "申请", "退款", "？"]]'
        # get_data() 使用 json.loads() 反序列化后，恢复为下面对应的 Python 列表结构。

        # 优先从 Redis 获取问题及其分词结果
        # 数据结构：list[str] | None
        # 示例：["如何修改密码？", "如何申请退款？"]
        # {"qa_original_questions":""}
        # {"qa_tokenized_questions":""}
        self.original_questions = self.redis_client.get_data(original_key)

        # 数据结构：list[list[str]] | None
        # 示例：[["如何", "修改密码", "？"], ["如何", "申请", "退款", "？"]]
        tokenized_questions = self.redis_client.get_data(tokenized_key)

        # Redis 中没有数据时，从 MySQL 加载
        if not self.original_questions or not tokenized_questions:
            # 默认 PyMySQL Cursor.fetchall() 的数据结构：tuple[tuple[str], ...]
            # 因为 SQL 只查询 question 一列，所以每一行都是仅包含一个问题字符串的元组
            # 示例：rows = (("如何修改密码？",), ("如何申请退款？",), ("忘记密码怎么办？",))
            rows = self.mysql_client.fetch_questions()
            if not rows:
                self.logger.warning("未加载到问题")
                return

            # 逐行取出第 1 列的字符串，将“元组的集合”转换成“问题字符串列表”
            # row[0]：从 ("如何修改密码？",) 中取出 "如何修改密码？"
            # isinstance 判断用于兼容 row 为 tuple、list 或已经是字符串的情况
            self.original_questions = [
                row[0] if isinstance(row, (tuple, list)) else row for row in rows
            ]
            # 转换后的数据结构：list[str]
            # 示例：self.original_questions = ["如何修改密码？", "如何申请退款？", "忘记密码怎么办？"]

            # 遍历 list[str] 类型的原始问题，每个问题经 preprocess_text() 转成 list[str]
            # 最终数据结构：list[list[str]]，外层每个元素对应一个问题的分词结果
            # 示例（具体分词结果由 Jieba 词典决定）：
            # self.original_questions = ["如何修改密码？", "如何申请退款？"]
            # tokenized_questions = [["如何", "修改密码", "？"], ["如何", "申请", "退款", "？"]]
            tokenized_questions = [
                preprocess_text(question) for question in self.original_questions
            ]
            self.redis_client.set_data(original_key, self.original_questions)
            self.redis_client.set_data(tokenized_key, tokenized_questions)

        # 数据结构：list[list[str]]，外层列表对应所有问题，内层列表是单个问题的分词结果
        # 示例：[["如何", "修改", "密码"], ["如何", "申请", "退款"], ["忘记", "密码", "怎么办"]]
        self.questions = tokenized_questions
        # 使用全部分词问题构建 BM25 检索模型，供后续计算查询与各问题的相关性分数
        self.bm25 = BM25Okapi(self.questions)
        self.logger.info("BM25 模型初始化完成")

    def _softmax(self, scores):
        # scores：numpy.ndarray，形状为 (问题总数,)，保存每个问题的 BM25 原始分数
        # 示例：scores = np.array([0.3, 1.2, 3.5])

        # 先减去最大值，避免指数运算时数值过大：[-3.2, -2.3, 0.0]
        # np.exp() 会对数组中的每个元素计算自然指数 e^x，其中 e ≈ 2.71828
        # 公式：exp_scores[i] = e ** (scores[i] - max(scores))
        # 示例：[e^-3.2, e^-2.3, e^0] ≈ [0.041, 0.100, 1.000]
        # exp_scores 的数据结构仍为 numpy.ndarray，形状与 scores 相同
        exp_scores = np.exp(scores - np.max(scores))

        # 返回相同形状的 numpy.ndarray；各元素之和为 1
        # 示例：[0.036, 0.088, 0.876]
        return exp_scores / exp_scores.sum()

    def search(self, query, threshold=0.85):
        if not query or not isinstance(query, str):
            self.logger.error("无效查询")
            return None, False

        cached_answer = self.redis_client.get_answer(query)
        if cached_answer:
            return cached_answer, False

        if self.bm25 is None or not self.original_questions:
            self.logger.warning("问题库为空，无法执行检索")
            return None, True

        try:
            # 数据结构：list[str]，表示查询语句的分词结果
            # 示例：query = "忘记密码怎么办" -> query_tokens = ["忘记", "密码", "怎么办"]
            query_tokens = preprocess_text(query)

            # 数据结构：numpy.ndarray，形状为 (问题总数,)
            # 每个元素是查询与对应问题的 BM25 原始相关性分数
            # 示例：scores = np.array([0.3, 1.2, 3.5])
            scores = self.bm25.get_scores(query_tokens)

            # 数据结构：numpy.ndarray，形状不变；所有元素之和为 1
            # 与上面的示例对应：softmax_scores ≈ np.array([0.036, 0.088, 0.876])
            softmax_scores = self._softmax(scores)

            # 数据结构：numpy 整数标量（通常为 numpy.int64）
            # argmax() 返回最大分数的位置；示例中 best_idx = 2，表示第 3 个问题
            best_idx = softmax_scores.argmax()

            # 数据结构：numpy 浮点标量（通常为 numpy.float64）
            # 根据最佳位置取出最高分；示例中 best_score ≈ 0.876
            best_score = softmax_scores[best_idx]

            if best_score >= threshold:
                original_question = self.original_questions[best_idx]
                answer = self.redis_client.get_answer(original_question)
                if answer:
                    self.logger.info(f"从 Redis 中获取到了问题的答案: {answer}")
                else:
                    answer = self.mysql_client.fetch_answer(original_question)
                    self.logger.info(f"从 MySQL 中获取到了问题的答案: {answer}")

                if answer:
                    self.redis_client.set_data(f"answer:{query}", answer)
                    self.logger.info(
                        f"搜索成功，Softmax 相似度: {best_score:.3f}"
                    )
                    return answer, False

            self.logger.info(
                f"未找到可靠答案，最高 Softmax 相似度: {best_score:.3f}"
            )
            return None, True
        except Exception as exc:
            self.logger.error(f"搜索失败: {exc}")
            return None, True