pip install langchain langchain-community \
             chromadb sentence-transformers \
             pypdf unstructured \
             openai tiktoken

from langchain_community.document_loaders import PyPDFLoader

def load_pdf(file_path: str):
    """加載 PDF 文件，返回 Document 列表"""
    loader = PyPDFLoader(file_path)
    pages = loader.load()
    print(f"成功加載 {len(pages)} 頁(yè)")
    return pages

# 每個(gè)Document對(duì)象包含:
# - page_content: 文本內(nèi)容
# - metadata: 元數(shù)據(jù)（頁(yè)碼、來(lái)源等）

from langchain_community.document_loaders import UnstructuredMarkdownLoader

def load_markdown(file_path: str):
    """加載 Markdown 文件"""
    loader = UnstructuredMarkdownLoader(file_path)
    docs = loader.load()
    print(f"成功加載 Markdown: {len(docs)} 段")
    return docs

from langchain_community.document_loaders import DirectoryLoader

def load_directory(dir_path: str, glob_pattern: str = "**/*.pdf"):
    """批量加載目錄中的文檔"""
    loader = DirectoryLoader(
        dir_path,
        glob=glob_pattern,
        show_progress=True,
        use_multithreading=True
    )
    documents = loader.load()
    print(f"從 {dir_path} 加載了 {len(documents)} 個(gè)文檔")
    return documents

from langchain.text_splitter import RecursiveCharacterTextSplitter

def split_documents(documents, chunk_size=500, chunk_overlap=50):
    """
    遞歸字符分塊器 —— 按段落、句子邊界智能切分

    參數(shù):
        chunk_size: 每塊最大字符數(shù)
        chunk_overlap: 相鄰塊重疊字符數(shù)（保持上下文連貫）
    """
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", "。", "！", "？", ".", " ", ""],
        length_function=len
    )
    chunks = text_splitter.split_documents(documents)
    print(f"文檔被切分為 {len(chunks)} 個(gè)文本塊")
    return chunks

from langchain_community.embeddings import HuggingFaceEmbeddings

def get_embedding_model():
    """使用本地 Sentence Transformer 模型"""
    model = HuggingFaceEmbeddings(
        model_name="BAAI/bge-small-zh-v1.5",  # 中文優(yōu)秀模型
        model_kwargs={"device": "cpu"},
        encode_kwargs={"normalize_embeddings": True}
    )
    return model

# 測(cè)試嵌入效果
model = get_embedding_model()
vector = model.embed_query("什么是機(jī)器學(xué)習(xí)？")
print(f"向量維度: {len(vector)}")  # 輸出: 向量維度: 384

from langchain_community.embeddings import OpenAIEmbeddings

def get_openai_embedding():
    """使用 OpenAI text-embedding-3-small"""
    return OpenAIEmbeddings(
        model="text-embedding-3-small",
        openai_api_key="your-api-key"
    )

from langchain_community.vectorstores import Chroma

def build_vector_store(chunks, embedding_model, persist_directory="./vector_db"):
    """
    構(gòu)建向量數(shù)據(jù)庫(kù)并持久化

    流程: 文本塊 → Embedding → 存入 Chroma
    """
    vectorstore = Chroma.from_documents(
        documents=chunks,
        embedding=embedding_model,
        persist_directory=persist_directory
    )
    vectorstore.persist()
    print(f"向量數(shù)據(jù)庫(kù)構(gòu)建完成，共 {vectorstore._collection.count()} 條記錄")
    return vectorstore

def load_vector_store(embedding_model, persist_directory="./vector_db"):
    """加載已有的向量數(shù)據(jù)庫(kù)"""
    vectorstore = Chroma(
        persist_directory=persist_directory,
        embedding_function=embedding_model
    )
    return vectorstore

def similarity_search(vectorstore, query: str, k: int = 4):
    """
    基礎(chǔ)語(yǔ)義檢索 —— 返回最相似的 k 個(gè)文本塊

    原理: 將 query 向量化 → 與數(shù)據(jù)庫(kù)中所有向量計(jì)算余弦相似度 → 返回 Top-K
    """
    results = vectorstore.similarity_search(
        query=query,
        k=k
    )
    for i, doc in enumerate(results, 1):
        print(f"\n--- 檢索結(jié)果 {i} ---")
        print(f"內(nèi)容: {doc.page_content[:200]}...")
        print(f"來(lái)源: {doc.metadata.get('source', '未知')}")
    return results

