找回密码
 立即注册
首页 业界区 安全 给大模型通过RAG挂上知识库

给大模型通过RAG挂上知识库

类饲冰 7 天前
前言

因为大模型的知识库存在于训练期间,因此对于一些最新发生的事或者是专业性问题可能会出现不准确或者是幻觉,因此可以使用RAG技术给大模型外挂知识库来达到精准回答的目的。
实操

gpt4all

可以参考之前的文章:Llama模型私有化教程
[img=720,377.3333333333333]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609664.png[/img]

1.png

他的优点就是通过UI在线下载模型和导入知识库,操作都比较一站式、傻瓜式。注意的是gpt4all的模型文件和ollama不通用。
open-webui

安装可以参考Llama模型私有化教程,也比较简单就不多赘述。
先看下在没有知识库的情况下,咨询相关问题时得到的结果是错误的:
[img=720,326.36182902584494]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609668.png[/img]

可以通过如下方式进行知识库的构建:
  1. 右上角-工作空间-知识库-新增知识库空间-上传知识库文件
复制代码
[img=720,123.6774193548387]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609671.png[/img]

这个时候再咨询知识库中存在的内容时就可以得到满意的结果(引用的方式是在输入框中输入#):
[img=720,306.9980879541109]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609672.png[/img]

ima

https://ima.qq.com/
ima是腾讯出品的AI+知识库的软件。创建知识库的流程为:
[img=720,294.33376455368693]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609673.png[/img]

首先有个缺点,它竟然不能上传markdown。还有些其他BUG,比如明明存在知识库,但是却选择不了:
[img=720,369.35593220338984]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609674.png[/img]

[img=720,354.38961038961037]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609675.png[/img]

因为没法设置prompt,如果你想让大模型每次都只从知识库中搜索不要联想,那么就就需要每次在输入框中输入特定prompt告知不要胡乱回答,结果发现又是混元问题,问答模型改成deepseek后好点:
[img=720,489.46889226100154]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609676.png[/img]

[img=720,618.4615384615385]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609677.png[/img]

终于明白这些公司为什么要接deepseek了,因此自己公司的太差。
【----帮助网安学习,以下所有学习资料免费领!加vx:yj520400,备注 “博客园” 获取!】
 ① 网安学习成长路径思维导图
 ② 60+网安经典常用工具包
 ③ 100+SRC漏洞分析报告
 ④ 150+网安攻防实战技术电子书
 ⑤ 最权威CISSP 认证考试指南+题库
 ⑥ 超1800页CTF实战技巧手册
 ⑦ 最新网安大厂面试题合集(含答案)
 ⑧ APP客户端安全检测指南(安卓+IOS)
langchain+chroma

上面介绍的都是通过图形化的方式进行,但是在一些工程化的地方可能没法进行图形化操作,接下来介绍使用代码的方式来进行让大模型外挂知识库。把文档投喂给大模型时需要先对文档进行向量转换,这里以chroma 官方代码为例:
  1. import chromadb
  2. # setup Chroma in-memory, for easy prototyping. Can add persistence easily!
  3. client = chromadb.Client()
  4.  
  5. # Create collection. get_collection, get_or_create_collection, delete_collection also available!
  6. collection = client.create_collection("all-my-documents")
  7.  
  8. # Add docs to the collection. Can also update and delete. Row-based API coming soon!
  9. collection.add(
  10.    documents=["This is document1", "This is document2"], # we handle tokenization, embedding, and indexing automatically. You can skip that and add your own embeddings as well
  11.    metadatas=[{"source": "notion"}, {"source": "google-docs"}], # filter on these!
  12.    ids=["doc1", "doc2"], # unique for each doc
  13. )
  14.  
  15. # Query/search 2 most similar results. You can also .get by id
  16. results = collection.query(
  17.    query_texts=["This is document1"],
  18.    n_results=2,
  19.    # where={"metadata_field": "is_equal_to_this"}, # optional filter
  20.    # where_document={"$contains":"search_string"}  # optional filter
  21. )
  22. print(results)
复制代码
上述代码含义是创建了一个集合,并且往集合中添加知识库,每个知识库都必须有自己的独立id。注意,chroma只支持传入文本不支持直接引用文件,因此想要把文件转成向量需要先把文件读取出内容给到chroma才行。
得到的内容如下:
  1. {'ids': [['doc1', 'doc2']], 'embeddings': None, 'documents': [['This is document1', 'This is document2']], 'uris': None, 'data': None, 'metadatas': [[{'source': 'notion'}, {'source': 'google-docs'}]], 'distances': [[0.0, 0.2221483439207077]], 'included': [<IncludeEnum.distances: 'distances'>, <IncludeEnum.documents: 'documents'>, <IncludeEnum.metadatas: 'metadatas'>]}
复制代码
其中distances代表是距离,笔者特地把搜索的问题和id为doc1的内容一致,因此可以看到得到的距离为0(距离越小,相似度越高),代表问题和文档一模一样,因此在后续投喂给大模型时,可以选择小于多少距离的投喂给大模型来解决token过长的问题。
接下来介绍langchain,langchain功能和它的名字一样,简单理解就是它可以把各个东西和大模型串在一起,比如可以把上面chroma生成的文档向量投喂给大模型进行知识库问答。langchain牛逼的点是他做了很多第三方工具的集成,比如以langchains调用chroma生成向量数据库为例:
  1. from langchain_ollama import OllamaEmbeddings
  2. from langchain_chroma import Chroma
  3. from uuid import uuid4
  4. from langchain_core.documents import Document
  5.  
  6. embeddings = OllamaEmbeddings(model="nomic-embed-text:latest")
  7.  
  8.  
  9. vector_store = Chroma(
  10.    collection_name="example_collection",
  11.    embedding_function=embeddings,
  12.    persist_directory="./chroma_langchain_db",  # Where to save data locally, remove if not necessary
  13. )
  14.  
  15. document_1 = Document(
  16.    page_content="I had chocolate chip pancakes and scrambled eggs for breakfast this morning.",
  17.    metadata={"source": "tweet"},
  18.    id=1,
  19. )
  20.  
  21. document_2 = Document(
  22.    page_content="The weather forecast for tomorrow is cloudy and overcast, with a high of 62 degrees.",
  23.    metadata={"source": "news"},
  24.    id=2,
  25. )
  26.  
  27. document_3 = Document(
  28.    page_content="Building an exciting new project with LangChain - come check it out!",
  29.    metadata={"source": "tweet"},
  30.    id=3,
  31. )
  32.  
  33. document_4 = Document(
  34.    page_content="Robbers broke into the city bank and stole $1 million in cash.",
  35.    metadata={"source": "news"},
  36.    id=4,
  37. )
  38.  
  39. document_5 = Document(
  40.    page_content="Wow! That was an amazing movie. I can't wait to see it again.",
  41.    metadata={"source": "tweet"},
  42.    id=5,
  43. )
  44.  
  45. document_6 = Document(
  46.    page_content="Is the new iPhone worth the price? Read this review to find out.",
  47.    metadata={"source": "website"},
  48.    id=6,
  49. )
  50.  
  51. document_7 = Document(
  52.    page_content="The top 10 soccer players in the world right now.",
  53.    metadata={"source": "website"},
  54.    id=7,
  55. )
  56.  
  57. document_8 = Document(
  58.    page_content="LangGraph is the best framework for building stateful, agentic applications!",
  59.    metadata={"source": "tweet"},
  60.    id=8,
  61. )
  62.  
  63. document_9 = Document(
  64.    page_content="The stock market is down 500 points today due to fears of a recession.",
  65.    metadata={"source": "news"},
  66.    id=9,
  67. )
  68.  
  69. document_10 = Document(
  70.    page_content="I have a bad feeling I am going to get deleted :(",
  71.    metadata={"source": "tweet"},
  72.    id=10,
  73. )
  74.  
  75. documents = [
  76.    document_1,
  77.    document_2,
  78.    document_3,
  79.    document_4,
  80.    document_5,
  81.    document_6,
  82.    document_7,
  83.    document_8,
  84.    document_9,
  85.    document_10,
  86. ]
  87. uuids = [str(uuid4()) for _ in range(len(documents))]
  88.  
  89. vector_store.add_documents(documents=documents, ids=uuids)
  90.  
  91. results = vector_store.similarity_search_with_score(
  92.    "Will it be hot tomorrow?", k=1, filter={"source": "news"}
  93. )
  94. print("-----")
  95. print(results)
  96. print("-----")
  97. for res, score in results:
  98.    print(f"* [SIM={score:3f}] {res.page_content} [{res.metadata}]")
  99. print("-----")
复制代码
上述代码意思是指生成10个文档,然后通过langchain内置的第三方模块能力把这10个文档写入到了example_collection集合中,且向量数据库持久化,保存的路径为chroma_langchain_db目录中,最后在向量数据库中以source为news、最接近的1个为条件文档中搜索问题:
[img=720,248.27586206896552]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609678.png[/img]

[img=720,66.1706783369803]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609679.png[/img]

接下来尝试使用langchain调用ollama进行与本地大模型进行沟通:
  1. from langchain_ollama import ChatOllama
  2.  
  3. llm = ChatOllama(
  4.    model="deepseek-r1:latest",
  5.    temperature=0.5,
  6. )
  7. messages = [
  8.    (
  9.        "system",
  10.        "角色:你是IT小助手,你只回答IT相关问题,其他问题不回答。当别人问你是谁时,你回答:我是IT小助手。",
  11.    ),
  12.    ("human", "你是谁"),
  13. ]
  14. ai_msg = llm.invoke(messages)
  15. print(ai_msg)
复制代码
上述代码通过设置system prompt来约束了大模型的输出:
[img=720,86.92363636363636]https://m-1254331109.cos.ap-guangzhou.myqcloud.com/202502281609680.png[/img]

上面提到chroma无法直接传入文件,因此langchian提供了文档加载器来实现读取不同类型的文件并输入给chroma。为了解决嵌入模型和大语言模型输入的的token限制,需要对文档进行分割,下面以读取txt文件为例,通过对内容进行分割,然后提供给嵌入模型转成向量并搜索相似度后,带入到大语言模型的上下文中进行提问:
  1. from typing import Dict
  2. import logging
  3. from pathlib import Path
  4.  
  5. from langchain_ollama import ChatOllama
  6. from langchain_core.prompts import ChatPromptTemplate
  7. from langchain_ollama import OllamaEmbeddings
  8. from langchain_text_splitters import RecursiveCharacterTextSplitter
  9. from langchain_community.document_loaders import TextLoader
  10. from langchain.chains import RetrievalQA
  11. from langchain_chroma import Chroma
  12. class VectorStoreQA:
  13.    def __init__(self,
  14.                 model_name: str = "deepseek-r1:latest",
  15.                 embedding_model: str = "nomic-embed-text:latest",
  16.                 temperature: float = 0.5,
  17.                 k: int = 4):
  18.        """
  19.        初始化 QA 系统
  20.        
  21.        Args:
  22.            model_name: LLM 模型名称
  23.            embedding_model: 嵌入模型名称
  24.            temperature: LLM 温度参数
  25.            k: 检索返回的文档数量
  26.        """
  27.        # 配置日志
  28.        logging.basicConfig(
  29.            level=logging.INFO,
  30.            format='%(asctime)s - %(levelname)s - %(message)s'
  31.        )
  32.        self.logger = logging.getLogger(__name__)
  33.        self.k = k
  34.        # 初始化 LLM
  35.        self.llm = ChatOllama(
  36.            model=model_name,
  37.            temperature=temperature,
  38.        )
  39.        
  40.        # 初始化 embeddings
  41.        self.embeddings = OllamaEmbeddings(model=embedding_model)
  42.        
  43.        # 初始化向量存储
  44.        self.vector_store = Chroma(embedding_function=self.embeddings)
  45.        
  46.        # 初始化 prompt 模板
  47.        # self.prompt = ChatPromptTemplate.from_messages([
  48.        #     ("system", """你的任务是且只基于提供的上下文信息回答用户问题。要求:1. 回答要准确、完整,并严格基于上下文信息2. 如果上下文信息不足以回答问题,不要编造信息和联想,直接说:在知识库中我找不到相关答案3. 采用结构化的格式组织回答,便于阅读"""),
  49.        #     ("user", """上下文信息:
  50.        #     {context}
  51.            
  52.        #     用户问题:{question}
  53.            
  54.        #     请提供你的回答:""")
  55.        # ])
  56.        self.prompt = ChatPromptTemplate.from_messages([
  57.            ("system", """上下文中没有相关资料的不要编造信息、不要从你历史库中搜索,直接说:在知识库中我找不到相关答案。"""),
  58.            ("user", """上下文信息:{context}
  59.            用户问题:{question}
  60.            请提供你的回答:""")
  61.        ])
  62.            
  63.  
  64.    def load_documents(self, file_path: str, chunk_size: int = 1000, chunk_overlap: int = 200) -> None:
  65.        """
  66.        加载并处理文本文档
  67.        
  68.        Args:
  69.            file_path: 文本文件路径
  70.            chunk_size: 文档分块大小
  71.            chunk_overlap: 分块重叠大小
  72.        """
  73.        try:
  74.            # 验证文件
  75.            path = Path(file_path)
  76.            if not path.exists():
  77.                raise FileNotFoundError(f"文件不存在: {file_path}")
  78.            
  79.            # 加载文档
  80.            loader = TextLoader(str(path))
  81.            docs = loader.load()
  82.            
  83.            # 文档分块
  84.            text_splitter = RecursiveCharacterTextSplitter(
  85.                chunk_size=chunk_size,
  86.                chunk_overlap=chunk_overlap
  87.            )
  88.            splits = text_splitter.split_documents(docs)
  89.            
  90.            # 添加到向量存储
  91.            self.vector_store.add_documents(documents=splits)
  92.            self.logger.info(f"成功加载文档: {file_path}")
  93.            
  94.        except Exception as e:
  95.            self.logger.error(f"文档处理错误: {str(e)}")
  96.            raise
  97.  
  98.    def get_answer(self, question: str) -> Dict:
  99.        """
  100.        获取问题的答案
  101.        Args:
  102.            question: 用户问题
  103.        Returns:
  104.            包含答案的字典
  105.        """
  106.        # 使用similarity_search_with_score方法获取文档和分数  
  107.        docs_and_scores = self.vector_store.similarity_search_with_score(  
  108.            query=question,  
  109.            k=self.k
  110.        )  
  111.        
  112.        # 打印每个文档的内容和相似度分数  
  113.        print("\n=== 检索到的相关文档 ===")  
  114.        for doc, score in docs_and_scores:  
  115.            print(f"\n相似度分数: {score:.4f}")  # 保留4位小数  
  116.            print(f"文档内容: {doc.page_content}")  
  117.            print(f"元数据: {doc.metadata}")  # 如果需要查看文档元数据  
  118.            print("-" * 50)  # 分隔线  
  119.  
  120.        # 提取文档内容用于后续处理  
  121.        context = "\n\n".join(doc.page_content for doc, _ in docs_and_scores)  
  122.        # 打印完整的prompt内容  
  123.        print("\n=== 实际发送给模型的Prompt ===")  
  124.        formatted_prompt = self.prompt.format(  
  125.            question=question,  
  126.            context=context  
  127.        )  
  128.        print(formatted_prompt)  
  129.        print("=" * 50)  
  130.        # docs = self.retriever.get_relevant_documents(question)  
  131.        # 将文档内容合并为上下文  
  132.        # context = "\n\n".join(doc.page_content for doc in docs)  
  133.        # print(context)
  134.        # 创建chain并调用
  135.        chain = self.prompt | self.llm  
  136.        response = chain.invoke({  
  137.            "question": question,  
  138.            "context": context  
  139.        })  
  140.        return response
  141.    def clear_vector_store(self):
  142.        """清空向量存储"""
  143.        try:
  144.            self.vector_store.delete_collection()
  145.            self.vector_store = Chroma(embedding_function=self.embeddings)
  146.            self.logger.info("已清空向量存储")
  147.        except Exception as e:
  148.            self.logger.error(f"清空向量存储时发生错误: {str(e)}")
  149.            raise
  150.  
  151. # 使用示例
  152. if __name__ == "__main__":
  153.    # 初始化 QA 系统
  154.    qa_system = VectorStoreQA(
  155.        model_name="deepseek-r1:latest",
  156.        k=4
  157.    )
  158.    
  159.    # 加载文档
  160.    qa_system.load_documents("/tmp/1.txt")
  161.    
  162.    # 提问
  163.    question = "猪八戒是谁?"
  164.    result = qa_system.get_answer(question)
  165.    print(result)
复制代码
总结

如果只是想简单尝试下大模型+知识库,那么gpt4all和ima都可以,毕竟都是图形化点点点就行,如果想要去自定义一些模型或者本身依赖ollama运行模型的话,可以选择open-webui,其可以有更多的自定义能力,如果想要在工程化中使用,建议使用langchain+chroma。
更多网安技能的在线实操练习,请点击这里>>
  

来源:程序园用户自行投稿发布,如果侵权,请联系站长删除
免责声明:如果侵犯了您的权益,请联系站长,我们会及时删除侵权内容,谢谢合作!
您需要登录后才可以回帖 登录 | 立即注册