检索前处理：查询构建、查询翻译、查询路由

检索前处理：查询构建、查询翻译、查询路由

自己搞的简单代码，便于理解

import logging
from typing import List, Tuple, Literal
from langchain_chroma import Chroma
from langchain_deepseek import ChatDeepSeek
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_core.prompts import PromptTemplate, ChatPromptTemplate
from langchain_core.output_parsers import LineListOutputParser, StrOutputParser
from langchain_core.pydantic_v1 import BaseModel, Field

# ===================== 基础日志与全局配置 =====================
logging.basicConfig(level=logging.INFO, format="%(message)s")
logger = logging.getLogger("retrieval_pipeline")
logging.getLogger("langchain.retrievers.multi_query").setLevel(logging.WARNING)

# 全局默认数据源（关闭路由时使用，与原有流水线完全兼容）
DEFAULT_DATASOURCE = "game_setting"

# ===================== 1. 基础组件初始化 =====================
llm = ChatDeepSeek(model="deepseek-chat", temperature=0)
line_parser = LineListOutputParser()
str_parser = StrOutputParser()
embed_model = HuggingFaceEmbeddings(model_name="BAAI/bge-small-zh")

# ===================== 2. 多数据源知识库构建 =====================
def build_vectorstore(docs_content: list, collection_name: str):
    """构建独立向量库集合"""
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=0)
    splits = text_splitter.create_documents(docs_content)
    return Chroma.from_documents(
        documents=splits,
        embedding=embed_model,
        collection_name=collection_name
    )

# 2.1 各领域独立知识库内容（可替换为你的真实文档）
setting_docs = [
    "剑客职业10级核心技能为基础剑式和格挡，普陀山关卡守关BOSS为黑风妖王，主打高物理单体伤害。",
    "法师职业主打远程法术输出，核心技能为火球术和冰墙，新手推荐先点满火球术提升清怪效率。",
    "游戏内装备分为白、绿、蓝、紫、橙五个品质，品质越高属性加成越高，可通过副本掉落获取。"
]

story_docs = [
    "主角猢狲生于花果山，自名孙悟空，拜师菩提祖师习得七十二变与筋斗云。",
    "孙悟空大闹天宫后被如来佛祖压于五行山下，五百年后被唐僧救出，踏上西天取经之路。",
    "普陀山一难中，唐僧被黑风妖王掳走，孙悟空前往黑风洞讨要，与黑熊精打斗数回合不分胜负。"
]

service_docs = [
    "游戏充值未到账请提供订单号和账号ID，联系客服24小时内处理补发。",
    "账号被盗请提供注册手机号和实名认证信息，提交申诉后3个工作日内回复。",
    "游戏闪退请更新显卡驱动，降低画质设置，若仍有问题请提交设备型号和闪退日志。"
]

# 2.2 构建向量库与对应检索器映射表
# 向量库映射（用于HyDE检索）
vectorstore_map = {
    "game_setting": build_vectorstore(setting_docs, "game_setting"),
    "game_story": build_vectorstore(story_docs, "game_story"),
    "customer_service": build_vectorstore(service_docs, "customer_service")
}

# MultiQuery检索器映射（用于多查询召回）
multiquery_retriever_map = {
    key: MultiQueryRetriever.from_llm(
        retriever=vs.as_retriever(search_kwargs={"k": 2}),
        llm=llm
    )
    for key, vs in vectorstore_map.items()
}

# ===================== 3. 模块零：逻辑路由（入口分发） =====================
# 3.1 路由结构化输出模型
class RouteQuery(BaseModel):
    """将用户问题路由到最适合回答的数据源"""
    datasource: Literal["game_setting", "game_story", "customer_service"] = Field(
        description="根据用户问题，选择最匹配的一个数据源，只能从给定选项中选择"
    )

# 3.2 路由提示词与路由链
route_system_prompt = """你是专业的查询路由专家，请根据用户问题的内容，路由到最合适的数据源：
- game_setting：游戏玩法、职业技能、关卡机制、装备系统、数值设定等游戏设定类问题
- game_story：游戏剧情、人物关系、故事背景、角色经历等剧情类问题
- customer_service：充值、账号异常、bug反馈、闪退、登录失败等客服与技术问题

只需选择一个最匹配的数据源即可。"""

route_prompt = ChatPromptTemplate.from_messages([
    ("system", route_system_prompt),
    ("human", "{question}"),
])

# 绑定结构化输出，自动解析为RouteQuery对象
structured_llm_router = llm.with_structured_output(RouteQuery)
router_chain = route_prompt | structured_llm_router

def run_routing(question: str) -> str:
    """执行路由分类，返回数据源key，失败时返回默认数据源兜底"""
    try:
        route_result = router_chain.invoke({"question": question})
        target = route_result.datasource
        logger.info(f"[路由阶段] 分类结果：→ {target} 知识库")
        return target
    except Exception as e:
        logger.warning(f"[路由阶段] 路由失败，使用默认数据源兜底：{str(e)}")
        return DEFAULT_DATASOURCE

