from langchain.text_splitter import CharacterTextSplitter

text_splitter = CharacterTextSplitter(
    separator="\n",
    chunk_size=500,
    chunk_overlap=50,
)
chunks = text_splitter.split_text(long_text)

from langchain_experimental.text_splitter import SemanticChunker
from langchain_openai.embeddings import OpenAIEmbeddings

embeddings = OpenAIEmbeddings()
semantic_splitter = SemanticChunker(
    embeddings=embeddings,
    breakpoint_threshold_type="percentile",  # 百分位阈值
    breakpoint_threshold_amount=95,          # 95分位作为分割点
)
chunks = semantic_splitter.create_documents([long_text])

from langchain.text_splitter import RecursiveCharacterTextSplitter

text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    separators=["\n\n", "\n", "。", "！", "？", "；", "，", " ", ""],
    length_function=len,
)
chunks = text_splitter.create_documents([long_text])

from langchain_text_splitters import MarkdownHeaderTextSplitter

headers_to_split_on = [
    ("#", "Header 1"),
    ("##", "Header 2"),
    ("###", "Header 3"),
]

markdown_splitter = MarkdownHeaderTextSplitter(
    headers_to_split_on=headers_to_split_on,
)
chunks = markdown_splitter.split_text(markdown_text)

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

# 父级切分（大 chunk）
parent_splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200)
# 子级切分（小 chunk，用于索引）
child_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)

# 向量存储（存子 chunk 的 embedding）
vectorstore = Chroma(embedding_function=OpenAIEmbeddings())
# 存储父文档
store = InMemoryStore()

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

# 添加文档
retriever.add_documents([document])

# 检索：用子 chunk 匹配，返回父 chunk
results = retriever.invoke("query")

from sentence_transformers import SentenceTransformer, SentencesDataset, losses
from sentence_transformers.readers import InputExample
from torch.utils.data import DataLoader

# 加载预训练模型
model = SentenceTransformer('BAAI/bge-base-zh-v1.5')

# 准备训练数据：(query, positive_doc, negative_doc) 三元组
train_examples = [
    InputExample(texts=["查询1", "相关文档1", "不相关文档1"]),
    InputExample(texts=["查询2", "相关文档2", "不相关文档2"]),
    # ... 更多三元组
]

train_dataset = SentencesDataset(train_examples, model)
train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=16)

# 使用 MultipleNegativesRankingLoss
train_loss = losses.MultipleNegativesRankingLoss(model)

# 微调
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=3,
    warmup_steps=100,
    output_path='./finetuned-bge-zh'
)

# BGE-M3 天然支持多语言，query 和 document 用同一模型
from FlagEmbedding import BGEM3FlagModel

model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)

# 中文 query 和英文文档都可以 embedding
query_emb = model.encode(['如何配置 CUDA 环境？'])['dense_vecs']
doc_emb = model.encode(['CUDA environment setup guide for Ubuntu'])['dense_vecs']

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

# === BGE-M3 ===
from FlagEmbedding import BGEM3FlagModel
bge_model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True)

# === OpenAI ===
from openai import OpenAI
openai_client = OpenAI()

def get_openai_embedding(text):
    resp = openai_client.embeddings.create(
        model="text-embedding-3-small",
        input=text
    )
    return np.array(resp.data[0].embedding)

# === 测试 ===
query = "RAG 中如何提升召回率？"
docs = [
    "本文介绍了多种提升 RAG 检索召回率的方法，包括优化切分策略、选择更好的 embedding 模型等。",
    "今天天气不错，适合出去散步。",
    "Retrieval-Augmented Generation combines retrieval with language model generation.",
]

# BGE-M3
q_bge = bge_model.encode([query])['dense_vecs']
d_bge = bge_model.encode(docs)['dense_vecs']
sim_bge = cosine_similarity(q_bge, d_bge)[0]

# OpenAI
q_openai = get_openai_embedding(query)
d_openai = np.array([get_openai_embedding(d) for d in docs])
sim_openai = cosine_similarity([q_openai], d_openai)[0]

print(f"{'文档':<70} | BGE-M3 | OpenAI")
print("-" * 95)
for i, doc in enumerate(docs):
    print(f"{doc:<70} | {sim_bge[i]:.4f} | {sim_openai[i]:.4f}")

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate

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

rewrite_prompt = ChatPromptTemplate.from_template("""
你是一位信息检索专家。请将以下用户问题改写为更适合搜索引擎检索的形式。
要求：
1. 提取关键实体和概念
2. 扩展同义词和相关术语
3. 保持语义不变

用户问题：{query}

改写后的查询：
""")

rewrite_chain = rewrite_prompt | llm

def rewrite_query(query: str) -> str:
    response = rewrite_chain.invoke({"query": query})
    return response.content.strip()

# 示例
original = "怎么配环境？"
rewritten = rewrite_query(original)
print(f"原始: {original}")
print(f"改写: {rewritten}")
# 输出示例: "开发环境配置步骤 CUDA PyTorch 安装教程"

def expand_with_synonyms(query: str) -> list[str]:
    """生成多个查询变体：原始 + 同义词扩展"""
    variants = [query]

    # 示例：简单同义词映射（实际可用 WordNet / HowNet / LLM）
    synonym_map = {
        "配置": ["设置", "安装", "部署", "环境搭建"],
        "提升": ["提高", "优化", "改善", "增强"],
        "召回率": ["recall", "检索准确率", "命中率"],
    }

    for word, synonyms in synonym_map.items():
        if word in query:
            for syn in synonyms[:2]:  # 取前2个同义词
                variants.append(query.replace(word, syn))

    return list(set(variants))

