第五篇：RAG高级篇 - 高级检索与优化#

# 用户查询:"Python 3.11的新特性"
# 向量检索可能返回:
# ❌ "Python 3.10的新特性"(语义相似,但版本不对)
# ❌ "Python的发展历史"(相关,但不精确)
# ✅ "Python 3.11 release notes"(精确匹配)

混合检索 = 向量检索(语义理解) + 全文检索(精确匹配)

from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader

# 步骤1: 加载和分块文档
loader = DirectoryLoader("./docs", glob="**/*.txt")
documents = loader.load()

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50
)
splits = text_splitter.split_documents(documents)

# 步骤2: 创建向量检索器
embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(
    documents=splits,
    embedding=embeddings
)
vector_retriever = vectorstore.as_retriever(
    search_kwargs={"k": 4}
)

# 步骤3: 创建BM25检索器
bm25_retriever = BM25Retriever.from_documents(splits)
bm25_retriever.k = 4

# 步骤4: 创建混合检索器(EnsembleRetriever)
ensemble_retriever = EnsembleRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    weights=[0.5, 0.5]  # 权重:向量50%,BM25 50%
)

# 步骤5: 使用混合检索
results = ensemble_retriever.invoke("Python 3.11的新特性有哪些?")
for doc in results:
    print(f"内容: {doc.page_content[:100]}...")
    print(f"来源: {doc.metadata.get('source', 'unknown')}\n")

内容: Python 3.11于2022年10月发布,主要新特性包括:
1. 性能提升(平均快25%)
2. 更好的错误提示...
来源: docs/python_releases.txt

内容: Python 3.11引入了异常组(Exception Groups)和except*语法...
来源: docs/python_311_features.txt

# 测试不同权重组合
test_weights = [
    ([0.7, 0.3], "向量为主"),
    ([0.5, 0.5], "平衡"),
    ([0.3, 0.7], "BM25为主")
]

query = "如何安装Python 3.11"

for weights, desc in test_weights:
    retriever = EnsembleRetriever(
        retrievers=[vector_retriever, bm25_retriever],
        weights=weights
    )
    results = retriever.invoke(query)
    print(f"\n{desc} (向量:{weights[0]}, BM25:{weights[1]})")
    print(f"Top-1: {results[0].page_content[:80]}...")

from langchain.agents import create_agent
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool

# 步骤1: 创建混合检索器(同上)
ensemble_retriever = EnsembleRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    weights=[0.6, 0.4]
)

# 步骤2: 包装为工具
@tool
def search_docs(query: str) -> str:
    """搜索文档库,返回相关信息。

    适用于:
    - 查找产品文档
    - 搜索技术资料
    - 获取配置说明
    """
    results = ensemble_retriever.invoke(query)
    # 格式化结果
    formatted = []
    for i, doc in enumerate(results[:3], 1):
        formatted.append(
            f"[文档{i}]\n"
            f"内容: {doc.page_content}\n"
            f"来源: {doc.metadata.get('source', 'unknown')}"
        )
    return "\n\n".join(formatted)

# 步骤3: 创建Agent
agent = create_agent(
    model="gpt-4",
    tools=[search_docs],
    system_prompt="""你是一个文档搜索助手,可以帮助用户查找技术文档。

使用 search_docs 工具搜索相关信息,然后给出准确的回答。"""
)

# 步骤4: 查询
result = agent.invoke({
    "messages": [{"role": "user", "content": "Python 3.11相比3.10有哪些性能提升?"}]
})
print(result["messages"][-1].content)

对于文档d,其RRF分数 = Σ (1 / (k + rank_i(d)))

其中:
- rank_i(d):文档d在第i个检索器中的排名
- k:常数(通常为60)

def reciprocal_rank_fusion(
    retriever_results: list[list],  # 多个检索器的结果列表
    k: int = 60
) -> list:
    """
    RRF算法实现

    Args:
        retriever_results: [[doc1, doc2, ...], [doc3, doc1, ...], ...]
        k: RRF常数

    Returns:
        融合后的文档列表(按RRF分数排序)
    """
    # 计算每个文档的RRF分数
    doc_scores = {}

    for retriever_docs in retriever_results:
        for rank, doc in enumerate(retriever_docs, start=1):
            doc_id = doc.metadata.get('id', id(doc))

            # RRF分数累加
            if doc_id not in doc_scores:
                doc_scores[doc_id] = {'doc': doc, 'score': 0}

            doc_scores[doc_id]['score'] += 1 / (k + rank)

    # 按分数排序
    sorted_docs = sorted(
        doc_scores.values(),
        key=lambda x: x['score'],
        reverse=True
    )

    return [item['doc'] for item in sorted_docs]

# 使用示例
vector_results = vector_retriever.invoke("Python 3.11")
bm25_results = bm25_retriever.invoke("Python 3.11")

