import re
from typing import List
# 医疗文本层级化语义Chunking工具类
# 针对医疗文本强结构化、高关联特性，设计"结构化拆分 + 语义补齐 + 边界保护"专属分段策略：
#   ① 结构化拆分：按医学文档层级标题、条款标识切分基础单元，保留原始逻辑结构
#   ② 语义补齐  ：过短且句子不完整的片段不独立成块，挂起并与后续内容按句补齐，避免残缺语义
#   ③ 边界保护  ：超长片段按句子边界回退软切分，不硬切句子、不拆散医学专业术语
class MedicalChunking:
    def __init__(self, max_chunk_size: int = 800, min_chunk_size: int = 50):
        self.max_chunk_size = max_chunk_size  # 医疗片段最大长度
        self.min_chunk_size = min_chunk_size  # 医疗片段最小长度
        # 匹配医学文档层级标题（如 "1." "3.1"）、条款（如 "①②③"）
        self.title_pattern = re.compile(r"^(\d+\.){0,3}\s*.+")
        self.clause_pattern = re.compile(r"^[①②③④⑤⑥⑦⑧⑨⑩]\s*.+")
        # 医学句子结束符（句号、分号、感叹号、问号）
        self.sentence_end_pattern = re.compile(r"[。；;！？]$")
    # ① 结构化拆分：按医学文档层级标识拆分基础单元
    def split_by_structure(self, text: str) -> List[str]:
        lines = text.split("\n")
        chunks = []
        temp_chunk = []
        for line in lines:
            line = line.strip()
            if not line:
                continue
            # 识别新标题/新条款，切割新片段
            if self.title_pattern.match(line) or self.clause_pattern.match(line):
                if temp_chunk:
                    chunks.append(" ".join(temp_chunk))
                    temp_chunk = []
                temp_chunk.append(line)
            else:
                temp_chunk.append(line)
        # 加入最后一个片段
        if temp_chunk:
            chunks.append(" ".join(temp_chunk))
        return chunks
    # 按句子边界切分（边界保护的基础能力）
    def _split_sentences(self, chunk: str) -> List[str]:
        sentences = re.split(r"(?<=[。；;！？])\s*", chunk.strip())
        return [s for s in sentences if s]
    # ② 语义补齐：过短且句子不完整的片段挂起，与后续内容合并补齐
    def semantic_completion(self, chunks: List[str]) -> List[str]:
        completed = []
        pending = ""  # 待补齐的短片段（语义挂起区）
        for chunk in chunks:
            if pending:
                chunk = pending + chunk  # 补齐合并，恢复完整语义
                pending = ""
            # 片段过短且未以句子结束符收尾 → 语义不完整，挂起等待补齐
            if len(chunk) < self.min_chunk_size and not self.sentence_end_pattern.search(chunk):
                pending = chunk
            else:
                completed.append(chunk)
        if pending:
            completed.append(pending)
        return completed
    # ③ 边界保护：超长片段按句子边界软切分，不硬切句子、不拆散术语
    def boundary_protect(self, chunks: List[str]) -> List[str]:
        protected = []
        for chunk in chunks:
            if len(chunk) <= self.max_chunk_size:
                protected.append(chunk)
                continue
            # 超长片段：退回到句子边界再切分，避免语义割裂
            sentences = self._split_sentences(chunk)
            buffer = ""
            for sent in sentences:
                if len(buffer) + len(sent) <= self.max_chunk_size:
                    buffer += sent
                else:
                    if buffer:
                        protected.append(buffer.strip())
                    buffer = sent
            if buffer:
                protected.append(buffer.strip())
        return protected
    # 完整分段流程：结构化拆分 → 语义补齐 → 边界保护
    def chunk(self, medical_text: str) -> List[str]:
        struct_chunks = self.split_by_structure(medical_text)
        completed_chunks = self.semantic_completion(struct_chunks)
        final_chunks = self.boundary_protect(completed_chunks)
        return final_chunks