variants = expand_with_synonyms("如何提升 RAG 召回率")
# ['如何提升 RAG 召回率', '如何提高 RAG 召回率', '如何优化 RAG 召回率']

from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain.vectorstores import Chroma

# 假设已有 vectorstore
vectorstore = Chroma(...)

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

retriever = MultiQueryRetriever.from_llm(
    retriever=vectorstore.as_retriever(),
    llm=llm,
    parser_key="line",  # 按行解析生成的多个查询
)

# 一次调用，多路检索
docs = retriever.invoke("RAG 召回率怎么提升？")

from langchain.prompts import PromptTemplate
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import LLMChain

# 步骤 1：用 LLM 生成假设性文档
hyde_prompt = PromptTemplate(
    input_variables=["query"],
    template="请根据以下问题，生成一个假设性的参考答案（不需要真实，只需要看起来合理）：\n\n问题：{query}\n\n假设性答案："
)

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.7)
hyde_chain = LLMChain(llm=llm, prompt=hyde_prompt)

# 步骤 2：用假设文档的 embedding 检索真实文档
def hyde_retrieve(query: str, vectorstore, top_k: int = 5):
    # 生成假设文档
    hypothetical_doc = hyde_chain.run(query=query)

    # 用假设文档做 embedding 检索
    docs = vectorstore.similarity_search(hypothetical_doc, k=top_k)

    return docs

# 示例
query = "Transformer 中 Self-Attention 的计算公式是什么？"
docs = hyde_retrieve(query, vectorstore)

from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI

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

decompose_prompt = ChatPromptTemplate.from_template("""
将以下复杂问题分解为多个独立的子问题，每个子问题可以单独检索回答。

问题：{query}

子问题列表（每行一个）：
""")

decompose_chain = decompose_prompt | llm

def decompose_and_retrieve(query: str, vectorstore):
    # 分解子问题
    result = decompose_chain.invoke({"query": query})
    sub_queries = [line.strip() for line in result.content.split("\n") if line.strip()]

    # 每个子问题独立检索
    all_docs = []
    for sq in sub_queries:
        docs = vectorstore.similarity_search(sq, k=3)
        all_docs.extend(docs)

    # 去重
    unique_docs = list({doc.page_content: doc for doc in all_docs}.values())
    return unique_docs

# 示例
complex_query = "对比 BERT 和 GPT 的架构差异，以及它们各自适合什么场景？"
docs = decompose_and_retrieve(complex_query, vectorstore)
# 分解为:
# 1. "BERT 的架构特点是什么？"
# 2. "GPT 的架构特点是什么？"
# 3. "BERT 和 GPT 的架构差异对比"
# 4. "BERT 适合什么应用场景？"
# 5. "GPT 适合什么应用场景？"

class EnhancedRetriever:
    """组合 Multi-Query + HyDE 的增强检索器"""

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

    def retrieve(self, query: str, top_k: int = 10) -> list:
        all_docs = []

        # 策略 1：原始 query 直接检索
        docs_original = self.vectorstore.similarity_search(query, k=top_k)
        all_docs.extend(docs_original)

        # 策略 2：Multi-Query
        multi_query_prompt = f"""
        Generate 3 different versions of the following search query.
        Each version should use different wording but preserve the original intent.

        Original query: {query}

        Query versions (one per line):
        """
        response = self.llm.invoke(multi_query_prompt)
        variants = [line.strip() for line in response.content.split("\n") if line.strip() and line.strip()[0].isdigit()]

        for variant in variants[:3]:
            docs = self.vectorstore.similarity_search(variant, k=5)
            all_docs.extend(docs)

        # 策略 3：HyDE
        hyde_prompt = f"Generate a hypothetical document that would answer this query:\n{query}"
        hyde_response = self.llm.invoke(hyde_prompt)
        docs_hyde = self.vectorstore.similarity_search(hyde_response.content, k=5)
        all_docs.extend(docs_hyde)

        # 去重
        unique_docs = list({doc.page_content: doc for doc in all_docs}.values())
        return unique_docs[:top_k]

from rank_bm25 import BM25Okapi
import jieba

# 中文需要分词
corpus = [
    "RAG 是一种结合检索和生成的架构",
    "BM25 是经典的关键词匹配算法",
    "向量检索擅长语义匹配，BM25 擅长精确匹配",
]
tokenized_corpus = [list(jieba.cut(doc)) for doc in corpus]

bm25 = BM25Okapi(tokenized_corpus)

query = "向量检索和 BM25 的区别"
tokenized_query = list(jieba.cut(query))

scores = bm25.get_scores(tokenized_query)
for i, score in enumerate(scores):
    print(f"文档{i}: {score:.4f} - {corpus[i]}")

from rank_bm25 import BM25Okapi
import jieba
import numpy as np

