GH-2137: Add getNativeClient API to VectorStore interface

Fixes: #2137

Adds getNativeClient API to VectorStore interface allowing access to the underlying native client implementation.

This change:
- Adds getNativeClient() default method to VectorStore interface returning Optional<T>
- Implements getNativeClient() in all vector store implementations exposing their respective native clients
- Adds integration tests verifying native client access for all implementations

Signed-off-by: Soby Chacko <[email protected]>

Author: sobychacko

spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java

@@ -108,6 +108,24 @@ default List<Document> similaritySearch(String query) {
 		return this.similaritySearch(SearchRequest.builder().query(query).build());
 	}
 
+	/**
+	 * Returns the native client if available in this vector store implementation.
+	 *
+	 * Note on usage: 1. Returns empty Optional when no native client is available 2. Due
+	 * to Java type erasure, runtime type checking is not possible
+	 *
+	 * Example usage: When working with implementation with known native client:
+	 * Optional<NativeClientType> client = vectorStore.getNativeClient();
+	 *
+	 * Note: Using Optional<?> will return the native client if one exists, rather than an
+	 * empty Optional. For type safety, prefer using the specific client type.
+	 * @return Optional containing native client if available, empty Optional otherwise
+	 * @param <T> The type of the native client
+	 */
+	default <T> Optional<T> getNativeClient() {
+		return Optional.empty();
+	}
+
 	/**
 	 * Builder interface for creating VectorStore instances. Implements a fluent builder
 	 * pattern for configuring observation-related settings.

vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStore.java

@@ -372,6 +372,13 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str
 			.similarityMetric("cosine");
 	}
 
+	@Override
+	public <T> Optional<T> getNativeClient() {
+		@SuppressWarnings("unchecked")
+		T client = (T) this.container;
+		return Optional.of(client);
+	}
+
 	/**
 	 * Builder class for creating {@link CosmosDBVectorStore} instances.
 	 * <p>

vector-stores/spring-ai-azure-cosmos-db-store/src/test/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStoreIT.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 the original author or authors.
+ * Copyright 2023-2025 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -19,9 +19,11 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.UUID;
 
 import com.azure.cosmos.CosmosAsyncClient;
+import com.azure.cosmos.CosmosAsyncContainer;
 import com.azure.cosmos.CosmosClientBuilder;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -169,6 +171,15 @@ void testSimilaritySearchWithFilter() {
 		assertThat(results4).isEmpty();
 	}
 
+	@Test
+	void getNativeClientTest() {
+		this.contextRunner.run(context -> {
+			CosmosDBVectorStore vectorStore = context.getBean(CosmosDBVectorStore.class);
+			Optional<CosmosAsyncContainer> nativeClient = vectorStore.getNativeClient();
+			assertThat(nativeClient).isPresent();
+		});
+	}
+
 	@SpringBootConfiguration
 	@EnableAutoConfiguration
 	public static class TestApplication {

vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java

@@ -332,6 +332,13 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str
 			.similarityMetric(this.initializeSchema ? VectorStoreSimilarityMetric.COSINE.value() : null);
 	}
 
+	@Override
+	public <T> Optional<T> getNativeClient() {
+		@SuppressWarnings("unchecked")
+		T client = (T) this.searchClient;
+		return Optional.of(client);
+	}
+
 	public record MetadataField(String name, SearchFieldDataType fieldType) {
 
 		public static MetadataField text(String name) {

vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreIT.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2023-2024 the original author or authors.
+ * Copyright 2023-2025 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -23,10 +23,12 @@
 import java.util.Date;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
 
 import com.azure.core.credential.AzureKeyCredential;
+import com.azure.search.documents.SearchClient;
 import com.azure.search.documents.indexes.SearchIndexClient;
 import com.azure.search.documents.indexes.SearchIndexClientBuilder;
 import org.awaitility.Awaitility;
@@ -318,6 +320,15 @@ public void searchThresholdTest() {
 		});
 	}
 
+	@Test
+	void getNativeClientTest() {
+		this.contextRunner.run(context -> {
+			AzureVectorStore vectorStore = context.getBean(AzureVectorStore.class);
+			Optional<SearchClient> nativeClient = vectorStore.getNativeClient();
+			assertThat(nativeClient).isPresent();
+		});
+	}
+
 	@SpringBootConfiguration
 	@EnableAutoConfiguration
 	public static class Config {

vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStore.java

@@ -712,6 +712,13 @@ private void ensureTableColumnsExist(int vectorDimension) {
 		}
 	}
 
+	@Override
+	public <T> Optional<T> getNativeClient() {
+		@SuppressWarnings("unchecked")
+		T client = (T) this.session;
+		return Optional.of(client);
+	}
+
 	/**
 	 * Indexes are automatically created with COSINE. This can be changed manually via
 	 * cqlsh

vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java

@@ -21,6 +21,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.UUID;
 import java.util.stream.Collectors;
 
@@ -524,6 +525,15 @@ void deleteWithComplexFilterExpression() {
 		});
 	}
 
+	@Test
+	void getNativeClientTest() {
+		this.contextRunner.run(context -> {
+			CassandraVectorStore vectorStore = context.getBean(CassandraVectorStore.class);
+			Optional<CqlSession> nativeClient = vectorStore.getNativeClient();
+			assertThat(nativeClient).isPresent();
+		});
+	}
+
 	@SpringBootConfiguration
 	@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
 	public static class TestApplication {

vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStore.java

@@ -275,6 +275,13 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str
 			.dimensions(this.embeddingModel.dimensions());
 	}
 
+	@Override
+	public <T> Optional<T> getNativeClient() {
+		@SuppressWarnings("unchecked")
+		T client = (T) this.session;
+		return Optional.of(client);
+	}
+
 	/**
 	 * Builder class for creating {@link CoherenceVectorStore} instances.
 	 * <p>

vector-stores/spring-ai-coherence-store/src/test/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStoreIT.java

@@ -23,6 +23,7 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.UUID;
 import java.util.stream.Stream;
 
@@ -290,6 +291,15 @@ public void searchWithThreshold() {
 		});
 	}
 
+	@Test
+	void getNativeClientTest() {
+		this.contextRunner.run(context -> {
+			CoherenceVectorStore vectorStore = context.getBean(CoherenceVectorStore.class);
+			Optional<Session> nativeClient = vectorStore.getNativeClient();
+			assertThat(nativeClient).isPresent();
+		});
+	}
+
 	private static boolean isSortedByDistance(final List<Document> documents) {
 		final List<Double> distances = documents.stream()
 			.map(doc -> (Double) doc.getMetadata().get(DocumentMetadata.DISTANCE.value()))

vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java

@@ -354,6 +354,13 @@ private String getSimilarityMetric() {
 		return SIMILARITY_TYPE_MAPPING.get(this.options.getSimilarity()).value();
 	}
 
+	@Override
+	public <T> Optional<T> getNativeClient() {
+		@SuppressWarnings("unchecked")
+		T client = (T) this.elasticsearchClient;
+		return Optional.of(client);
+	}
+
 	/**
 	 * Creates a new builder instance for ElasticsearchVectorStore.
 	 * @return a new ElasticsearchBuilder instance