使用Azure Cosmos DB MongoDB vCore进行向量搜索

最新推荐文章于 2025-06-15 17:36:11 发布

qahaj

最新推荐文章于 2025-06-15 17:36:11 发布

阅读量344

点赞数 9

文章标签：数据库 azure mongodb python

本文链接：https://blog.youkuaiyun.com/qahaj/article/details/146436426

版权

在这篇文章中，我们将探讨如何使用Azure Cosmos DB MongoDB vCore作为向量数据库来存储文档，并在集合中创建索引，通过近似最近邻算法（如余弦距离、欧氏距离和内积）执行向量搜索查询以定位与查询向量接近的文档。

技术背景介绍

Azure Cosmos DB是一个全球分布的多模型数据库服务，支持单位数毫秒响应时间，自动和即时可扩展性，以及在任何规模下的速度保证。它支持多种API，包括MongoDB vCore API，这是一个完全托管的MongoDB兼容数据库服务，允许开发者使用熟悉的MongoDB架构来构建现代应用程序。

核心原理解析

向量数据库是一种专门处理向量数据的数据库，用于执行例如向量相似性搜索的任务。这对于需要找到与某个查询向量最相似的文档这一类应用非常有用，常见算法包括余弦相似度、欧氏距离等。

代码实现演示

以下代码展示了如何将文档加载到Azure Cosmos DB的集合中，创建索引并进行向量搜索查询。

# 安装必要的库
%pip install --upgrade --quiet pymongo langchain-openai langchain-community

import os
from pymongo import MongoClient
from langchain_community.document_loaders import TextLoader
from langchain_community.vectorstores.azure_cosmos_db import AzureCosmosDBVectorSearch, CosmosDBSimilarityType, CosmosDBVectorSearchType
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import CharacterTextSplitter

# 设置MongoDB连接字符串和其他配置
CONNECTION_STRING = "YOUR_CONNECTION_STRING"
INDEX_NAME = "izzy-test-index"
NAMESPACE = "izzy_test_db.izzy_test_collection"
DB_NAME, COLLECTION_NAME = NAMESPACE.split(".")

# 配置OpenAI环境变量
os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2023-05-15"
os.environ["OPENAI_API_BASE"] = "YOUR_OPEN_AI_ENDPOINT"
os.environ["OPENAI_API_KEY"] = "YOUR_OPENAI_API_KEY"
os.environ["OPENAI_EMBEDDINGS_DEPLOYMENT"] = "smart-agent-embedding-ada"
os.environ["OPENAI_EMBEDDINGS_MODEL_NAME"] = "text-embedding-ada-002"

# 加载文档
SOURCE_FILE_NAME = "../../how_to/state_of_the_union.txt"
loader = TextLoader(SOURCE_FILE_NAME)
documents = loader.load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

# 初始化OpenAI Embeddings
openai_embeddings = OpenAIEmbeddings(
    deployment=os.getenv("OPENAI_EMBEDDINGS_DEPLOYMENT"),
    model=os.getenv("OPENAI_EMBEDDINGS_MODEL_NAME"),
    chunk_size=1
)

# 初始化MongoDB客户端
client = MongoClient(CONNECTION_STRING)
collection = client[DB_NAME][COLLECTION_NAME]

# 创建Azure Cosmos DB向量搜索索引
vectorstore = AzureCosmosDBVectorSearch.from_documents(
    docs,
    openai_embeddings,
    collection=collection,
    index_name=INDEX_NAME,
)

# 设置索引创建参数
num_lists = 100
dimensions = 1536
similarity_algorithm = CosmosDBSimilarityType.COS
kind = CosmosDBVectorSearchType.VECTOR_IVF
m = 16
ef_construction = 64
ef_search = 40

# 创建索引
vectorstore.create_index(num_lists, dimensions, similarity_algorithm, kind, m, ef_construction)

# 执行相似性搜索
query = "What did the president say about Ketanji Brown Jackson"
docs = vectorstore.similarity_search(query)

# 输出结果
print(docs[0].page_content)