用于与多个数据库聊天的智能 SQL 代理问答和 RAG 系统（3） —— 基于 LangChain 框架的文档检索与问答功能以及RAG Tool的使用

介绍基于 LangChain 框架的文档检索与问答功能，目标是通过查询存储的向量数据库（VectorDB），为用户的问题检索相关内容，并生成自然语言的答案。以下是代码逻辑的详细解析：

代码结构与功能

1. 初始化环境与加载配置

   from dotenv import load_dotenv
   load_dotenv()
   
   os.environ['OPENAI_API_KEY'] = os.getenv("OPENAI_API_KEY")
   

   • 使用 dotenv 加载 .env 文件中的环境变量（如 OPENAI_API_KEY 和 OPENAI_API_BASE）。
   • 将 OpenAI API 密钥设置为环境变量，供后续使用。

2. 加载向量数据库

   VECTORDB_DIR = "data/airline_policy_vectordb"
   K = 2
   
   vectordb = Chroma(
       collection_name="rag-chroma",
       persist_directory=str(here(VECTORDB_DIR)),
       embedding_function=OpenAIEmbeddings()
   )
   print("Number of vectors in vectordb:", vectordb._collection.count(), "\n\n")
   

   • VECTORDB_DIR：向量数据库的存储目录。
   • K=2：设置检索时返回的文档数。
   • 使用 Chroma 加载向量数据库：
     • collection_name 指定集合名称。
     • persist_directory 指向存储数据库的目录。
     • embedding_function 使用 OpenAIEmbeddings 生成向量嵌入。
   • 打印向量数据库中存储的向量数量。

3. 检索文档并生成提示

   message = "What is the cancelation rule for a flight ticket at swiss airline policy?"
   
   docs = vectordb.similarity_search(message, k=K)
   
   question = "# User new question:\n" + message
   retrieved_content = ""
   for doc in docs:
       retrieved_content += f"{doc.page_content}\n\n"
   prompt = f"# Content:\n{retrieved_content}\n\n{question}"
   
   pprint(prompt)
   

   • 定义用户问题 message，这是需要检索和回答的问题。
   • 检索文档：
     • 使用 vectordb.similarity_search 在向量数据库中查找与问题最相似的文档，返回 K 个相关文档。
   • 生成提示：
     • 将用户问题和检索到的文档内容整合为一个提示（prompt），格式如下：

       # Content:
       （检索到的文档内容）
       
       # User new question:
       （用户问题）
       

部分代码执行结果如下：

4. 调用聊天模型生成答案

   from langchain_openai import ChatOpenAI
   
   chat = ChatOpenAI(
       openai_api_base=os.getenv("OPENAI_API_BASE"),
       openai_api_key=os.getenv("OPENAI_API_KEY"),
       temperature=0
   )
   
   messages = [
       {"role": "system", "content": "You will receive a user's query and possible content where the answer might be. If the answer is found, provide it, if not, state that the answer does not exist."},
       {"role": "user", "content": prompt}
   ]
   
   response = chat.invoke(messages)
   print(response)
   

   • 使用 ChatOpenAI 调用 OpenAI 的聊天模型。
   • 定义对话上下文：
     • 系统消息：告诉模型用户会提供问题和可能的内容，要求模型判断答案是否存在。
     • 用户消息：将生成的提示作为输入。
   • 调用 chat.invoke 生成回答并打印。

5. 定义工具函数

   from langchain_core.tools import tool
   
   @tool
   def lookup_swiss_airline_policy(query: str) -> str:
       """Search within the Swiss Airline's company policies to check whether certain options are permitted. Input should be a search query."""
       vectordb = Chroma(
           collection_name="rag-chroma",
           persist_directory=str(here(VECTORDB_DIR)),
           embedding_function=OpenAIEmbeddings()
       )
       docs = vectordb.similarity_search(query, k=K)
       return "\n\n".join([doc.page_content for doc in docs])
   

   • 使用 @tool 装饰器创建一个可复用工具函数。
   • 功能：
     • 接受查询 query。
     • 从向量数据库中检索相关文档。
     • 返回检索到的文档内容。
   • 说明信息：提供关于工具功能的描述，供其他程序或用户调用时参考。

6. 工具函数的调用

   print(lookup_swiss_airline_policy.name)
   print(lookup_swiss_airline_policy.args)
   print(lookup_swiss_airline_policy.description)
   
   pprint(lookup_swiss_airline_policy.invoke("can I cancel my ticket?"))
   

   • 打印工具的元信息（名称、参数、描述）。
   • 调用 lookup_swiss_airline_policy.invoke，检索 “can I cancel my ticket?” 的相关内容并打印结果。

工具函数的执行结果：

运行流程总结

1. 加载向量数据库，并初始化嵌入模型。
2. 用户输入问题。
3. 从数据库中检索与问题相关的文档。
4. 将问题和文档内容发送至聊天模型，生成答案。
5. 定义工具函数供后续复用。

用途与适用场景

• 用途：实现基于文档的问答系统，用于快速查询特定文档的内容。
• 适用场景：
  • 公司政策文档检索（如航空政策）。
  • 产品说明文档或技术支持文档查询。
  • 法律条款、合同内容等信息的检索与问答。