fused_results = reciprocal_rank_fusion([vector_results, bm25_results])
print(f"融合后top-3文档:")
for i, doc in enumerate(fused_results[:3], 1):
    print(f"{i}. {doc.page_content[:80]}...")

from langchain_community.document_loaders import JSONLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_chroma import Chroma
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from langchain.agents import create_agent
from langchain_core.tools import tool
import json

# 步骤1: 准备产品数据
products_data = [
    {
        "id": "1",
        "name": "iPhone 14 Pro 256GB 深空黑",
        "category": "手机",
        "brand": "Apple",
        "model": "iPhone 14 Pro",
        "storage": "256GB",
        "color": "深空黑",
        "price": 8999,
        "description": "Apple 最新旗舰手机,搭载A16芯片,4800万像素主摄,支持灵动岛交互"
    },
    {
        "id": "2",
        "name": "iPhone 14 Pro 256GB 紫色",
        "category": "手机",
        "brand": "Apple",
        "model": "iPhone 14 Pro",
        "storage": "256GB",
        "color": "紫色",
        "price": 8999,
        "description": "Apple 旗舰机型,紫色配色,256GB大容量存储,专业级摄影系统"
    },
    {
        "id": "3",
        "name": "iPhone 13 Pro 256GB 远峰蓝",
        "category": "手机",
        "brand": "Apple",
        "model": "iPhone 13 Pro",
        "storage": "256GB",
        "color": "远峰蓝",
        "price": 7499,
        "description": "上一代旗舰,A15芯片,三摄系统,性价比之选"
    }
]

# 转换为Document对象
from langchain_core.documents import Document

documents = []
for product in products_data:
    # 构建富文本描述(便于检索)
    text = f"""
产品名称:{product['name']}
品牌:{product['brand']}
型号:{product['model']}
容量:{product['storage']}
颜色:{product['color']}
价格:¥{product['price']}
描述:{product['description']}
"""
    doc = Document(
        page_content=text,
        metadata=product
    )
    documents.append(doc)

# 步骤2: 创建混合检索器
embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(documents, embeddings)
vector_retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

bm25_retriever = BM25Retriever.from_documents(documents)
bm25_retriever.k = 3

# 产品搜索:BM25权重更高(精确匹配型号、容量、颜色)
product_retriever = EnsembleRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    weights=[0.4, 0.6]  # 更重视精确匹配
)

# 步骤3: 创建搜索工具
@tool
def search_products(query: str) -> str:
    """搜索产品库,返回匹配的商品信息。

    适用于查询:
    - 手机型号搜索(如"iPhone 14 Pro")
    - 特定配置搜索(如"256GB 紫色")
    - 价格范围查询
    """
    results = product_retriever.invoke(query)

    formatted = []
    for i, doc in enumerate(results, 1):
        meta = doc.metadata
        formatted.append(
            f"【商品{i}】{meta['name']}\n"
            f"价格:¥{meta['price']}\n"
            f"描述:{meta['description']}"
        )

    return "\n\n".join(formatted)

# 步骤4: 创建购物助手Agent
shopping_agent = create_agent(
    model="gpt-4",
    tools=[search_products],
    system_prompt="""你是一个专业的购物助手,帮助用户查找和推荐商品。

使用 search_products 工具搜索商品信息,然后给出准确的推荐。

注意:
- 优先推荐完全匹配用户需求的商品
- 如果没有完全匹配,推荐相近的替代品
- 清晰说明商品的价格和主要特点"""
)

# 步骤5: 测试查询
test_queries = [
    "我想买iPhone 14 Pro 256GB 紫色",
    "有没有256GB的紫色手机",
    "8000元左右的Apple旗舰机"
]

for query in test_queries:
    print(f"\n{'='*50}")
    print(f"用户查询:{query}")
    print('='*50)

    result = shopping_agent.invoke({
        "messages": [{"role": "user", "content": query}]
    })
    print(result["messages"][-1].content)

from typing import List
from langchain_core.documents import Document

def evaluate_retrieval(
    retriever,
    test_queries: List[str],
    ground_truth: List[List[str]]  # 每个查询的正确文档ID列表
) -> dict:
    """评估检索器性能"""
    precisions = []
    recalls = []

    for query, truth_ids in zip(test_queries, ground_truth):
        results = retriever.invoke(query)
        retrieved_ids = [doc.metadata.get('id') for doc in results]

        # 计算Precision@K
        hits = len(set(retrieved_ids) & set(truth_ids))
        precision = hits / len(retrieved_ids) if retrieved_ids else 0
        recall = hits / len(truth_ids) if truth_ids else 0

        precisions.append(precision)
        recalls.append(recall)

    return {
        'precision': sum(precisions) / len(precisions),
        'recall': sum(recalls) / len(recalls)
    }

# 测试不同权重
test_queries = [
    "iPhone 14 Pro 256GB 紫色",
    "8000元左右的旗舰机",
    "Apple最新手机"
]

ground_truth = [
    ["2"],  # 精确匹配
    ["1", "2", "3"],  # 价格范围
    ["1", "2"]  # 最新型号
]

