from langchain.llms import OpenAI
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain.vectorstores import Chroma
from langchain.embeddings import OpenAIEmbeddings
from langchain.chat_models import ChatOpenAI

# 初始化基础组件
llm = ChatOpenAI(temperature=0)
embeddings = OpenAIEmbeddings()
vectorstore = Chroma(embedding_function=embeddings, persist_directory="./chroma_db")

# 创建基础检索器
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 5})

# 包装成 MultiQueryRetriever
# 其内部会使用 LLM 根据原问题生成 3 个变体问题
retriever = MultiQueryRetriever.from_llm(
    retriever=base_retriever,
    llm=llm,
    include_original=True # 包含原始查询
)

# 检索示例
question = “Java 中如何实现一个线程安全的单例模式？”
docs = retriever.get_relevant_documents(question)
print(f"检索到 {len(docs)} 个文档片段")
for i, doc in enumerate(docs):
    print(f"片段 {i+1} 来源: {doc.metadata.get('source', 'N/A')}")

from langchain.llms import OpenAI
from langchain.embeddings import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import LLMChain
from langchain.prompts import PromptTemplate

class HyDERetriever:
    def __init__(self, vectorstore, llm, embedding_model):
        self.vectorstore = vectorstore
        self.llm = llm
        self.embedding_model = embedding_model
        # 提示词用于生成假设文档
        self.hyde_prompt = PromptTemplate(
            input_variables=["question"],
            template="请基于以下问题，生成一个假设性的答案段落。问题：{question}\n假设性答案："
        )
        self.hyde_chain = LLMChain(llm=llm, prompt=self.hyde_prompt)

    def get_relevant_documents(self, question, k=5):
        # 1. 生成假设文档
        hypothetical_doc = self.hyde_chain.run(question=question)
        print(f"生成的假设文档：{hypothetical_doc[:200]}...")

        # 2. 将假设文档向量化
        hypothetical_embedding = self.embedding_model.embed_query(hypothetical_doc)

        # 3. 用假设文档的向量进行相似度检索
        docs = self.vectorstore.similarity_search_by_vector(hypothetical_embedding, k=k)
        return docs

# 使用示例
llm = ChatOpenAI(temperature=0.7) # HyDE 可以稍高一些创造性
embeddings = OpenAIEmbeddings()
vectorstore = Chroma(embedding_function=embeddings, persist_directory="./chroma_db")

hyde_retriever = HyDERetriever(vectorstore, llm, embeddings)
question = “在微服务架构中，如何设计一个保证最终一致性的分布式事务？”
hyde_docs = hyde_retriever.get_relevant_documents(question)

from sentence_transformers import CrossEncoder
import numpy as np

class RerankRetriever:
    def __init__(self, base_retriever, reranker_model_name='BAAI/bge-reranker-large'):
        self.base_retriever = base_retriever
        # 加载交叉编码器重排序模型
        self.reranker = CrossEncoder(reranker_model_name, max_length=512)

    def get_relevant_documents(self, query, k=5, rerank_k=20):
        # 1. 首先用基础检索器召回更多文档 (e.g., 20个)
        initial_docs = self.base_retriever.get_relevant_documents(query, k=rerank_k)
        if not initial_docs:
            return []

        # 2. 准备用于重排序的 (query, doc) 对
        pairs = [(query, doc.page_content) for doc in initial_docs]

        # 3. 使用交叉编码器计算相关性分数
        scores = self.reranker.predict(pairs)

        # 4. 根据分数重新排序文档
        ranked_indices = np.argsort(scores)[::-1] # 降序排列
        reranked_docs = [initial_docs[i] for i in ranked_indices[:k]]

        # 可选：打印分数以供调试
        for i, idx in enumerate(ranked_indices[:3]):
            print(f"重排序后第 {i+1} 名，分数：{scores[idx]:.4f}， 内容：{reranked_docs[i].page_content[:100]}...")
        return reranked_docs

# 集成到检索链中
base_retriever = vectorstore.as_retriever(search_kwargs={"k": 20})
rerank_retriever = RerankRetriever(base_retriever)
final_docs = rerank_retriever.get_relevant_documents(question, k=5)

# 以简化的两跳迭代检索为例
def iterative_retrieval(initial_question, max_hops=2):
    llm = ChatOpenAI(temperature=0)
    context = ""
    current_question = initial_question

    for hop in range(max_hops):
        print(f"\n=== 第 {hop+1} 轮检索 ===")
        print(f"当前问题：{current_question}")

        # 检索
        docs = retriever.get_relevant_documents(current_question)
        retrieved_text = "\n\n".join([d.page_content for d in docs[:3]])
        context += f"\n\n[第 {hop+1} 轮检索结果]:\n{retrieved_text}"

        # 判断是否需要进一步检索
        prompt = f"""
        你已获得以下信息：
        {context}

        原始问题是：{initial_question}
        基于当前已有信息，能否直接、准确地回答原始问题？
        如果**不能**，请分析还缺少什么关键信息，并生成一个最有助于获取该信息的新搜索问题。
        如果**能**，请输出 'ANSWER_READY'。

        请只输出新搜索问题或 'ANSWER_READY'。
        """
        decision = llm.predict(prompt)

        if "ANSWER_READY" in decision:
            print("信息已充足，准备生成最终答案。")
            break
        else:
            current_question = decision.strip()
            print(f"生成的新检索问题：{current_question}")

    return context

# 使用
complex_question = “Apache Kafka 的 Exactly-Once 语义是如何实现的？与 Apache Pulsar 的实现方式相比有什么优劣？”
accumulated_context = iterative_retrieval(complex_question)

# 使用 RAGAS (https://github.com/explodinggradients/ragas) 进行量化评估
from ragas.metrics import faithfulness, answer_relevancy, context_recall, context_precision
from ragas.metrics.critique import harmfulness
from ragas import evaluate
from datasets import Dataset
import os

# 准备评估数据：问题、参考答案、检索到的上下文、生成的答案
questions = ["Java 中 volatile 关键字的作用是什么？"]
ground_truths = [["volatile 关键字确保变量的可见性，禁止指令重排序，但不保证原子性。"]]
contexts = [["在Java中，volatile关键字的主要作用是保证变量的可见性..."]]
answers = ["volatile 关键字保证了多线程环境下变量的可见性和有序性。"]

dataset_dict = {
    "question": questions,
    "answer": answers,
    "contexts": contexts,
    "ground_truth": ground_truths
}
dataset = Dataset.from_dict(dataset_dict)

# 选择要评估的指标
metrics = [
    faithfulness, # 答案是否基于上下文，有无幻觉
    answer_relevancy, # 答案与问题的相关程度
    context_recall, # 所有必要的参考信息是否都被检索到
    # context_precision, # 检索到的上下文是否都相关
]

# 执行评估
result = evaluate(dataset, metrics)
print(result)
df = result.to_pandas()
print(df[['faithfulness', 'answer_relevancy', 'context_recall']])