MongoDB RAG

使用 MongoDB 向量搜索实现 RAG

🌐 查看英文原文 | 源码 Notebook

如何使用 Claude 3 和 MongoDB 构建 RAG 系统

本教程实现了一个被提示扮演风险投资技术分析师角色的聊天机器人。该聊天机器人是一个基础的 RAG 系统，以一组科技新闻文章作为其知识源。 本教程涵盖以下内容：

  1. 按照全面的教程设置你的开发环境，从安装必要的库到配置 MongoDB 数据库。
  2. 学习高效的数据处理方法，包括创建向量搜索索引和准备数据以进行数据导入和查询处理。
  3. 了解如何在 RAG 系统中使用 Claude 3 模型，基于从数据库检索的上下文信息生成精确的回复。

你需要准备以下内容： - Claude API Key - VoyageAI API Key - Hugging Face Access Token

第一步：库安装、数据加载和准备

以下是实现代码中使用的工具和库的简要说明： - anthropic：Anthropic 的官方 Python 库，可访问最先进的语言模型。该库提供对 Claude 3 系列模型的访问，能够理解文本和图像。 - datasets：该库是 Hugging Face 生态系统的一部分。通过安装 datasets，我们可以访问许多预处理好的即用型数据集，这对于训练和微调机器学习模型或基准测试其性能至关重要。 - pandas：这个数据科学库提供了强大的数据结构和方法，用于数据处理、加工和分析。 - voyageai：这是访问 VoyageAI embedding 模型套件的官方 Python 客户端库。 - pymongo：PyMongo 是 MongoDB 的 Python 工具包，它实现了与 MongoDB 数据库的交互。

!pip install pymongo datasets pandas anthropic voyageai

下面的代码片段执行以下步骤： 1. 导入必要的库： - os 用于与操作系统交互， - requests 用于发起 HTTP 请求， - BytesIO 来自 io 模块，用于在内存中像文件一样处理字节对象， - pandas（简写为 pd）用于数据操作和分析， - userdata 来自 google.colab，可访问存储在 Google Colab 密钥中的环境变量。 2. 函数定义：download_and_combine_parquet_files 函数定义了两个参数： - parquet_file_urls：URL 字符串列表，每个 URL 指向一个包含 tech-news-embedding 数据集子集合的 Parquet 文件。 - hf_token 是表示 Hugging Face 授权令牌的字符串。访问令牌可以从 Hugging Face 平台创建或复制。 3. 下载和读取 Parquet 文件：该函数遍历 parquet_file_urls 中的每个 URL。对于每个 URL，它会： - 使用 requests.get 方法发起 GET 请求，传入 URL 和用于授权的 headers。 - 检查响应状态码是否为 200（OK），表示请求成功。 - 如果成功，将响应内容读取到 BytesIO 对象中（在内存中作为文件处理），然后使用 pandas.read_parquet 从该对象读取 Parquet 文件到 Pandas DataFrame。 - 将 DataFrame 追加到 all_dataframes 列表中。 4. 合并 DataFrame：在将所有 Parquet 文件下载并读取为 DataFrame 后，会检查 all_dataframes 是否为空。如果有可用的 DataFrame，则使用 pd.concat 将所有 DataFrame 连接成一个，ignore_index=True 用于重新索引合并后的 DataFrame。这个合并的 DataFrame 是 download_and_combine_parquet_files 函数的最终输出。

from io import BytesIO

import pandas as pd
import requests
from google.colab import userdata


def download_and_combine_parquet_files(parquet_file_urls, hf_token):
    """
    Downloads Parquet files from the provided URLs using the given Hugging Face token,
    and returns a combined DataFrame.

    Parameters:
    - parquet_file_urls: List of strings, URLs to the Parquet files.
    - hf_token: String, Hugging Face authorization token.

    Returns:
    - combined_df: A pandas DataFrame containing the combined data from all Parquet files.
    """
    headers = {"Authorization": f"Bearer {hf_token}"}
    all_dataframes = []

    for parquet_file_url in parquet_file_urls:
        response = requests.get(parquet_file_url, headers=headers, timeout=60)
        if response.status_code == 200:
            parquet_bytes = BytesIO(response.content)
            df = pd.read_parquet(parquet_bytes)
            all_dataframes.append(df)
        else:
            print(
                f"Failed to download Parquet file from {parquet_file_url}: {response.status_code}"
            )

    if all_dataframes:
        combined_df = pd.concat(all_dataframes, ignore_index=True)
        return combined_df
    else:
        print("No dataframes to concatenate.")
        return None