# ===================== 4. 模块一：查询澄清 =====================
CLARIFY_GENERATE_PROMPT = PromptTemplate(
    input_variables=["question", "max_clarify_num"],
    template="""你是专业的查询澄清助手。请判断用户的问题是否存在语义模糊、信息缺失、范围不明确、存在歧义的情况。
规则：
1. 如果问题足够明确，不需要补充信息，请直接输出空行，不要输出任何其他内容。
2. 如果问题需要澄清，请生成 {max_clarify_num} 个针对性的澄清问题，用于补全信息、缩小范围、消除歧义。
3. 每个澄清问题只询问一个信息点，表述简洁中立。
4. 每个问题单独占一行，不要编号，不要多余解释。

用户原始问题：{question}"""
)

QUERY_INTEGRATE_PROMPT = PromptTemplate(
    input_variables=["original_question", "clarify_history"],
    template="""请根据用户的原始问题和澄清问答记录，整合出一个信息完整、表述精准的最终查询语句。
要求：
1. 保留原始问题的核心诉求，补充所有澄清得到的限定条件。
2. 语句通顺自然，是一个完整的问题，不要分点。
3. 只输出最终的查询语句，不要输出其他任何内容。

原始问题：{original_question}
澄清问答记录：
{clarify_history}

最终查询："""
)

def generate_clarify_questions(question: str, max_num: int = 3) -> List[str]:
    chain = CLARIFY_GENERATE_PROMPT | llm | line_parser
    result = chain.invoke({"question": question, "max_clarify_num": max_num})
    return [q.strip() for q in result if q.strip()]

def integrate_final_query(original_question: str, clarify_history: List[Tuple[str, str]]) -> str:
    history_text = "\n".join([f"问：{q}\n答：{a}" for q, a in clarify_history])
    chain = QUERY_INTEGRATE_PROMPT | llm
    return chain.invoke({
        "original_question": original_question,
        "clarify_history": history_text
    }).content.strip()

def run_clarification(
    original_question: str,
    max_rounds: int = 2,
    questions_per_round: int = 3
) -> str:
    current_question = original_question
    clarify_history = []

    for round_idx in range(max_rounds):
        clarify_questions = generate_clarify_questions(current_question, questions_per_round)
        if not clarify_questions:
            logger.info(f"\n[澄清阶段] 第 {round_idx+1} 轮：问题已足够明确，终止澄清")
            break

        logger.info(f"\n[澄清阶段] 第 {round_idx+1} 轮澄清问题：")
        for i, q in enumerate(clarify_questions, 1):
            logger.info(f"  {i}. {q}")

        print("\n请依次回答以下问题：")
        user_answers = []
        for q in clarify_questions:
            ans = input(f"Q: {q}\nA: ")
            user_answers.append(ans)

        for q, a in zip(clarify_questions, user_answers):
            clarify_history.append((q, a))

        current_question = integrate_final_query(original_question, clarify_history)
        logger.info(f"\n[澄清阶段] 本轮整合后查询：{current_question}")

    return current_question

# ===================== 5. 模块二：复合问题拆分 =====================
def split_complex_question(question: str) -> List[str]:
    prompt = f"""请将用户的问题拆分成若干个独立的单一子问题，每个子问题只包含一个疑问点。
要求：
1. 每个子问题单独占一行
2. 不要编号、不要多余解释
3. 保留原问题的所有限定条件和语义

用户问题：{question}"""
    response = llm.invoke(prompt)
    sub_questions = [line.strip() for line in response.content.splitlines() if line.strip()]
    return sub_questions

# ===================== 6. 模块三：HyDE 假设文档召回 =====================
HYDE_PROMPT = ChatPromptTemplate.from_template("""请撰写一段与以下问题相关的游戏设定内容，要求表述风格和游戏官方设定文档一致，语义连贯，包含具体细节。
问题：{question}
内容：""")

generate_hyde_doc = HYDE_PROMPT | llm | str_parser

def hyde_retrieve(query: str, vectorstore, top_k: int = 2):
    """HyDE单路检索，传入指定向量库，生成失败自动跳过"""
    try:
        hypo_doc = generate_hyde_doc.invoke({"question": query})
        if not hypo_doc or len(hypo_doc.strip()) < 20:
            logger.warning(f"  [HyDE] 生成内容无效，跳过该路召回")
            return []
        
        logger.info(f"  [HyDE] 生成假设文档：{hypo_doc[:80]}...")
        hyde_retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
        docs = hyde_retriever.invoke(hypo_doc)
        logger.info(f"  [HyDE] 召回文档数：{len(docs)}")
        return docs
    except Exception as e:
        logger.warning(f"  [HyDE] 执行异常，跳过：{str(e)}")
        return []

