用户提问 → 检索相关文档 → 将[文档 + 问题]提供给 LLM → 生成答案

from sentence_transformers import SentenceTransformer
import numpy as np

# 加载嵌入模型
model = SentenceTransformer('all-MiniLM-L6-v2')

# 文档嵌入
documents = [
    "Python 是广泛用于数据科学的编程语言。",
    "机器学习是从数据中学习模式的人工智能领域。",
    "巴黎是法国的首都。",
    "深度学习是使用神经网络的机器学习方法。",
]

doc_embeddings = model.encode(documents)
print(f"嵌入形状: {doc_embeddings.shape}")  # (4, 384)

# 查询嵌入
query = "人工智能与数据分析"
query_embedding = model.encode([query])[0]

# 计算余弦相似度
def cosine_similarity(a, b):
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

similarities = [cosine_similarity(query_embedding, doc_emb)
                for doc_emb in doc_embeddings]

# 按结果排序
ranked = sorted(zip(similarities, documents), reverse=True)
print("\n按相似度排序:")
for score, doc in ranked:
    print(f"  {score:.4f}: {doc}")

from langchain.text_splitter import CharacterTextSplitter

text = """
机器学习是人工智能的一个分支，是让计算机无需
显式编程即可从数据中学习的技术。它包括监督学习、
无监督学习、强化学习等方法论，各自解决不同类型的问题。
"""

splitter = CharacterTextSplitter(
    chunk_size=100,
    chunk_overlap=20,
    separator="\n"
)
chunks = splitter.split_text(text)
for i, chunk in enumerate(chunks):
    print(f"分块 {i}: {chunk[:50]}...")

from langchain.text_splitter import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    separators=["\n\n", "\n", ".", "!", "?", ",", " ", ""]
)

with open("document.txt", "r", encoding="utf-8") as f:
    text = f.read()

chunks = splitter.split_text(text)
print(f"总分块数: {len(chunks)}")
print(f"平均分块长度: {sum(len(c) for c in chunks) / len(chunks):.0f} 字符")

from langchain_experimental.text_splitter import SemanticChunker
from langchain_openai import OpenAIEmbeddings

# 语义分块器(基于嵌入相似度切分)
semantic_splitter = SemanticChunker(
    embeddings=OpenAIEmbeddings(),
    breakpoint_threshold_type="percentile",  # 或 "standard_deviation"
    breakpoint_threshold_amount=95           # 在相似度变化最大的前 5% 处切分
)

chunks = semantic_splitter.create_documents([text])
print(f"语义分块结果: {len(chunks)} 个分块")

from langchain.retrievers import ParentDocumentRetriever
from langchain.storage import InMemoryStore
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter

# 子分块器：小分块(用于检索)
child_splitter = RecursiveCharacterTextSplitter(chunk_size=200)

# 父分块器：大分块(用于上下文)
parent_splitter = RecursiveCharacterTextSplitter(chunk_size=1000)

# 向量存储(存储子分块嵌入)
vectorstore = Chroma(
    collection_name="full_documents",
    embedding_function=OpenAIEmbeddings()
)

# 父文档存储(存储完整的父分块)
store = InMemoryStore()

retriever = ParentDocumentRetriever(
    vectorstore=vectorstore,
    docstore=store,
    child_splitter=child_splitter,
    parent_splitter=parent_splitter,
)

import numpy as np
import faiss
from langchain_community.vectorstores import FAISS, Chroma
from langchain_openai import OpenAIEmbeddings

# ===== 直接使用 FAISS =====
# 随机向量数据(实际使用中应为嵌入)
d = 384          # 向量维度
n = 10000        # 文档数量
vectors = np.random.randn(n, d).astype('float32')

# 创建索引(L2 距离)
index_flat = faiss.IndexFlatL2(d)
index_flat.add(vectors)
print(f"FAISS 索引大小: {index_flat.ntotal}")

# 检索相似向量
query = np.random.randn(1, d).astype('float32')
k = 5
distances, indices = index_flat.search(query, k)
print(f"前 {k} 个结果: {indices[0]}")

