RRF_score(d) = Σ 1 / (k + rank_i(d))

# hybrid_retrieval.py
# 混合检索示例

from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings
from langchain.schema import Document

# 准备测试文档
docs = [
    Document(page_content="苹果公司发布了新款 iPhone 16，售价 5999 元起。"),
    Document(page_content="订单号 20240315001 于 3 月 15 日下单，商品为蓝牙耳机，状态：已发货。"),
    Document(page_content="退货政策：电子产品 7 天无理由退货，需保持包装完整。"),
    Document(page_content="iPhone 的设置方法：进入「设置」App，选择「通用」，可以调整语言、字体大小等。"),
    Document(page_content="订单查询方式：登录 App，进入「我的」页面，点击「订单管理」。"),
]

# 稀疏检索器：BM25，基于关键词
bm25_retriever = BM25Retriever.from_documents(docs)
bm25_retriever.k = 3

# 稠密检索器：向量相似度
embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(docs, embeddings)
dense_retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

# 融合检索器：weights 控制两路的权重
# 0.5/0.5 是等权，可以根据场景调整
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, dense_retriever],
    weights=[0.5, 0.5]
)

# 测试
test_queries = [
    "订单号 20240315001 的发货状态",   # 精确字符串，BM25 应该更有帮助
    "苹果手机怎么调字体",             # 语义查询，向量检索更有帮助
]

for query in test_queries:
    print(f"\n查询：{query}")
    results = ensemble_retriever.invoke(query)
    for i, doc in enumerate(results):
        print(f"  [{i+1}] {doc.page_content[:60]}...")

# reranking_demo.py
# 用 BGE Reranker 做重排序（本地运行，免费）

from langchain.schema import Document
from sentence_transformers import CrossEncoder

# 加载 BGE Reranker（第一次运行会自动下载模型，约 1GB）
reranker = CrossEncoder("BAAI/bge-reranker-base")  # base 版本小一些，large 效果更好

# 模拟第一阶段召回的文档（实际是混合检索的结果）
query = "iPhone 16 多少钱？"
candidate_docs = [
    "苹果公司发布了新款 iPhone 16，售价 5999 元起，提供黑白粉三种颜色。",
    "退货政策：电子产品 7 天无理由退货，需保持包装完整。",
    "iPhone 的设置方法：进入设置 App，选择通用，可以调整语言。",
    "iPhone 16 Pro 版售价 8999 元起，搭载 A18 Pro 芯片，支持摄影风格功能。",
    "订单查询方式：登录 App，进入我的页面，点击订单管理。",
]

# Cross-encoder 打分：输入 (query, doc) 对，输出相关性分数
pairs = [(query, doc) for doc in candidate_docs]
scores = reranker.predict(pairs)

# 按分数降序排列
ranked = sorted(zip(scores, candidate_docs), key=lambda x: x[0], reverse=True)

print(f"查询：{query}\n")
print("重排序结果：")
for score, doc in ranked:
    print(f"  [{score:.4f}] {doc[:60]}...")

# 取 top-3 送给 LLM
top_docs = [doc for _, doc in ranked[:3]]

# query_rewriting.py
from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

rewrite_prompt = ChatPromptTemplate.from_template("""
你是一个搜索查询优化助手。请将用户的口语化问题改写成更适合文档检索的查询语句。
要求：保留关键信息，去掉口语词，扩展同义词。

用户问题：{question}
改写后的查询（只输出查询语句，不要解释）：
""")

def rewrite_query(question: str) -> str:
    chain = rewrite_prompt | llm
    result = chain.invoke({"question": question})
    return result.content.strip()

# 测试
test_questions = [
    "刚买的东西能退吗",
    "我在哪里看订单",
    "手机坏了咋整",
]

for q in test_questions:
    rewritten = rewrite_query(q)
    print(f"原始：{q}")
    print(f"改写：{rewritten}\n")

# HyDE 示例
hyde_prompt = ChatPromptTemplate.from_template("""
请根据以下问题，生成一段可能的答案文档。这个答案不需要正确，只需要格式和风格
像真实文档一样，用于辅助检索。

问题：{question}
假设答案（100字以内）：
""")

def hyde_retrieval(question: str, retriever) -> list:
    # 第一步：生成假设答案
    chain = hyde_prompt | llm
    hypothetical_answer = chain.invoke({"question": question}).content

    # 第二步：用假设答案去检索（而不是用原始问题）
    docs = retriever.invoke(hypothetical_answer)
    return docs