class HybridRetriever:
    """BM25 + 向量混合检索器"""

    def __init__(self, corpus, vectorstore, top_k=10):
        self.corpus = corpus
        self.vectorstore = vectorstore
        self.top_k = top_k

        # 构建 BM25 索引
        tokenized = [list(jieba.cut(doc)) for doc in corpus]
        self.bm25 = BM25Okapi(tokenized)

    def retrieve(self, query: str) -> list:
        # 1. 向量检索
        vector_docs = self.vectorstore.similarity_search_with_score(query, k=self.top_k)

        # 2. BM25 检索
        tokenized_query = list(jieba.cut(query))
        bm25_scores = self.bm25.get_scores(tokenized_query)

        # 3. RRF 融合
        results = self._rrf_fusion(vector_docs, bm25_scores)
        return results

    def _rrf_fusion(self, vector_docs, bm25_scores, k_rrf=60):
        """Reciprocal Rank Fusion"""
        rrf_scores = {}

        # 向量检索结果的 RRF 贡献
        for rank, (doc, vscore) in enumerate(vector_docs):
            content = doc.page_content
            rrf_scores[content] = rrf_scores.get(content, 0) + 1 / (k_rrf + rank + 1)

        # BM25 结果的 RRF 贡献
        bm25_ranked = sorted(enumerate(bm25_scores), key=lambda x: -x[1])
        for rank, (idx, bscore) in enumerate(bm25_ranked[:self.top_k]):
            if bscore > 0:
                content = self.corpus[idx]
                rrf_scores[content] = rrf_scores.get(content, 0) + 1 / (k_rrf + rank + 1)

        # 按 RRF 分数排序
        sorted_results = sorted(rrf_scores.items(), key=lambda x: -x[1])
        return [(content, score) for content, score in sorted_results[:self.top_k]]

from elasticsearch import Elasticsearch

es = Elasticsearch("http://localhost:9200")

# 创建索引（同时支持 BM25 和 向量）
index_name = "rag-docs"

es.indices.create(index=index_name, body={
    "mappings": {
        "properties": {
            "content": {
                "type": "text",
                "analyzer": "ik_max_word"  # 中文分词
            },
            "embedding": {
                "type": "dense_vector",
                "dims": 1024,
                "index": True,
                "similarity": "cosine"
            }
        }
    }
})

# 索引文档
def index_document(doc_id, content, embedding):
    es.index(index=index_name, id=doc_id, body={
        "content": content,
        "embedding": embedding
    })

# 混合检索
def hybrid_search(query, query_embedding, top_k=10):
    response = es.search(index=index_name, body={
        "retriever": {
            "rrf": {
                "retrievers": [
                    # BM25 检索
                    {
                        "standard": {
                            "query": {
                                "match": {
                                    "content": {
                                        "query": query,
                                        "analyzer": "ik_max_word"
                                    }
                                }
                            }
                        }
                    },
                    # 向量检索（kNN）
                    {
                        "knn": {
                            "field": "embedding",
                            "query_vector": query_embedding,
                            "k": top_k,
                            "num_candidates": top_k * 2
                        }
                    }
                ],
                "rank_constant": 60,  # RRF 的 k 值
                "window_size": top_k
            }
        },
        "size": top_k
    })

    return [hit["_source"]["content"] for hit in response["hits"]["hits"]]

import os
import re
from datetime import datetime

def extract_metadata(file_path: str, content: str) -> dict:
    """从文件和文档内容中提取元数据"""
    metadata = {}

    # 文件级元数据
    metadata["source"] = file_path
    metadata["filename"] = os.path.basename(file_path)
    metadata["extension"] = os.path.splitext(file_path)[1]
    metadata["modified_time"] = os.path.getmtime(file_path)

    # 从内容中提取时间（如 "2024-03-15", "March 2024"）
    date_patterns = [
        r'(\d{4}-\d{2}-\d{2})',          # 2024-03-15
        r'(\d{4}/\d{2}/\d{2})',          # 2024/03/15
        r'((?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)\w* \d{4})',
    ]
    for pattern in date_patterns:
        match = re.search(pattern, content)
        if match:
            metadata["doc_date"] = match.group(1)
            break

    # 从 Markdown 标题提取主题
    headings = re.findall(r'^#{1,3}\s+(.+)$', content, re.MULTILINE)
    metadata["headings"] = headings

    # 用 LLM 提取标签
    # （此处省略，见 7.1 节进阶）

    return metadata

from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import JsonOutputParser

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

tagging_prompt = """
你是一个文档分类专家。请分析以下文档内容，提取以下信息：

1. 主题类别（如：教程、API文档、FAQ、最佳实践等）
2. 技术领域（如：机器学习、Web开发、数据库等）
3. 关键词（3-5个）
4. 难度等级（入门/中级/高级）

文档内容（前1000字）：
{text}

以 JSON 格式输出。
"""

tagging_chain = tagging_prompt | llm | JsonOutputParser()

metadata = tagging_chain.invoke({"text": content[:1000]})
# 输出: {"category": "教程", "domain": "机器学习", "keywords": ["RAG", "检索", "召回率"], "difficulty": "中级"}

from langchain_chroma import Chroma

# 只检索特定类型和时间范围的文档
docs = vectorstore.similarity_search(
    query="RAG 优化方法",
    k=10,
    filter={
        "doc_type": "tutorial",
        "doc_date": {"$gte": "2024-01-01"}
    }
)

# 先检索更多结果
docs = vectorstore.similarity_search(query="RAG 优化方法", k=50)

# 再按元数据过滤
filtered_docs = [
    doc for doc in docs
    if doc.metadata.get("doc_type") in ["tutorial", "guide"]
]

# 取 Top-K
final_docs = filtered_docs[:10]

from langchain.retrievers import SelfQueryRetriever
from langchain_chroma import Chroma
from langchain_openai import OpenAIEmbeddings, ChatOpenAI

