使用智谱的GLM-4.6构建RAG

1.前提条件

1.1 python 开发环境准备

python>=3.12

1.1.1 依赖解释

==sentence_transformers==
专门用于计算句子、文本和图像嵌入的 Python 框架。简单来说，它能将一段文字转换成一个高维度的数学向量（也叫嵌入），这个向量能够捕捉这段文字的语义信息。
==chromadb==
ChromaDB 是一个开源的嵌入式数据库，专门用于存储、管理和查询嵌入。它让你可以轻松地将文本（或其他数据）转换成向量，并根据这些向量的语义相似性进行高效搜索。
==zai-sdk==
智谱的 python-SDK，用于调用智谱 AI 开放平台的大模型。
==python-dotenv==
主要作用是从 .env 文件中读取环境变量，并将它们加载到 os.environ 中，这样你就可以在 Python 代码中像读取系统环境变量一样读取这些配置了。

1.1.2 安装依赖

pip install sentence_transformers chromadb zai-sdk python-dotenv

1.2 智谱 API_KEY

1.2.1 进入官网

智谱AI开放平台

1.2.2 注册后创建 API_KEY

1.3 文件准备

创建目录RAG,在目录下创建.env文件，和doc.md以及main.py文件

.env 内容

ZHIPU_API_KEY="你的API_KEY"
HF_ENDPOINT = "https://hf-mirror.com"

doc.md 内容

# 哆啦A梦与超级赛亚人：时空之战  
​  
    
​  
在一个寻常的午后，大雄依旧坐在书桌前发呆，作业堆得像山，连第一页都没动。哆啦A梦在一旁翻着漫画，时不时叹口气，觉得这孩子还是一如既往的不靠谱。正当他们的生活照常进行时，一道强光突然从天而降，整个房间震动不已。光芒中走出一名金发少年，身披战甲、气势惊人，他就是来自未来的超级赛亚人——特兰克斯。他一出现便说出了惊人的话：未来的地球即将被黑暗势力摧毁，他来此是为了寻求哆啦A梦的帮助。  
​  
    
​  
哆啦A梦与大雄听后大惊，但也从特兰克斯坚定的眼神中读出了不容拒绝的决心。特兰克斯解释说，未来的敌人并非普通反派，而是一个名叫“黑暗赛亚人”的存在，他由邪恶科学家复制了贝吉塔的基因并加以改造，实力超乎想象。这个敌人不仅拥有赛亚人战斗力，还能操纵扭曲的时间能量，几乎无人可敌。特兰克斯已经独自战斗多年，但每一次都以惨败告终。他说：“科技，是我那个时代唯一缺失的武器，而你们，正好拥有它。”  
​  
    
​  
于是，哆啦A梦带着特兰克斯与大雄启动时光机，穿越到了那个即将崩溃的未来世界。眼前的景象令人震撼：城市沦为废墟，大地裂痕纵横，天空中浮动着压抑的黑雾。特兰克斯说，这正是黑暗赛亚人带来的结果，一切生命几乎都被抹杀，只剩他在苦苦支撑。大雄虽感到恐惧，但看到无辜的人类遭殃，内心逐渐燃起斗志。哆啦A梦则冷静地分析局势，决定使用他最强的三样秘密道具来对抗黑暗势力。  
​  
    
​  
三件秘密道具分别是：可以临时赋予超级战力的“复制斗篷”，能暂停时间五秒的“时间停止手表”，以及可在一分钟中完成一年修行的“精神与时光屋便携版”。大雄被推进精神屋内，在其中接受密集的训练，虽然只有几分钟现实时间，他却经历了整整一年的苦修。刚开始他依旧软弱，想放弃、想逃跑，但当他想起静香、父母，还有哆啦A梦那坚定的眼神时，他终于咬牙坚持了下来。出来之后，他的身体与精神都焕然一新，眼神中多了一份成熟与自信。  
​  
    
​  
最终战在黑暗赛亚人的空中要塞前爆发，特兰克斯率先出击，释放全力与敌人正面对决。哆啦A梦则用任意门和道具支援，从各个方向制造混乱，尽量压制敌人的时空能力。但黑暗赛亚人太过强大，仅凭特兰克斯一人根本无法压制，更别说击败。就在特兰克斯即将被击倒之际，大雄披上复制斗篷、冲破恐惧从高空跃下。他的拳头燃烧着金色光焰，目标直指敌人心脏。  
​  
    
    
​  
时间停止装置在关键时刻启动，世界陷入静止，大雄用这个短短五秒接近了敌人的盲点。他集中全力，一记重拳击穿了黑暗赛亚人的能量核心，引发巨大的能量反冲。黑暗赛亚人尖叫着化为碎光，天空中的黑雾瞬间散去，阳光重新洒落大地。特兰克斯倒在地上，看着眼前这个曾经懦弱的少年，露出了欣慰的笑容。他知道，这一次，是大雄救了世界。  
​  
    