以下是本教程所需的 Parquet 文件列表。所有文件的完整列表位于这里。每个 Parquet 文件大约包含 45,000 个数据点。

在下面的代码片段中，tech-news-embeddings 数据集的一个子集被组合成单个 DataFrame，然后赋值给变量 combined_df。

# Uncomment the links below to load more data
# For the full list of data visit: https://huggingface.co/datasets/MongoDB/tech-news-embeddings/tree/refs%2Fconvert%2Fparquet/default/train
parquet_files = [
    "https://huggingface.co/api/datasets/AIatMongoDB/tech-news-embeddings/parquet/default/train/0000.parquet",
    # "https://huggingface.co/api/datasets/AIatMongoDB/tech-news-embeddings/parquet/default/train/0001.parquet",
    # "https://huggingface.co/api/datasets/AIatMongoDB/tech-news-embeddings/parquet/default/train/0002.parquet",
    # "https://huggingface.co/api/datasets/AIatMongoDB/tech-news-embeddings/parquet/default/train/0003.parquet",
    # "https://huggingface.co/api/datasets/AIatMongoDB/tech-news-embeddings/parquet/default/train/0004.parquet",
    # "https://huggingface.co/api/datasets/AIatMongoDB/tech-news-embeddings/parquet/default/train/0005.parquet",
]

hf_token = userdata.get("HF_TOKEN")
combined_df = download_and_combine_parquet_files(parquet_files, hf_token)

作为数据准备的最后阶段，下面的代码片段展示了从分组数据集中移除 _id 列的步骤，因为它在本教程的后续步骤中不需要。此外，每个数据点的 embedding 列中的数据从 numpy 数组转换为 Python 列表，以防止在数据导入过程中出现与不兼容数据类型相关的错误。

# Remove the _id coloum from the intital dataset
combined_df = combined_df.drop(columns=["_id"])

# Remove the initial embedding coloumn as we are going to create new embeddings with VoyageAI embedding model
combined_df = combined_df.drop(columns=["embedding"])
combined_df.head()
# Limiting the amount of document used to 500 for this demo due to the rate limit on VoyageAI API
# Read more on VoyageAI rate limits: https://docs.voyageai.com/docs/rate-limits
max_documents = 500

if len(combined_df) > max_documents:
    combined_df = combined_df[:max_documents]
import voyageai

vo = voyageai.Client(api_key=userdata.get("VOYAGE_API_KEY"))


def get_embedding(text: str) -> list[float]:
    if not text.strip():
        print("Attempted to get embedding for empty text.")
        return []

    embedding = vo.embed(text, model="voyage-large-2", input_type="document")

    return embedding.embeddings[0]


combined_df["embedding"] = combined_df["description"].apply(get_embedding)

combined_df.head()

第二步：数据库和集合创建

要创建新的 MongoDB 数据库，请设置数据库集群： 1. 注册免费的 MongoDB Atlas 账户，或已有用户登录 MongoDB Atlas 2. 在左侧面板选择”Database”选项，将导航到 Database Deployment 页面，其中显示任何现有集群的部署规格。点击”+Create”按钮创建新的数据库集群。 3. 如需数据库集群设置和获取 URI 的帮助，请参阅我们关于设置 MongoDB 集群和获取连接字符串的指南。 注意：在创建概念验证时，不要忘记将 Python 主机的 IP 或 0.0.0.0/0 加入白名单。 4. 成功创建和部署集群后，集群可以在’Database Deployment’页面上访问。 5. 点击集群的”Connect”按钮，查看通过各种语言驱动程序设置集群连接的选项。 6. 本教程只需要集群的 URI（唯一资源标识符）。获取 URI 并将其复制到 Google Colab 的 Secrets 环境中的 MONGO_URI 变量中，或放入 .env 文件或等效位置。

创建集群后，导航到集群页面，点击 + Create Database 在 MongoDB Atlas 集群中创建数据库和集合。 数据库将命名为 tech_news，集合将命名为 hacker_noon_tech_news。

第三步：向量搜索索引创建

到此为止，你已经创建了集群、数据库和集合。

本节中的步骤对于确保能够使用输入到聊天机器人的查询来搜索 hacker_noon_tech_news 集合中的记录至关重要。此步骤的目标是创建向量搜索索引。要实现这一点，请参阅官方的向量搜索索引创建指南。

