# RAG全链路极简入门Demo
import re
from typing import List
# 模拟文档解析、清洗
def parse_and_clean_doc(raw_text: str) -> str:
    # 剔除空行、特殊符号、冗余空格
    text = re.sub(r'\s+', ' ', raw_text)
    text = re.sub(r'页码|页眉|页脚|版权', '', text)
    return text.strip()

# 模拟固定长度分块
def chunk_text(text: str, chunk_size: int = 600, overlap: int = 80) -> List[str]:
    chunks = []
    start = 0
    while start < len(text):
        end = start + chunk_size
        chunk = text[start:end]
        chunks.append(chunk)
        start = end - overlap
    return chunks

# 模拟双路检索占位（后续章节完善）
if __name__ == "__main__":
    # 原始文档文本
    raw_doc = "这里是业务知识库原始文档内容，包含各类业务规则、参数说明、流程方案..."
    clean_text = parse_and_clean_doc(raw_doc)
    chunk_list = chunk_text(clean_text)
    print(f"文档分块完成，总块数：{len(chunk_list)}")

# 文档解析与标准化清洗实战代码
import re
from pathlib import Path

# 通用文本清洗函数
def clean_document_text(text: str) -> str:
    # 1. 剔除页眉页脚、水印、版权等无效关键词
    invalid_words = ["页眉", "页脚", "页码", "版权所有", "水印", "广告", "免责声明"]
    for word in invalid_words:
        text = text.replace(word, "")
    # 2. 修复文本错乱，统一换行和空格
    text = re.sub(r'\n+', '\n', text)
    text = re.sub(r'\s+', ' ', text)
    # 3. 剔除过短无意义文本
    text_lines = [line.strip() for line in text.split("\n") if len(line.strip()) > 5]
    return "\n".join(text_lines)

# TXT文档解析
def parse_txt(file_path: str) -> str:
    with open(file_path, "r", encoding="utf-8") as f:
        raw_text = f.read()
    return clean_document_text(raw_text)

# PDF文档解析（可复制PDF）
def parse_pdf(file_path: str) -> str:
    import pdfplumber
    raw_text = ""
    with pdfplumber.open(file_path) as pdf:
        for page in pdf.pages:
            page_text = page.extract_text() or ""
            raw_text += page_text + "\n"
    return clean_document_text(raw_text)

# 批量解析文档入口
if __name__ == "__main__":
    txt_content = parse_txt("test.txt")
    pdf_content = parse_pdf("test.pdf")
    print("TXT解析清洗后内容：\n", txt_content[:500])

# Chunk分块三大策略完整实战代码
import re
from typing import List
from sentence_transformers import SentenceTransformer

# 1. 固定长度重叠分块（基础版）
def fixed_chunk(text: str, chunk_size: int = 600, overlap: int = 80) -> List[str]:
    chunks = []
    start = 0
    text_len = len(text)
    while start < text_len:
        end = min(start + chunk_size, text_len)
        chunk = text[start:end]
        chunks.append(chunk.strip())
        start = end - overlap
    return chunks

# 2. 层级结构分块（进阶生产版）
def hierarchy_chunk(text: str) -> List[str]:
    # 按标题、段落层级拆分
    title_pattern = re.compile(r'第.+[章节]|[\d]+\. ')
    raw_blocks = title_pattern.split(text)
    chunks = []
    for block in raw_blocks:
        if len(block.strip()) > 100:
            # 超长段落二次拆分
            if len(block) > 800:
                chunks.extend(fixed_chunk(block))
            else:
                chunks.append(block.strip())
    return chunks