# 测试医疗文本分段（高血压诊疗指南片段）
if __name__ == "__main__":
    test_text = """1. 高血压病因
    高血压的病因分为原发性与继发性两类。原发性高血压无明确病因，与遗传、饮食、作息、情绪密切相关。
    2. 高血压临床表现
    早期高血压无明显症状，部分患者会出现头晕、头痛、乏力等症状，长期高血压可引发心脑血管并发症。
    3. 基础治疗方案
    3.1 生活干预：低盐饮食、规律运动、戒烟限酒、控制体重。
    3.2 药物治疗：常用降压药包括利尿剂、钙通道阻滞剂等，需遵医嘱服用，禁止擅自停药。"""
    chunker = MedicalChunking()
    print("① 结构化拆分：按层级标题切分基础单元")
    struct_chunks = chunker.split_by_structure(test_text)
    for i, c in enumerate(struct_chunks):
        print(f"   基础单元{i+1}（{len(c)}字）：{c[:45]}...")
    print("\n② 语义补齐：过短片段挂起合并，保证语义完整")
    completed = chunker.semantic_completion(struct_chunks)
    for i, c in enumerate(completed):
        print(f"   补齐片段{i+1}（{len(c)}字）：{c[:45]}...")
    print("\n③ 边界保护 + 最终结果：超长片段按句回退切分")
    result = chunker.chunk(test_text)
    for idx, res in enumerate(result):
        print(f"   片段{idx+1}（{len(res)}字）：{res}")
        print(f"   片段长度：{len(res)}\n")

import numpy as np
from rank_bm25 import BM25Okapi
from sentence_transformers import SentenceTransformer
from modelscope import snapshot_download
from typing import List, Dict
# 医疗混合检索工具类
class MedicalHybridSearch:
    def __init__(self):
        cache_dir = "D:\\modelscope\\hub"
        # 加载Embedding模型用于语义匹配
        embedding_model_dir = snapshot_download(
            model_id="sentence-transformers/all-MiniLM-L6-v2",
            cache_dir=cache_dir,
            revision="master"
        )
        # 加载医疗适配向量模型
        self.vec_model = SentenceTransformer(embedding_model_dir)
        self.bm25 = None
        self.corpus = []
        self.tokenized_corpus = []
    def build_corpus(self, chunks: List[str]):
        # 构建检索语料库
        self.corpus = chunks
        # BM25分词处理（简易医疗分词）
        self.tokenized_corpus = [chunk.split() for chunk in chunks]
        self.bm25 = BM25Okapi(self.tokenized_corpus)
        # 预计算向量库
        self.corpus_embeds = self.vec_model.encode(chunks)
    def bm25_search(self, query: str, top_n: int = 10) -> List[Dict]:
        # BM25稀疏检索
        tokenized_query = query.split()
        scores = self.bm25.get_scores(tokenized_query)
        top_idx = np.argsort(scores)[::-1][:top_n]
        return [{"text": self.corpus[i], "score": scores[i]} for i in top_idx]
    def vector_search(self, query: str, top_n: int = 10) -> List[Dict]:
        # 稠密向量检索
        query_embed = self.vec_model.encode(query)
        scores = np.dot(self.corpus_embeds, query_embed)
        top_idx = np.argsort(scores)[::-1][:top_n]
        return [{"text": self.corpus[i], "score": scores[i]} for i in top_idx]
    def normalize_score(self, score_list: List[Dict]) -> List[Dict]:
        # 分数归一化处理
        if not score_list:
            return score_list
        scores = [x["score"] for x in score_list]
        min_s, max_s = min(scores), max(scores)
        if max_s == min_s:
            for item in score_list:
                item["norm_score"] = 1.0
        else:
            for item in score_list:
                item["norm_score"] = (item["score"] - min_s) / (max_s - min_s)
        return score_list
    def hybrid_search(self, query: str, top_n: int = 8) -> List[Dict]:
        # 双路检索
        bm25_res = self.normalize_score(self.bm25_search(query))
        vec_res = self.normalize_score(self.vector_search(query))
        
        # 结果合并去重
        all_res = {}
        for item in bm25_res:
            all_res[item["text"]] = 0.6 * item["norm_score"]
        for item in vec_res:
            if item["text"] in all_res:
                all_res[item["text"]] += 0.4 * item["norm_score"]
            else:
                all_res[item["text"]] = 0.4 * item["norm_score"]
        
        # 排序取TopN
        sorted_res = sorted(all_res.items(), key=lambda x: x[1], reverse=True)
        return [{"text": k, "final_score": v} for k, v in sorted_res[:top_n]]