weights_to_test = [
    (0.3, 0.7),
    (0.4, 0.6),
    (0.5, 0.5),
    (0.6, 0.4),
    (0.7, 0.3)
]

print("权重调优结果:\n")
for vec_weight, bm25_weight in weights_to_test:
    retriever = EnsembleRetriever(
        retrievers=[vector_retriever, bm25_retriever],
        weights=[vec_weight, bm25_weight]
    )

    metrics = evaluate_retrieval(retriever, test_queries, ground_truth)
    print(f"向量:{vec_weight:.1f}, BM25:{bm25_weight:.1f} => "
          f"Precision: {metrics['precision']:.2%}, "
          f"Recall: {metrics['recall']:.2%}")

查询类型
├── 包含精确关键词(型号、版本、规格)
│   └── 使用混合检索(BM25权重0.5-0.7)
├── 自然语言问句
│   └── 使用混合检索(向量权重0.6-0.7)
├── 纯语义查询(概念、主题)
│   └── 仅使用向量检索
└── 代码/命令搜索
    └── 使用混合检索(BM25权重0.7-0.8)

# 用户查询:"如何在Python中实现多线程安全的单例模式?"

# 向量检索top-5结果:
# 1. "Python单例模式实现"(相关度:0.85)
# 2. "多线程编程基础"(相关度:0.82)
# 3. "Python设计模式大全"(相关度:0.80)← 内容太泛化
# 4. "线程安全的单例模式"(相关度:0.78)← 应该排更前
# 5. "Python线程锁使用"(相关度:0.75)

# 问题:
# - 第3个文档太泛化,但相关度分数高
# - 第4个文档最相关,但排序靠后

原始检索(快速,粗排)→ top-100候选
         ↓
重排序模型(精细,慢)→ top-5最终结果
         ↓
LLM生成

from langchain.retrievers import ContextualCompressionRetriever
from langchain_community.retrievers.document_compressors import LLMChainExtractor
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader

# 步骤1: 准备文档
loader = DirectoryLoader("./docs", glob="**/*.txt")
documents = loader.load()

splitter = RecursiveCharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50
)
splits = splitter.split_documents(documents)

# 步骤2: 创建基础检索器(粗排)
embeddings = OpenAIEmbeddings()
vectorstore = Chroma.from_documents(splits, embeddings)
base_retriever = vectorstore.as_retriever(
    search_kwargs={"k": 20}  # 先检索20个候选
)

# 步骤3: 创建LLM压缩器(精排)
llm = ChatOpenAI(model="gpt-4", temperature=0)
compressor = LLMChainExtractor.from_llm(llm)

# 步骤4: 组合为压缩检索器
compression_retriever = ContextualCompressionRetriever(
    base_compressor=compressor,
    base_retriever=base_retriever
)

# 步骤5: 使用
query = "如何在Python中实现线程安全的单例模式?"
compressed_docs = compression_retriever.invoke(query)

print(f"原始检索器返回:20个文档")
print(f"重排序后返回:{len(compressed_docs)}个文档\n")

for i, doc in enumerate(compressed_docs, 1):
    print(f"【文档{i}】")
    print(f"内容:{doc.page_content[:150]}...")
    print(f"来源:{doc.metadata.get('source', 'unknown')}\n")

from langchain_community.retrievers.document_compressors import EmbeddingsFilter

# 创建embedding过滤器
embeddings_filter = EmbeddingsFilter(
    embeddings=OpenAIEmbeddings(),
    similarity_threshold=0.75  # 相似度阈值
)

# 组合为压缩检索器
compression_retriever = ContextualCompressionRetriever(
    base_compressor=embeddings_filter,
    base_retriever=base_retriever
)

# 使用
query = "Python单例模式实现"
filtered_docs = compression_retriever.invoke(query)

print(f"过滤前:{20}个文档")
print(f"过滤后:{len(filtered_docs)}个文档(相似度≥0.75)")

from sentence_transformers import CrossEncoder
from langchain_core.documents import Document
from typing import List

class CrossEncoderReranker:
    """基于Cross-Encoder的重排序器"""

    def __init__(self, model_name: str = "BAAI/bge-reranker-large"):
        """
        初始化Cross-Encoder重排序器

        Args:
            model_name: HuggingFace模型名称
                - BAAI/bge-reranker-base(英文,速度快)
                - BAAI/bge-reranker-large(英文,精度高)
                - BAAI/bge-reranker-v2-m3(多语言,推荐中文)
        """
        self.model = CrossEncoder(model_name)

    def rerank(
        self,
        query: str,
        documents: List[Document],
        top_n: int = 5
    ) -> List[Document]:
        """
        重排序文档列表

        Args:
            query: 用户查询
            documents: 候选文档列表
            top_n: 返回top-N个文档

        Returns:
            重排序后的文档列表
        """
        # 构建查询-文档对
        pairs = [[query, doc.page_content] for doc in documents]

        # 计算相关性分数
        scores = self.model.predict(pairs)

        # 按分数排序
        doc_scores = list(zip(documents, scores))
        doc_scores.sort(key=lambda x: x[1], reverse=True)

        # 返回top-N
        return [doc for doc, score in doc_scores[:top_n]]

