AI之RAG系列-问题增强

# 问题增强示例import loggingfrom typing import List, Optional, Dict, Any, Protocolfrom dataclasses import dataclass, fieldimport numpy as npfrom openai import OpenAI as _OpenAIfrom rag_embedding import RAGEmbedding# 调用向量化模块的RAGEmbedding类logging.basicConfig(    level=logging.INFO,    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",)logger = logging.getLogger(__name__)# ==================== embedding模型与查询单元 ====================class EmbedderProtocol(Protocol):    """Embedding 模型协议，支持语义去重时注入"""    def __init__(self,        model_name: str = "BAAI/bge-small-zh-v1.5",        device: str = "auto",        cache_dir: Optional[str] = None        ):        self.embedding=RAGEmbedding(            model_name=model_name,            device=device,            normalize=True,            cache_dir=cache_dir            )    def embed_query(self, text: str) -> List[float]:        return self.embedding.encode(text).tolist()@dataclassclass EnhancedQuery:    """增强后的查询单元"""    text: str    query_type: str = "original"  # original / rewrite / hyde    def get_search_text(self, query_instruction: Optional[str] = None) -> str:        """        获取用于检索的文本。        规则：            - original / rewrite: 加 query_instruction 前缀（BGE 需要）            - hyde: 直接返回，不加前缀（已是文档形态）        """        if self.query_type == "hyde":            return self.text        if query_instruction:            return query_instruction + self.text        return self.text# ==================== 配置类 ====================@dataclassclass QueryEnhancementConfig:    """问题增强配置"""    enabled: bool = False    method: str = "rewrite"                    # "rewrite" / "hyde" / "both"    num_variants: int = 3    temperature_rewrite: float = 0.7    temperature_hyde: float = 0.3    max_tokens_rewrite: int = 256    max_tokens_hyde: int = 300    # 语义去重配置    semantic_dedup_enabled: bool = True          # 是否启用向量语义去重    semantic_dedup_threshold: float = 0.95       # 余弦相似度阈值，>= 则丢弃    prompt_rewrite: str = """你是一个问题改写专家。请基于用户的原始问题，生成 {num_variants} 个语义相近但表述不同的问题，用于帮助检索系统更全面地召回相关信息。原始问题: {query}请直接输出改写后的问题，每行一个，不要添加编号或额外说明："""    prompt_hyde: str = """你是一位知识库专家。请根据用户的问题，撰写一段详细的、包含关键信息的答案片段（100~200字）。这段文字不需要完全准确，但应尽可能包含与用户问题相关的专业术语和核心概念，以便用于从知识库中检索相关的真实文档。用户问题: {query}假设答案："""@dataclassclass LLMConfig:    """LLM 客户端配置"""    base_url: str = "http://localhost:8000/v1"    api_key: str = "EMPTY"    model_name: str = "Qwen2.5-7B-Instruct"    top_p: float = 0.9# ==================== 余弦相似函数 ====================def _cosine_similarity(a: List[float], b: List[float]) -> float:    """计算两个向量的余弦相似度"""    a_arr = np.array(a, dtype=np.float32)    b_arr = np.array(b, dtype=np.float32)    norm_a = np.linalg.norm(a_arr)    norm_b = np.linalg.norm(b_arr)    if norm_a == 0 or norm_b == 0:        return 0.0    return float(np.dot(a_arr, b_arr) / (norm_a * norm_b))# ==================== 核心增强类 ====================class QueryEnhancer:    """    问题增强器（含向量语义去重）    去重策略（按顺序执行）：        1. 字符串去重：前100字符精确匹配，快速过滤完全重复项        2. 语义去重（可选）：基于 Embedding 余弦相似度，过滤语义等价项    """    def __init__(        self,        config: Optional[QueryEnhancementConfig] = None,        llm_config: Optional[LLMConfig] = None,        llm_client: Optional[Any] = None,        embedding_model: Optional[EmbedderProtocol] = None,  # 用于语义去重    ):        self.config = config or QueryEnhancementConfig()        self.llm_config = llm_config or LLMConfig()        self.embedding_model = embedding_model        if llm_client is not None:            self.llm_client = llm_client        else:            self.llm_client = _OpenAI(                base_url=self.llm_config.base_url,                api_key=self.llm_config.api_key,            )    def enhance(self, query: str) -> List[EnhancedQuery]:        """        执行问题增强，返回增强后的查询列表（首项始终为原始查询）        """        results: List[EnhancedQuery] = [            EnhancedQuery(text=query, query_type="original")        ]        if not self.config.enabled:            logger.info("[Enhance] 问题增强已禁用，返回原始查询")            return results        method = self.config.method.lower()        if method in ("rewrite", "both"):            rewrite_results = self._rewrite(query)            results.extend(rewrite_results)            logger.info(f"[Enhance] 查询改写生成 {len(rewrite_results)} 个变体")        if method in ("hyde", "both"):            hyde_results = self._hyde(query)            results.extend(hyde_results)            logger.info(f"[Enhance] HyDE 生成 {len(hyde_results)} 个假设文档")        # 双层去重        results = self._deduplicate(results)        logger.info(f"[Enhance] 最终返回 {len(results)} 个查询（含原始查询）")        return results    def _rewrite(self, query: str) -> List[EnhancedQuery]:        """查询改写：调用 LLM 生成语义相近的查询变体"""        cfg = self.config        if cfg.num_variants <= 0:            return []        try:            prompt = cfg.prompt_rewrite.format(                query=query, num_variants=cfg.num_variants            )            response = self.llm_client.chat.completions.create(                model=self.llm_config.model_name,                messages=[                    {"role": "system", "content": "你是一个专业的问题改写助手。"},                    {"role": "user", "content": prompt},                ],                temperature=cfg.temperature_rewrite,                max_tokens=cfg.max_tokens_rewrite,                top_p=self.llm_config.top_p,            )            generated = response.choices[0].message.content.strip()            variants = [line.strip() for line in generated.split('\n') if line.strip()]            variants = variants[: cfg.num_variants]            results = []            for v in variants:                if v != query:                    results.append(EnhancedQuery(text=v, query_type="rewrite"))            return results        except Exception as e:            logger.warning(f"[Enhance] 查询改写失败: {e}")            return []    def _hyde(self, query: str) -> List[EnhancedQuery]:        """HyDE：调用 LLM 生成假设答案文档"""        cfg = self.config        try:            prompt = cfg.prompt_hyde.format(query=query)            response = self.llm_client.chat.completions.create(                model=self.llm_config.model_name,                messages=[                    {"role": "system", "content": "你是一位知识库专家，擅长撰写用于检索的假设答案。"},                    {"role": "user", "content": prompt},                ],                temperature=cfg.temperature_hyde,                max_tokens=cfg.max_tokens_hyde,                top_p=self.llm_config.top_p,            )            hypothetical_doc = response.choices[0].message.content.strip()            if not hypothetical_doc:                logger.warning("[Enhance] HyDE 生成空文档")                return []            return [EnhancedQuery(text=hypothetical_doc, query_type="hyde")]        except Exception as e:            logger.warning(f"[Enhance] HyDE 生成失败: {e}")            return []    # ==================== 双层去重 ====================    def _deduplicate(self, queries: List[EnhancedQuery]) -> List[EnhancedQuery]:        """        双层去重：            1. 字符串去重（快速过滤完全重复项）            2. 语义去重（基于 Embedding 余弦相似度，过滤语义等价项）        """        # 第一层：字符串去重        unique = self._deduplicate_by_text(queries)        # 第二层：语义去重（仅在启用且有 embedding_model 时执行）        if self.config.semantic_dedup_enabled and self.embedding_model is not None:            try:                unique = self._deduplicate_by_semantic(unique)            except Exception as e:                logger.warning(f"[Enhance] 语义去重失败，回退到字符串去重: {e}")        return unique    def _deduplicate_by_text(self, queries: List[EnhancedQuery]) -> List[EnhancedQuery]:        """基于文本前100字符精确去重"""        seen = set()        unique = []        for q in queries:            key = q.text[:100].strip()            if key not in seen:                seen.add(key)                unique.append(q)        return unique    def _deduplicate_by_semantic(self, queries: List[EnhancedQuery]) -> List[EnhancedQuery]:        """        基于向量余弦相似度的语义去重。        规则：            - 以原始查询（queries[0]）为锚点            - 对每个改写变体（rewrite），计算与所有已保留项的最大相似度            - 若相似度 >= threshold，则认为语义等价，丢弃该变体            - HyDE 文档不参与语义去重（形态差异大，与原始查询的相似度天然较低）            - 变体之间也互相计算相似度，避免生成多个雷同改写        """        if len(queries) <= 1:            return queries        threshold = self.config.semantic_dedup_threshold        # 编码所有查询文本（统一加指令前缀，确保与检索时一致）        vectors = []        for q in queries:            if q.query_type == "hyde":                text = q.text            else:                # 语义去重时，对 original/rewrite 统一加前缀编码，保证可比性                text = q.get_search_text("为这个句子生成表示以用于检索相关文章：")            vec = self.embedding_model.embed_query(text)            vectors.append(vec)        kept = [queries[0]]  # 保留原始查询        kept_vectors = [vectors[0]]        for i in range(1, len(queries)):            q = queries[i]            current_vec = vectors[i]            # 计算与所有已保留项的最大相似度            max_sim = 0.0            for kept_vec in kept_vectors:                sim = _cosine_similarity(current_vec, kept_vec)                if sim > max_sim:                    max_sim = sim            if max_sim >= threshold:                logger.info(                    f"[SemanticDedup] 丢弃语义等价变体: '{q.text[:40]}...' "                    f"(max_sim={max_sim:.4f} >= threshold={threshold})"                )            else:                kept.append(q)                kept_vectors.append(current_vec)        removed = len(queries) - len(kept)        if removed > 0:            logger.info(f"[SemanticDedup] 语义去重移除 {removed} 个变体，保留 {len(kept)} 个")        return kept