# HNSW 索引(快速近似检索)
index_hnsw = faiss.IndexHNSWFlat(d, 32)  # M=32 条连接
index_hnsw.add(vectors)
distances_hnsw, indices_hnsw = index_hnsw.search(query, k)
print(f"HNSW 结果: {indices_hnsw[0]}")

# ===== LangChain + Chroma =====
from langchain.schema import Document

documents = [
    Document(page_content="Python 被用于数据科学。", metadata={"source": "doc1"}),
    Document(page_content="机器学习从数据中学习模式。", metadata={"source": "doc2"}),
    Document(page_content="深度学习是基于神经网络的机器学习。", metadata={"source": "doc3"}),
    Document(page_content="自然语言处理用于分析文本。", metadata={"source": "doc4"}),
]

# 创建 Chroma 向量存储
embeddings = OpenAIEmbeddings()
chroma_db = Chroma.from_documents(
    documents,
    embeddings,
    persist_directory="./chroma_db"
)

# 相似度检索
results = chroma_db.similarity_search("AI 与机器学习", k=2)
for doc in results:
    print(f"来源: {doc.metadata['source']}, 内容: {doc.page_content}")

# 附带分数一起检索
results_with_score = chroma_db.similarity_search_with_score("深度学习", k=2)
for doc, score in results_with_score:
    print(f"分数: {score:.4f}, 内容: {doc.page_content}")

from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
import os

# ===== 1. 加载文档 =====
# PDF 加载器
loader = PyPDFLoader("company_handbook.pdf")
# 或者加载整个目录
# loader = DirectoryLoader("./docs/", glob="**/*.pdf", loader_cls=PyPDFLoader)
pages = loader.load()
print(f"已加载页数: {len(pages)}")

# ===== 2. 切分文本 =====
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,
    length_function=len,
    separators=["\n\n", "\n", ".", "!", "?", ",", " "]
)
chunks = text_splitter.split_documents(pages)
print(f"生成的分块数: {len(chunks)}")

# ===== 3. 生成嵌入并存入向量数据库 =====
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma.from_documents(
    chunks,
    embeddings,
    persist_directory="./rag_db"
)
print("向量数据库存储完成")

# ===== 4. 构建 RAG 链 =====
# 自定义提示词模板
prompt_template = """你是一个乐于助人的 AI 助手。
请仅使用给定的上下文来回答问题。
如果上下文中没有相关信息，请回答"在提供的文档中找不到相关信息"。

上下文:
{context}

问题: {question}

回答:"""

PROMPT = PromptTemplate(
    template=prompt_template,
    input_variables=["context", "question"]
)

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
retriever = vectorstore.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 4}
)

# RetrievalQA 链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    chain_type_kwargs={"prompt": PROMPT},
    return_source_documents=True
)

# ===== 5. 问答 =====
query = "公司的休假政策是怎样的？"
result = qa_chain.invoke({"query": query})

print(f"\n问题: {query}")
print(f"回答: {result['result']}")
print(f"\n参考文档:")
for doc in result['source_documents']:
    print(f"  - 第 {doc.metadata.get('page', '?')} 页: {doc.page_content[:100]}...")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.core.node_parser import SentenceSplitter
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

# 配置
Settings.llm = OpenAI(model="gpt-4o-mini", temperature=0)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
Settings.node_parser = SentenceSplitter(chunk_size=512, chunk_overlap=20)

# 加载文档
documents = SimpleDirectoryReader("./docs/").load_data()

# 创建索引
index = VectorStoreIndex.from_documents(documents)

# 查询引擎
query_engine = index.as_query_engine(
    similarity_top_k=4,
    response_mode="compact"  # "tree_summarize"、"refine" 等
)

# 查询
response = query_engine.query("What is the main topic of the documents?")
print(f"回答: {response}")
print(f"\n来源节点:")
for node in response.source_nodes:
    print(f"  - Score: {node.score:.4f}")
    print(f"    Text: {node.text[:100]}...")

from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from langchain_community.vectorstores import Chroma

# BM25 检索(基于关键词)
bm25_retriever = BM25Retriever.from_documents(chunks)
bm25_retriever.k = 4

# 向量检索(基于语义)
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 4})