# 使用示例
reranker = CrossEncoderReranker(model_name="BAAI/bge-reranker-v2-m3")

# 先用向量检索获取候选
base_results = base_retriever.invoke("Python单例模式线程安全实现")
print(f"初筛文档数:{len(base_results)}")

# 重排序
reranked_results = reranker.rerank(
    query="Python单例模式线程安全实现",
    documents=base_results,
    top_n=5
)

print(f"\n重排序后top-5:")
for i, doc in enumerate(reranked_results, 1):
    print(f"{i}. {doc.page_content[:80]}...")

from langchain.agents import create_agent
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool

# 步骤1: 创建带重排序的检索器
class RerankedRetriever:
    """带重排序的检索器"""

    def __init__(
        self,
        base_retriever,
        reranker,
        initial_k: int = 20,
        final_k: int = 5
    ):
        self.base_retriever = base_retriever
        self.reranker = reranker
        self.initial_k = initial_k
        self.final_k = final_k

    def invoke(self, query: str) -> List[Document]:
        """执行检索和重排序"""
        # 粗排:检索initial_k个候选
        candidates = self.base_retriever.invoke(query)

        # 精排:重排序并返回final_k个
        reranked = self.reranker.rerank(
            query=query,
            documents=candidates,
            top_n=self.final_k
        )

        return reranked

# 创建检索器实例
reranked_retriever = RerankedRetriever(
    base_retriever=base_retriever,  # 粗排检索器
    reranker=CrossEncoderReranker("BAAI/bge-reranker-v2-m3"),  # 重排序器
    initial_k=20,  # 粗排top-20
    final_k=5      # 精排top-5
)

# 步骤2: 包装为工具
@tool
def search_docs_with_rerank(query: str) -> str:
    """搜索文档库并重排序,返回最相关的内容。

    使用两阶段检索:
    1. 向量检索:快速筛选top-20候选
    2. Cross-Encoder重排序:精确排序top-5结果
    """
    results = reranked_retriever.invoke(query)

    formatted = []
    for i, doc in enumerate(results, 1):
        formatted.append(
            f"【文档{i}】\n"
            f"{doc.page_content}\n"
            f"来源:{doc.metadata.get('source', 'unknown')}"
        )

    return "\n\n".join(formatted)

# 步骤3: 创建Agent
agent = create_agent(
    model="gpt-4",
    tools=[search_docs_with_rerank],
    system_prompt="""你是一个技术文档助手,帮助用户查找和理解技术文档。

使用 search_docs_with_rerank 工具搜索相关文档,然后给出详细的技术解答。

注意:
- 优先使用检索到的文档内容
- 如果文档中有代码示例,请引用
- 保持回答的专业性和准确性"""
)

# 步骤4: 测试
result = agent.invoke({
    "messages": [{"role": "user", "content": "Python中如何实现线程安全的单例模式?请给出完整代码示例。"}]
})
print(result["messages"][-1].content)

# 原始查询:"怎么让Python快一点"
# 问题:
# - 表达模糊("快一点"指什么?)
# - 可能遗漏相关文档

# 解决方案:生成多个改写查询
# 1. "如何优化Python代码性能"
# 2. "Python程序加速方法"
# 3. "提升Python执行效率的技巧"

from langchain.retrievers import MultiQueryRetriever
from langchain_openai import ChatOpenAI

# 创建Multi-Query检索器
llm = ChatOpenAI(model="gpt-4", temperature=0)

multi_query_retriever = MultiQueryRetriever.from_llm(
    retriever=vectorstore.as_retriever(search_kwargs={"k": 5}),
    llm=llm
)

# 使用(自动生成3-5个改写查询并检索)
query = "怎么让Python快一点"
results = multi_query_retriever.invoke(query)

print(f"检索到{len(results)}个文档(去重后)")
for doc in results[:3]:
    print(f"- {doc.page_content[:80]}...")

用户查询:"怎么让Python快一点"
    ↓
LLM生成改写查询:
    ├── "如何优化Python代码性能"
    ├── "Python程序加速方法"
    └── "提升Python执行效率的技巧"
    ↓
并行检索3个查询
    ↓
合并 + 去重
    ↓
返回结果

查询:"张三的领导的领导是谁?"

向量RAG:
├── 检索到:"张三的直接领导是李四"
├── 检索到:"李四的绩效评估为优秀"
└── ❌ 无法推理出"张三的领导的领导"

知识图谱RAG:
├── 实体:张三 --[reports_to]--> 李四
├── 实体:李四 --[reports_to]--> 王五
└── ✅ 推理:张三 --> 李四 --> 王五(多跳查询)