# 定义元数据字段
metadata_field_info = [
    {"name": "doc_type", "type": "string", "description": "文档类型"},
    {"name": "domain", "type": "string", "description": "技术领域"},
    {"name": "difficulty", "type": "string", "description": "难度等级"},
    {"name": "doc_date", "type": "string", "description": "文档日期"},
]

# SelfQueryRetriever：LLM 自动从 query 中提取过滤条件
retriever = SelfQueryRetriever.from_llm(
    llm=ChatOpenAI(model="gpt-4o-mini", temperature=0),
    vectorstore=vectorstore,
    document_contents="技术文档集合，涵盖 AI、编程、数据库等领域",
    metadata_field_info=metadata_field_info,
)

# LLM 自动解析：
# "找一下最近半年关于 RAG 优化的中级教程"
# → 自动提取 filter: {"difficulty": "中级", "doc_date": ">= 2024-01-01"}
docs = retriever.invoke("找一下最近半年关于 RAG 优化的中级教程")

from langchain_community.graphs import NetworkxEntityGraph
import networkx as nx

# 构建简单的知识图谱
graph = NetworkxEntityGraph()
graph.add_node("PyTorch", {"type": "framework"})
graph.add_node("RAG", {"type": "architecture"})
graph.add_node("向量数据库", {"type": "component"})

# 添加关系
graph.add_edge("PyTorch", "RAG", relation="used_in")
graph.add_edge("向量数据库", "RAG", relation="component_of")

# 查询：找出与 RAG 直接相关的所有组件
related = list(graph._graph.neighbors("RAG"))
# ['PyTorch', '向量数据库']

# 用这些相关实体扩展检索查询
expanded_query = "RAG " + " ".join(related)

import chromadb
from chromadb.utils import embedding_functions

# 初始化
client = chromadb.PersistentClient(path="./chroma_db")
embed_fn = embedding_functions.SentenceTransformerEmbeddingFunction(
    model_name="BAAI/bge-m3"
)
collection = client.get_or_create_collection(
    name="tech-docs",
    embedding_function=embed_fn,
)

# 索引带元数据的文档
collection.add(
    documents=[
        "RAG 是一种结合检索和生成的架构...",
        "PyTorch 2.0 性能优化指南...",
        "如何部署 LLM 到生产环境...",
    ],
    metadatas=[
        {"doc_type": "guide", "domain": "LLM", "year": 2024},
        {"doc_type": "tutorial", "domain": "PyTorch", "year": 2024},
        {"doc_type": "ops", "domain": "LLM", "year": 2023},
    ],
    ids=["doc1", "doc2", "doc3"],
)

# 带元数据过滤的检索
results = collection.query(
    query_texts=["LLM 部署最佳实践"],
    n_results=5,
    where={"year": {"$gte": 2024}},        # 过滤条件
    where_document={"$contains": "部署"},  # 内容关键词过滤
)

for i, (doc, metadata, distance) in enumerate(zip(
    results["documents"][0], results["metadatas"][0], results["distances"][0]
)):
    print(f"[{i}] dist={distance:.4f} type={metadata['doc_type']} - {doc[:80]}...")

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