# 集成(混合)
ensemble_retriever = EnsembleRetriever(
    retrievers=[bm25_retriever, vector_retriever],
    weights=[0.5, 0.5]  # 各方式的权重
)

results = ensemble_retriever.invoke("Python 编程教程")
print(f"混合检索结果数: {len(results)}")

from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(temperature=0)
multi_query_retriever = MultiQueryRetriever.from_llm(
    retriever=vectorstore.as_retriever(),
    llm=llm
)

# 内部由 LLM 将问题改写为多个版本
# 例如: "RAG 的优点是什么?"
# → "RAG 技术的优势是什么?"
# → "检索增强生成相比普通 LLM 好在哪里?"
# → "为什么应该使用 RAG?"

results = multi_query_retriever.invoke("RAG 的优点是什么?")
print(f"多查询检索结果: {len(results)} 个")

# MMR 检索
mmr_retriever = vectorstore.as_retriever(
    search_type="mmr",
    search_kwargs={
        "k": 4,           # 最终返回数量
        "fetch_k": 20,    # 初始候选数量
        "lambda_mult": 0.5  # 相似度 vs 多样性的平衡(0=多样性, 1=相似度)
    }
)

results = mmr_retriever.invoke("机器学习算法")
print(f"MMR 结果: {len(results)} 个")

from langchain.schema import Document

# 带元数据的文档
docs_with_metadata = [
    Document(
        page_content="2024年第一季度营收为100亿韩元。",
        metadata={"year": 2024, "quarter": "Q1", "category": "financial"}
    ),
    Document(
        page_content="2024年第二季度营收为120亿韩元。",
        metadata={"year": 2024, "quarter": "Q2", "category": "financial"}
    ),
    Document(
        page_content="技术路线图：计划强化 AI 功能。",
        metadata={"year": 2024, "quarter": "Q1", "category": "strategy"}
    ),
]

# 用元数据过滤缩小检索范围
filtered_results = vectorstore.similarity_search(
    "营收表现",
    k=2,
    filter={"category": "financial", "year": 2024}
)

from sentence_transformers import CrossEncoder
import numpy as np

# 加载 Cross-Encoder 模型
cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

# 初始检索结果(基于双编码器)
query = "机器学习算法的种类"
initial_results = vectorstore.similarity_search(query, k=20)  # 多检索一些

# 用交叉编码器重新排序
pairs = [[query, doc.page_content] for doc in initial_results]
scores = cross_encoder.predict(pairs)

# 按分数排序
ranked = sorted(zip(scores, initial_results), reverse=True)
top_k = [doc for _, doc in ranked[:5]]  # 只使用前 5 个

print("重排序后的头部结果:")
for score, doc in ranked[:3]:
    print(f"  分数 {score:.4f}: {doc.page_content[:80]}...")

import cohere
from langchain.retrievers.document_compressors import CohereRerank
from langchain.retrievers import ContextualCompressionRetriever

co = cohere.Client("your-api-key")

# Cohere 重排序压缩器
compressor = CohereRerank(
    cohere_api_key="your-api-key",
    top_n=3,
    model="rerank-multilingual-v3.0"  # 支持韩语
)

# 压缩检索器(检索 + 重排序)
compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=vectorstore.as_retriever(search_kwargs={"k": 20})
)

results = compression_retriever.invoke("请告诉我公司的政策")
print(f"重排序后的文档数: {len(results)}")

from FlagEmbedding import FlagReranker

# BGE 重排序器(开源，本地运行)
reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True)

query = "What is RAG?"
passages = [
    "RAG stands for Retrieval-Augmented Generation.",
    "A rag is a piece of cloth used for cleaning.",
    "RAG systems combine retrieval with generation for better LLM responses.",
]

# 逐对计算分数
scores = reranker.compute_score([[query, p] for p in passages])
ranked = sorted(zip(scores, passages), reverse=True)

for score, passage in ranked:
    print(f"  {score:.4f}: {passage}")

问题 → LLM 生成假想答案 → 对假想答案做嵌入 → 检索真实文档