# ===================== 7. 模块四：多路召回 + 全局去重 =====================
def single_query_retrieve(query: str, base_retriever, vectorstore, enable_hyde: bool = False):
    """单个查询执行多路召回：MultiQuery + 可选HyDE"""
    # 第1路：MultiQuery多角度检索
    docs = base_retriever.invoke(query)
    
    # 第2路：HyDE假设文档检索
    if enable_hyde:
        hyde_docs = hyde_retrieve(query, vectorstore)
        docs.extend(hyde_docs)
    
    return docs

def batch_retrieve(sub_questions: List[str], base_retriever, vectorstore, enable_hyde: bool = False):
    """批量子问题检索 + 全局去重"""
    all_docs = []
    for idx, q in enumerate(sub_questions, 1):
        logger.info(f"\n[检索阶段] 正在处理子问题 {idx}/{len(sub_questions)}：{q}")
        docs = single_query_retrieve(q, base_retriever, vectorstore, enable_hyde)
        all_docs.extend(docs)

    # 按文档内容全局去重
    unique_docs = []
    seen_contents = set()
    for doc in all_docs:
        if doc.page_content not in seen_contents:
            seen_contents.add(doc.page_content)
            unique_docs.append(doc)
    return unique_docs

# ===================== 8. 主流水线：全模块串联 =====================
def retrieval_pipeline(
    original_question: str,
    enable_routing: bool = False,
    enable_clarification: bool = True,
    enable_split: bool = True,
    enable_hyde: bool = False,
    max_clarify_rounds: int = 2,
    questions_per_round: int = 3
):
    """
    完整检索前处理流水线
    :param original_question: 用户原始问题
    :param enable_routing: 是否开启逻辑路由（多数据源分发）
    :param enable_clarification: 是否开启查询澄清
    :param enable_split: 是否开启复合问题拆分
    :param enable_hyde: 是否开启HyDE假设文档召回
    :param max_clarify_rounds: 最大澄清轮数
    :param questions_per_round: 单轮澄清问题数
    :return: 最终去重后的文档列表
    """
    logger.info("=" * 60)
    logger.info(f"原始问题：{original_question}")
    logger.info(f"功能开关：路由={enable_routing} | 澄清={enable_clarification} | 拆分={enable_split} | HyDE={enable_hyde}")
    logger.info("=" * 60)

    # 第0步：逻辑路由，选择对应数据源的检索器与向量库
    target_datasource = DEFAULT_DATASOURCE
    if enable_routing:
        target_datasource = run_routing(original_question)
    
    # 获取当前数据源对应的检索组件
    current_retriever = multiquery_retriever_map.get(target_datasource, multiquery_retriever_map[DEFAULT_DATASOURCE])
    current_vectorstore = vectorstore_map.get(target_datasource, vectorstore_map[DEFAULT_DATASOURCE])

    # 第一步：查询澄清
    current_question = original_question
    if enable_clarification:
        current_question = run_clarification(
            current_question,
            max_rounds=max_clarify_rounds,
            questions_per_round=questions_per_round
        )
        logger.info(f"\n[澄清完成] 最终明确问题：{current_question}")

    # 第二步：复合问题拆分
    sub_questions = [current_question]
    if enable_split:
        sub_questions = split_complex_question(current_question)
        logger.info(f"\n[拆分完成] 拆分为 {len(sub_questions)} 个子问题：")
        for i, q in enumerate(sub_questions, 1):
            logger.info(f"  {i}. {q}")

    # 第三步：多路检索 + 全局去重
    final_docs = batch_retrieve(sub_questions, current_retriever, current_vectorstore, enable_hyde)

    logger.info("\n" + "=" * 60)
    logger.info(f"流水线执行完成，共召回 {len(final_docs)} 条去重文档")
    logger.info("=" * 60)
    return final_docs

# ===================== 运行示例 =====================
if __name__ == "__main__":
    # 测试1：游戏设定类问题（开启路由+HyDE，关闭澄清方便演示）
    print("\n" + "#"*60)
    print("测试1：游戏设定类问题")
    print("#"*60)
    docs = retrieval_pipeline(
        original_question="普陀山的BOSS怎么打？",
        enable_routing=True,
        enable_clarification=False,
        enable_split=False,
        enable_hyde=True
    )
    print("最终检索结果：")
    for doc in docs:
        print(f"- {doc.page_content}")

    # 测试2：剧情类问题
    print("\n" + "#"*60)
    print("测试2：剧情类问题")
    print("#"*60)
    docs = retrieval_pipeline(
        original_question="孙悟空和黑熊精认识吗？",
        enable_routing=True,
        enable_clarification=False,
        enable_split=False,
        enable_hyde=False
    )
    print("最终检索结果：")
    for doc in docs:
        print(f"- {doc.page_content}")

    # 测试3：客服类问题
    print("\n" + "#"*60)
    print("测试3：客服类问题")
    print("#"*60)
    docs = retrieval_pipeline(
        original_question="充了钱没到账要找谁？",
        enable_routing=True,
        enable_clarification=False,
        enable_split=False,
        enable_hyde=False
    )
    print("最终检索结果：")
    for doc in docs:
        print(f"- {doc.page_content}")