# 3. 简易语义分块（高阶精准版）
model = SentenceTransformer('all-MiniLM-L6-v2')
def semantic_chunk(text: str, threshold: float = 0.75) -> List[str]:
    sentences = [s.strip() for s in re.split(r'[。！？]', text) if s.strip()]
    chunks = []
    temp_chunk = [sentences[0]]
    for sent in sentences[1:]:
        # 计算句子语义相似度
        vec1 = model.encode("".join(temp_chunk))
        vec2 = model.encode(sent)
        sim = float(vec1 @ vec2.T)
        if sim > threshold:
            temp_chunk.append(sent)
        else:
            chunks.append("".join(temp_chunk))
            temp_chunk = [sent]
    chunks.append("".join(temp_chunk))
    return chunks

# 场景化调用示例
if __name__ == "__main__":
    doc_text = "此处为清洗后的完整文档文本..."
    # 通用场景
    print(fixed_chunk(doc_text))
    # 结构化文档场景
    print(hierarchy_chunk(doc_text))
    # 高精度问答场景
    print(semantic_chunk(doc_text))

# 向量检索 + BM25双路检索实战代码
import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer

# 初始化模型
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')

# 1. 向量检索实现
def vector_search(query: str, chunk_list: List[str], top_k: int = 10) -> List[tuple]:
    query_vec = embedding_model.encode(query)
    chunk_vecs = embedding_model.encode(chunk_list)
    # 余弦相似度计算
    sim_scores = np.dot(chunk_vecs, query_vec) / (np.linalg.norm(chunk_vecs, axis=1) * np.linalg.norm(query_vec))
    # 排序取TopK
    rank_results = sorted(zip(chunk_list, sim_scores), key=lambda x: x[1], reverse=True)
    return rank_results[:top_k]

# 2. BM25关键词检索实现
def bm25_search(query: str, chunk_list: List[str], top_k: int = 10) -> List[tuple]:
    # 文本分词（简易分词，中文场景可替换jieba分词）
    tokenized_chunks = [list(chunk) for chunk in chunk_list]
    bm25 = BM25Okapi(tokenized_chunks)
    tokenized_query = list(query)
    scores = bm25.get_scores(tokenized_query)
    rank_results = sorted(zip(chunk_list, scores), key=lambda x: x[1], reverse=True)
    return rank_results[:top_k]

# 双路检索调用
if __name__ == "__main__":
    chunks = ["业务故障码TS-999排查方案...", "系统通用故障解决流程...", "参数配置规范说明..."]
    user_query = "TS-999故障怎么解决"
    vec_res = vector_search(user_query, chunks)
    bm25_res = bm25_search(user_query, chunks)
    print("向量检索结果：", vec_res)
    print("BM25检索结果：", bm25_res)

# RRF倒数排序融合 实战代码
def rrf_fusion(vec_results: list, bm25_results: list, k: int = 60) -> list:
    """
    :param vec_results: 向量检索结果（有序列表）
    :param bm25_results: BM25检索结果（有序列表）
    :param k: RRF平滑系数，默认60
    :return: 融合后排序结果
    """
    score_dict = {}
    # 累加向量检索排名分数
    for idx, (chunk, _) in enumerate(vec_results):
        score_dict[chunk] = score_dict.get(chunk, 0) + 1 / (idx + 1 + k)
    # 累加BM25检索排名分数
    for idx, (chunk, _) in enumerate(bm25_results):
        score_dict[chunk] = score_dict.get(chunk, 0) + 1 / (idx + 1 + k)
    # 综合排序
    sorted_res = sorted(score_dict.items(), key=lambda x: x[1], reverse=True)
    return sorted_res

# 场景化权重微调封装
def scene_fusion(vec_res, bm25_res, scene_type="common"):
    if scene_type == "precision":
        # 精准场景：加重BM25权重
        fusion_res = rrf_fusion([], bm25_res) + rrf_fusion(vec_res, [])
    elif scene_type == "semantic":
        # 语义场景：加重向量权重
        fusion_res = rrf_fusion(vec_res, []) + rrf_fusion([], bm25_res)
    else:
        # 通用场景：均等融合
        fusion_res = rrf_fusion(vec_res, bm25_res)
    return fusion_res