from langchain.chains import HypotheticalDocumentEmbedder
from langchain_openai import OpenAI, OpenAIEmbeddings
from langchain_community.vectorstores import Chroma

# HyDE 提示词
hyde_prompt = """你是一位专家。请针对以下问题撰写一份详细的答案文档。
请让这份文档写得像是真实存在的一样。

问题: {question}

假想答案文档:"""

# HyDE 嵌入生成器
llm = OpenAI()
embeddings = OpenAIEmbeddings()

hyde_embeddings = HypotheticalDocumentEmbedder.from_llm(
    llm=llm,
    embeddings=embeddings,
    prompt_key="web_search"  # 或使用自定义提示词
)

# 使用 HyDE 的检索
query = "量子计算对加密的影响"
# HyDE 内部会先生成假想文档，再做嵌入
hyde_retriever = vectorstore.as_retriever(embedding_function=hyde_embeddings)
results = hyde_retriever.invoke(query)

# 手动实现 HyDE
def manual_hyde(query, llm, embeddings, vectorstore, k=4):
    # 1. 生成假想文档
    hypothetical_doc = llm.invoke(
        f"请针对以下问题撰写详细答案: {query}"
    )

    # 2. 对假想文档做嵌入
    hyp_embedding = embeddings.embed_query(hypothetical_doc.content)

    # 3. 用假想文档的嵌入进行检索
    results = vectorstore.similarity_search_by_vector(hyp_embedding, k=k)

    return results, hypothetical_doc.content

from langchain_openai import ChatOpenAI
chat_llm = ChatOpenAI(temperature=0.7)
results, hyp_doc = manual_hyde(
    "深度学习的历史", chat_llm, embeddings, vectorstore
)
print(f"生成的假想文档: {hyp_doc[:200]}...")
print(f"检索到的实际文档: {len(results)} 个")

from langchain_openai import ChatOpenAI
from langchain.prompts import ChatPromptTemplate

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

# Self-RAG 模拟(真正的 Self-RAG 需要专门训练的模型)
class SelfRAGSimulator:
    def __init__(self, retriever, llm):
        self.retriever = retriever
        self.llm = llm

    def should_retrieve(self, query: str) -> bool:
        """判断是否需要检索"""
        prompt = ChatPromptTemplate.from_template("""
请判断回答以下问题是否需要检索外部文档。
如果仅凭常识或推理即可回答，请只回答 'NO'；
如果需要特定事实或专业知识，请只回答 'YES'。

问题: {query}
判断 (YES/NO):""")
        response = self.llm.invoke(prompt.format_messages(query=query))
        return "YES" in response.content.upper()

    def is_relevant(self, query: str, doc_content: str) -> bool:
        """判断文档是否与问题相关"""
        prompt = ChatPromptTemplate.from_template("""
请判断以下文档是否与问题相关。
只回答 'RELEVANT' 或 'IRRELEVANT'。

问题: {query}
文档: {doc}
判断:""")
        response = self.llm.invoke(
            prompt.format_messages(query=query, doc=doc_content[:500])
        )
        return "RELEVANT" in response.content.upper()

    def generate_with_reflection(self, query: str) -> str:
        """以 Self-RAG 方式生成应答"""
        # 1. 判断是否需要检索
        need_retrieve = self.should_retrieve(query)
        print(f"是否需要检索: {need_retrieve}")

        if not need_retrieve:
            # 不检索，直接回答
            response = self.llm.invoke(query)
            return response.content

        # 2. 检索文档
        docs = self.retriever.invoke(query)

        # 3. 按相关性过滤
        relevant_docs = [d for d in docs if self.is_relevant(query, d.page_content)]
        print(f"相关文档: {len(relevant_docs)}/{len(docs)} 个")

        if not relevant_docs:
            return "未能找到相关文档，将以常识回答: " + \
                   self.llm.invoke(query).content

        # 4. 基于上下文生成答案
        context = "\n\n".join([d.page_content for d in relevant_docs[:3]])
        prompt = f"""请使用上下文回答问题。
上下文: {context}
问题: {query}
回答:"""
        return self.llm.invoke(prompt).content

from langchain_community.tools.tavily_search import TavilySearchResults
from typing import List, Tuple

