# 最基础的固定分块（使用 LangChain）
from langchain.text_splitter import CharacterTextSplitter

splitter = CharacterTextSplitter(
    chunk_size=500,
    chunk_overlap=50,  # 块与块之间的重叠，防止关键信息被切到边界
    separator="\n"
)
chunks = splitter.split_text(document)

from langchain.text_splitter import RecursiveCharacterTextSplitter

splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    separators=["\n\n", "\n", "。", "，", " ", ""]  # 优先级从高到低
)

import spacy

nlp = spacy.load("zh_core_web_sm")  # 中文可用 spacy 或 jieba

def sentence_chunk(text, max_sentences=5):
    doc = nlp(text)
    sentences = [sent.text for sent in doc.sents]
    chunks = []
    for i in range(0, len(sentences), max_sentences):
        chunks.append(" ".join(sentences[i:i+max_sentences]))
    return chunks

# langchain 的 SemanticChunker
from langchain_experimental.text_splitter import SemanticChunker
from langchain_openai import OpenAIEmbeddings

splitter = SemanticChunker(
    embeddings=OpenAIEmbeddings(),
    breakpoint_threshold_type="percentile",  # 或 "standard_deviation"
    breakpoint_threshold_amount=95  # 低于第95百分位的相似度就切
)

# LlamaIndex 的 SentenceWindowNodeParser
from llama_index.core.node_parser import SentenceWindowNodeParser

parser = SentenceWindowNodeParser(
    window_size=3,  # 前后各取 3 句
    window_metadata_key="window"
)

import re

def clean_text(text: str) -> str:
    # 1. 移除多余的空白行
    text = re.sub(r'\n\s*\n\s*\n+', '\n\n', text)
    # 2. 替换特殊 Unicode 字符
    text = text.replace('​', '')   # 零宽空格
    text = text.replace('\xa0', ' ')    # 非断空格
    # 3. 统一全角/半角
    # （根据你的语料决定是否做）
    return text.strip()

# 使用 BGE-M3 给中文文档生成 Embedding
from sentence_transformers import SentenceTransformer

# BGE-M3 支持中英双语，且能输出 Dense + Sparse 两种向量
model = SentenceTransformer("BAAI/bge-m3", device="cuda")

# 示例文档块
chunks = [
    "RAG 是检索增强生成的缩写...",
    "向量数据库用于存储和检索高维向量...",
    "苹果公司推出了新款 iPhone..."
]

# 生成 Dense Embedding（常规用途）
dense_embeddings = model.encode(
    chunks,
    normalize_embeddings=True,  # L2 归一化，使得内积 = 余弦相似度
    show_progress_bar=True
)

# BGE-M3 还能同时输出 Sparse Embedding（用于关键词匹配）
sparse_embeddings = model.encode(
    chunks,
    return_dense=False,
    return_sparse=True  # 输出词袋向量
)
print(f"向量维度：{dense_embeddings.shape}")  # (3, 1024)

# 混合检索示例（使用 RRF: Reciprocal Rank Fusion）
from typing import List

def reciprocal_rank_fusion(
    dense_results: List[str],   # 向量检索排名 (rank 1, 2, 3...)
    sparse_results: List[str],  # BM25 检索排名
    k: int = 60  # RRF 平滑参数
) -> List[str]:
    scores = {}
    
    # 向量检索得分
    for rank, doc_id in enumerate(dense_results):
        scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)
    
    # BM25 检索得分
    for rank, doc_id in enumerate(sparse_results):
        scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)
    
    # 按融合分降序排列
    return sorted(scores.keys(), key=lambda x: scores[x], reverse=True)

# 使用 BGE-Reranker-v2 进行精排
from FlagEmbedding import FlagReranker

reranker = FlagReranker("BAAI/bge-reranker-v2-m3", use_fp16=True)

# 候选文档列表
query = "RAG 系统如何减少幻觉？"
candidates = [
    "RAG 通过引入外部知识库来约束 LLM 输出，从而减少幻觉现象",
    "GPT-4 在数学推理方面的准确率达到了 97%",
    "增加检索到的相关文档数量可以显著降低幻觉率",
    "RAG 系统的延迟主要取决于向量检索的速度"
]

# 逐条打分
pairs = [[query, doc] for doc in candidates]
scores = reranker.compute_score(pairs, normalize=True)

# 按分数排序
ranked = sorted(zip(candidates, scores), key=lambda x: x[1], reverse=True)
for doc, score in ranked:
    print(f"[{score:.3f}] {doc}")