# 调用示例
if __name__ == "__main__":
    # 承接上一章双路检索结果
    vec_res = [("语义匹配文本1", 0.89), ("普通文本2", 0.76)]
    bm25_res = [("精准关键词文本3", 25.6), ("普通文本2", 18.2)]
    final_res = scene_fusion(vec_res, bm25_res)
    print("RRF融合最终结果：", final_res)

# Reranker重排序实战代码
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# 加载轻量开源重排序模型
rerank_model_path = "cross-encoder/ms-marco-MiniLM-L-6-v2"
tokenizer = AutoTokenizer.from_pretrained(rerank_model_path)
rerank_model = AutoModelForSequenceClassification.from_pretrained(rerank_model_path)
rerank_model.eval()

def rerank(query: str, candidate_chunks: list, top_n: int = 8) -> list:
    """
    对融合后的候选文本进行精准重排序
    :param query: 用户提问
    :param candidate_chunks: 多路融合后的候选文本列表
    :param top_n: 重排序后保留最优条数
    :return: 排序后的优质文本
    """
    score_list = []
    with torch.no_grad():
        for chunk in candidate_chunks:
            # 双向交互编码，拼接query和文本
            inputs = tokenizer(query, chunk, return_tensors="pt", truncation=True, max_length=512)
            score = rerank_model(**inputs).logits[0][0].item()
            score_list.append((chunk, score))
    # 按相关性分数倒序排序
    sorted_res = sorted(score_list, key=lambda x: x[1], reverse=True)
    return [item[0] for item in sorted_res[:top_n]]

# 完整链路调用示例
if __name__ == "__main__":
    # 多路融合后的粗召回结果（20条）
    candidate_texts = ["候选文本1", "候选文本2", "候选文本3"]
    user_query = "TS-999故障详细排查步骤"
    # 重排序筛选Top8最优内容
    final_chunks = rerank(user_query, candidate_texts)
    print("重排序后最优检索内容：", final_chunks)

# 场景化RAG完整流水线整合代码
from typing import List

class SceneRAGPipeline:
    def __init__(self):
        self.model = SentenceTransformer('all-MiniLM-L6-v2')

    # 通用问答场景流水线
    def common_qa_pipeline(self, raw_text: str, query: str) -> List[str]:
        text = clean_document_text(raw_text)
        chunks = fixed_chunk(text, chunk_size=700, overlap=80)
        vec_res = vector_search(query, chunks, top_k=20)
        bm25_res = bm25_search(query, chunks, top_k=20)
        fuse_res = scene_fusion(vec_res, bm25_res, scene_type="common")
        return rerank(query, [i[0] for i in fuse_res])

    # 精准查询场景流水线
    def precision_qa_pipeline(self, raw_text: str, query: str) -> List[str]:
        text = clean_document_text(raw_text)
        chunks = fixed_chunk(text, chunk_size=400, overlap=50)
        vec_res = vector_search(query, chunks, top_k=20)
        bm25_res = bm25_search(query, chunks, top_k=20)
        fuse_res = scene_fusion(vec_res, bm25_res, scene_type="precision")
        return rerank(query, [i[0] for i in fuse_res])

    # 长逻辑推理场景流水线
    def logic_qa_pipeline(self, raw_text: str, query: str) -> List[str]:
        text = clean_document_text(raw_text)
        chunks = hierarchy_chunk(text)
        vec_res = vector_search(query, chunks, top_k=20)
        bm25_res = bm25_search(query, chunks, top_k=20)
        fuse_res = scene_fusion(vec_res, bm25_res, scene_type="common")
        return rerank(query, [i[0] for i in fuse_res])

# 场景化调用
if __name__ == "__main__":
    rag_pipe = SceneRAGPipeline()
    doc = "业务完整知识库文档内容..."
    # 根据业务场景切换流水线
    res = rag_pipe.common_qa_pipeline(doc, "产品使用常见问题")
    print("最终检索结果：", res)