class CrossEncoderReranker:
    """Cross-Encoder 重排序器"""

    def __init__(self, model_name="BAAI/bge-reranker-v2-m3", device="cuda"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.model.to(device)
        self.model.eval()

    def rerank(self, query: str, documents: list[str], top_k: int = 5) -> list:
        """
        对候选文档进行重排序
        :param query: 用户查询
        :param documents: 候选文档列表（通常来自向量检索的 Top-30~Top-50）
        :param top_k: 返回的 top 结果数
        :return: 排序后的 (文档, 分数) 列表
        """
        # 构造 (query, document) 对
        pairs = [[query, doc] for doc in documents]

        # 批量推理
        with torch.no_grad():
            inputs = self.tokenizer(
                pairs,
                padding=True,
                truncation=True,
                max_length=512,
                return_tensors="pt"
            ).to(self.model.device)

            scores = self.model(**inputs, return_dict=True).logits.view(-1).float()

        # 排序
        scored_docs = sorted(
            zip(documents, scores.cpu().tolist()),
            key=lambda x: x[1],
            reverse=True
        )

        return scored_docs[:top_k]

from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.vectorstores import Chroma
from langchain_core.documents import Document

class TwoStageRetriever:
    """两阶段检索：向量召回 + Cross-Encoder 精排"""

    def __init__(self, vectorstore, reranker, recall_k=50, final_k=5):
        self.vectorstore = vectorstore
        self.reranker = reranker
        self.recall_k = recall_k    # 阶段1召回数量
        self.final_k = final_k      # 阶段2精排后返回数量

    def retrieve(self, query: str) -> list:
        # === 阶段1：向量召回（快速、大规模）===
        docs = self.vectorstore.similarity_search(query, k=self.recall_k)
        candidates = [doc.page_content for doc in docs]

        if len(candidates) == 0:
            return []

        # === 阶段2：Cross-Encoder 重排（精确、小规模）===
        ranked = self.reranker.rerank(query, candidates, top_k=self.final_k)

        # 返回结果
        return [Document(page_content=doc, metadata={"rerank_score": score})
                for doc, score in ranked]

# === 使用示例 ===
# 1. 初始化组件
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma(
    collection_name="tech-docs",
    embedding_function=embeddings,
    persist_directory="./chroma_db"
)

reranker = CrossEncoderReranker(model_name="BAAI/bge-reranker-v2-m3")

retriever = TwoStageRetriever(
    vectorstore=vectorstore,
    reranker=reranker,
    recall_k=50,
    final_k=5
)

# 2. 检索
query = "如何提升 RAG 的召回率？"
results = retriever.retrieve(query)

for i, doc in enumerate(results):
    print(f"[{i}] score={doc.metadata['rerank_score']:.4f}")
    print(f"    {doc.page_content[:100]}...\n")

# Batch 推理优化
def rerank_batch(self, query: str, documents: list[str], batch_size: int = 32):
    """分批量推理，避免 OOM"""
    all_scores = []
    for i in range(0, len(documents), batch_size):
        batch_docs = documents[i:i + batch_size]
        batch_pairs = [[query, doc] for doc in batch_docs]

        with torch.no_grad():
            inputs = self.tokenizer(batch_pairs, padding=True, truncation=True,
                                    max_length=512, return_tensors="pt").to(self.model.device)
            batch_scores = self.model(**inputs, return_dict=True).logits.view(-1).float()
            all_scores.extend(batch_scores.cpu().tolist())

    return all_scores

import numpy as np

def recall_at_k(retrieved_docs, relevant_docs, k):
    """Recall@K"""
    retrieved_top_k = retrieved_docs[:k]
    hits = len(set(retrieved_top_k) & set(relevant_docs))
    return hits / len(relevant_docs) if relevant_docs else 0

def precision_at_k(retrieved_docs, relevant_docs, k):
    """Precision@K"""
    retrieved_top_k = retrieved_docs[:k]
    hits = len(set(retrieved_top_k) & set(relevant_docs))
    return hits / k

def mrr(retrieved_docs, relevant_docs):
    """Mean Reciprocal Rank（单条 query）"""
    for i, doc in enumerate(retrieved_docs):
        if doc in relevant_docs:
            return 1.0 / (i + 1)
    return 0.0

def ndcg_at_k(retrieved_docs, relevant_docs, k, rel_map):
    """NDCG@K，rel_map 是 {doc_id: relevance_score} 的映射"""
    dcg = 0.0
    for i, doc in enumerate(retrieved_docs[:k]):
        rel = rel_map.get(doc, 0)
        dcg += (2**rel - 1) / np.log2(i + 2)

    # Ideal DCG
    ideal_rels = sorted([rel_map.get(d, 0) for d in retrieved_docs], reverse=True)[:k]
    idcg = sum((2**r - 1) / np.log2(i + 2) for i, r in enumerate(ideal_rels))

    return dcg / idcg if idcg > 0 else 0.0

from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import JsonOutputParser

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

qa_gen_prompt = """
你是评估数据集构建专家。请根据以下文档内容，生成 3-5 个问答对。

要求：
1. 问题应该是用户真实会问的形式
2. 标注出回答该问题需要的文档编号
3. 问题应覆盖不同难度

文档：
{document}

以 JSON 格式输出：
[
  {{"question": "...", "relevant_doc_ids": [0, 2], "difficulty": "easy"}}
]
"""

qa_gen_chain = qa_gen_prompt | llm | JsonOutputParser()

# 为每个文档生成 QA 对
def build_eval_dataset(documents):
    eval_data = []
    for i, doc in enumerate(documents):
        qa_pairs = qa_gen_chain.invoke({"document": doc})
        for qa in qa_pairs:
            eval_data.append({
                "question": qa["question"],
                "relevant_doc_ids": [i],
                "difficulty": qa.get("difficulty", "medium"),
            })
    return eval_data

[
  {
    "question": "如何提升 RAG 的召回率？",
    "relevant_doc_ids": ["doc_chunk_3", "doc_chunk_7", "doc_chunk_12"],
    "difficulty": "hard"
  },
  {
    "question": "BM25 算法是什么？",
    "relevant_doc_ids": ["doc_chunk_15"],
    "difficulty": "easy"
  }
]

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

# 准备评估数据
data_samples = {
    "question": [
        "如何提升 RAG 召回率？",
        "BM25 和向量检索的区别是什么？",
        "什么是 RRF 算法？",
    ],
    "contexts": [
        ["本文介绍了7种提升RAG召回率的方案，包括优化切分、混合检索、重排序等。"],
        ["BM25是关键词匹配算法，向量检索是语义匹配算法。两者互补。"],
        ["RRF（Reciprocal Rank Fusion）是一种结果融合算法，公式为 Σ 1/(k+rank)。"],
    ],
    "answer": [
        "可以通过优化文档切分、选择更好的 embedding 模型、查询重写、混合检索、元数据过滤、重排序等方案提升召回率。",
        "BM25 基于关键词频率和逆文档频率做匹配；向量检索将文本映射为向量后用余弦相似度计算。前者精确匹配好，后者语义理解强。",
        "RRF 是 Reciprocal Rank Fusion 的缩写，用于融合多个检索结果。核心公式是每个检索器中 1/(k+排名) 的累加和。",
    ],
    "ground_truth": [
        "主要方案有：优化文档切分策略、embedding 模型选型、查询重写与增强、混合检索（BM25+向量）、元数据过滤、重排序。",
        "BM25 是传统关键词匹配算法，擅长精确匹配；向量检索将文本转为向量后计算语义相似度，擅长语义理解。两者结合效果最佳。",
        "RRF 即 Reciprocal Rank Fusion，通过累加各检索器中 1/(60+排名) 来融合结果，无需分数归一化。",
    ],
}

dataset = Dataset.from_dict(data_samples)

# 运行评估
llm = ChatOpenAI(model="gpt-4o")
embeddings = OpenAIEmbeddings()

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

print(result)
# 输出:
# context_recall:      0.85   ← 上下文是否包含了正确答案
# context_precision:   0.80   ← 上下文中有多少是相关的
# faithfulness:        0.90   ← 答案是否忠实于上下文
# answer_relevancy:    0.88   ← 答案是否切题

class RagMetricsTracker:
    """RAG 管线指标追踪"""

    def __init__(self):
        self.metrics = {
            "total_queries": 0,
            "latency_p50": [],
            "latency_p99": [],
            "user_feedback_thumbs_up": 0,
            "user_feedback_thumbs_down": 0,
            "retrieval_recall_at_5": [],
        }

    def log_query(self, query, response, latency_ms, user_feedback=None):
        self.metrics["total_queries"] += 1
        self.metrics["latency_p50"].append(latency_ms)

        if user_feedback == "up":
            self.metrics["user_feedback_thumbs_up"] += 1
        elif user_feedback == "down":
            self.metrics["user_feedback_thumbs_down"] += 1

    def get_satisfaction_rate(self):
        total = self.metrics["user_feedback_thumbs_up"] + self.metrics["user_feedback_thumbs_down"]
        if total == 0:
            return None
        return self.metrics["user_feedback_thumbs_up"] / total

    def get_p99_latency(self):
        if not self.metrics["latency_p50"]:
            return None
        sorted_latencies = sorted(self.metrics["latency_p50"])
        idx = int(len(sorted_latencies) * 0.99)
        return sorted_latencies[min(idx, len(sorted_latencies) - 1)]

import json
from typing import List, Dict

class RAGEvaluator:
    """完整的 RAG 评估管线"""

    def __init__(self, retriever, llm_evaluator=None):
        self.retriever = retriever
        self.llm_evaluator = llm_evaluator or ChatOpenAI(model="gpt-4o-mini")

    def evaluate(self, eval_dataset: List[Dict], k: int = 5) -> Dict:
        """
        评估 RAG 管线
        :param eval_dataset: [{"question": "...", "relevant_doc_ids": [...]}]
        :param k: Recall/Precision@K 的 K 值
        """
        results = {
            "recall_at_k": [],
            "precision_at_k": [],
            "mrr": [],
            "latency_ms": [],
        }

        for sample in eval_dataset:
            query = sample["question"]
            relevant_ids = set(sample["relevant_doc_ids"])

            # 检索
            import time
            start = time.time()
            retrieved = self.retriever.retrieve(query, top_k=k)
            latency = (time.time() - start) * 1000

            # 提取检索到的 doc_ids
            retrieved_ids = set(doc.metadata.get("doc_id", "") for doc in retrieved)

            # 计算指标
            hits = len(retrieved_ids & relevant_ids)
            recall = hits / len(relevant_ids) if relevant_ids else 0
            precision = hits / k if k > 0 else 0

            # MRR
            mrr_score = 0
            for i, doc in enumerate(retrieved):
                if doc.metadata.get("doc_id") in relevant_ids:
                    mrr_score = 1.0 / (i + 1)
                    break

            results["recall_at_k"].append(recall)
            results["precision_at_k"].append(precision)
            results["mrr"].append(mrr_score)
            results["latency_ms"].append(latency)

        # 汇总
        summary = {
            "num_queries": len(eval_dataset),
            "mean_recall_at_k": np.mean(results["recall_at_k"]),
            "mean_precision_at_k": np.mean(results["precision_at_k"]),
            "mean_mrr": np.mean(results["mrr"]),
            "p50_latency_ms": np.percentile(results["latency_ms"], 50),
            "p99_latency_ms": np.percentile(results["latency_ms"], 99),
        }

        return summary

# === 使用 ===
evaluator = RAGEvaluator(retriever=TwoStageRetriever(...))
summary = evaluator.evaluate(eval_dataset, k=5)

print(json.dumps(summary, indent=2))
# {
#   "num_queries": 100,
#   "mean_recall_at_5": 0.72,
#   "mean_precision_at_5": 0.64,
#   "mean_mrr": 0.58,
#   "p50_latency_ms": 180,
#   "p99_latency_ms": 450
# }

"""
高召回率 RAG 管线 - 完整实现
整合方案：智能切分 + Embedding 优化 + 查询重写 + 混合检索 + 元数据过滤 + 重排序
"""

import os
import json
import re
import time
import jieba
from typing import List, Dict, Optional
from dataclasses import dataclass, field

# === LangChain 组件 ===
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.documents import Document
from langchain.retrievers import ParentDocumentRetriever, SelfQueryRetriever

# === 向量存储 ===
from langchain_chroma import Chroma

# === BM25 ===
from rank_bm25 import BM25Okapi

# === Reranker ===
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# === 评估 ===
from ragas.metrics import context_recall, context_precision
from ragas import evaluate


# ============================================
# 组件1：文档处理
# ============================================

@dataclass
class DocumentProcessor:
    """文档处理管线：解析 → 切分 → 元数据提取"""

    chunk_size: int = 1000
    chunk_overlap: int = 200
    parent_chunk_size: int = 2000
    enable_parent_child: bool = True

    def process(self, raw_text: str, source: str = "") -> List[Document]:
        """处理原始文本为带元数据的 Document 列表"""

        # 1. 提取元数据
        metadata = self._extract_metadata(raw_text, source)

        # 2. 切分
        if self.enable_parent_child:
            docs = self._parent_child_split(raw_text, metadata)
        else:
            docs = self._recursive_split(raw_text, metadata)

        return docs

    def _extract_metadata(self, text: str, source: str) -> Dict:
        """从文档提取元数据"""
        metadata = {"source": source}

        # 提取标题
        title_match = re.search(r'^#\s+(.+)$', text, re.MULTILINE)
        if title_match:
            metadata["title"] = title_match.group(1).strip()

        # 提取日期
        date_match = re.search(r'(\d{4}-\d{2}-\d{2})', text)
        if date_match:
            metadata["date"] = date_match.group(1)

        # 提取标签
        tags_match = re.findall(r'tags?:\s*\[?(.+?)\]?', text)
        if tags_match:
            metadata["tags"] = [t.strip() for t in tags_match[0].split(",")]

        # 提取层级标题
        metadata["headings"] = re.findall(r'^#{1,3}\s+(.+)$', text, re.MULTILINE)

        return metadata

    def _parent_child_split(self, text: str, metadata: Dict) -> List[Document]:
        """Parent-Child 切分"""
        parent_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.parent_chunk_size,
            chunk_overlap=self.chunk_overlap,
            separators=["\n\n", "\n", "。", "！", "？", "；", "，", " "],
        )
        child_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap // 4,
            separators=["\n\n", "\n", "。", "！", "？", "；", "，", " "],
        )

        parent_docs = parent_splitter.create_documents([text], metadatas=[metadata])
        child_docs = child_splitter.create_documents([text], metadatas=[metadata])

        # 关联父子文档（通过 source + 位置信息）
        for i, child in enumerate(child_docs):
            child.metadata["parent_idx"] = i // 3  # 简化关联

        # 实际检索时返回 parent，索引时用 child
        return child_docs  # 索引用 child

    def _recursive_split(self, text: str, metadata: Dict) -> List[Document]:
        """递归切分"""
        splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.chunk_size,
            chunk_overlap=self.chunk_overlap,
            separators=["\n\n", "\n", "。", "！", "？", "；", "，", " "],
        )
        return splitter.create_documents([text], metadatas=[metadata])