# 安装Neo4j(使用Docker)
docker run \
    --name neo4j \
    -p 7474:7474 -p 7687:7687 \
    -e NEO4J_AUTH=neo4j/password \
    neo4j:latest

# 安装Python依赖
pip install langchain-community langchain-neo4j neo4j

from langchain_community.graphs import Neo4jGraph
from langchain_community.chains.graph_qa.cypher import GraphCypherQAChain
from langchain_openai import ChatOpenAI

# 步骤1: 连接Neo4j
graph = Neo4jGraph(
    url="bolt://localhost:7687",
    username="neo4j",
    password="password"
)

# 步骤2: 构建示例知识图谱
# 创建组织架构图
graph.query("""
// 创建员工节点
CREATE (张三:Employee {name: '张三', position: '工程师', department: '研发部'})
CREATE (李四:Employee {name: '李四', position: '技术经理', department: '研发部'})
CREATE (王五:Employee {name: '王五', position: '技术总监', department: '研发部'})
CREATE (赵六:Employee {name: '赵六', position: '产品经理', department: '产品部'})

// 创建汇报关系
CREATE (张三)-[:REPORTS_TO]->(李四)
CREATE (李四)-[:REPORTS_TO]->(王五)
CREATE (赵六)-[:REPORTS_TO]->(王五)

// 创建项目节点和参与关系
CREATE (项目A:Project {name: '项目A', status: '进行中'})
CREATE (张三)-[:WORKS_ON]->(项目A)
CREATE (赵六)-[:WORKS_ON]->(项目A)
""")

# 步骤3: 创建Cypher QA链
cypher_chain = GraphCypherQAChain.from_llm(
    llm=ChatOpenAI(model="gpt-4", temperature=0),
    graph=graph,
    verbose=True
)

# 步骤4: 查询(自动生成Cypher语句)
queries = [
    "张三的领导是谁?",
    "张三的领导的领导是谁?",  # 多跳查询
    "研发部有哪些人?",
    "项目A有哪些参与者?"
]

for query in queries:
    print(f"\n查询:{query}")
    result = cypher_chain.invoke({"query": query})
    print(f"答案:{result['result']}")

查询:张三的领导是谁?
生成Cypher:
MATCH (e:Employee {name: '张三'})-[:REPORTS_TO]->(manager)
RETURN manager.name

答案:张三的领导是李四

查询:张三的领导的领导是谁?
生成Cypher:
MATCH (e:Employee {name: '张三'})-[:REPORTS_TO*2]->(manager)
RETURN manager.name

答案:张三的领导的领导是王五

from langchain_community.vectorstores import Neo4jVector
from langchain_openai import OpenAIEmbeddings

# 使用Neo4jVector实现"向量入口 + 图遍历"的检索
# 假设Graph中已经存在 Employee 节点, 且包含 "name", "position" 等属性

# 步骤1: 创建向量检索器 (连接现有的图)
vector_store = Neo4jVector.from_existing_graph(
    embedding=OpenAIEmbeddings(),
    url="bolt://localhost:7687",
    username="neo4j",
    password="password",
    index_name="employee_index",
    node_label="Employee",
    text_node_properties=["name", "position", "department"], # 这些属性内容会被向量化
    embedding_node_property="embedding", # 向量存放在节点的embedding属性中

    # 【核心Magic】retrieval_query: 向量检索找到节点后, 执行此Cypher获取上下文
    # 这里的 'node' 是向量匹配到的节点
    retrieval_query="""
    // 找到该员工的直接下属 (1跳关系)
    MATCH (node)<-[:REPORTS_TO]-(subordinate)
    RETURN "员工: " + node.name + " (" + node.position + ")" +
           " 管理着: " + subordinate.name + " (" + subordinate.position + ")" AS text,
           score,
           {} AS metadata
    """
)

# 步骤2: 执行检索
# 查询: "谁是李四的下属?"
# 1. 向量检索找到 "李四" 节点
# 2. 执行 retrieval_query 找到李四的下属
results = vector_store.similarity_search("李四", k=1)

print(f"GraphRAG检索结果:")
for doc in results:
    print(doc.page_content)

# 输出示例:
# 员工: 李四 (技术经理) 管理着: 张三 (工程师)

传统RAG:
查询 → 检索 → 生成

Self-RAG:
查询 → 检索 → 评估相关性 → 生成 → 验证答案 → (重新检索)

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