# 示例：多轮对话中的查询重写
chat_history = [
    {"role": "user", "content": "介绍一下 Transformer"},
    {"role": "assistant", "content": "Transformer 是 Google 在 2017 年提出的..."},
    {"role": "user", "content": "它的注意力机制怎么工作？"}  # ← "它"指的是？
]

# 使用 LLM 将模糊问题 → 独立完整的问题
rewritten_query = llm.rewrite(
    user_query="它的注意力机制怎么工作？",
    chat_history=chat_history
)
# → "Transformer 模型中的自注意力机制是如何工作的？"

# 多跳检索的简化实现
def multi_hop_retrieval(query: str, retriever, llm, max_hops=3):
    current_query = query
    all_contexts = []
    
    for hop in range(max_hops):
        # 本轮检索
        docs = retriever.search(current_query, top_k=5)
        all_contexts.extend(docs)
        
        # 让 LLM 判断：还缺什么信息？
        response = llm.check_completeness(
            original_query=query,
            current_context=all_contexts
        )
        
        if response["is_complete"]:
            break  # 信息够了，停止检索
        
        current_query = response["next_query"]  # 下一轮要搜的问题
    
    return all_contexts

def allocate_context_budget(
    ranked_docs: list,  # 按相关性排好序的文档
    max_tokens: int = 4000,  # 最大上下文预算
    min_per_doc: int = 200   # 每篇文档最少保留
) -> list:
    """智能分配上下文预算：高分文档多给 tokens"""
    selected = []
    remaining = max_tokens
    
    # 按分数加权分配
    total_score = sum(doc["score"] for doc in ranked_docs)
    
    for doc in ranked_docs:
        if remaining < min_per_doc:
            break
        # 按相关性比例分配 tokens
        allocated = max(
            min_per_doc,
            int(doc["score"] / total_score * max_tokens)
        )
        allocated = min(allocated, remaining)  # 不能超预算
        
        selected.append({
            "content": doc["content"][:allocated * 4],  # 粗略估算
            "source": doc["source"],
            "score": doc["score"]
        })
        remaining -= allocated
    
    return selected

def format_with_citation(docs: list, answer: str) -> str:
    """给答案添加引用标记"""
    citation_map = {}
    
    # 在答案中找到与各文档匹配的句子，添加引用
    for i, doc in enumerate(docs):
        source_id = f"[{i+1}]"
        # 简化做法：追加引用列表
        citation_map[source_id] = doc["source"]
    
    # 在答案末尾添加引用列表
    citations = "\n\n---\n**参考资料：**\n"
    for i, doc in enumerate(docs):
        citations += f"[{i+1}] {doc['title']} — {doc['source']}\n"
    
    return answer + citations

# 简化版 GraphRAG 实现思路
import networkx as nx

# Step 1: 用 LLM 从文档中提取实体和关系
entities_relations = llm.extract_graph(text)  # [(实体1, 关系, 实体2), ...]
# e.g., [("Transformer", "提出于", "2017"), ("Transformer", "用于", "NLP"), ...]

# Step 2: 构建 NetworkX 图谱
G = nx.Graph()
for e1, rel, e2 in entities_relations:
    G.add_edge(e1, e2, relation=rel)

# Step 3: 社区检测
from networkx.algorithms import community
communities = community.louvain_communities(G)
# → [{节点集1}, {节点集2}, ...]

# Step 4: 对每个社区用 LLM 生成摘要
community_summaries = []
for nodes in communities:
    summary = llm.summarize(f"这个社区包含: {nodes}")
    community_summaries.append(summary)

# LightRAG 的简化核心概念
# 不使用 NetworkX 全局分析，而是把图谱三元组向量化
class LightRAG:
    def retrieve(self, query):
        # 1. 向量检索相关实体
        entities = self.entity_index.search(query, top_k=20)
        
        # 2. 沿图谱边扩展 1-2 跳
        extended = self.graph.expand(entities, hops=2)
        
        # 3. 检索局部子图的关系
        relations = self.get_relations(extended)
        
        # 4. 直接用 "实体+关系" 文本拼接 Prompt
        return self.format_context(extended, relations)

# Self-RAG 的简化决策流程
def self_rag_pipeline(query: str) -> str:
    # Step 1: 判断是否需要检索
    need_retrieval = reflection_llm.decide(
        f"这个问题需要检索外部知识吗？\n问题：{query}"
    )
    
    if need_retrieval == "不需要":
        return llm.generate(query)  # 直接回答
    
    # Step 2: 检索后逐段评估相关性
    docs = retriever.search(query)
    relevant_docs = []
    for doc in docs:
        is_relevant = reflection_llm.judge(
            f"这段内容与问题相关吗？\n问题：{query}\n内容：{doc}"
        )
        if is_relevant:
            relevant_docs.append(doc)
    
    # Step 3: 生成答案
    answer = llm.generate_with_context(query, relevant_docs)
    
    # Step 4: 自我评估
    supported = reflection_llm.verify(
        f"答案的每个声明是否都得到了参考资料的支持？\n答案：{answer}"
    )
    
    if not supported:
        answer = llm.regenerate_with_constraint(
            "重新生成，这次只能使用参考资料中的信息"
        )
    
    return answer

