# 轻量级RAG + SKILL 深度融合 完整可运行示例
# 环境安装：pip install langchain chromadb python-dotenv modelscope sentence-transformers matplotlib
import re
import os
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from dotenv import load_dotenv
from modelscope import snapshot_download
from sentence_transformers import SentenceTransformer
from langchain_text_splitters import CharacterTextSplitter
# 加载环境变量
load_dotenv()
# ====================== 本地嵌入模型配置 ======================
print("步骤1: 正在下载/校验嵌入模型缓存...")
embedding_model_dir = snapshot_download(
    model_id=os.getenv("MODEL_ID"),
    cache_dir=os.getenv("CACHE_DIR"),
    revision="master"
)
print(f"嵌入模型路径: {embedding_model_dir}")
print("\n步骤2: 加载嵌入模型...")
embedding_model = SentenceTransformer(embedding_model_dir)
print("✓ 嵌入模型加载成功")
# ====================== 模块1：SKILL 解析器 ======================
class SkillParser:
    """解析SKILL.md文件，提取触发规则、逻辑、边界"""
    def __init__(self, skill_path):
        self.skill_path = skill_path
        self.content = self._read_file()
        
    def _read_file(self):
        with open(self.skill_path, 'r', encoding='utf-8') as f:
            return f.read()
    
    def get_trigger_keywords(self):
        match = re.search(r'关键词：(.*)', self.content)
        return match.group(1).split('、') if match else []
    
    def get_skill_logic(self):
        match = re.search(r'## 3. 核心逻辑\n(.*?)\n##', self.content, re.DOTALL)
        return match.group(1).strip() if match else ""
    
    def get_knowledge_boundary(self):
        match = re.search(r'## 4. 知识边界\n(.*?)\n##', self.content, re.DOTALL)
        return match.group(1).strip() if match else ""
# ====================== 模块2：轻量级RAG微型知识库 ======================
class MiniKnowledgeBase:
    """单SKILL专属向量库，本地零成本部署"""
    def __init__(self, skill_name, knowledge_path):
        self.skill_name = skill_name
        self.knowledge_path = knowledge_path
        self.chunks = self._split_text()
        self.chunk_embeddings = self._build_embeddings()
    def _read_knowledge(self):
        with open(self.knowledge_path, 'r', encoding='utf-8') as f:
            return f.read()
    def _split_text(self):
        # 按行分割，每行作为一个独立的chunk
        text = self._read_knowledge()
        lines = [line.strip() for line in text.split('\n') if line.strip()]
        return lines
    def _build_embeddings(self):
        global embedding_model
        return embedding_model.encode(self.chunks, convert_to_numpy=True)
    def search(self, query, top_k=3):
        global embedding_model
        query_embedding = embedding_model.encode([query], convert_to_numpy=True)
        similarities = np.dot(self.chunk_embeddings, query_embedding.T).flatten()
        top_indices = similarities.argsort()[-top_k:][::-1]
        # 显示相似度分数
        print(f"📊 检索结果相似度：")
        for i, idx in enumerate(top_indices):
            print(f"   [{i+1}] {similarities[idx]:.4f} - {self.chunks[idx][:30]}...")
        class SimpleDoc:
            def __init__(self, page_content, similarity):
                self.page_content = page_content
                self.similarity = similarity
        return [SimpleDoc(self.chunks[idx], similarities[idx]) for idx in top_indices]
# ====================== 模块3：SKILL 调度中心 ======================
class SkillDispatcher:
    """自动匹配用户问题对应的SKILL"""
    def __init__(self, skill_list):
        self.skills = skill_list
        # 为每个技能提取关键词并计算向量
        self.skill_keywords = []
        for skill in skill_list:
            parser = skill['parser']
            keywords = parser.get_trigger_keywords()
            # 技能名称 + 关键词作为匹配依据
            match_text = f"{skill['name']} {', '.join(keywords)}"
            self.skill_keywords.append({
                'skill': skill,
                'text': match_text,
                'embedding': embedding_model.encode(match_text, convert_to_numpy=True)
            })
    def match_skill(self, query):
        global embedding_model
        # 计算问题的向量
        query_embedding = embedding_model.encode(query, convert_to_numpy=True)
        # 计算与每个技能的相似度
        max_similarity = -1
        matched_skill = None
        for skill_info in self.skill_keywords:
            # 使用余弦相似度
            similarity = np.dot(query_embedding, skill_info['embedding']) / (
                np.linalg.norm(query_embedding) * np.linalg.norm(skill_info['embedding'])
            )
            if similarity > max_similarity:
                max_similarity = similarity
                matched_skill = skill_info['skill']
        print(f"📊 技能匹配相似度: {max_similarity:.4f}")
        return matched_skill
# ====================== 模块5：完整执行引擎 ======================
def run_agent(query, skill_list):
    # 1. 匹配SKILL
    dispatcher = SkillDispatcher(skill_list)
    matched_skill = dispatcher.match_skill(query)
    print(f"\n🎯 匹配技能：{matched_skill['name']}")
    
    # 2. 检索专属知识库
    kb = matched_skill['kb']
    docs = kb.search(query)
    print(f"✅ 精准检索法律知识完成")
    # 3. 读取SKILL规则
    parser = matched_skill['parser']
    boundary = parser.get_knowledge_boundary()
    logic = parser.get_skill_logic()
    # 4. 智能组织答案
    answer_parts = []
    # 添加最相关的检索结果
    if docs and docs[0].similarity > 0.3:  # 设置相似度阈值
        answer_parts.append(f"根据相关法律规定：\n{docs[0].page_content}")
    # 添加辅助信息
    if len(docs) > 1 and docs[1].similarity > 0.2:
        answer_parts.append(f"\n补充信息：\n{docs[1].page_content}")
    # 添加知识边界说明
    answer_parts.append(f"\n（服务范围：{boundary}）")
    return "\n".join(answer_parts)
# ====================== 主函数：一键运行 ======================
if __name__ == "__main__":
    print("="*60)
    print("🚀 轻量级RAG + SKILL 系统启动中...")
    print("="*60)
    
    # 初始化法律咨询SKILL（所有文件已内置）
    legal_parser = SkillParser("legal_skill.md")
    legal_kb = MiniKnowledgeBase("legal_skill", "legal_knowledge.txt")
    
    # 注册SKILL
    skills = [
        {
            "name": "法律咨询助手",
            "parser": legal_parser,
            "kb": legal_kb
        }
    ]
    
    # 测试问题1
    query1 = "民间借贷的合法利率上限是多少？"
    answer1 = run_agent(query1, skills)
    print("\n📝 用户问题：", query1)
    print("💡 系统回答：\n", answer1)
    
    print("\n" + "-"*60)
    
    # 测试问题2
    query2 = "试用期不签劳动合同违法吗？"
    answer2 = run_agent(query2, skills)
    print("\n📝 用户问题：", query2)
    print("💡 系统回答：\n", answer2)
    
    print("\n" + "="*60)
    print("✅ 系统运行完成！所有功能正常执行")
    print("="*60)