# ============================================
# 组件2：混合检索器
# ============================================

class HybridRetriever:
    """BM25 + 向量混合检索 + RRF 融合"""

    def __init__(self, vectorstore, corpus: List[str], top_k: int = 30):
        self.vectorstore = vectorstore
        self.top_k = top_k
        self.corpus = corpus

        # 构建 BM25 索引
        tokenized_corpus = [list(jieba.cut(doc)) for doc in corpus]
        self.bm25 = BM25Okapi(tokenized_corpus)

    def retrieve(self, query: str) -> List[Document]:
        # 向量检索
        vector_docs = self.vectorstore.similarity_search_with_score(query, k=self.top_k)

        # BM25 检索
        tokenized_query = list(jieba.cut(query))
        bm25_scores = self.bm25.get_scores(tokenized_query)

        # RRF 融合
        return self._rrf_fusion(vector_docs, bm25_scores)

    def _rrf_fusion(self, vector_docs, bm25_scores, k_rrf: int = 60) -> List[Document]:
        rrf_scores = {}

        # 向量结果
        for rank, (doc, vscore) in enumerate(vector_docs):
            content = doc.page_content
            rrf_scores[content] = rrf_scores.get(content, 0) + 1 / (k_rrf + rank + 1)

        # BM25 结果
        bm25_ranked = sorted(enumerate(bm25_scores), key=lambda x: -x[1])
        for rank, (idx, bscore) in enumerate(bm25_ranked[:self.top_k]):
            if bscore > 0:
                content = self.corpus[idx]
                rrf_scores[content] = rrf_scores.get(content, 0) + 1 / (k_rrf + rank + 1)

        sorted_results = sorted(rrf_scores.items(), key=lambda x: -x[1])

        return [Document(page_content=content, metadata={"rrf_score": score})
                for content, score in sorted_results[:self.top_k]]