class CorrectiveRAG:
    def __init__(self, retriever, llm):
        self.retriever = retriever
        self.llm = llm
        self.web_search = TavilySearchResults(max_results=3)

    def evaluate_documents(self, query: str, docs: list) -> Tuple[str, List]:
        """
        评估文档相关性
        返回: ("CORRECT"|"INCORRECT"|"AMBIGUOUS", 过滤后的文档)
        """
        evaluation_prompt = """请评估检索到的文档与问题的相关性。
- CORRECT: 文档可以直接回答问题
- INCORRECT: 文档与问题无关
- AMBIGUOUS: 部分相关但不完整

问题: {query}
文档:
{docs}

评估 (CORRECT/INCORRECT/AMBIGUOUS):"""

        docs_text = "\n---\n".join([d.page_content[:300] for d in docs[:4]])
        response = self.llm.invoke(
            evaluation_prompt.format(query=query, docs=docs_text)
        )

        evaluation = response.content.strip().upper()
        if "CORRECT" in evaluation:
            return "CORRECT", docs
        elif "INCORRECT" in evaluation:
            return "INCORRECT", []
        else:
            return "AMBIGUOUS", docs

    def run(self, query: str) -> str:
        # 1. 初始检索
        docs = self.retriever.invoke(query)

        # 2. 评估文档质量
        status, filtered_docs = self.evaluate_documents(query, docs)
        print(f"文档评估结果: {status}")

        # 3. 按情况处理
        if status == "INCORRECT":
            # 用网络搜索补充
            print("正在用网络搜索补充...")
            web_results = self.web_search.invoke(query)
            context = "\n".join([r['content'] for r in web_results])
        elif status == "AMBIGUOUS":
            # 知识精炼 + 网络搜索相结合
            web_results = self.web_search.invoke(query)
            web_context = "\n".join([r['content'] for r in web_results])
            doc_context = "\n".join([d.page_content for d in filtered_docs[:2]])
            context = doc_context + "\n\n[网络搜索补充]\n" + web_context
        else:
            # 直接使用检索到的文档
            context = "\n\n".join([d.page_content for d in filtered_docs[:4]])

        # 4. 生成最终应答
        response = self.llm.invoke(
            f"上下文:\n{context}\n\n问题: {query}\n回答:"
        )
        return response.content

class AdaptiveRAG:
    def __init__(self, simple_retriever, advanced_retriever, llm):
        self.simple_retriever = simple_retriever   # 简单向量检索
        self.advanced_retriever = advanced_retriever  # 混合检索 + 重排序
        self.llm = llm

    def classify_query(self, query: str) -> str:
        """对查询复杂度分类"""
        prompt = f"""请对以下问题的复杂度进行分类:
- simple: 单纯的事实确认，或可直接回答
- complex: 需要综合多个来源、多步推理

问题: {query}
分类 (simple/complex):"""

        response = self.llm.invoke(prompt)
        return "complex" if "complex" in response.content.lower() else "simple"

    def run(self, query: str) -> str:
        query_type = self.classify_query(query)
        print(f"查询类型: {query_type}")

        if query_type == "simple":
            docs = self.simple_retriever.invoke(query)
        else:
            docs = self.advanced_retriever.invoke(query)

        context = "\n\n".join([d.page_content for d in docs])
        return self.llm.invoke(
            f"上下文:\n{context}\n\n问题: {query}\n回答:"
        ).content

# 安装并初始化 GraphRAG
pip install graphrag

# 初始化项目
python -m graphrag.index --init --root ./ragtest

# 修改配置文件后建立索引
python -m graphrag.index --root ./ragtest

# 全局检索(需要理解整体文档)
python -m graphrag.query --root ./ragtest --method global "What are the main themes?"

# 局部检索(围绕特定实体)
python -m graphrag.query --root ./ragtest --method local "Tell me about company X"

from ragas import evaluate
from ragas.metrics import (
    faithfulness,
    answer_relevancy,
    context_recall,
    context_precision
)
from datasets import Dataset