class SelfRAG:
    """Self-RAG实现"""

    def __init__(self, retriever, llm):
        self.retriever = retriever
        self.llm = llm

        # 相关性评估提示词
        self.relevance_prompt = ChatPromptTemplate.from_template("""
评估以下文档与查询的相关性。

查询:{query}

文档:
{document}

相关性(0-10分):
理由:
""")

        # 答案验证提示词
        self.verification_prompt = ChatPromptTemplate.from_template("""
验证答案是否被文档支持。

查询:{query}
答案:{answer}

文档:
{documents}

验证结果(支持/不支持):
理由:
""")

    def invoke(self, query: str, max_iterations: int = 2) -> dict:
        """Self-RAG查询"""
        iteration = 0
        while iteration < max_iterations:
            iteration += 1
            print(f"\n第{iteration}轮检索...")

            # 步骤1: 检索文档
            docs = self.retriever.invoke(query)

            # 步骤2: 评估相关性
            relevant_docs = []
            for doc in docs[:5]:
                relevance_result = (self.relevance_prompt | self.llm).invoke({
                    "query": query,
                    "document": doc.page_content
                })

                # 简单解析评分(实际应用中可用structured output)
                if "相关性" in relevance_result.content:
                    score_line = [
                        line for line in relevance_result.content.split('\n')
                        if '相关性' in line or '分' in line
                    ][0]

                    try:
                        score = int(''.join(filter(str.isdigit, score_line))[:2])
                        if score >= 7:
                            relevant_docs.append(doc)
                    except:
                        continue

            if not relevant_docs:
                print("  未找到相关文档,重新检索...")
                continue

            print(f"  找到{len(relevant_docs)}个相关文档")

            # 步骤3: 生成答案
            context = "\n\n".join([doc.page_content for doc in relevant_docs])
            answer_prompt = f"""
基于以下文档回答问题:

{context}

问题:{query}

答案:
"""
            answer = self.llm.invoke(answer_prompt).content

            # 步骤4: 验证答案
            verification = (self.verification_prompt | self.llm).invoke({
                "query": query,
                "answer": answer,
                "documents": context
            })

            if "支持" in verification.content:
                return {
                    "query": query,
                    "answer": answer,
                    "documents": relevant_docs,
                    "iterations": iteration,
                    "verified": True
                }

            print("  答案未通过验证,重新检索...")

        # 最大迭代次数后仍未验证通过
        return {
            "query": query,
            "answer": answer,
            "documents": relevant_docs,
            "iterations": iteration,
            "verified": False
        }

# 使用示例
self_rag = SelfRAG(
    retriever=vectorstore.as_retriever(),
    llm=ChatOpenAI(model="gpt-4", temperature=0)
)

result = self_rag.invoke("Python 3.11的主要新特性是什么?")
print(f"\n最终答案:\n{result['answer']}")
print(f"\n迭代次数:{result['iterations']}")
print(f"验证通过:{result['verified']}")

from langchain_community.tools import DuckDuckGoSearchRun

class CorrectiveRAG:
    """Corrective RAG实现"""

    def __init__(self, retriever, llm):
        self.retriever = retriever
        self.llm = llm
        self.web_search = DuckDuckGoSearchRun()

    def invoke(self, query: str) -> str:
        """CRAG查询"""
        # 步骤1: 本地检索
        docs = self.retriever.invoke(query)

        # 步骤2: 评估相关性
        relevance_score = self._evaluate_relevance(query, docs)

        if relevance_score >= 0.8:
            # 高相关:直接使用
            context = "\n\n".join([doc.page_content for doc in docs[:3]])
            return self._generate_answer(query, context)

        elif relevance_score < 0.4:
            # 低相关:网络搜索
            print("本地文档相关性低,启动网络搜索...")
            web_results = self.web_search.invoke(query)
            return self._generate_answer(query, web_results)

        else:
            # 中等相关:知识精炼
            print("文档需要精炼...")
            refined_context = self._refine_knowledge(query, docs)
            return self._generate_answer(query, refined_context)

    def _evaluate_relevance(self, query: str, docs: list) -> float:
        """评估文档相关性(简化版)"""
        # 实际应用中可使用专门的评估模型
        prompt = f"""
评估以下文档与查询的平均相关性(0-1之间的分数)。

查询:{query}

文档:
{docs[0].page_content[:200]}...

相关性分数(0-1):
"""
        response = self.llm.invoke(prompt).content
        try:
            score = float(''.join(filter(lambda x: x.isdigit() or x == '.', response)))
            return min(1.0, score)
        except:
            return 0.5

    def _refine_knowledge(self, query: str, docs: list) -> str:
        """知识精炼:提取关键信息"""
        context = "\n\n".join([doc.page_content for doc in docs[:5]])
        prompt = f"""
从以下文档中提取与查询相关的关键信息(去除无关内容)。

查询:{query}

文档:
{context}

关键信息:
"""
        return self.llm.invoke(prompt).content

    def _generate_answer(self, query: str, context: str) -> str:
        """生成答案"""
        prompt = f"""
基于以下信息回答问题:

{context}

问题:{query}

答案:
"""
        return self.llm.invoke(prompt).content

# 使用
crag = CorrectiveRAG(
    retriever=vectorstore.as_retriever(),
    llm=ChatOpenAI(model="gpt-4", temperature=0)
)

answer = crag.invoke("量子计算机的最新进展")
print(answer)

from langchain.agents import create_agent
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool

# 定义多个检索工具
@tool
def search_technical_docs(query: str) -> str:
    """搜索技术文档"""
    # 使用向量检索器
    results = tech_retriever.invoke(query)
    return "\n".join([doc.page_content for doc in results[:3]])

@tool
def search_company_policies(query: str) -> str:
    """搜索公司政策文档"""
    results = policy_retriever.invoke(query)
    return "\n".join([doc.page_content for doc in results[:3]])

@tool
def query_database(sql_query: str) -> str:
    """查询数据库获取统计数据"""
    # 简化示例
    return "查询结果:..."

# 创建Agentic RAG
agentic_rag = create_agent(
    model="gpt-4",
    tools=[search_technical_docs, search_company_policies, query_database],
    system_prompt="""你是一个智能助手,可以查询技术文档、公司政策和数据库。

根据用户问题自主决定:
1. 需要使用哪些工具
2. 以什么顺序使用工具
3. 如何组合不同来源的信息

保持回答准确、完整、专业。"""
)

# 使用
result = agentic_rag.invoke({
    "messages": [{"role": "user", "content": "公司的远程办公政策是什么?需要提交哪些申请?"}]
})
print(result["messages"][-1].content)

用户查询
    ↓
┌───────────────────────────────────┐
│  查询处理层 (Query Processing)      │
│  - 查询改写                         │
│  - 意图识别                         │
│  - 参数提取                         │
└───────────┬───────────────────────┘
            ↓
┌───────────────────────────────────┐
│  并发检索层 (Parallel Retrieval)    │
│  ┌─────────┐      ┌─────────┐     │
│  │ 向量检索 │      │ BM25检索 │     │
│  └─────────┘      └─────────┘     │
└───────────┬───────────────────────┘
            ↓
┌───────────────────────────────────┐
│  融合排序层 (Hybrid Ranking)        │
│  - RRF融合                          │
│  - 权重调整                         │
│  - 去重                             │
└───────────┬───────────────────────┘
            ↓
┌───────────────────────────────────┐
│  重排序层 (Reranking) [可选]        │
│  - Cross-Encoder精排                │
│  - 多样性优化                        │
└───────────┬───────────────────────┘
            ↓
┌───────────────────────────────────┐
│  缓存层 (Caching)                   │
│  - 查询缓存                         │
│  - 结果缓存                         │
└───────────┬───────────────────────┘
            ↓
┌───────────────────────────────────┐
│  监控层 (Monitoring)                │
│  - 延迟监控                         │
│  - 质量监控                         │
│  - 成本监控                         │
└───────────────────────────────────┘

# production_hybrid_retriever.py
from typing import List, Dict, Any, Optional
from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma
from langchain_community.retrievers import BM25Retriever
from langchain.retrievers import EnsembleRetriever
from functools import lru_cache
import asyncio
import time
import logging
from dataclasses import dataclass

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

@dataclass
class RetrievalMetrics:
    """检索指标"""
    latency_ms: float
    num_results: int
    cache_hit: bool
    retriever_used: str

