大模型RAG实现智能知识库，仅几十行代码(openai接口+faiss)

检索增强生成 (RAG) 是一种使用来自私有或专有数据源的信息来辅助文本生成的技术。它将检索模型（设计用于搜索大型数据集或知识库）和生成模型（例如大型语言模型 (LLM)，此类模型会使用检索到的信息生成可供阅读的文本回复）结合在一起。
RAG的主要使用场景，就是基于自有知识库的大模型。本文通过python来实现最简答的一个rag搜索

实现思路

1. 对知识库进行分段
2. 对分段后的内容进行Embedding向量计算，并存入faiss数据库
3. 对搜索的问题进行Embedding向量计算
4. 查询近似度topN的段落
5. 把topN段落和问题一起发给gpt
6. 得到基于知识库的一个回答

faiss数据库

faiss是一个用于向量搜索的向量数据库，可以基于内存和cpu来运行，python只需要几行代码就能实现数据的存储和搜索，具体的用法或者深入的学习可以自行百度一下

代码实现

话不多说直接上代码,我不会写python,基本都是chatgpt写的，代码丑陋但是能跑。
引入依赖

import faiss
import numpy as np
from openai import OpenAI

定义一个分割函数，用来分割大段落成小段落的，这里随便分的，就5句话分成一段，也可以根据实际需要或者用一些比较好的库去进行分段

def split_string(x: str) -> list:
  punctuation = [',', '.']
  segments = []
  current = []
  count = 0
  for char in x:
    current.append(char)
    if char in punctuation:
      count += 1
      if count == 5:
        segments.append(''.join(current).strip())
        current = []
        count = 0

  if current:
    segments.append(''.join(current).strip())

  return segments

定义一个embedding函数，这里用openai的模型api直接调用

def getEmbeddings(x: str, client: OpenAI):
  response = client.embeddings.create(
    input=x,
    model="text-embedding-3-small"
  )
  return response.data[0].embedding

定义搜索方法,这里就是搜索和问题最相关的topk个段落。这里写成一次性方法，如果是会一直提问，分好的段落就不用重复embedding和插入数据库了

def get_similar_documents(client: OpenAI, query: str, documents: list,
    top_k: int = 5) -> list:
  queryEmbedding = getEmbeddings(query, client)
  queryEmbedding = np.array(queryEmbedding).reshape(1, -1)
  docEmbeddings = np.array([getEmbeddings(i, client) for i in documents])
  index = faiss.IndexFlatL2(docEmbeddings.shape[1])
  index.add(docEmbeddings)
  distances, indices = index.search(queryEmbedding, top_k)
  doc = []
  for i in indices[0]:
    doc.append(documents[i])
  return doc

下面是主程序

client = OpenAI(api_key="sk-xxxxx",base_url="https://xxx.xxx.xxx")
prompt = "你是一个智能问答机器人，以下是你的知识库内容："
doc_text = ["xxxxxxxxxxxxxxxxxxxxx","xxxxxxxxxxxxxxxxxx","xxxxxxxxxxxxxxxxx"]
query = input("输入你想问的问题")          
doc = get_similar_documents(client,query, doc_text,10)
response = client.chat.completions.create(
  model="gpt-3.5-turbo",
  messages=[
    {"role": "system", "content": prompt +('\n'.join(doc))},
    {"role": "user", "content": query}
  ]
)
print(response.choices[0].message.content)    

至此，你就实现了一个基于私有知识库的大模型智能问答了。里面的提示词都是随便写的，如果正式使用的话，话术肯定是需要调整一下，然后数据最好是持久化多次重复使用。