​  
战后，未来世界开始恢复，植物重新生长，人类重建家园。特兰克斯告别时紧紧握住大雄的手，说：“你是我见过最特别的战士。”哆啦A梦也为大雄感到骄傲，说他终于真正成长了一次。三人站在山丘上，看着远方重新明亮的地平线，心中感受到从未有过的安宁。随后，哆啦A梦与大雄乘坐时光机返回了属于他们的那个年代，一切仿佛又恢复平静。  
​  
    
​  
回到现代后，大雄仿佛变了一个人，不再轻易抱怨、不再逃避责任。他认真写完作业，帮妈妈买菜，甚至主动练习体育，哆啦A梦惊讶得说不出话来。他知道，这不是一时兴起，而是大雄真正内心成长的结果。大雄有时会望着天空出神，仿佛还能看见未来世界的那一片废墟与重生的希望。他不会说出来，但他心中永远铭记那一战。  
​  
    
​  
几天后，电视新闻中突然出现一则画面：一位金发少年在街头击退了失控的机器人，引发市民围观与猜测。大雄放下手中的课本，望向哆啦A梦，两人心照不宣地笑了。也许，特兰克斯又回来了，也许，新的敌人正在逼近。冒险从未真正结束，而他们，早已准备好了。无论时空如何动荡，他们将永远并肩作战。

main.py 内容

import os  

import chromadb 

from typing import List  
from sentence_transformers import SentenceTransformer, CrossEncoder  
from dotenv import load_dotenv  
from zai import ZhipuAiClient  
​  
​  
def split_into_chunks(doc_file: str) -> List[str]:  
    """  
    将文档文件分割成多个文本块。  
​  
    参数:  
        doc_file (str): 文档文件的路径。  
​  
    返回:  
        List[str]: 由文档中的段落组成的列表，每个段落是一个文本块。  
​  
    说明:  
        - 该函数读取指定路径的文档文件，并将其内容按照两个连续的换行符("\n\n")进行分割。  
        - 分割后的每个部分作为一个文本块，返回包含所有文本块的列表。  
    """  
    with open(doc_file, 'r', encoding='utf-8') as file:  
        content = file.read()  
​  
    return [chunk for chunk in content.split("\n\n")]  
​  
​  
def embed_chunk(chunk: str) -> List[float]:  
    """  
    将输入的文本块转换为向量嵌入表示。  
​  
    Args:  
        chunk (str): 需要转换为向量嵌入的文本块。  
​  
    Returns:  
        List[float]: 文本块的向量嵌入表示，是一个浮点数列表。  
    """  
    embedding = embedding_model.encode(chunk, normalize_embeddings=True)  
    return embedding.tolist()  
​  
​  
def save_embeddings_to_db(  
        chunks: List[str], embeddings_list: List[List[float]]):  
    """  
    将文本块和对应的嵌入向量保存到ChromaDB数据库中。  
      
    Args:  
        chunks (List[str]): 文本块列表，每个元素是一个字符串文本块  
        embeddings_list (List[List[float]]): 嵌入向量列表，每个元素是一个浮点数列表，表示对应文本块的嵌入向量  
      
    Returns:  
        Collection: ChromaDB数据库集合对象，包含已添加的文档和嵌入向量  
      
    Notes:  
        - 使用临时客户端(EphemeralClient)创建数据库实例  
        - 可以通过取消注释PersistentClient行来使用持久化存储  
        - 为每个文本块生成唯一的字符串ID  
        - 创建或获取名为"default"的集合  
        - 将文档、嵌入向量和ID批量添加到集合中  
    """  
    ids = [str(i) for i in range(len(chunks))]  
    chromadb_client = chromadb.EphemeralClient()  
    # chromadb_client = chromadb.PersistentClient("./chroma.db")  
​  
    db_collection = chromadb_client.get_or_create_collection(name="default")  
    db_collection.add(documents=chunks, embeddings=embeddings_list, ids=ids)  
    return db_collection  
​  
​  
def retrieve(query: str, top_k: int, db_collection):  
    """返回查询结果中的第一个文档。  
​  
    Args:  
        无直接参数，依赖于包含'documents'键的results字典。  
​  
    Returns:  
        list: results['documents'][0] - 查询结果中的第一个文档列表。  
​  
    Note:  
        - 假设输入的results字典中存在'documents'键  
        - 'documents'对应的值是一个列表，且至少包含一个元素  
        - 返回的是'documents'列表的第一个元素  
    """  
    query_embedding = embed_chunk(query)  
    results = db_collection.query(  
        query_embeddings=[query_embedding],  
        n_results=top_k)  
      
    return results['documents'][0]  