# 测试混合检索（含执行过程说明）
if __name__ == "__main__":
    # 模拟医疗知识库片段
    med_chunks = [
        "高血压患者需低盐低脂饮食，每日食盐摄入量不超过5克，规律监测血压。",
        "高血压常用药物包括钙通道阻滞剂、利尿剂、ACEI类药物，需个体化用药。",
        "糖尿病患者需控制碳水摄入，规律运动，定期监测血糖，预防并发症。",
        "青霉素过敏患者禁止使用青霉素类药物，过敏反应需立即使用肾上腺素急救。"
    ]
    print("=" * 50)
    print("  医疗混合检索（BM25 + 向量双路）执行过程")
    print("=" * 50)
    # 步骤1：初始化检索器（加载Embedding模型）
    print("\n" + "-" * 50)
    print("[步骤1] 初始化检索器，加载Embedding模型（all-MiniLM-L6-v2）...")
    searcher = MedicalHybridSearch()
    print("  ✓ 模型加载完成")
    # 步骤2：构建语料库（BM25分词 + 向量预编码）
    print("\n" + "-" * 50)
    print("[步骤2] 构建检索语料库...")
    searcher.build_corpus(med_chunks)
    print(f"  ✓ 语料片段数：{len(searcher.corpus)}")
    print(f"  ✓ BM25分词索引构建完成 | 向量库编码完成")
    # 步骤3：发起医疗查询
    query = "高血压日常饮食和用药注意事项"
    print("\n" + "-" * 50)
    print(f"[步骤3] 发起医疗查询：{query}")
    # 步骤4：双路检索 + 归一化 + 加权融合 + 排序
    print("\n" + "-" * 50)
    print("[步骤4] 双路检索执行中...")
    print("  ├─ BM25稀疏检索（关键词匹配）→ 分数归一化，权重0.6")
    print("  ├─ 向量稠密检索（语义匹配）→ 分数归一化，权重0.4")
    print("  └─ 结果去重合并 → 按加权总分降序取TopN")
    result = searcher.hybrid_search(query)
    # 步骤5：输出检索结果
    print("\n" + "-" * 50)
    print("[步骤5] 检索结果输出：")
    for idx, res in enumerate(result):
        print(f"  Top{idx+1} 分数：{res['final_score']:.4f} 内容：{res['text']}")

from sentence_transformers import CrossEncoder
from typing import List, Dict
# 医疗专属重排工具类
class MedicalRerank:
    def __init__(self):
        # 加载医疗微调重排模型
        self.rerank_model = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
        # 医疗来源权威度权重映射
        self.authority_weight = {
            "诊疗指南": 1.0,
            "药典标准": 1.0,
            "三甲规范": 0.95,
            "核心期刊": 0.9,
            "医学教材": 0.85,
            "普通科普": 0.6,
            "未知来源": 0.5
        }
    def get_authority_score(self, text: str) -> float:
        # 识别文本来源权威分数
        for source, score in self.authority_weight.items():
            if source in text:
                return score
        return 0.5
    def rerank(self, query: str, search_res: List[Dict], top_k: int = 5) -> List[Dict]:
        # 构建重排输入对
        pair_list = [(query, item["text"]) for item in search_res]
        # 语义重排分数
        semantic_scores = self.rerank_model.predict(pair_list)
        
        # 综合打分：语义分数 + 权威度分数
        for idx, item in enumerate(search_res):
            auth_score = self.get_authority_score(item["text"])
            # 加权综合得分
            item["rerank_score"] = 0.7 * semantic_scores[idx] + 0.3 * auth_score
        
        # 排序筛选TopK
        sorted_res = sorted(search_res, key=lambda x: x["rerank_score"], reverse=True)
        return sorted_res[:top_k]
# 测试重排流程
if __name__ == "__main__":
    reranker = MedicalRerank()
    # 模拟混合检索召回结果
    test_res = [
        {"text": "【高血压诊疗指南】高血压患者每日食盐摄入≤5g，规律监测血压，适度有氧运动。", "final_score": 0.92},
        {"text": "高血压日常养生科普：少吃盐多运动，保持好心情有助于血压稳定。", "final_score": 0.90},
        {"text": "糖尿病患者需控制糖分摄入，避免高糖食物，预防血糖升高。", "final_score": 0.65}
    ]
    query = "高血压日常养护注意事项"
    final_res = reranker.rerank(query, test_res)
    for idx, res in enumerate(final_res):
        print(f"重排Top{idx+1} 综合分数：{res['rerank_score']:.4f}")
        print(f"内容：{res['text']}\n")