def mmr_search(vectorstore, query: str, k: int = 4, fetch_k: int = 20):
    """
    MMR（最大邊際相關(guān)性）檢索
    在保證相關(guān)性的同時(shí)，盡量減少結(jié)果之間的冗余

    相比普通檢索:
    - 普通檢索可能返回內(nèi)容高度重復(fù)的多個(gè)結(jié)果
    - MMR 檢索能返回既相關(guān)又多樣化的結(jié)果
    """
    results = vectorstore.max_marginal_relevance_search(
        query=query,
        k=k,
        fetch_k=fetch_k
    )
    return results

# ===== 三種檢索方式對(duì)比 =====

# 1. 純相似度檢索 —— 最簡(jiǎn)單，可能冗余
docs1 = vectorstore.similarity_search("什么是深度學(xué)習(xí)？", k=4)

# 2. 相似度 + 分?jǐn)?shù) —— 可按閾值過(guò)濾
docs2 = vectorstore.similarity_search_with_relevance_scores(
    "什么是深度學(xué)習(xí)？", k=4
)
for doc, score in docs2:
    print(f"分?jǐn)?shù): {score:.4f} | {doc.page_content[:80]}")

# 3. MMR 檢索 —— 相關(guān)性 + 多樣性兼顧（推薦）
docs3 = vectorstore.max_marginal_relevance_search(
    "什么是深度學(xué)習(xí)？", k=4, fetch_k=20
)

from langchain.prompts import ChatPromptTemplate

RAG_PROMPT_TEMPLATE = """你是一個(gè)專業(yè)的知識(shí)助手。請(qǐng)根據(jù)以下檢索到的參考資料來(lái)回答用戶的問(wèn)題。

要求：
- 只基于參考資料回答，不要編造信息
- 如果參考資料中沒有相關(guān)內(nèi)容，請(qǐng)誠(chéng)實(shí)說(shuō)明
- 回答要條理清晰，必要時(shí)使用列表或分點(diǎn)說(shuō)明

參考資料：
{context}

用戶問(wèn)題：{question}

請(qǐng)給出你的回答："""

def build_rag_prompt(query: str, retrieved_docs: list) -> str:
    """將檢索到的文檔拼接為 context，構(gòu)建完整 Prompt"""
    context = "\n\n".join([
        f"[來(lái)源 {i+1}]: {doc.page_content}"
        for i, doc in enumerate(retrieved_docs)
    ])
    prompt = ChatPromptTemplate.from_template(RAG_PROMPT_TEMPLATE)
    return prompt.format(context=context, question=query)

from langchain_community.chat_models import ChatOpenAI

def generate_answer(query: str, retrieved_docs: list, llm=None):
    """調(diào)用 LLM 生成最終回答"""
    if llm is None:
        llm = ChatOpenAI(
            model="gpt-4o-mini",
            temperature=0,
            openai_api_key="your-api-key"
        )

    prompt = build_rag_prompt(query, retrieved_docs)
    response = llm.invoke(prompt)
    return response.content

"""
完整的 RAG 系統(tǒng) —— 從文檔到問(wèn)答
"""
from langchain_community.document_loaders import DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.prompts import ChatPromptTemplate
from langchain_community.chat_models import ChatOpenAI


