From 21c2773efea2a627ca81ca9bd34d6efddb448453 Mon Sep 17 00:00:00 2001 From: Soby Chacko Date: Sun, 2 Feb 2025 15:20:33 -0500 Subject: [PATCH] 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 - 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 --- .../ai/vectorstore/VectorStore.java | 18 ++++++++++++++++ .../cosmosdb/CosmosDBVectorStore.java | 7 +++++++ .../cosmosdb/CosmosDBVectorStoreIT.java | 13 +++++++++++- .../vectorstore/azure/AzureVectorStore.java | 7 +++++++ .../vectorstore/azure/AzureVectorStoreIT.java | 13 +++++++++++- .../cassandra/CassandraVectorStore.java | 7 +++++++ .../cassandra/CassandraVectorStoreIT.java | 10 +++++++++ .../coherence/CoherenceVectorStore.java | 7 +++++++ .../coherence/CoherenceVectorStoreIT.java | 10 +++++++++ .../ElasticsearchVectorStore.java | 7 +++++++ .../ElasticsearchVectorStoreIT.java | 21 +++++++++++++++++++ .../mariadb/MariaDBVectorStore.java | 7 +++++++ .../vectorstore/mariadb/MariaDBStoreIT.java | 10 +++++++++ .../vectorstore/milvus/MilvusVectorStore.java | 7 +++++++ .../milvus/MilvusVectorStoreIT.java | 10 +++++++++ .../atlas/MongoDBAtlasVectorStore.java | 7 +++++++ .../atlas/MongoDBAtlasVectorStoreIT.java | 10 +++++++++ .../vectorstore/neo4j/Neo4jVectorStore.java | 7 +++++++ .../vectorstore/neo4j/Neo4jVectorStoreIT.java | 10 +++++++++ .../opensearch/OpenSearchVectorStore.java | 7 +++++++ .../opensearch/OpenSearchVectorStoreIT.java | 10 +++++++++ .../vectorstore/oracle/OracleVectorStore.java | 7 +++++++ .../oracle/OracleVectorStoreIT.java | 13 ++++++++++++ .../vectorstore/pgvector/PgVectorStore.java | 7 +++++++ .../vectorstore/pgvector/PgVectorStoreIT.java | 9 ++++++++ .../pinecone/PineconeVectorStore.java | 7 +++++++ .../pinecone/PineconeVectorStoreIT.java | 11 ++++++++++ .../vectorstore/qdrant/QdrantVectorStore.java | 7 +++++++ .../qdrant/QdrantVectorStoreIT.java | 10 +++++++++ .../vectorstore/redis/RedisVectorStore.java | 7 +++++++ .../vectorstore/redis/RedisVectorStoreIT.java | 10 +++++++++ .../typesense/TypesenseVectorStore.java | 7 +++++++ .../typesense/TypesenseVectorStoreIT.java | 10 +++++++++ .../weaviate/WeaviateVectorStore.java | 7 +++++++ .../weaviate/WeaviateVectorStoreIT.java | 10 +++++++++ 35 files changed, 325 insertions(+), 2 deletions(-) diff --git a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java index 2b1bc863d5d..814efc2eaaa 100644 --- a/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java +++ b/spring-ai-core/src/main/java/org/springframework/ai/vectorstore/VectorStore.java @@ -108,6 +108,24 @@ default List 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 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 The type of the native client + */ + default Optional getNativeClient() { + return Optional.empty(); + } + /** * Builder interface for creating VectorStore instances. Implements a fluent builder * pattern for configuring observation-related settings. diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStore.java b/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStore.java index 45f8e8d01bf..5763728cead 100644 --- a/vector-stores/spring-ai-azure-cosmos-db-store/src/main/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStore.java +++ b/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 Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.container; + return Optional.of(client); + } + /** * Builder class for creating {@link CosmosDBVectorStore} instances. *

diff --git a/vector-stores/spring-ai-azure-cosmos-db-store/src/test/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStoreIT.java b/vector-stores/spring-ai-azure-cosmos-db-store/src/test/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStoreIT.java index cecd3869454..39729c6806e 100644 --- a/vector-stores/spring-ai-azure-cosmos-db-store/src/test/java/org/springframework/ai/vectorstore/cosmosdb/CosmosDBVectorStoreIT.java +++ b/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 nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration public static class TestApplication { diff --git a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java b/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java index e7dfb531d1a..eff32d5db99 100644 --- a/vector-stores/spring-ai-azure-store/src/main/java/org/springframework/ai/vectorstore/azure/AzureVectorStore.java +++ b/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 Optional 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) { diff --git a/vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreIT.java b/vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreIT.java index 862a52979d7..1e4bd4aea04 100644 --- a/vector-stores/spring-ai-azure-store/src/test/java/org/springframework/ai/vectorstore/azure/AzureVectorStoreIT.java +++ b/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 nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration public static class Config { diff --git a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStore.java b/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStore.java index 9575ef5bcbf..6e9b1610cc7 100644 --- a/vector-stores/spring-ai-cassandra-store/src/main/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStore.java +++ b/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 Optional 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 diff --git a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java b/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java index 737131950be..e30ba900d44 100644 --- a/vector-stores/spring-ai-cassandra-store/src/test/java/org/springframework/ai/vectorstore/cassandra/CassandraVectorStoreIT.java +++ b/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 nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestApplication { diff --git a/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStore.java b/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStore.java index 928425299a4..32fdd60fcf7 100644 --- a/vector-stores/spring-ai-coherence-store/src/main/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStore.java +++ b/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 Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.session; + return Optional.of(client); + } + /** * Builder class for creating {@link CoherenceVectorStore} instances. *

diff --git a/vector-stores/spring-ai-coherence-store/src/test/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStoreIT.java b/vector-stores/spring-ai-coherence-store/src/test/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStoreIT.java index 32739a36152..526fb361a01 100644 --- a/vector-stores/spring-ai-coherence-store/src/test/java/org/springframework/ai/vectorstore/coherence/CoherenceVectorStoreIT.java +++ b/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 nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + private static boolean isSortedByDistance(final List documents) { final List distances = documents.stream() .map(doc -> (Double) doc.getMetadata().get(DocumentMetadata.DISTANCE.value())) diff --git a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java b/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java index 4d2a1924209..71958f72b83 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/main/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStore.java +++ b/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 Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.elasticsearchClient; + return Optional.of(client); + } + /** * Creates a new builder instance for ElasticsearchVectorStore. * @return a new ElasticsearchBuilder instance diff --git a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java index 5c8b74d6dd6..1d8a2521d1f 100644 --- a/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java +++ b/vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java @@ -24,6 +24,7 @@ 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; @@ -529,6 +530,26 @@ public void overDefaultSizeTest() { }); } + @Test + public void getNativeClientTest() { + getContextRunner().run(context -> { + ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine", + ElasticsearchVectorStore.class); + + // Test successful native client retrieval + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + + // Verify client functionality + ElasticsearchClient client = nativeClient.get(); + IndicesStats stats = client.indices() + .stats(s -> s.index("spring-ai-document-index")) + .indices() + .get("spring-ai-document-index"); + assertThat(stats).isNotNull(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestApplication { diff --git a/vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBVectorStore.java b/vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBVectorStore.java index 4efc6554f91..c04e3dcd80e 100644 --- a/vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBVectorStore.java +++ b/vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBVectorStore.java @@ -460,6 +460,13 @@ private String getSimilarityMetric() { return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value(); } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.jdbcTemplate; + return Optional.of(client); + } + public enum MariaDBDistanceType { EUCLIDEAN, COSINE diff --git a/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBStoreIT.java b/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBStoreIT.java index b2b4109b1bc..315254e4b4d 100644 --- a/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBStoreIT.java +++ b/vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBStoreIT.java @@ -22,6 +22,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.Collectors; import java.util.stream.Stream; @@ -461,6 +462,15 @@ public void deleteWithComplexFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner.run(context -> { + MariaDBVectorStore vectorStore = context.getBean(MariaDBVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestApplication { diff --git a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStore.java b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStore.java index a430a951997..8da7c515ce4 100644 --- a/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStore.java +++ b/vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStore.java @@ -559,6 +559,13 @@ private String getSimilarityMetric() { return SIMILARITY_TYPE_MAPPING.get(this.metricType).value(); } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.milvusClient; + return Optional.of(client); + } + public static class Builder extends AbstractVectorStoreBuilder { private final MilvusServiceClient milvusClient; diff --git a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.java b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.java index ddad8ff4424..956e13819c5 100644 --- a/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.java +++ b/vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.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; @@ -378,6 +379,15 @@ public void deleteWithComplexFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=COSINE").run(context -> { + MilvusVectorStore vectorStore = context.getBean(MilvusVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestApplication { diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStore.java b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStore.java index e03294b4f29..cea8dafdeaa 100644 --- a/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStore.java +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStore.java @@ -330,6 +330,13 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str .fieldName(this.pathName); } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.mongoTemplate; + return Optional.of(client); + } + /** * Creates a new builder instance for MongoDBAtlasVectorStore. * @return a new MongoDBBuilder instance diff --git a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java index cdadd75b0bf..3e98bae372d 100644 --- a/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java +++ b/vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java @@ -22,6 +22,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; @@ -355,6 +356,15 @@ void deleteWithComplexFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner.run(context -> { + MongoDBAtlasVectorStore vectorStore = context.getBean(MongoDBAtlasVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + public static String getText(String uri) { var resource = new DefaultResourceLoader().getResource(uri); try { diff --git a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java index 6efec2d2b2f..f9146f95a6a 100644 --- a/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java +++ b/vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java @@ -368,6 +368,13 @@ private String getSimilarityMetric() { return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value(); } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.driver; + return Optional.of(client); + } + /** * An enum to configure the distance function used in the Neo4j vector index. */ diff --git a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java index 3d74e561920..29974c8009d 100644 --- a/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java +++ b/vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java @@ -19,6 +19,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; @@ -397,6 +398,15 @@ void deleteWithComplexFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner.run(context -> { + Neo4jVectorStore vectorStore = context.getBean(Neo4jVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestApplication { diff --git a/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStore.java b/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStore.java index 360a8c046c9..67ebd4bbc06 100644 --- a/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStore.java +++ b/vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStore.java @@ -380,6 +380,13 @@ else if ("l2".equalsIgnoreCase(this.similarityFunction)) { return this.similarityFunction; } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.openSearchClient; + return Optional.of(client); + } + /** * The representation of {@link Document} along with its embedding. * diff --git a/vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java b/vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java index f0ad4aedc28..bdb68b0887b 100644 --- a/vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java +++ b/vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java @@ -25,6 +25,7 @@ import java.util.Date; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.TimeZone; import java.util.UUID; import java.util.concurrent.TimeUnit; @@ -524,6 +525,15 @@ void deleteWithComplexFilterExpression() { }); } + @Test + void getNativeClientTest() { + getContextRunner().run(context -> { + OpenSearchVectorStore vectorStore = context.getBean("vectorStore", OpenSearchVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestApplication { diff --git a/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/oracle/OracleVectorStore.java b/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/oracle/OracleVectorStore.java index d624ef809e4..08c8604d96c 100644 --- a/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/oracle/OracleVectorStore.java +++ b/vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/oracle/OracleVectorStore.java @@ -530,6 +530,13 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str .similarityMetric(getSimilarityMetric()); } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.jdbcTemplate; + return Optional.of(client); + } + private String getSimilarityMetric() { if (!SIMILARITY_TYPE_MAPPING.containsKey(this.distanceType)) { return this.distanceType.name(); diff --git a/vector-stores/spring-ai-oracle-store/src/test/java/org/springframework/ai/vectorstore/oracle/OracleVectorStoreIT.java b/vector-stores/spring-ai-oracle-store/src/test/java/org/springframework/ai/vectorstore/oracle/OracleVectorStoreIT.java index e4cb92aabfe..aaa7fb9a98d 100644 --- a/vector-stores/spring-ai-oracle-store/src/test/java/org/springframework/ai/vectorstore/oracle/OracleVectorStoreIT.java +++ b/vector-stores/spring-ai-oracle-store/src/test/java/org/springframework/ai/vectorstore/oracle/OracleVectorStoreIT.java @@ -22,6 +22,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.Collectors; @@ -420,6 +421,18 @@ void deleteWithComplexFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner + .withPropertyValues("test.spring.ai.vectorstore.oracle.distanceType=COSINE", + "test.spring.ai.vectorstore.oracle.searchAccuracy=" + OracleVectorStore.DEFAULT_SEARCH_ACCURACY) + .run(context -> { + OracleVectorStore vectorStore = context.getBean(OracleVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestClient { diff --git a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java index 00a32457147..6fa057258b7 100644 --- a/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java +++ b/vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java @@ -476,6 +476,13 @@ private String getSimilarityMetric() { return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value(); } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.jdbcTemplate; + return Optional.of(client); + } + /** * By default, pgvector performs exact nearest neighbor search, which provides perfect * recall. You can add an index to use approximate nearest neighbor search, which diff --git a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java index 4cd00d3b800..59d6ed50c27 100644 --- a/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java +++ b/vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java @@ -482,6 +482,15 @@ public void deleteWithStringFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner.run(context -> { + PgVectorStore vectorStore = context.getBean(PgVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestApplication { diff --git a/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStore.java b/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStore.java index d8a494786ef..596b4803b09 100644 --- a/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStore.java +++ b/vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStore.java @@ -374,6 +374,13 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str .fieldName(this.pineconeContentFieldName); } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.pineconeConnection; + return Optional.of(client); + } + /** * Builder class for creating {@link PineconeVectorStore} instances. This implements a * type-safe step builder pattern to ensure all required fields are provided in a diff --git a/vector-stores/spring-ai-pinecone-store/src/test/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStoreIT.java b/vector-stores/spring-ai-pinecone-store/src/test/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStoreIT.java index 32f934a9539..1e183ecc417 100644 --- a/vector-stores/spring-ai-pinecone-store/src/test/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStoreIT.java +++ b/vector-stores/spring-ai-pinecone-store/src/test/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStoreIT.java @@ -21,10 +21,12 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.UUID; import java.util.concurrent.TimeUnit; import java.util.stream.Collectors; +import io.pinecone.PineconeConnection; import org.awaitility.Awaitility; import org.awaitility.Duration; import org.junit.jupiter.api.BeforeAll; @@ -369,6 +371,15 @@ void deleteWithComplexFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner.run(context -> { + PineconeVectorStore vectorStore = context.getBean(PineconeVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + private void cleanupExistingDocuments(VectorStore vectorStore, String query) { List existingDocs = searchDocuments(vectorStore, query, DEFAULT_TOP_K); if (!existingDocs.isEmpty()) { diff --git a/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java b/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java index 24202c6be2d..865ce3253b3 100644 --- a/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java +++ b/vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java @@ -345,6 +345,13 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.qdrantClient; + return Optional.of(client); + } + /** * Builder for creating instances of {@link QdrantVectorStore}. This builder provides * a fluent API for configuring all aspects of the vector store. diff --git a/vector-stores/spring-ai-qdrant-store/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java b/vector-stores/spring-ai-qdrant-store/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java index 0c4b06d43f6..cf73709b329 100644 --- a/vector-stores/spring-ai-qdrant-store/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java +++ b/vector-stores/spring-ai-qdrant-store/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java @@ -19,6 +19,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.UUID; import java.util.concurrent.ExecutionException; import java.util.stream.Collectors; @@ -344,6 +345,15 @@ void deleteWithComplexFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner.run(context -> { + QdrantVectorStore vectorStore = context.getBean(QdrantVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration public static class TestApplication { diff --git a/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/redis/RedisVectorStore.java b/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/redis/RedisVectorStore.java index f19eb89994e..ff1f949630f 100644 --- a/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/redis/RedisVectorStore.java +++ b/vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/redis/RedisVectorStore.java @@ -466,6 +466,13 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.jedis; + return Optional.of(client); + } + public static Builder builder(JedisPooled jedis, EmbeddingModel embeddingModel) { return new Builder(jedis, embeddingModel); } diff --git a/vector-stores/spring-ai-redis-store/src/test/java/org/springframework/ai/vectorstore/redis/RedisVectorStoreIT.java b/vector-stores/spring-ai-redis-store/src/test/java/org/springframework/ai/vectorstore/redis/RedisVectorStoreIT.java index d43f0aa4f94..a2e7dd53efe 100644 --- a/vector-stores/spring-ai-redis-store/src/test/java/org/springframework/ai/vectorstore/redis/RedisVectorStoreIT.java +++ b/vector-stores/spring-ai-redis-store/src/test/java/org/springframework/ai/vectorstore/redis/RedisVectorStoreIT.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; @@ -347,6 +348,15 @@ void deleteWithComplexFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner.run(context -> { + RedisVectorStore vectorStore = context.getBean(RedisVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestApplication { diff --git a/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStore.java b/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStore.java index dd3ff09efb0..d6aec99f09f 100644 --- a/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStore.java +++ b/vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStore.java @@ -378,6 +378,13 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str .similarityMetric(VectorStoreSimilarityMetric.COSINE.value()); } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.client; + return Optional.of(client); + } + public static class Builder extends AbstractVectorStoreBuilder { private String collectionName = DEFAULT_COLLECTION_NAME; diff --git a/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java b/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java index 583382cfe4f..ed4f444701e 100644 --- a/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java +++ b/vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java @@ -23,6 +23,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; @@ -335,6 +336,15 @@ void deleteWithComplexFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner.run(context -> { + TypesenseVectorStore vectorStore = context.getBean(TypesenseVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class }) public static class TestApplication { diff --git a/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStore.java b/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStore.java index 92f3aab4630..9aface670a3 100644 --- a/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStore.java +++ b/vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStore.java @@ -443,6 +443,13 @@ public VectorStoreObservationContext.Builder createObservationContextBuilder(Str .collectionName(this.weaviateObjectClass); } + @Override + public Optional getNativeClient() { + @SuppressWarnings("unchecked") + T client = (T) this.weaviateClient; + return Optional.of(client); + } + /** * Defines the consistency levels for Weaviate operations. * diff --git a/vector-stores/spring-ai-weaviate-store/src/test/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStoreIT.java b/vector-stores/spring-ai-weaviate-store/src/test/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStoreIT.java index 28e329e8028..459d48ed8d4 100644 --- a/vector-stores/spring-ai-weaviate-store/src/test/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStoreIT.java +++ b/vector-stores/spring-ai-weaviate-store/src/test/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStoreIT.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 io.weaviate.client.Config; @@ -311,6 +312,15 @@ void deleteWithStringFilterExpression() { }); } + @Test + void getNativeClientTest() { + this.contextRunner.run(context -> { + WeaviateVectorStore vectorStore = context.getBean(WeaviateVectorStore.class); + Optional nativeClient = vectorStore.getNativeClient(); + assertThat(nativeClient).isPresent(); + }); + } + @SpringBootConfiguration @EnableAutoConfiguration public static class TestApplication {