# ============================================
# 组件3：重排序器
# ============================================

class CrossEncoderReranker:
    """Cross-Encoder 重排序器"""

    def __init__(self, model_name: str = "BAAI/bge-reranker-v2-m3",
                 device: str = None):
        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.model.to(self.device)
        self.model.eval()

    def rerank(self, query: str, documents: List[Document],
               top_k: int = 5) -> List[Document]:
        if not documents:
            return []

        pairs = [[query, doc.page_content] for doc in documents]

        with torch.no_grad():
            inputs = self.tokenizer(
                pairs, padding=True, truncation=True,
                max_length=512, return_tensors="pt"
            ).to(self.device)

            scores = self.model(**inputs, return_dict=True).logits.view(-1).float()

        scored_docs = sorted(
            zip(documents, scores.cpu().tolist()),
            key=lambda x: x[1],
            reverse=True
        )

        return [Document(page_content=doc.page_content,
                         metadata={**doc.metadata, "rerank_score": float(score)})
                for doc, score in scored_docs[:top_k]]


# ============================================
# 组件4：查询重写器
# ============================================

class QueryRewriter:
    """LLM 驱动的查询重写"""

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

    def multi_query(self, query: str, num_variants: int = 3) -> List[str]:
        """生成多个查询变体"""
        prompt = f"""Generate {num_variants} different versions of this search query.
Each version should use different wording but preserve the original intent.

Original query: {query}

Query versions (one per line, no numbering):"""

        response = self.llm.invoke(prompt)
        variants = [line.strip() for line in response.content.split("\n")
                    if line.strip()]

        return [query] + variants[:num_variants]

    def hyde(self, query: str) -> str:
        """生成假设性文档"""
        prompt = f"Generate a hypothetical document that would answer this query:\n{query}"
        response = self.llm.invoke(prompt)
        return response.content.strip()


