Azure Cosmos DB
本节将指导您如何设置 CosmosDBVectorStore 来存储文档嵌入向量并执行相似性搜索。
什么是 Azure Cosmos DB?
Azure Cosmos DB 是微软面向关键任务应用程序设计的全球分布式云原生数据库服务。它提供高可用性、低延迟以及水平扩展能力,以满足现代应用程序的需求。其核心设计理念自底向上就融入了全球分布、细粒度多租户和水平可扩展性。它是 Azure 中的一项基础服务,被微软大多数全球范围的关键任务应用程序所使用,包括 Teams、Skype、Xbox Live、Office 365、Bing、Azure Active Directory、Azure 门户、Microsoft Store 等。它也被数千家外部客户使用,包括 OpenAI 的 ChatGPT 以及其他需要弹性扩展、开箱即用的全球分布、全球范围内的低延迟和高可用性的关键任务 AI 应用程序。
什么是 DiskANN?
DiskANN(基于磁盘的近似最近邻搜索)是一项创新技术,应用于Azure Cosmos DB中,旨在提升向量搜索的性能。它通过索引存储在Cosmos DB中的嵌入向量,实现对高维数据的高效且可扩展的相似性搜索。
DiskANN 具有以下优势:
-
高效性:通过利用基于磁盘的数据结构,DiskANN与传统方法相比,显著缩短了查找最近邻所需的时间。
-
可扩展性:它能处理超出内存容量的大型数据集,适用于包括机器学习和人工智能驱动解决方案在内的多种应用场景。
-
低延迟:DiskANN 在搜索操作期间最大限度地减少了延迟,确保即使在处理大量数据时,应用程序也能快速检索结果。
在Spring AI for Azure Cosmos DB的场景中,向量搜索将创建并利用DiskANN索引,以确保相似性查询的最佳性能。
使用自动配置设置 Azure Cosmos DB 向量存储
以下代码演示了如何通过自动配置设置 CosmosDBVectorStore:
package com.example.demo;
import io.micrometer.observation.ObservationRegistry;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
import org.springframework.boot.CommandLineRunner;
import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Lazy;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import static org.assertj.core.api.Assertions.assertThat;
@SpringBootApplication
@EnableAutoConfiguration
public class DemoApplication implements CommandLineRunner {
private static final Logger log = LoggerFactory.getLogger(DemoApplication.class);
@Lazy
@Autowired
private VectorStore vectorStore;
public static void main(String[] args) {
SpringApplication.run(DemoApplication.class, args);
}
@Override
public void run(String... args) throws Exception {
Document document1 = new Document(UUID.randomUUID().toString(), "Sample content1", Map.of("key1", "value1"));
Document document2 = new Document(UUID.randomUUID().toString(), "Sample content2", Map.of("key2", "value2"));
this.vectorStore.add(List.of(document1, document2));
List<Document> results = this.vectorStore.similaritySearch(SearchRequest.builder().query("Sample content").topK(1).build());
log.info("Search results: {}", results);
// Remove the documents from the vector store
this.vectorStore.delete(List.of(document1.getId(), document2.getId()));
}
@Bean
public ObservationRegistry observationRegistry() {
return ObservationRegistry.create();
}
}
自动配置
Spring AI 的自动配置和 starter 模块的 artifact 名称发生了重大变化。更多信息请参阅升级说明。
将以下依赖添加到你的 Maven 项目中:
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-starter-vector-store-azure-cosmos-db</artifactId>
</dependency>
配置属性
以下是 Cosmos DB 向量存储可用的配置属性:
| 属性 | 描述 |
|---|---|
| spring.ai.vectorstore.cosmosdb.databaseName | 要使用的 Cosmos DB 数据库名称。 |
| spring.ai.vectorstore.cosmosdb.containerName | 要使用的 Cosmos DB 容器名称。 |
| spring.ai.vectorstore.cosmosdb.partitionKeyPath | 分区键的路径。 |
| spring.ai.vectorstore.cosmosdb.metadataFields | 元数据字段的逗号分隔列表。 |
| spring.ai.vectorstore.cosmosdb.vectorStoreThroughput | 向量存储的吞吐量。 |
| spring.ai.vectorstore.cosmosdb.vectorDimensions | 向量的维度数。 |
| spring.ai.vectorstore.cosmosdb.endpoint | Cosmos DB 的端点。 |
| spring.ai.vectorstore.cosmosdb.key | Cosmos DB 的密钥(如果密钥不存在,将使用 [DefaultAzureCredential](learn.microsoft.com/azure/developer/java/sdk/authentication/credential-chains#defaultazurecredential-overview))。 |
复杂搜索与过滤器
您可以使用 Cosmos DB 向量存储中的过滤器执行更复杂的搜索。以下示例展示了如何在搜索查询中使用过滤器。
Map<String, Object> metadata1 = new HashMap<>();
metadata1.put("country", "UK");
metadata1.put("year", 2021);
metadata1.put("city", "London");
Map<String, Object> metadata2 = new HashMap<>();
metadata2.put("country", "NL");
metadata2.put("year", 2022);
metadata2.put("city", "Amsterdam");
Document document1 = new Document("1", "A document about the UK", this.metadata1);
Document document2 = new Document("2", "A document about the Netherlands", this.metadata2);
vectorStore.add(List.of(document1, document2));
FilterExpressionBuilder builder = new FilterExpressionBuilder();
List<Document> results = vectorStore.similaritySearch(SearchRequest.builder().query("The World")
.topK(10)
.filterExpression((this.builder.in("country", "UK", "NL")).build()).build());
在不使用自动配置的情况下设置 Azure Cosmos DB 向量存储
以下代码演示了如何在不依赖自动配置的情况下设置 CosmosDBVectorStore。建议使用 DefaultAzureCredential 对 Azure Cosmos DB 进行身份验证。
@Bean
public VectorStore vectorStore(ObservationRegistry observationRegistry) {
// Create the Cosmos DB client
CosmosAsyncClient cosmosClient = new CosmosClientBuilder()
.endpoint(System.getenv("COSMOSDB_AI_ENDPOINT"))
.credential(new DefaultAzureCredentialBuilder().build())
.userAgentSuffix("SpringAI-CDBNoSQL-VectorStore")
.gatewayMode()
.buildAsyncClient();
// Create and configure the vector store
return CosmosDBVectorStore.builder(cosmosClient, embeddingModel)
.databaseName("test-database")
.containerName("test-container")
// Configure metadata fields for filtering
.metadataFields(List.of("country", "year", "city"))
// Set the partition key path (optional)
.partitionKeyPath("/id")
// Configure performance settings
.vectorStoreThroughput(1000)
.vectorDimensions(1536) // Match your embedding model's dimensions
// Add custom batching strategy (optional)
.batchingStrategy(new TokenCountBatchingStrategy())
// Add observation registry for metrics
.observationRegistry(observationRegistry)
.build();
}
@Bean
public EmbeddingModel embeddingModel() {
return new TransformersEmbeddingModel();
}
此配置展示了所有可用的构建器选项:
-
databaseName:您的 Cosmos DB 数据库名称 -
containerName:数据库中容器的名称 -
partitionKeyPath:分区键的路径(例如:"/id") -
metadataFields:用于筛选的元数据字段列表 -
vectorStoreThroughput:向量存储容器的吞吐量(RU/s) -
vectorDimensions:向量的维度数量(应与您的嵌入模型匹配) -
batchingStrategy:文档操作的分批策略(可选)
手动依赖设置
在你的Maven项目中添加以下依赖:
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-azure-cosmos-db-store</artifactId>
</dependency>
访问原生客户端
Azure Cosmos DB 向量存储实现通过 getNativeClient() 方法提供对底层原生 Azure Cosmos DB 客户端 (CosmosClient) 的访问:
CosmosDBVectorStore vectorStore = context.getBean(CosmosDBVectorStore.class);
Optional<CosmosClient> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
CosmosClient client = nativeClient.get();
// Use the native client for Azure Cosmos DB-specific operations
}
原生客户端让您能够访问 Azure Cosmos DB 特有的功能和操作,这些可能无法通过 VectorStore 接口直接暴露。