​  
​  
def rerank(query: str, retrieved_chunks: List[str], top_k: int) -> List[str]:  
    """  
    对检索到的文本块进行重排序，返回最相关的top_k个文本块。  
      
    Args:  
        query (str): 查询字符串  
        retrieved_chunks (List[str]): 待排序的文本块列表  
        top_k (int): 返回的文本块数量  
          
    Returns:  
        List[str]: 按相关性排序后的前top_k个文本块  
          
    Note:  
        使用CrossEncoder模型计算query和每个chunk的相关性分数，  
        然后根据分数对chunks进行降序排序。  
        CrossEncoder模型路径为：'C:\\Users\\HUAWEI\\.cache\\huggingface\\hub\\models--cross-encoder--mmarco-mMiniLMv2-L12-H384-v1\\snapshots\\1427fd652930e4ba29e8149678df786c240d8825'  
    """  
    cross_encoder = CrossEncoder(  
        r'C:\Users\HUAWEI\.cache\huggingface\hub\models--cross-encoder--mmarco-mMiniLMv2-L12-H384-v1\snapshots\1427fd652930e4ba29e8149678df786c240d8825')  
    pairs = [(query, chunk) for chunk in retrieved_chunks]  
    scores = cross_encoder.predict(pairs)  
​  
    chunk_with_score_list = [(chunk, score)  
                             for chunk, score in zip(retrieved_chunks, scores)]  
    chunk_with_score_list.sort(key=lambda pair: pair[1], reverse=True)  
​  
    return [chunk for chunk, _ in chunk_with_score_list][:top_k]  
​  
​  
def generate(query: str, chunks: List[str]) -> str:  
    """\n  
    基于给定的问题和相关文本片段生成回答。\n  
    \n  
    Args:\n  
        query (str): 用户的问题\n  
        chunks (List[str]): 相关文本片段列表\n  
    \n  
    Returns:\n  
        str: 生成的回答内容\n  
    \n  
    Process:\n  
        1. 从环境变量获取智谱AI的API密钥\n  
        2. 创建智谱AI客户端\n  
        3. 构建系统提示和用户提示\n  
        4. 调用chat completion API生成回答\n  
        5. 返回生成的回答内容\n  
    """  
    api_key = os.getenv("ZHIPU_API_KEY")  
    client = ZhipuAiClient(api_key=api_key)  
    prompt = [  
        {"role": "system",  
            "content": "你是一位知识助手，请根据用户的问题和下列片段生成准确的回答。"  
         },  
        {  
            "role": "user",  
            "content": f"""用户问题: {query}相关片段:{"\n\n".join(chunks)}请基于上述内容作答，不要编造信息。"""  
        }]  
​  
    print(f"{prompt}\n\n---\n")  
​  
    response = client.chat.completions.create(  
        model="glm-4.6",  
        messages=prompt,  
        temperature=0.6  
    )  
​  
    return response.choices[0].message.content  
​  
​  
if __name__ == "__main__":  
    load_dotenv()  
    print('分片开始！')  
    chunks = split_into_chunks(rf"D:\code_space\demo\RAG\doc.md")  
    print('分片完成！')  
    print('向量化开始！')  
    embedding_model = SentenceTransformer(  
        r"C:\Users\HUAWEI\.cache\huggingface\hub\models--shibing624--text2vec-base-chinese\snapshots\183bb99aa7af74355fb58d16edf8c13ae7c5433e")  
    embeddings = [embed_chunk(chunk) for chunk in chunks]  
​  
    print('向量化完成！')  
    print('索引开始！')  
    db_collection = save_embeddings_to_db(chunks, embeddings)  
    print('索引完成！')  
​  
    query = '哆啦A梦使用的3个秘密道具分别是什么？'  
    retrieved_chunks = retrieve(query, 5,db_collection)  
    reranked_chunks = rerank(query, retrieved_chunks, 3)  
    answer = generate(query, reranked_chunks)  
    print(answer)  

2. 运行 main. py

3 .说明

3.1~3.3 完成了 RAG 的建立，3. 4 及后续构成了通过大模型检索 RAG 进行回答。

3.1 将文档 doc.md 中的内容进行切片。

3.2 将分片进行向量化

3.3 将向量临时存储到 chromdb 中，并生成索引

3.4 根据用户问题检索 RAG

3 .5 相关性排序

3.6 构建 prompt (提示词+相关片段)

3.7 整体流程实现