# Agentic RAG 的简化实现（使用 LangGraph）
from langgraph.graph import StateGraph, END
from langgraph.prebuilt import ToolNode

tools = [
    vector_search_tool,   # 向量检索工具
    web_search_tool,      # 网络搜索工具
    sql_query_tool,       # 数据库查询工具
    calculator_tool,      # 计算工具
    graph_search_tool,    # 知识图谱检索工具
]

# Agent 自主循环：执行 → 观察 → 思考 → 执行...
# 直到 Agent 认为"信息足够了"才生成最终答案

graph = StateGraph(AgentState)
graph.add_node("agent", agent_node)  # Agent 思考
graph.add_node("tools", ToolNode(tools))  # 工具执行
graph.add_edge("tools", "agent")  # 工具结果 → Agent 继续思考
graph.add_conditional_edges(
    "agent",
    should_continue,  # Agent 自主决定：继续搜 or 结束
    {"continue": "tools", "end": END}
)

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

# 准备评估数据
eval_data = Dataset.from_dict({
    "question": ["RAG 的工作原理是什么？", "什么是 GraphRAG？"],
    "answer": ["RAG 通过检索外部知识库来增强...", "GraphRAG 是微软提出的..."],
    "contexts": [
        ["RAG 是检索增强生成技术的缩写...", "LLM 存在幻觉问题..."],
        ["GraphRAG 使用知识图谱...", "Leiden 社区检测算法..."]
    ],
    "ground_truth": [
        "RAG 结合了检索和生成两个阶段，先从知识库检索相关内容，再交给 LLM 生成答案。",
        "GraphRAG 是微软在 2024 年提出的 RAG 改进方案，通过构建知识图谱实现全局理解。"
    ]
})

# 一键评估
result = evaluate(eval_data, metrics=[
    faithfulness,
    answer_relevancy,
    context_precision,
    context_recall
])
print(result)

"""
mini-rag: A production-ready RAG system in ~300 lines
依赖: pip install qdrant-client sentence-transformers llama-index-core
"""
from pathlib import Path
from typing import List, Optional, Dict, Any
import hashlib

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct, Filter
from sentence_transformers import CrossEncoder
from llama_index.core import Settings
from llama_index.core.node_parser import SentenceSplitter
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.readers.file import PyMuPDFReader

# ─── 配置 ─────────────────────────────────────────
EMBED_MODEL_NAME = "BAAI/bge-m3"
RERANKER_MODEL_NAME = "BAAI/bge-reranker-v2-m3"
COLLECTION_NAME = "knowledge_base"
VECTOR_SIZE = 1024
CHUNK_SIZE = 1000
CHUNK_OVERLAP = 200
RETRIEVAL_TOP_K = 20
FINAL_TOP_K = 5