# 准备评估数据
evaluation_data = {
    "question": [
        "公司的年假制度是怎样的？",
        "远程办公政策是什么？",
    ],
    "answer": [
        "入职一年后可获得15天年假，此后每年增加1天。",
        "每周可远程办公3天。",
    ],
    "contexts": [
        ["员工在入职一年后获得15天年假。此后每年增加1天。"],
        ["员工每周可远程办公2天。经特殊批准可增加。"],
    ],
    "ground_truth": [
        "入职一年后15天，每年增加1天",
        "每周远程办公2天为基本，经批准可增加",
    ]
}

dataset = Dataset.from_dict(evaluation_data)

result = evaluate(
    dataset,
    metrics=[
        faithfulness,
        answer_relevancy,
        context_recall,
        context_precision
    ]
)

print(result)
# faithfulness: 0.75 (检测到答案与上下文不一致的部分)
# answer_relevancy: 0.92
# context_recall: 0.85
# context_precision: 0.78

import json
from typing import List, Dict

class RAGEvaluator:
    def __init__(self, rag_chain, llm):
        self.rag_chain = rag_chain
        self.llm = llm

    def evaluate_faithfulness(self, answer: str, context: str) -> float:
        """评估答案是否忠实于上下文 (0-1)"""
        prompt = f"""请评估以下答案是否仅基于给定上下文中的信息写成。
0.0 (完全不是) ~ 1.0 (完全忠实)

上下文: {context}
答案: {answer}

忠实度分数 (仅数字):"""
        response = self.llm.invoke(prompt)
        try:
            return float(response.content.strip())
        except:
            return 0.5

    def evaluate_answer_relevancy(self, question: str, answer: str) -> float:
        """评估答案与问题的相关性"""
        prompt = f"""请以 0.0~1.0 评估以下答案与问题的相关程度。

问题: {question}
答案: {answer}

相关度分数 (仅数字):"""
        response = self.llm.invoke(prompt)
        try:
            return float(response.content.strip())
        except:
            return 0.5

    def run_evaluation(self, test_cases: List[Dict]) -> Dict:
        results = []
        for case in test_cases:
            question = case["question"]
            expected = case.get("expected_answer", "")

            # 生成 RAG 应答
            result = self.rag_chain.invoke({"query": question})
            answer = result["result"]
            context = "\n".join([d.page_content for d in result["source_documents"]])

            # 评估
            faithfulness_score = self.evaluate_faithfulness(answer, context)
            relevancy_score = self.evaluate_answer_relevancy(question, answer)

            results.append({
                "question": question,
                "answer": answer,
                "faithfulness": faithfulness_score,
                "relevancy": relevancy_score,
            })

        # 汇总
        avg_faithfulness = sum(r["faithfulness"] for r in results) / len(results)
        avg_relevancy = sum(r["relevancy"] for r in results) / len(results)

        return {
            "results": results,
            "avg_faithfulness": avg_faithfulness,
            "avg_relevancy": avg_relevancy,
            "overall_score": (avg_faithfulness + avg_relevancy) / 2
        }

import hashlib
import json
from functools import lru_cache
import redis

class CachedRAGSystem:
    def __init__(self, rag_chain, redis_client=None, ttl=3600):
        self.rag_chain = rag_chain
        self.redis = redis_client
        self.ttl = ttl

    def _get_cache_key(self, query: str) -> str:
        return f"rag:{hashlib.md5(query.encode()).hexdigest()}"

    def query(self, query: str) -> dict:
        cache_key = self._get_cache_key(query)

        # 检查缓存
        if self.redis:
            cached = self.redis.get(cache_key)
            if cached:
                print("缓存命中！")
                return json.loads(cached)

        # 执行 RAG
        result = self.rag_chain.invoke({"query": query})
        response = {
            "answer": result["result"],
            "sources": [d.metadata for d in result["source_documents"]]
        }

        # 存储缓存
        if self.redis:
            self.redis.setex(cache_key, self.ttl, json.dumps(response))

        return response

# 嵌入缓存
from langchain.globals import set_llm_cache
from langchain_community.cache import InMemoryCache, SQLiteCache

# 内存缓存(用于开发)
set_llm_cache(InMemoryCache())