# ============================================
# 组件5：端到端 RAG 管线
# ============================================

@dataclass
class HighRecallRAGPipeline:
    """高召回率 RAG 管线"""

    vectorstore: Chroma
    corpus: List[str]
    llm: ChatOpenAI
    reranker: Optional[CrossEncoderReranker] = None
    enable_multi_query: bool = True
    enable_hyde: bool = False
    recall_k: int = 30
    final_k: int = 5

    def __post_init__(self):
        self.hybrid_retriever = HybridRetriever(
            self.vectorstore, self.corpus, top_k=self.recall_k
        )
        self.query_rewriter = QueryRewriter(self.llm)

    def retrieve(self, query: str) -> List[Document]:
        """完整检索流程"""
        all_candidates = []

        # 策略1：原始 query 直接检索
        docs = self.hybrid_retriever.retrieve(query)
        all_candidates.extend(docs)

        # 策略2：Multi-Query
        if self.enable_multi_query:
            variants = self.query_rewriter.multi_query(query)
            for variant in variants[1:]:  # 跳过原始 query
                docs = self.hybrid_retriever.retrieve(variant)
                all_candidates.extend(docs)

        # 策略3：HyDE
        if self.enable_hyde:
            hyde_doc = self.query_rewriter.hyde(query)
            docs = self.vectorstore.similarity_search(hyde_doc, k=self.recall_k)
            all_candidates.extend(docs)

        # 去重
        unique = self._deduplicate(all_candidates)

        # 重排序
        if self.reranker and len(unique) > self.final_k:
            unique = self.reranker.rerank(query, unique, top_k=self.final_k)

        return unique[:self.final_k]

    def _deduplicate(self, docs: List[Document]) -> List[Document]:
        seen = set()
        unique = []
        for doc in docs:
            if doc.page_content not in seen:
                seen.add(doc.page_content)
                unique.append(doc)
        return unique

    def generate(self, query: str) -> Dict:
        """检索 + 生成"""
        start_time = time.time()

        # 检索
        docs = self.retrieve(query)
        retrieval_time = time.time() - start_time

        # 构建上下文
        context = "\n\n".join([
            f"[文档{i+1}]\n{doc.page_content}"
            for i, doc in enumerate(docs)
        ])

        # 生成
        prompt = f"""基于以下参考文档回答问题。如果参考文档中没有相关信息，请直接说明。

参考文档：
{context}

问题：{query}

回答："""

        response = self.llm.invoke(prompt)
        generation_time = time.time() - start_time - retrieval_time

        return {
            "query": query,
            "answer": response.content,
            "sources": [
                {
                    "content": doc.page_content[:200] + "...",
                    "rerank_score": doc.metadata.get("rerank_score", None),
                    "rrf_score": doc.metadata.get("rrf_score", None),
                }
                for doc in docs
            ],
            "retrieval_time_ms": round(retrieval_time * 1000, 1),
            "generation_time_ms": round(generation_time * 1000, 1),
            "total_time_ms": round((time.time() - start_time) * 1000, 1),
        }


# ============================================
# 使用示例
# ============================================

def main():
    """完整使用示例"""

    # 1. 初始化组件
    llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)
    embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
    vectorstore = Chroma(
        collection_name="tech-docs",
        embedding_function=embeddings,
        persist_directory="./chroma_db"
    )

    # 语料库（实际应从文档处理管线获取）
    corpus = [
        "RAG 是一种结合检索和生成的架构，召回率是关键指标。",
        "BM25 是经典的关键词匹配算法，与向量检索互补。",
        "RRF（Reciprocal Rank Fusion）是结果融合算法，公式为 Σ 1/(k+rank)。",
        "Cross-Encoder 重排序可以大幅提升检索精度。",
        "BGE-M3 是当前中文场景最强的开源 embedding 模型。",
    ]

    reranker = CrossEncoderReranker()

    # 2. 创建管线
    pipeline = HighRecallRAGPipeline(
        vectorstore=vectorstore,
        corpus=corpus,
        llm=llm,
        reranker=reranker,
        enable_multi_query=True,
        enable_hyde=False,
        recall_k=30,
        final_k=3,
    )

    # 3. 检索 + 生成
    query = "如何提升 RAG 的召回率？"
    result = pipeline.generate(query)

    print(f"问题: {result['query']}")
    print(f"回答: {result['answer']}")
    print(f"耗时: {result['total_time_ms']}ms (检索 {result['retrieval_time_ms']}ms)")
    print(f"来源数: {len(result['sources'])}")


if __name__ == "__main__":
    main()

# 概念性示例（非实际可运行代码）
# 训练数据: (query, ground_truth_context, ground_truth_answer)
# 损失函数: LLM生成loss + 检索相关性loss
# 目标: 让模型内化检索策略