import os
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS  # 使用FAISS替代Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_core.prompts import ChatPromptTemplate
from langchain.schema import Document
from typing import List
import requests
import json
import dashscope
from dashscope import Generation
from modelscope import snapshot_download
from modelscope.models import Model
from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
# 配置阿里千问模型
class QwenLLM:
    def __init__(self):
        # 设置阿里云API密钥
        self.api_key = os.environ["DASHSCOPE_API_KEY"]  # 从环境变量获取API密钥
        dashscope.api_key = self.api_key  # 设置dashscope的API密钥
    def invoke(self, prompt: str) -> str:
        try:
            # 调用阿里千问大模型API生成回答
            response = Generation.call(
                model="qwen-max",  # 使用千问最大模型
                prompt=prompt,  # 传入提示词
                temperature=0.1,  # 低温度值使输出更确定性
                top_p=0.9,  # 核采样参数，控制生成多样性
                result_format='message'  # 返回消息格式
            )
            # 检查响应状态
            if response.status_code == 200:
                return response.output.choices[0].message.content  # 返回模型生成的内容
            else:
                return f"请求失败: {response.message}"  # 返回错误信息
        except Exception as e:
            return f"请求异常: {str(e)}"  # 返回异常信息
# 从ModelScope下载并使用嵌入模型
class ModelScopeEmbeddings:
    def __init__(self, model_name="damo/nlp_corom_sentence-embedding_chinese-base"):
        # 从ModelScope下载中文句子嵌入模型
        self.model_dir = snapshot_download(model_name, cache_dir='D:\modelscope\hub\models')
        print(f"模型已下载到: {self.model_dir}")
        # 使用HuggingFaceEmbeddings包装ModelScope模型，使其兼容LangChain
        self.embeddings = HuggingFaceEmbeddings(
            model_name=self.model_dir,  # 使用下载的模型路径
            model_kwargs={'device': 'cpu'},  # 指定使用CPU设备
            encode_kwargs={'normalize_embeddings': True}  # 对嵌入向量进行归一化
        )
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """将文档列表转换为嵌入向量"""
        return self.embeddings.embed_documents(texts)
    def embed_query(self, text: str) -> List[float]:
        """将查询文本转换为嵌入向量"""
        return self.embeddings.embed_query(text)
# 初始化RAG系统
def setup_rag_system(knowledge_texts: List[str]):
    # 1. 文本分割 - 将长文本分割成适合处理的小块
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=500,  # 每个文本块的最大字符数
        chunk_overlap=50,  # 块之间的重叠字符数，保持上下文连贯
        length_function=len  # 使用Python内置len函数计算长度
    )
    # 创建文档对象
    documents = []
    for i, text in enumerate(knowledge_texts):
        docs = text_splitter.split_text(text)  # 分割文本
        for doc in docs:
            # 创建LangChain文档对象，包含内容和元数据
            documents.append(Document(page_content=doc, metadata={"source": f"doc_{i}"}))
    # 2. 初始化ModelScope嵌入模型
    embeddings = ModelScopeEmbeddings()
    # 3. 创建FAISS向量数据库 - 高效的相似性搜索库
    vectorstore = FAISS.from_documents(
        documents=documents,  # 文档列表
        embedding=embeddings.embeddings  # 嵌入模型
    )
    return vectorstore, embeddings  # 返回向量存储和嵌入模型
# 构建RAG链
def create_rag_chain(vectorstore, embeddings):
    # 创建检索器 - 用于从向量库中搜索相关文档
    retriever = vectorstore.as_retriever(
        search_type="similarity",  # 使用相似性搜索
        search_kwargs={"k": 3}  # 返回最相关的3个文档
    )
    # 定义提示模板 - 指导模型如何基于上下文回答问题
    template = """
    你是一个专业的助手，请根据以下背景信息回答问题。
    背景信息：
    {context}
    问题：{question}
    请根据背景信息提供准确、简洁的回答。如果背景信息中没有相关答案，请如实告知。
    """
    prompt = ChatPromptTemplate.from_template(template)  # 创建提示模板
    # 创建RAG处理函数
    def rag_chain(question: str) -> str:
        # 检索相关文档 - 基于问题查找最相关的文档片段
        relevant_docs = retriever.invoke(question)
        # 合并检索到的文档内容作为上下文
        context = "\n\n".join([doc.page_content for doc in relevant_docs])
        # 构建提示 - 将上下文和问题填入模板
        formatted_prompt = prompt.invoke({
            "context": context, 
            "question": question
        })
        # 调用千问模型生成回答
        qwen_llm = QwenLLM()
        response = qwen_llm.invoke(str(formatted_prompt))
        return response, context  # 返回模型回答和使用的上下文
    return rag_chain  # 返回RAG处理函数
# 示例知识库数据
knowledge_data = [
    """
    人工智能发展趋势：2024年，生成式AI技术将继续快速发展，多模态能力成为重点。
    企业应用AI的主要方向包括智能客服、内容生成、数据分析等。深度学习和神经网络
    技术不断突破，大模型参数规模持续增长。
    """,
    """
    公司规章制度：员工工作时间为每周5天，每天8小时。年假根据工龄计算，1-3年员工
    享有5天年假，3-5年员工享有10天年假。所有员工需要遵守保密协议，不得泄露公司
    商业秘密和技术信息。
    """,
    """
    产品介绍：我司主要产品包括智能客服系统、企业知识管理平台和AI内容生成工具。
    智能客服系统支持多渠道接入，能够自动回答常见问题，提升客户服务效率。知识管理
    平台帮助企业构建专属知识库，实现知识的集中管理和智能检索。
    """
]
# 初始化RAG系统
print("正在初始化RAG系统...")
vectorstore, embeddings = setup_rag_system(knowledge_data)
rag_chain = create_rag_chain(vectorstore, embeddings)
# 进行问答测试
questions = [
    "公司的年假政策是怎样的？",
    "智能客服系统有哪些功能？",
    "2024年AI发展趋势是什么？"
]
for question in questions:
    print(f"\n问题：{question}")
    answer, context = rag_chain(question)
    print(f"答案：{answer}")
    print(f"检索到的上下文：{context[:100]}...")  # 只显示前100个字符