cosmos db entra id support and fixes

Author: TheovanKraay

auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure-cosmos-db/pom.xml

@@ -45,6 +45,11 @@
 			<version>${project.parent.version}</version>
 			<optional>true</optional>
 		</dependency>
+		<dependency>
+			<groupId>com.azure</groupId>
+			<artifactId>azure-identity</artifactId>
+			<version>1.15.4</version> <!-- or the latest version -->
+		</dependency>
 		<dependency>
 			<groupId>org.springframework.boot</groupId>
 			<artifactId>spring-boot-starter</artifactId>

auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure-cosmos-db/src/main/java/org/springframework/ai/vectorstore/cosmosdb/autoconfigure/CosmosDBVectorStoreAutoConfiguration.java

@@ -18,6 +18,7 @@
 
 import com.azure.cosmos.CosmosAsyncClient;
 import com.azure.cosmos.CosmosClientBuilder;
+import com.azure.identity.DefaultAzureCredentialBuilder;
 import io.micrometer.observation.ObservationRegistry;
 
 import org.springframework.ai.embedding.BatchingStrategy;
@@ -33,6 +34,7 @@
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.context.annotation.Bean;
+
 import java.util.List;
 
 /**
@@ -43,19 +45,37 @@
  * @author Soby Chacko
  * @since 1.0.0
  */
+
 @AutoConfiguration
 @ConditionalOnClass({ CosmosDBVectorStore.class, EmbeddingModel.class, CosmosAsyncClient.class })
 @EnableConfigurationProperties(CosmosDBVectorStoreProperties.class)
 @ConditionalOnProperty(name = SpringAIVectorStoreTypes.TYPE, havingValue = SpringAIVectorStoreTypes.AZURE_COSMOS_DB,
 		matchIfMissing = true)
 public class CosmosDBVectorStoreAutoConfiguration {
 
+	private final String agentSuffix = "SpringAI-CDBNoSQL-VectorStore";
+
 	@Bean
 	public CosmosAsyncClient cosmosClient(CosmosDBVectorStoreProperties properties) {
-		return new CosmosClientBuilder().endpoint(properties.getEndpoint())
-			.userAgentSuffix("SpringAI-CDBNoSQL-VectorStore")
-			.key(properties.getKey())
-			.gatewayMode()
+		String mode = properties.getConnectionMode();
+		if (mode == null) {
+			properties.setConnectionMode("gateway");
+		}
+		else if (!mode.equals("direct") && !mode.equals("gateway")) {
+			throw new IllegalArgumentException("Connection mode must be either 'direct' or 'gateway'");
+		}
+
+		CosmosClientBuilder builder = new CosmosClientBuilder().endpoint(properties.getEndpoint())
+			.userAgentSuffix(agentSuffix);
+
+		if (properties.getKey() == null || properties.getKey().isEmpty()) {
+			builder.credential(new DefaultAzureCredentialBuilder().build());
+		}
+		else {
+			builder.key(properties.getKey());
+		}
+
+		return ("direct".equals(properties.getConnectionMode()) ? builder.directMode() : builder.gatewayMode())
 			.buildAsyncClient();
 	}
 
@@ -75,12 +95,11 @@ public CosmosDBVectorStore cosmosDBVectorStore(ObservationRegistry observationRe
 		return CosmosDBVectorStore.builder(cosmosAsyncClient, embeddingModel)
 			.databaseName(properties.getDatabaseName())
 			.containerName(properties.getContainerName())
-			.metadataFields(List.of(properties.getMetadataFields()))
+			.metadataFields(properties.getMetadataFieldList())
 			.vectorStoreThroughput(properties.getVectorStoreThroughput())
 			.vectorDimensions(properties.getVectorDimensions())
 			.partitionKeyPath(properties.getPartitionKeyPath())
 			.build();
-
 	}
 
 }

auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure-cosmos-db/src/main/java/org/springframework/ai/vectorstore/cosmosdb/autoconfigure/CosmosDBVectorStoreProperties.java

@@ -19,13 +19,15 @@
 import org.springframework.ai.vectorstore.properties.CommonVectorStoreProperties;
 import org.springframework.boot.context.properties.ConfigurationProperties;
 
+import java.util.Arrays;
+import java.util.List;
+
 /**
  * Configuration properties for CosmosDB Vector Store.
  *
  * @author Theo van Kraay
  * @since 1.0.0
  */