class MiniRAG:
    """一个简洁但功能完整的 RAG 系统"""
    
    def __init__(self, qdrant_url: str = "http://localhost:6333"):
        # ─── 向量数据库 ───
        self.qdrant = QdrantClient(url=qdrant_url)
        
        # ─── Embedding 模型 ───
        self.embed_model = HuggingFaceEmbedding(
            model_name=EMBED_MODEL_NAME,
            device="cuda"
        )
        
        # ─── Reranker ───
        self.reranker = CrossEncoder(RERANKER_MODEL_NAME)
        
        # ─── 确保集合存在 ───
        self._ensure_collection()
    
    def _ensure_collection(self):
        """创建集合（如果不存在）"""
        collections = [c.name for c in self.qdrant.get_collections().collections]
        if COLLECTION_NAME not in collections:
            self.qdrant.create_collection(
                collection_name=COLLECTION_NAME,
                vectors_config=VectorParams(
                    size=VECTOR_SIZE,
                    distance=Distance.COSINE
                )
            )
            print("✨ 向量集合已创建")
    
    # ═══════════════════════════════════════════════
    #  数据摄入
    # ═══════════════════════════════════════════════
    
    def ingest_directory(self, dir_path: str) -> int:
        """批量导入目录中的所有 PDF 和 TXT 文件"""
        count = 0
        for ext in ["*.pdf", "*.txt", "*.md"]:
            for file_path in Path(dir_path).rglob(ext):
                count += self.ingest_file(str(file_path))
        return count
    
    def ingest_file(self, file_path: str) -> int:
        """导入单个文件并返回摄入的 chunk 数"""
        # 读取文件
        if file_path.endswith(".pdf"):
            reader = PyMuPDFReader()
            documents = reader.load(file_path)
            text = "\n\n".join([doc.text for doc in documents])
        else:
            with open(file_path, "r", encoding="utf-8") as f:
                text = f.read()
        
        # 分块
        splitter = SentenceSplitter(
            chunk_size=CHUNK_SIZE,
            chunk_overlap=CHUNK_OVERLAP
        )
        chunks = splitter.split_text(text)
        
        # 生成向量并存储
        points = []
        for i, chunk in enumerate(chunks):
            if len(chunk.strip()) < 50:  # 跳过太短的块
                continue
            
            embedding = self.embed_model.get_text_embedding(chunk)
            chunk_id = hashlib.md5(
                f"{file_path}_{i}".encode()
            ).hexdigest()
            
            points.append(PointStruct(
                id=chunk_id,
                vector=embedding,
                payload={
                    "text": chunk,
                    "source": file_path,
                    "chunk_index": i,
                    "char_count": len(chunk)
                }
            ))
        
        if points:
            self.qdrant.upsert(
                collection_name=COLLECTION_NAME,
                points=points
            )
        
        print(f"📥 已摄入: {file_path} ({len(points)} 块)")
        return len(points)
    
    # ═══════════════════════════════════════════════
    #  检索
    # ═══════════════════════════════════════════════
    
    def retrieve(self, query: str, top_k: int = RETRIEVAL_TOP_K) -> List[Dict]:
        """向量检索 + Reranker 精排"""
        # Step 1: 向量检索（粗召回）
        query_vec = self.embed_model.get_text_embedding(query)
        results = self.qdrant.search(
            collection_name=COLLECTION_NAME,
            query_vector=query_vec,
            limit=top_k
        )
        
        if not results:
            return []
        
        # Step 2: Reranker 精排
        pairs = [[query, hit.payload["text"]] for hit in results]
        rerank_scores = self.reranker.predict(pairs)
        
        # Step 3: 按 Reranker 分数排序
        ranked = sorted(
            zip(results, rerank_scores),
            key=lambda x: x[1],
            reverse=True
        )
        
        return [
            {
                "text": hit.payload["text"],
                "source": hit.payload["source"],
                "score": float(score)
            }
            for hit, score in ranked[:FINAL_TOP_K]
        ]
    
    # ═══════════════════════════════════════════════
    #  生成回答（调用 LLM）
    # ═══════════════════════════════════════════════
    
    def build_prompt(self, query: str, contexts: List[Dict]) -> str:
        """构建 RAG Prompt"""
        context_text = "\n\n".join([
            f"[{i+1}] (来源: {ctx['source']}, 相关度: {ctx['score']:.2f})\n{ctx['text']}"
            for i, ctx in enumerate(contexts)
        ])
        
        return f"""你是一个知识丰富、回答精准的 AI 助手。

## 规则
- 严格基于以下参考资料回答问题
- 资料不足时，明确说明"现有资料无法回答该问题"
- 引用资料时注明编号 [1]、[2] 等
- 中文回答，保持友好专业的语气

## 参考资料
{context_text}

## 用户问题
{query}

## 回答"""
    
    def query(self, llm, user_query: str) -> str:
        """端到端查询"""
        contexts = self.retrieve(user_query)
        
        if not contexts:
            return "抱歉，我在知识库中没有找到相关信息 (´;ω;`)" 
        
        prompt = self.build_prompt(user_query, contexts)
        answer = llm.generate(prompt)  # 你需要实现自己的 LLM 调用
        
        return answer


# ─── 使用示例 ─────────────────────────────────────
if __name__ == "__main__":
    # 初始化
    rag = MiniRAG(qdrant_url="http://localhost:6333")
    
    # 导入文档
    count = rag.ingest_directory("./my_documents/")
    print(f"总共摄入 {count} 个文档块 (๑•̀ㅂ•́)و✧")
    
    # 查询
    contexts = rag.retrieve("RAG 系统如何减少大模型的幻觉问题？")
    for i, ctx in enumerate(contexts):
        print(f"\n[{i+1}] 相关度: {ctx['score']:.3f}")
        print(f"    来源: {ctx['source']}")
        print(f"    内容: {ctx['text'][:200]}...")