import uuid
from datetime import datetime
from typing import List, Dict
# 医疗RAG溯源管理工具类
class MedicalTraceSource:
    def __init__(self):
        # 存储知识库片段与元数据
        self.knowledge_base = {}
        # 溯源日志
        self.trace_log = []
    def add_knowledge(self, chunk_list: List[str], source_name: str, 
                      source_level: str, publish_time: str):
        # 知识库入库打标
        for chunk in chunk_list:
            chunk_id = str(uuid.uuid4())[:8]
            self.knowledge_base[chunk_id] = {
                "content": chunk,
                "source": source_name,
                "level": source_level,
                "publish_time": publish_time
            }
    def match_trace(self, answer: str, rerank_res: List[Dict]) -> List[str]:
        # 匹配回答对应的溯源片段
        trace_ids = []
        for chunk_id, meta in self.knowledge_base.items():
            chunk_content = meta["content"]
            # 简单文本匹配（生产环境可改用相似度匹配提升精度）
            if chunk_content in answer and chunk_id not in trace_ids:
                trace_ids.append(chunk_id)
        return trace_ids
    def generate_trace_report(self, trace_ids: List[str]) -> str:
        # 生成标准化溯源清单
        if not trace_ids:
            return "暂无有效溯源依据"
        report = "\n【权威溯源依据】\n"
        for idx, tid in enumerate(trace_ids):
            meta = self.knowledge_base[tid]
            report += f"{idx+1}. 来源：{meta['source']}（{meta['level']}）\n"
            report += f"   发布时间：{meta['publish_time']}\n"
            report += f"   核心依据：{meta['content'][:100]}...\n"
        return report
    def save_trace_log(self, query: str, answer: str, trace_ids: List[str]):
        # 留存溯源日志
        log = {
            "time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
            "query": query,
            "answer": answer,
            "trace_ids": trace_ids
        }
        self.trace_log.append(log)
# 测试溯源体系（含执行过程说明）
if __name__ == "__main__":
    print("=" * 50)
    print("  医疗RAG引用溯源体系 — 执行过程")
    print("=" * 50)
    tracer = MedicalTraceSource()
    # 步骤1：权威医学知识库入库打标
    print("\n" + "-" * 50)
    print("[步骤1] 知识库入库打标...")
    chunk_data = ["高血压患者每日食盐摄入量不超过5克，需规律监测血压，适度进行有氧运动。"]
    tracer.add_knowledge(
        chunk_list=chunk_data,
        source_name="《中国高血压防治指南2024版》",
        source_level="国家级权威指南",
        publish_time="2024-03-15"
    )
    print("  ✓ 已入库片段，自动生成唯一ID（uuid前8位）")
    print("  ✓ 元数据打标：来源 + 权威级别 + 发布时间")
    # 步骤2：模拟模型回答
    print("\n" + "-" * 50)
    print("[步骤2] 模拟模型生成回答...")
    user_query = "高血压日常饮食注意什么"
    model_answer = "高血压患者在日常饮食中需严格控制盐分摄入，每日食盐摄入量不宜超过5克，同时配合规律运动、定期监测血压，维持血压稳定。"
    print(f"  ✓ 用户问题：{user_query}")
    print(f"  ✓ 模型回答：{model_answer[:40]}...")
    # 步骤3：溯源匹配（回答内容 ↔ 知识库片段）
    print("\n" + "-" * 50)
    print("[步骤3] 溯源匹配：回答内容与知识库片段比对...")
    trace_ids = tracer.match_trace(model_answer, [])
    print(f"  ✓ 匹配到溯源片段ID：{trace_ids if trace_ids else '无'}")
    # 步骤4：生成标准化溯源清单
    print("\n" + "-" * 50)
    print("[步骤4] 生成标准化溯源报告...")
    trace_report = tracer.generate_trace_report(trace_ids)
    # 步骤5：留存溯源日志
    tracer.save_trace_log(user_query, model_answer, trace_ids)
    print("  ✓ 溯源日志已留存（时间 + 问题 + 回答 + 溯源ID）")
    # 输出带溯源的完整回答
    print("\n" + "-" * 50)
    print("[最终输出] 带权威溯源的回答")
    print(model_answer)
    print(trace_report)