class SimpleRAG:
    """一個(gè)簡(jiǎn)單但完整的 RAG 系統(tǒng)"""

    def __init__(self, docs_dir: str = "./data", db_dir: str = "./vector_db"):
        self.docs_dir = docs_dir
        self.db_dir = db_dir
        self.embedding_model = HuggingFaceEmbeddings(
            model_name="BAAI/bge-small-zh-v1.5",
            encode_kwargs={"normalize_embeddings": True}
        )
        self.vectorstore = None
        self.llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

    # ---------- 離線階段：構(gòu)建知識(shí)庫(kù) ----------

    def build_index(self):
        """Step 1~4: 加載文檔 → 分塊 → 嵌入 → 存入向量數(shù)據(jù)庫(kù)"""
        print("=" * 50)
        print("Step 1: 加載文檔...")
        loader = DirectoryLoader(self.docs_dir, glob="**/*.pdf", show_progress=True)
        documents = loader.load()
        print(f"  加載了 {len(documents)} 個(gè)文檔")

        print("Step 2: 文本分塊...")
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=500, chunk_overlap=50,
            separators=["\n\n", "\n", "。", ".", " ", ""]
        )
        chunks = splitter.split_documents(documents)
        print(f"  切分為 {len(chunks)} 個(gè)文本塊")

        print("Step 3~4: 向量嵌入并存儲(chǔ)...")
        self.vectorstore = Chroma.from_documents(
            documents=chunks,
            embedding=self.embedding_model,
            persist_directory=self.db_dir
        )
        self.vectorstore.persist()
        print(f"  向量數(shù)據(jù)庫(kù)構(gòu)建完成！")
        print("=" * 50)

    def load_index(self):
        """加載已有的向量數(shù)據(jù)庫(kù)"""
        self.vectorstore = Chroma(
            persist_directory=self.db_dir,
            embedding_function=self.embedding_model
        )
        print("已加載向量數(shù)據(jù)庫(kù)")

    # ---------- 在線階段：檢索 + 生成 ----------

    def query(self, question: str, k: int = 4) -> str:
        """完整 RAG 問(wèn)答流程"""
        # Step 5: 語(yǔ)義檢索
        retrieved = self.vectorstore.max_marginal_relevance_search(
            query=question, k=k, fetch_k=k*5
        )

        # Step 6: 拼接 Prompt + LLM 生成
        context = "\n\n".join([
            f"[來(lái)源 {i+1}]: {doc.page_content}"
            for i, doc in enumerate(retrieved)
        ])

        prompt = f"""基于以下參考資料回答問(wèn)題。如果資料中沒有答案，請(qǐng)說(shuō)明。

參考資料：
{context}

問(wèn)題：{question}"""

        response = self.llm.invoke(prompt)
        return response.content


# ===== 使用示例 =====
if __name__ == "__main__":
    rag = SimpleRAG(docs_dir="./data", db_dir="./vector_db")

    # 首次運(yùn)行：構(gòu)建索引
    # rag.build_index()

    # 后續(xù)運(yùn)行：直接加載
    rag.load_index()

    # 提問(wèn)
    answer = rag.query("什么是深度學(xué)習(xí)？它有哪些主要應(yīng)用？")
    print(f"\n回答:\n{answer}")

from langchain.retrievers import BM25Retriever, EnsembleRetriever

def build_hybrid_retriever(chunks, vectorstore, k=4):
    """
    混合檢索 = 向量檢索(語(yǔ)義) + BM25檢索(關(guān)鍵詞)
    取長(zhǎng)補(bǔ)短，效果優(yōu)于單一檢索
    """
    # 關(guān)鍵詞檢索（BM25）
    bm25_retriever = BM25Retriever.from_documents(chunks)
    bm25_retriever.k = k

    # 語(yǔ)義檢索（向量）
    vector_retriever = vectorstore.as_retriever(
        search_kwargs={"k": k}
    )

    # 混合檢索器（各占 50% 權(quán)重）
    ensemble_retriever = EnsembleRetriever(
        retrievers=[bm25_retriever, vector_retriever],
        weights=[0.5, 0.5]
    )
    return ensemble_retriever

from sentence_transformers import CrossEncoder

def rerank_results(query: str, documents: list, top_k: int = 4) -> list:
    """
    使用 Cross-Encoder 對(duì)檢索結(jié)果重排序
    Cross-Encoder 比 Bi-Encoder 更精確，但速度較慢
    適合對(duì) Top-K 候選做精排
    """
    reranker = CrossEncoder("BAAI/bge-reranker-base")

    pairs = [[query, doc.page_content] for doc in documents]
    scores = reranker.predict(pairs)

    # 按分?jǐn)?shù)降序排列
    ranked = sorted(
        zip(documents, scores),
        key=lambda x: x[1],
        reverse=True
    )
    return [doc for doc, score in ranked[:top_k]]