在使用 MongoDB Atlas 上的 JSON 编辑器创建向量搜索索引时，请确保你的向量搜索索引命名为 vector_index，向量搜索索引定义如下：

{
 "fields": [{
     "numDimensions": 1536,
     "path": "embedding",
     "similarity": "cosine",
     "type": "vector"
   }]
}

第四步：数据导入

要将数据导入到前面步骤中创建的 MongoDB 数据库中，需要执行以下操作： - 连接到数据库和集合 - 清除集合中的任何现有记录 - 在导入前将数据集的 Pandas DataFrame 转换为字典 - 使用批量操作将字典导入 MongoDB

本教程需要集群的 URI（唯一资源标识符）。获取 URI 并将其复制到 Google Colab 的 Secrets 环境中的 MONGO_URI 变量中，或放入 .env 文件或等效位置。

import pymongo
from google.colab import userdata


def get_mongo_client(mongo_uri):
    """Establish connection to the MongoDB."""
    try:
        client = pymongo.MongoClient(mongo_uri)
        print("Connection to MongoDB successful")
        return client
    except pymongo.errors.ConnectionFailure as e:
        print(f"Connection failed: {e}")
        return None


mongo_uri = userdata.get("MONGO_URI")
if not mongo_uri:
    print("MONGO_URI not set in environment variables")

mongo_client = get_mongo_client(mongo_uri)

DB_NAME = "tech_news"
COLLECTION_NAME = "hacker_noon_tech_news"

db = mongo_client[DB_NAME]
collection = db[COLLECTION_NAME]
Connection to MongoDB successful
# To ensure we are working with a fresh collection
# delete any existing records in the collection
collection.delete_many({})
DeleteResult({'n': 228012, 'electionId': ObjectId('7fffffff000000000000000e'), 'opTime': {'ts': Timestamp(1709660559, 7341), 't': 14}, 'ok': 1.0, '$clusterTime': {'clusterTime': Timestamp(1709660559, 7341), 'signature': {'hash': b'jT\xf1\xb4\xa9\xd3\xe3suu\x03`\x15(}\x8f\x00\x9f\xe9\x8a', 'keyId': 7320226449804230661}}, 'operationTime': Timestamp(1709660559, 7341)}, acknowledged=True)
# Data Ingestion
combined_df_json = combined_df.to_dict(orient="records")
collection.insert_many(combined_df_json)

第五步：向量搜索

本节展示了如何创建一个自定义的向量搜索函数，该函数接受用户查询，对应于聊天机器人的输入。该函数还接受第二个参数 collection，指向包含记录的数据库集合，向量搜索操作将针对这些记录进行。

vector_search 函数根据 MongoDB 聚合管道中定义的一系列操作生成向量搜索结果。该管道包含 $vectorSearch 和 $project 阶段，并根据用户查询的向量 embedding 执行查询。然后格式化结果，省略后续流程不需要的记录属性。

下面的代码片段执行以下操作以实现语义搜索： 1. 定义 vector_search 函数，该函数接受用户的查询字符串和 MongoDB 集合作为输入，并返回基于向量相似性搜索匹配查询的文档列表。 2. 通过调用之前定义的 get_embedding 函数为用户查询生成 embedding，该函数将查询字符串转换为向量表示。 3. 为 MongoDB 的聚合函数构建管道，包含两个主要阶段：$vectorSearch 和 $project。 4. $vectorSearch 阶段执行实际的向量搜索。index 字段指定用于向量搜索的向量索引，这应与前序步骤中向量搜索索引定义中输入的名称一致。queryVector 字段接收用户查询的 embedding 表示。path 字段对应于包含 embedding 的文档字段。numCandidates 指定要考虑的候选文档数量以及返回结果数量的限制。 5. $project 阶段格式化结果，排除 _id 和 embedding 字段。 6. aggregate 执行定义的管道以获取向量搜索结果。最终操作将从数据库返回的游标转换为列表。