-
 @ConfigurationProperties(CosmosDBVectorStoreProperties.CONFIG_PREFIX)
 public class CosmosDBVectorStoreProperties extends CommonVectorStoreProperties {
 
@@ -47,6 +49,8 @@ public class CosmosDBVectorStoreProperties extends CommonVectorStoreProperties {
 
 	private String key;
 
+	private String connectionMode;
+
 	public int getVectorStoreThroughput() {
 		return this.vectorStoreThroughput;
 	}
@@ -63,6 +67,12 @@ public void setMetadataFields(String metadataFields) {
 		this.metadataFields = metadataFields;
 	}
 
+	public List<String> getMetadataFieldList() {
+		return this.metadataFields != null
+				? Arrays.stream(this.metadataFields.split(",")).map(String::trim).filter(s -> !s.isEmpty()).toList()
+				: List.of();
+	}
+
 	public String getEndpoint() {
 		return this.endpoint;
 	}
@@ -79,6 +89,14 @@ public void setKey(String key) {
 		this.key = key;
 	}
 
+	public void setConnectionMode(String connectionMode) {
+		this.connectionMode = connectionMode;
+	}
+
+	public String getConnectionMode() {
+		return this.connectionMode;
+	}
+
 	public String getDatabaseName() {
 		return this.databaseName;
 	}

auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-azure-cosmos-db/src/test/java/org/springframework/ai/vectorstore/cosmosdb/autoconfigure/CosmosDBVectorStoreAutoConfigurationIT.java

@@ -44,22 +44,35 @@
  * @author Theo van Kraay
  * @since 1.0.0
  */
-
 @EnabledIfEnvironmentVariable(named = "AZURE_COSMOSDB_ENDPOINT", matches = ".+")
 @EnabledIfEnvironmentVariable(named = "AZURE_COSMOSDB_KEY", matches = ".+")
 public class CosmosDBVectorStoreAutoConfigurationIT {
 
-	private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
-		.withConfiguration(AutoConfigurations.of(CosmosDBVectorStoreAutoConfiguration.class))
-		.withPropertyValues("spring.ai.vectorstore.cosmosdb.databaseName=test-database")
-		.withPropertyValues("spring.ai.vectorstore.cosmosdb.containerName=test-container")
-		.withPropertyValues("spring.ai.vectorstore.cosmosdb.partitionKeyPath=/id")
-		.withPropertyValues("spring.ai.vectorstore.cosmosdb.metadataFields=country,year,city")
-		.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorStoreThroughput=1000")
-		.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorDimensions=384")
-		.withPropertyValues("spring.ai.vectorstore.cosmosdb.endpoint=" + System.getenv("AZURE_COSMOSDB_ENDPOINT"))
-		.withPropertyValues("spring.ai.vectorstore.cosmosdb.key=" + System.getenv("AZURE_COSMOSDB_KEY"))
-		.withUserConfiguration(Config.class);
+	private final ApplicationContextRunner contextRunner;
+
+	public CosmosDBVectorStoreAutoConfigurationIT() {
+		String endpoint = System.getenv("AZURE_COSMOSDB_ENDPOINT");
+		String key = System.getenv("AZURE_COSMOSDB_KEY");
+
+		ApplicationContextRunner contextRunner = new ApplicationContextRunner()
+			.withConfiguration(AutoConfigurations.of(CosmosDBVectorStoreAutoConfiguration.class))
+			.withPropertyValues("spring.ai.vectorstore.cosmosdb.databaseName=test-database")
+			.withPropertyValues("spring.ai.vectorstore.cosmosdb.containerName=test-container")
+			.withPropertyValues("spring.ai.vectorstore.cosmosdb.partitionKeyPath=/id")
+			.withPropertyValues("spring.ai.vectorstore.cosmosdb.metadataFields=country,year,city")
+			.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorStoreThroughput=1000")
+			.withPropertyValues("spring.ai.vectorstore.cosmosdb.vectorDimensions=384");
+
+		if (endpoint != null && !"null".equalsIgnoreCase(endpoint)) {
+			contextRunner = contextRunner.withPropertyValues("spring.ai.vectorstore.cosmosdb.endpoint=" + endpoint);
+		}
+
+		if (key != null && !"null".equalsIgnoreCase(key)) {
+			contextRunner = contextRunner.withPropertyValues("spring.ai.vectorstore.cosmosdb.key=" + key);
+		}
+
+		this.contextRunner = contextRunner.withUserConfiguration(Config.class);
+	}
 
 	private VectorStore vectorStore;
 
@@ -124,14 +137,15 @@ void testSimilaritySearchWithFilter() {
 		metadata4.put("country", "US");
 		metadata4.put("year", 2020);
 		metadata4.put("city", "Sofia");
-
 		Document document1 = new Document("1", "A document about the UK", metadata1);
 		Document document2 = new Document("2", "A document about the Netherlands", metadata2);
 		Document document3 = new Document("3", "A document about the US", metadata3);
 		Document document4 = new Document("4", "A document about the US", metadata4);
 
 		this.vectorStore.add(List.of(document1, document2, document3, document4));
+
 		FilterExpressionBuilder b = new FilterExpressionBuilder();
+
 		List<Document> results = this.vectorStore.similaritySearch(SearchRequest.builder()
 			.query("The World")
 			.topK(10)
@@ -190,7 +204,7 @@ public void autoConfigurationEnabledByDefault() {
 
 	@Test
 	public void autoConfigurationEnabledWhenTypeIsAzureCosmosDB() {
-		this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=azure-cosmmos-db").run(context -> {
+		this.contextRunner.withPropertyValues("spring.ai.vectorstore.type=azure-cosmos-db").run(context -> {
 			assertThat(context.getBeansOfType(CosmosDBVectorStoreProperties.class)).isNotEmpty();
 			assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
 			assertThat(context.getBean(VectorStore.class)).isInstanceOf(CosmosDBVectorStore.class);

vector-stores/spring-ai-azure-cosmos-db-store/pom.xml

@@ -47,6 +47,11 @@
 			<artifactId>azure-spring-data-cosmos</artifactId>
 			<version>${azure-cosmos.version}</version>
 		</dependency>
+		<dependency>
+			<groupId>com.azure</groupId>
+			<artifactId>azure-identity</artifactId>
+			<version>1.15.4</version> <!-- or the latest version -->
+		</dependency>
 		<dependency>
 			<groupId>org.springframework.ai</groupId>
 			<artifactId>spring-ai-vector-store</artifactId>