# SQLite 缓存(用于生产)
set_llm_cache(SQLiteCache(database_path=".langchain.db"))

from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain_openai import ChatOpenAI

# 流式 LLM
streaming_llm = ChatOpenAI(
    model="gpt-4o-mini",
    streaming=True,
    callbacks=[StreamingStdOutCallbackHandler()]
)

# 通过 FastAPI 实现的流式端点
from fastapi import FastAPI
from fastapi.responses import StreamingResponse
import asyncio

app = FastAPI()

async def generate_rag_stream(query: str):
    docs = retriever.invoke(query)
    context = "\n\n".join([d.page_content for d in docs])

    async for chunk in streaming_llm.astream(
        f"上下文:\n{context}\n\n问题: {query}\n回答:"
    ):
        if chunk.content:
            yield f"data: {chunk.content}\n\n"

@app.get("/rag/stream")
async def rag_stream_endpoint(query: str):
    return StreamingResponse(
        generate_rag_stream(query),
        media_type="text/event-stream"
    )

# 追踪 token 用量
from langchain.callbacks import get_openai_callback

with get_openai_callback() as cb:
    result = qa_chain.invoke({"query": "问题文本"})
    print(f"总 token 数: {cb.total_tokens}")
    print(f"提示词 token 数: {cb.prompt_tokens}")
    print(f"补全 token 数: {cb.completion_tokens}")
    print(f"成本: ${cb.total_cost:.6f}")

# 通过压缩分块节省上下文 token
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain.retrievers import ContextualCompressionRetriever

# 只抽取相关部分以压缩上下文
compressor = LLMChainExtractor.from_llm(llm)
compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=vectorstore.as_retriever(search_kwargs={"k": 8})
)

# 只把压缩后的文档传给 LLM(节省 token)
compressed_docs = compression_retriever.invoke("问题")
total_tokens = sum(len(d.page_content.split()) for d in compressed_docs)
print(f"压缩后的上下文 token 数(估算): {total_tokens}")

import time
import logging
from dataclasses import dataclass, field
from typing import Optional

@dataclass
class RAGMetrics:
    query: str
    retrieval_time: float = 0.0
    generation_time: float = 0.0
    num_docs_retrieved: int = 0
    answer_length: int = 0
    error: Optional[str] = None

class MonitoredRAGSystem:
    def __init__(self, rag_chain, logger=None):
        self.rag_chain = rag_chain
        self.logger = logger or logging.getLogger(__name__)
        self.metrics_history = []

    def query(self, query: str) -> dict:
        metrics = RAGMetrics(query=query)
        start_total = time.time()

        try:
            # 检索计时
            retrieval_start = time.time()
            docs = retriever.invoke(query)
            metrics.retrieval_time = time.time() - retrieval_start
            metrics.num_docs_retrieved = len(docs)

            # 生成计时
            gen_start = time.time()
            result = self.rag_chain.invoke({"query": query})
            metrics.generation_time = time.time() - gen_start
            metrics.answer_length = len(result["result"])

        except Exception as e:
            metrics.error = str(e)
            self.logger.error(f"RAG 错误: {e}")
            raise

        finally:
            total_time = time.time() - start_total
            self.metrics_history.append(metrics)
            self.logger.info(
                f"查询处理完成 | "
                f"检索: {metrics.retrieval_time:.2f}s | "
                f"生成: {metrics.generation_time:.2f}s | "
                f"总计: {total_time:.2f}s | "
                f"文档: {metrics.num_docs_retrieved} 个"
            )

        return result

    def get_stats(self) -> dict:
        if not self.metrics_history:
            return {}
        retrieval_times = [m.retrieval_time for m in self.metrics_history if not m.error]
        gen_times = [m.generation_time for m in self.metrics_history if not m.error]
        return {
            "total_queries": len(self.metrics_history),
            "error_rate": sum(1 for m in self.metrics_history if m.error) / len(self.metrics_history),
            "avg_retrieval_time": sum(retrieval_times) / len(retrieval_times) if retrieval_times else 0,
            "avg_generation_time": sum(gen_times) / len(gen_times) if gen_times else 0,
        }