def vector_search(user_query, collection):
    """
    Perform a vector search in the MongoDB collection based on the user query.

    Args:
    user_query (str): The user's query string.
    collection (MongoCollection): The MongoDB collection to search.

    Returns:
    list: A list of matching documents.
    """

    # Generate embedding for the user query
    query_embedding = get_embedding(user_query)

    if query_embedding is None:
        return "Invalid query or embedding generation failed."

    # Define the vector search pipeline
    pipeline = [
        {
            "$vectorSearch": {
                "index": "vector_index",
                "queryVector": query_embedding,
                "path": "embedding",
                "numCandidates": 150,  # Number of candidate matches to consider
                "limit": 5,  # Return top 5 matches
            }
        },
        {
            "$project": {
                "_id": 0,  # Exclude the _id field
                "embedding": 0,  # Exclude the embedding field
                "score": {
                    "$meta": "vectorSearchScore"  # Include the search score
                },
            }
        },
    ]

    # Execute the search
    results = collection.aggregate(pipeline)
    return list(results)

第六步：使用 Claude 3 模型处理用户查询

本教程的最后一部分概述了以下操作序列：

  • 接受字符串形式的用户查询。
  • 使用 VoyageAI embedding 模型为用户查询生成 embedding。
  • 加载 Anthropic Claude 3，特别是 ‘claude-opus-4-1’ 模型，作为 RAG 系统的基础模型。
  • 使用用户查询的 embedding 执行向量搜索，从知识库中获取相关信息，为基础模型提供额外的上下文。
  • 将用户查询和收集到的额外信息一起提交给基础模型以生成回复。

一个重要的注意事项是，用户查询 embedding 的维度必须与 MongoDB Atlas 上向量搜索索引定义中设置的维度一致。

本节的下一步是导入 anthropic 库并加载客户端，以访问 anthropic 的消息处理方法和 Claude 模型。确保你从Anthropic 官方网站的设置页面获取了 Claude API key。

import anthropic

client = anthropic.Client(api_key=userdata.get("ANTHROPIC_API_KEY"))

下面是代码片段中操作的更详细说明：

  1. 向量搜索执行：函数首先使用用户的查询和指定的集合作为参数调用 vector_search。这会在集合内执行搜索，利用向量 embedding 查找与查询相关的信息。
  2. 编译搜索结果：search_result 被初始化为空字符串，用于聚合搜索中的信息。搜索结果通过遍历 vector_search 函数返回的结果来编译，将每个条目的详情（标题、公司名称、URL、发布日期、文章 URL 和描述）格式化为人类可读的字符串，并将这些信息以换行符 追加到 search_result 中。
  3. 使用 Anthropic 客户端生成回复：然后该函数构建对 Claude API 的请求（通过一个客户端对象，应该是之前创建的 anthropic.Client 类的实例）。它指定了：
  • 要使用的模型（“claude-opus-4-1”）表示 Claude 3 模型的特定版本。
  • 生成回复的最大 token 限制（max_tokens=1024）。
  • 一个系统描述引导模型表现为”风险投资技术分析师”，可以访问科技公司文章和信息，并使用此上下文来提供建议。
  • 模型要处理的消息将用户查询与聚合的搜索结果作为上下文结合在一起。
  1. 返回生成的回复和搜索结果：它从回复内容中的第一项提取并返回回复文本，以及编译好的搜索结果。
def handle_user_query(query, collection):
    get_knowledge = vector_search(query, collection)

    search_result = ""
    for result in get_knowledge:
        search_result += (
            f"Title: {result.get('title', 'N/A')}, "
            f"Company Name: {result.get('companyName', 'N/A')}, "
            f"Company URL: {result.get('companyUrl', 'N/A')}, "
            f"Date Published: {result.get('published_at', 'N/A')}, "
            f"Article URL: {result.get('url', 'N/A')}, "
            f"Description: {result.get('description', 'N/A')}, \n"
        )

    response = client.messages.create(
        model="claude-opus-4-1",
        max_tokens=1024,
        system="You are Venture Captital Tech Analyst with access to some tech company articles and information. You use the information you are given to provide advice.",
        messages=[
            {
                "role": "user",
                "content": "Answer this user query: "
                + query
                + " with the following context: "
                + search_result,
            }
        ],
    )

    return (response.content[0].text), search_result

本教程的最后一步是初始化查询，将其传递给 handle_user_query 函数并打印返回的回复。

# Conduct query with retrieval of sources
query = "Give me the best tech stock to invest in and tell me why"
response, source_information = handle_user_query(query, collection)

print(f"Response: {response}")
print(f"\\nSource Information: \\n{source_information}")
Response: Based on the information provided in the article titles and descriptions, Alibaba Group Holding Limited appears to be a top technology stock pick for 2023 according to renowned investor Ray Dalio. The article "Top 10 Technology Stocks to Buy in 2023 According to Ray Dalio" suggests that Alibaba is one of Dalio's favored tech investments for the year.