class ProductionHybridRetriever:
    """生产级混合检索器"""

    def __init__(
        self,
        documents: List[Document],
        vector_weight: float = 0.5,
        bm25_weight: float = 0.5,
        top_k: int = 4,
        enable_cache: bool = True,
        enable_rerank: bool = False,
        enable_monitoring: bool = True
    ):
        self.vector_weight = vector_weight
        self.bm25_weight = bm25_weight
        self.top_k = top_k
        self.enable_cache = enable_cache
        self.enable_rerank = enable_rerank
        self.enable_monitoring = enable_monitoring

        # 初始化检索器
        logger.info("初始化生产级混合检索器...")

        # 向量检索器
        embeddings = OpenAIEmbeddings()
        self.vectorstore = Chroma.from_documents(
            documents=documents,
            embedding=embeddings
        )
        self.vector_retriever = self.vectorstore.as_retriever(
            search_kwargs={"k": top_k}
        )

        # BM25检索器
        self.bm25_retriever = BM25Retriever.from_documents(documents)
        self.bm25_retriever.k = top_k

        # 混合检索器
        self.ensemble_retriever = EnsembleRetriever(
            retrievers=[self.vector_retriever, self.bm25_retriever],
            weights=[vector_weight, bm25_weight]
        )

        # 监控数据
        self.metrics_history: List[RetrievalMetrics] = []

        logger.info("✅ 混合检索器初始化完成")

    @lru_cache(maxsize=1000)
    def _cached_retrieve(self, query: str, k: int) -> tuple:
        """缓存的检索（使用tuple因为List不可hash）"""
        results = self._retrieve_internal(query, k)
        # 转为tuple以支持缓存
        return tuple((doc.page_content, doc.metadata) for doc in results)

    def _retrieve_internal(self, query: str, k: int) -> List[Document]:
        """内部检索实现"""
        return self.ensemble_retriever.invoke(query)

    async def _async_retrieve(self, query: str) -> List[Document]:
        """异步并发检索"""
        # 并发执行向量检索和BM25检索
        vector_task = asyncio.create_task(
            asyncio.to_thread(self.vector_retriever.invoke, query)
        )
        bm25_task = asyncio.create_task(
            asyncio.to_thread(self.bm25_retriever.invoke, query)
        )

        vector_results, bm25_results = await asyncio.gather(vector_task, bm25_task)

        # 手动融合（使用RRF）
        return self._rrf_fusion(vector_results, bm25_results)

    def _rrf_fusion(
        self,
        vector_results: List[Document],
        bm25_results: List[Document],
        k: int = 60
    ) -> List[Document]:
        """Reciprocal Rank Fusion融合"""
        scores = {}

        # 向量检索评分
        for rank, doc in enumerate(vector_results):
            doc_id = id(doc)
            rrf_score = 1 / (k + rank + 1)
            scores[doc_id] = scores.get(doc_id, 0) + self.vector_weight * rrf_score

        # BM25检索评分
        for rank, doc in enumerate(bm25_results):
            doc_id = id(doc)
            rrf_score = 1 / (k + rank + 1)
            scores[doc_id] = scores.get(doc_id, 0) + self.bm25_weight * rrf_score

        # 合并并去重
        all_docs = {id(doc): doc for doc in vector_results + bm25_results}

        # 按分数排序
        sorted_doc_ids = sorted(scores.keys(), key=lambda x: scores[x], reverse=True)

        return [all_docs[doc_id] for doc_id in sorted_doc_ids[:self.top_k]]

    def retrieve(
        self,
        query: str,
        use_async: bool = False,
        bypass_cache: bool = False
    ) -> Dict[str, Any]:
        """执行检索并返回结果+指标"""

        start_time = time.time()
        cache_hit = False

        # 缓存检索
        if self.enable_cache and not bypass_cache:
            try:
                cached_results = self._cached_retrieve(query, self.top_k)
                documents = [
                    Document(page_content=content, metadata=metadata)
                    for content, metadata in cached_results
                ]
                cache_hit = True
            except Exception as e:
                logger.warning(f"缓存失败，使用直接检索: {e}")
                documents = self._retrieve_internal(query, self.top_k)
        elif use_async:
            # 异步并发检索
            documents = asyncio.run(self._async_retrieve(query))
        else:
            # 同步检索
            documents = self._retrieve_internal(query, self.top_k)

        latency_ms = (time.time() - start_time) * 1000

        # 记录指标
        metrics = RetrievalMetrics(
            latency_ms=latency_ms,
            num_results=len(documents),
            cache_hit=cache_hit,
            retriever_used="async" if use_async else "sync"
        )

        if self.enable_monitoring:
            self.metrics_history.append(metrics)
            logger.info(
                f"检索完成: {latency_ms:.0f}ms | "
                f"结果数: {len(documents)} | "
                f"缓存: {'命中' if cache_hit else '未命中'}"
            )

        return {
            "documents": documents,
            "metrics": metrics
        }

    def get_performance_stats(self) -> Dict[str, Any]:
        """获取性能统计"""
        if not self.metrics_history:
            return {}

        latencies = [m.latency_ms for m in self.metrics_history]
        cache_hits = sum(1 for m in self.metrics_history if m.cache_hit)

        return {
            "total_queries": len(self.metrics_history),
            "avg_latency_ms": sum(latencies) / len(latencies),
            "p50_latency_ms": sorted(latencies)[len(latencies) // 2],
            "p95_latency_ms": sorted(latencies)[int(len(latencies) * 0.95)],
            "p99_latency_ms": sorted(latencies)[int(len(latencies) * 0.99)],
            "cache_hit_rate": cache_hits / len(self.metrics_history),
            "avg_results_per_query": sum(m.num_results for m in self.metrics_history) / len(self.metrics_history)
        }

from langchain_community.document_loaders import DirectoryLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter

# 加载文档
loader = DirectoryLoader("./docs", glob="**/*.txt")
documents = loader.load()

# 分块
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
splits = splitter.split_documents(documents)

# 创建生产级检索器
retriever = ProductionHybridRetriever(
    documents=splits,
    vector_weight=0.6,
    bm25_weight=0.4,
    top_k=5,
    enable_cache=True,
    enable_monitoring=True
)

# 执行检索
result = retriever.retrieve("Python 3.11的新特性")

print(f"检索结果: {len(result['documents'])}个")
print(f"延迟: {result['metrics'].latency_ms:.0f}ms")
print(f"缓存: {'命中' if result['metrics'].cache_hit else '未命中'}")

# 查看性能统计
stats = retriever.get_performance_stats()
print(f"\n性能统计:")
print(f"  平均延迟: {stats['avg_latency_ms']:.0f}ms")
print(f"  P95延迟: {stats['p95_latency_ms']:.0f}ms")
print(f"  缓存命中率: {stats['cache_hit_rate']:.1%}")