As a venture capital tech analyst, I would recommend considering an investment in Alibaba for the following reasons:

1. Endorsement from a respected investor: Ray Dalio, known for his successful investment strategies, has included Alibaba in his top 10 technology stock picks for 2023. His backing lends credibility to the investment potential of the company.

2. Strong market position: Alibaba is a leading e-commerce company in China with a significant market share. It has a diversified business model spanning e-commerce, cloud computing, digital media, and entertainment.

3. Growth potential: With China's large and growing middle class, Alibaba is well-positioned to benefit from increasing consumer spending and the shift towards online shopping.

However, it's essential to consider the following points as well:

1. Regulatory risks: Chinese tech companies, including Alibaba, have faced increased regulatory scrutiny in recent times. Changes in government policies could impact the company's growth and profitability.

2. Competition: While Alibaba is a dominant player, it faces competition from other tech giants like Tencent and JD.com in various business segments.

3. Geopolitical tensions: Ongoing tensions between the U.S. and China could lead to market volatility and impact investor sentiment towards Chinese stocks.

As with any investment, it's crucial to conduct thorough research, consider your risk tolerance, and diversify your portfolio. Keep in mind that the information provided here is based on limited data points, and stock prices can be influenced by various factors beyond the scope of this context.
Source Information: \nTitle: Top 10 Technology Stocks to Buy in 2023 According to Ray Dalio, Company Name: ALIBABA GROUP HOLDING LIMITED, Company URL: https://hackernoon.com/company/alibabagroupholdinglimited, Date Published: 2023-04-21 11:58:00, Article URL: https://uk.finance.yahoo.com/news/top-10-technology-stocks-buy-155830366.html, Description: In this article we discuss the top 10 technology stocks to buy in 2023 according to Ray Dalio. If you want to skip our detailed analysis of Dalio’s investment philosophy and portfolio construction, 
Title: 3 Tech Stocks I Love Right Now, Company Name: 10Clouds, Company URL: https://hackernoon.com/company/10clouds, Date Published: 2023-04-02 11:30:00, Article URL: https://www.msn.com/en-xl/money/other/3-tech-stocks-i-love-right-now/ar-AA19n9Ht, Description: These are tech giants but they''re also great investments., 
Title: 3 Millionaire-Maker Hydrogen Stocks to Buy Before the Window Closes, Company Name: Air Products & Chemicals, Company URL: https://hackernoon.com/company/airproductschemicals, Date Published: 2023-07-28 12:18:00, Article URL: https://www.msn.com/en-us/money/topstocks/3-millionaire-maker-hydrogen-stocks-to-buy-before-the-window-closes/ar-AA1etN8O, Description: These are the best hydrogen stocks to buy with multibagger returns potential.More From InvestorPlace Buy This $5 Stock BEFORE This Apple Project Goes Live Wall Street Titan: Here’s My #1 Stock for 2023 The $1 Investment You MUST Take Advantage of Right Now It doesn’t matter if you have $500 or $5 million., 
Title: Why it may be time to sell the pop in tech stocks: BlackRock, Company Name: BlackRock, Company URL: https://hackernoon.com/company/blackrock, Date Published: 2023-02-13 19:06:00, Article URL: https://news.yahoo.com/why-it-may-be-time-to-sell-the-pop-in-tech-stocks-blackrock-190606866.html, Description: Household tech names like Apple Meta and Netflix have soared so far in 2023 but one strategist says the gains aren''t likely to last. It''s time to take profits on tech stocks — the early sector winner of 2023 — as the Federal Reserve may soon dash hopes for a pivot on interest rates, 
Title: The harsh reality for investors eyeing tech stocks in 2023: Morning Brief, Company Name: 10Clouds, Company URL: https://hackernoon.com/company/10clouds, Date Published: 2023-01-02 11:18:00, Article URL: https://news.yahoo.com/the-harsh-reality-for-investors-eyeing-tech-stocks-in-2023-morning-brief-111854804.html, Description: Curious on how to buy battered tech stocks? Here''s a quick tip. More on that and what else to watch in business on Monday.,
Tip

中文读者提示：本章节的代码和输出为英文原文。如需理解具体实现细节，可参考上方中文导读和代码注释。如有疑问，欢迎在 GitHub Issues 讨论。