Spring/spring-ai
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Add getNativeClient API to VectorStore interface

Browse Source 
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

Fixes: #2137

Signed-off-by: Soby Chacko <soby.chacko@broadcom.com>
This commit is contained in:
Soby Chacko
2025-02-02 15:20:33 -05:00
committed by Mark Pollack
parent 54463e6622
commit 16a596f8b7
53 changed files with 618 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -108,6 +108,24 @@ public interface VectorStore extends DocumentWriter {
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.

4
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs.adoc
View File

@@ -44,6 +44,10 @@ public interface VectorStore extends DocumentWriter {
List<Document> similaritySearch(String query);
List<Document> similaritySearch(SearchRequest request);
default <T> Optional<T> getNativeClient() {
return Optional.empty();
}
}
```

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/apache-cassandra.adoc
View File

@@ -319,3 +319,20 @@ nodetool import wikidata articles ${CASSANDRA_DATA}/data/wikidata/articles-*/
* An alternative to `nodetool import` is to just restart Cassandra.
* If there are any failures in the indexes they will be rebuilt automatically.
====
== Accessing the Native Client
The Cassandra Vector Store implementation provides access to the underlying native Cassandra client (`CqlSession`) through the `getNativeClient()` method:
[source,java]
----
CassandraVectorStore vectorStore = context.getBean(CassandraVectorStore.class);
Optional<CqlSession> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
CqlSession session = nativeClient.get();
// Use the native client for Cassandra-specific operations
}
----
The native client gives you access to Cassandra-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/azure-cosmos-db.adoc
View File

@@ -205,3 +205,20 @@ Add the following dependency in your Maven project:
<artifactId>spring-ai-azure-cosmos-db-store</artifactId>
</dependency>
----
== Accessing the Native Client
The Azure Cosmos DB Vector Store implementation provides access to the underlying native Azure Cosmos DB client (`CosmosClient`) through the `getNativeClient()` method:
[source,java]
----
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
}
----
The native client gives you access to Azure Cosmos DB-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/azure.adoc
View File

@@ -226,3 +226,20 @@ is converted into the following Azure OData link:https://learn.microsoft.com/en-
----
$filter search.in(meta_country, 'UK,NL', ',') and meta_year ge 2020
----
== Accessing the Native Client
The Azure Vector Store implementation provides access to the underlying native Azure Search client (`SearchClient`) through the `getNativeClient()` method:
[source,java]
----
AzureVectorStore vectorStore = context.getBean(AzureVectorStore.class);
Optional<SearchClient> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
SearchClient client = nativeClient.get();
// Use the native client for Azure Search-specific operations
}
----
The native client gives you access to Azure Search-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/coherence.adoc Normal file
View File

@@ -0,0 +1,17 @@
== Accessing the Native Client
The Coherence Vector Store implementation provides access to the underlying native Coherence client (`Session`) through the `getNativeClient()` method:
[source,java]
----
CoherenceVectorStore vectorStore = context.getBean(CoherenceVectorStore.class);
Optional<Session> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
Session session = nativeClient.get();
// Use the native client for Coherence-specific operations
}
----
The native client gives you access to Coherence-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/elasticsearch.adoc
View File

@@ -269,3 +269,20 @@ public EmbeddingModel embeddingModel() {
return new OpenAiEmbeddingModel(new OpenAiApi(System.getenv("OPENAI_API_KEY")));
}
----
== Accessing the Native Client
The Elasticsearch Vector Store implementation provides access to the underlying native Elasticsearch client (`ElasticsearchClient`) through the `getNativeClient()` method:
[source,java]
----
ElasticsearchVectorStore vectorStore = context.getBean(ElasticsearchVectorStore.class);
Optional<ElasticsearchClient> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
ElasticsearchClient client = nativeClient.get();
// Use the native client for Elasticsearch-specific operations
}
----
The native client gives you access to Elasticsearch-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/mariadb.adoc
View File

@@ -204,3 +204,20 @@ vectorStore.similaritySearch(SearchRequest.builder()
----
NOTE: These filter expressions are automatically converted into the equivalent MariaDB JSON path expressions.
== Accessing the Native Client
The MariaDB Vector Store implementation provides access to the underlying native JDBC client (`JdbcTemplate`) through the `getNativeClient()` method:
[source,java]
----
MariaDBVectorStore vectorStore = context.getBean(MariaDBVectorStore.class);
Optional<JdbcTemplate> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
JdbcTemplate jdbc = nativeClient.get();
// Use the native client for MariaDB-specific operations
}
----
The native client gives you access to MariaDB-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/milvus.adoc
View File

@@ -229,3 +229,20 @@ If Docker complains about resources, then execute:
----
docker system prune --all --force --volumes
----
== Accessing the Native Client
The Milvus Vector Store implementation provides access to the underlying native Milvus client (`MilvusServiceClient`) through the `getNativeClient()` method:
[source,java]
----
MilvusVectorStore vectorStore = context.getBean(MilvusVectorStore.class);
Optional<MilvusServiceClient> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
MilvusServiceClient client = nativeClient.get();
// Use the native client for Milvus-specific operations
}
----
The native client gives you access to Milvus-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/mongodb.adoc
View File

@@ -228,3 +228,20 @@ To get started with Spring AI and MongoDB:
* See the https://www.mongodb.com/docs/atlas/atlas-vector-search/ai-integrations/spring-ai/#std-label-spring-ai[Getting Started guide for Spring AI Integration].
* For a comprehensive code example demonstrating Retrieval Augmented Generation (RAG) with Spring AI and MongoDB, refer to this https://www.mongodb.com/developer/languages/java/retrieval-augmented-generation-spring-ai/[detailed tutorial].
== Accessing the Native Client
The MongoDB Atlas Vector Store implementation provides access to the underlying native MongoDB client (`MongoClient`) through the `getNativeClient()` method:
[source,java]
----
MongoDBAtlasVectorStore vectorStore = context.getBean(MongoDBAtlasVectorStore.class);
Optional<MongoClient> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
MongoClient client = nativeClient.get();
// Use the native client for MongoDB-specific operations
}
----
The native client gives you access to MongoDB-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/neo4j.adoc
View File

@@ -237,3 +237,20 @@ is converted into the proprietary Neo4j filter format:
----
node.`metadata.author` IN ["john","jill"] AND node.`metadata.'article_type'` = "blog"
----
== Accessing the Native Client
The Neo4j Vector Store implementation provides access to the underlying native Neo4j client (`Driver`) through the `getNativeClient()` method:
[source,java]
----
Neo4jVectorStore vectorStore = context.getBean(Neo4jVectorStore.class);
Optional<Driver> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
Driver driver = nativeClient.get();
// Use the native client for Neo4j-specific operations
}
----
The native client gives you access to Neo4j-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/opensearch.adoc
View File

@@ -238,3 +238,20 @@ is converted into the proprietary OpenSearch filter format:
----
(metadata.author:john OR jill) AND metadata.article_type:blog
----
== Accessing the Native Client
The OpenSearch Vector Store implementation provides access to the underlying native OpenSearch client (`OpenSearchClient`) through the `getNativeClient()` method:
[source,java]
----
OpenSearchVectorStore vectorStore = context.getBean(OpenSearchVectorStore.class);
Optional<OpenSearchClient> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
OpenSearchClient client = nativeClient.get();
// Use the native client for OpenSearch-specific operations
}
----
The native client gives you access to OpenSearch-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/oracle.adoc
View File

@@ -206,3 +206,20 @@ You can then connect to the database using:
----
sql mlops/mlops@localhost/freepdb1
----
== Accessing the Native Client
The Oracle Vector Store implementation provides access to the underlying native Oracle client (`OracleConnection`) through the `getNativeClient()` method:
[source,java]
----
OracleVectorStore vectorStore = context.getBean(OracleVectorStore.class);
Optional<OracleConnection> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
OracleConnection connection = nativeClient.get();
// Use the native client for Oracle-specific operations
}
----
The native client gives you access to Oracle-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/pgvector.adoc
View File

@@ -243,3 +243,20 @@ You can connect to this server like this:
----
psql -U postgres -h localhost -p 5432
----
== Accessing the Native Client
The PGVector Store implementation provides access to the underlying native JDBC client (`JdbcTemplate`) through the `getNativeClient()` method:
[source,java]
----
PgVectorStore vectorStore = context.getBean(PgVectorStore.class);
Optional<JdbcTemplate> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
JdbcTemplate jdbc = nativeClient.get();
// Use the native client for PostgreSQL-specific operations
}
----
The native client gives you access to PostgreSQL-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/pinecone.adoc
View File

@@ -222,3 +222,20 @@ List<Document> results = vectorStore.similaritySearch(SearchRequest.query("Sprin
----
If all goes well, you should retrieve the document containing the text "Spring AI rocks!!".
== Accessing the Native Client
The Pinecone Vector Store implementation provides access to the underlying native Pinecone client (`PineconeConnection`) through the `getNativeClient()` method:
[source,java]
----
PineconeVectorStore vectorStore = context.getBean(PineconeVectorStore.class);
Optional<PineconeConnection> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
PineconeConnection client = nativeClient.get();
// Use the native client for Pinecone-specific operations
}
----
The native client gives you access to Pinecone-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/qdrant.adoc
View File

@@ -194,3 +194,20 @@ vectorStore.similaritySearch(SearchRequest.builder()
----
NOTE: These (portable) filter expressions get automatically converted into the proprietary Qdrant link:https://qdrant.tech/documentation/concepts/filtering/[filter expressions].
== Accessing the Native Client
The Qdrant Vector Store implementation provides access to the underlying native Qdrant client (`QdrantClient`) through the `getNativeClient()` method:
[source,java]
----
QdrantVectorStore vectorStore = context.getBean(QdrantVectorStore.class);
Optional<QdrantClient> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
QdrantClient client = nativeClient.get();
// Use the native client for Qdrant-specific operations
}
----
The native client gives you access to Qdrant-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/redis.adoc
View File

@@ -212,3 +212,20 @@ public EmbeddingModel embeddingModel() {
You must list explicitly all metadata field names and types (`TAG`, `TEXT`, or `NUMERIC`) for any metadata field used in filter expressions.
The `metadataFields` above registers filterable metadata fields: `country` of type `TAG`, `year` of type `NUMERIC`.
====
== Accessing the Native Client
The Redis Vector Store implementation provides access to the underlying native Redis client (`JedisPooled`) through the `getNativeClient()` method:
[source,java]
----
RedisVectorStore vectorStore = context.getBean(RedisVectorStore.class);
Optional<JedisPooled> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
JedisPooled jedis = nativeClient.get();
// Use the native client for Redis-specific operations
}
----
The native client gives you access to Redis-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/typesense.adoc
View File

@@ -229,3 +229,20 @@ If you are not retrieving the documents in the expected order or the search resu
Embedding models can have a significant impact on the search results (i.e. make sure if your data is in Spanish to use a Spanish or multilingual embedding model).
====
== Accessing the Native Client
The Typesense Vector Store implementation provides access to the underlying native Typesense client (`Client`) through the `getNativeClient()` method:
[source,java]
----
TypesenseVectorStore vectorStore = context.getBean(TypesenseVectorStore.class);
Optional<Client> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
Client client = nativeClient.get();
// Use the native client for Typesense-specific operations
}
----
The native client gives you access to Typesense-specific features and operations that might not be exposed through the `VectorStore` interface.

17
spring-ai-docs/src/main/antora/modules/ROOT/pages/api/vectordbs/weaviate.adoc
View File

@@ -230,3 +230,20 @@ You can use the following properties in your Spring Boot configuration to custom
|`spring.ai.vectorstore.weaviate.consistency-level`|Desired tradeoff between consistency and speed|ConsistentLevel.ONE
|`spring.ai.vectorstore.weaviate.filter-field`|Configures metadata fields that can be used in filters. Format: spring.ai.vectorstore.weaviate.filter-field.<field-name>=<field-type>|
|===
== Accessing the Native Client
The Weaviate Vector Store implementation provides access to the underlying native Weaviate client (`WeaviateClient`) through the `getNativeClient()` method:
[source,java]
----
WeaviateVectorStore vectorStore = context.getBean(WeaviateVectorStore.class);
Optional<WeaviateClient> nativeClient = vectorStore.getNativeClient();
if (nativeClient.isPresent()) {
WeaviateClient client = nativeClient.get();
// Use the native client for Weaviate-specific operations
}
----
The native client gives you access to Weaviate-specific features and operations that might not be exposed through the `VectorStore` interface.

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

@@ -372,6 +372,13 @@ public class CosmosDBVectorStore extends AbstractObservationVectorStore implemen
.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>

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

@@ -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 @@ package org.springframework.ai.vectorstore.cosmosdb;
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 @@ public class CosmosDBVectorStoreIT {
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 {

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

@@ -332,6 +332,13 @@ public class AzureVectorStore extends AbstractObservationVectorStore implements
.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) {

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

@@ -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.Collections;
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 class AzureVectorStoreIT {
});
}
@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 {

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

@@ -712,6 +712,13 @@ public class CassandraVectorStore extends AbstractObservationVectorStore impleme
}
}
@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

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

@@ -21,6 +21,7 @@ import java.nio.charset.StandardCharsets;
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 @@ class CassandraVectorStoreIT {
});
}
@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 {

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

@@ -275,6 +275,13 @@ public class CoherenceVectorStore extends AbstractObservationVectorStore impleme
.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>

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

@@ -23,6 +23,7 @@ import java.util.Collections;
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 class CoherenceVectorStoreIT {
});
}
@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()))

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

@@ -354,6 +354,13 @@ public class ElasticsearchVectorStore extends AbstractObservationVectorStore imp
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

21
vector-stores/spring-ai-elasticsearch-store/src/test/java/org/springframework/ai/vectorstore/elasticsearch/ElasticsearchVectorStoreIT.java
View File

@@ -24,6 +24,7 @@ import java.util.ArrayList;
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 @@ class ElasticsearchVectorStoreIT {
});
}
@Test
public void getNativeClientTest() {
getContextRunner().run(context -> {
ElasticsearchVectorStore vectorStore = context.getBean("vectorStore_cosine",
ElasticsearchVectorStore.class);
// Test successful native client retrieval
Optional<ElasticsearchClient> 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 {

7
vector-stores/spring-ai-mariadb-store/src/main/java/org/springframework/ai/vectorstore/mariadb/MariaDBVectorStore.java
View File

@@ -460,6 +460,13 @@ public class MariaDBVectorStore extends AbstractObservationVectorStore implement
return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value();
}
@Override
public <T> Optional<T> getNativeClient() {
@SuppressWarnings("unchecked")
T client = (T) this.jdbcTemplate;
return Optional.of(client);
}
public enum MariaDBDistanceType {
EUCLIDEAN, COSINE

10
vector-stores/spring-ai-mariadb-store/src/test/java/org/springframework/ai/vectorstore/mariadb/MariaDBStoreIT.java
View File

@@ -22,6 +22,7 @@ import java.util.Collections;
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 class MariaDBStoreIT {
});
}
@Test
void getNativeClientTest() {
this.contextRunner.run(context -> {
MariaDBVectorStore vectorStore = context.getBean(MariaDBVectorStore.class);
Optional<JdbcTemplate> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {

7
vector-stores/spring-ai-milvus-store/src/main/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStore.java
View File

@@ -559,6 +559,13 @@ public class MilvusVectorStore extends AbstractObservationVectorStore implements
return SIMILARITY_TYPE_MAPPING.get(this.metricType).value();
}
@Override
public <T> Optional<T> getNativeClient() {
@SuppressWarnings("unchecked")
T client = (T) this.milvusClient;
return Optional.of(client);
}
public static class Builder extends AbstractVectorStoreBuilder<Builder> {
private final MilvusServiceClient milvusClient;

10
vector-stores/spring-ai-milvus-store/src/test/java/org/springframework/ai/vectorstore/milvus/MilvusVectorStoreIT.java
View File

@@ -21,6 +21,7 @@ import java.nio.charset.StandardCharsets;
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 class MilvusVectorStoreIT {
});
}
@Test
void getNativeClientTest() {
this.contextRunner.withPropertyValues("test.spring.ai.vectorstore.milvus.metricType=COSINE").run(context -> {
MilvusVectorStore vectorStore = context.getBean(MilvusVectorStore.class);
Optional<MilvusServiceClient> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {

7
vector-stores/spring-ai-mongodb-atlas-store/src/main/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStore.java
View File

@@ -330,6 +330,13 @@ public class MongoDBAtlasVectorStore extends AbstractObservationVectorStore impl
.fieldName(this.pathName);
}
@Override
public <T> Optional<T> getNativeClient() {
@SuppressWarnings("unchecked")
T client = (T) this.mongoTemplate;
return Optional.of(client);
}
/**
* Creates a new builder instance for MongoDBAtlasVectorStore.
* @return a new MongoDBBuilder instance

10
vector-stores/spring-ai-mongodb-atlas-store/src/test/java/org/springframework/ai/vectorstore/mongodb/atlas/MongoDBAtlasVectorStoreIT.java
View File

@@ -22,6 +22,7 @@ import java.util.Arrays;
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 @@ class MongoDBAtlasVectorStoreIT {
});
}
@Test
void getNativeClientTest() {
this.contextRunner.run(context -> {
MongoDBAtlasVectorStore vectorStore = context.getBean(MongoDBAtlasVectorStore.class);
Optional<MongoTemplate> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
public static String getText(String uri) {
var resource = new DefaultResourceLoader().getResource(uri);
try {

7
vector-stores/spring-ai-neo4j-store/src/main/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStore.java
View File

@@ -368,6 +368,13 @@ public class Neo4jVectorStore extends AbstractObservationVectorStore implements
return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value();
}
@Override
public <T> Optional<T> 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.
*/

10
vector-stores/spring-ai-neo4j-store/src/test/java/org/springframework/ai/vectorstore/neo4j/Neo4jVectorStoreIT.java
View File

@@ -19,6 +19,7 @@ package org.springframework.ai.vectorstore.neo4j;
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 @@ class Neo4jVectorStoreIT {
});
}
@Test
void getNativeClientTest() {
this.contextRunner.run(context -> {
Neo4jVectorStore vectorStore = context.getBean(Neo4jVectorStore.class);
Optional<Driver> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {

7
vector-stores/spring-ai-opensearch-store/src/main/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStore.java
View File

@@ -380,6 +380,13 @@ public class OpenSearchVectorStore extends AbstractObservationVectorStore implem
return this.similarityFunction;
}
@Override
public <T> Optional<T> getNativeClient() {
@SuppressWarnings("unchecked")
T client = (T) this.openSearchClient;
return Optional.of(client);
}
/**
* The representation of {@link Document} along with its embedding.
*

10
vector-stores/spring-ai-opensearch-store/src/test/java/org/springframework/ai/vectorstore/opensearch/OpenSearchVectorStoreIT.java
View File

@@ -25,6 +25,7 @@ import java.time.ZonedDateTime;
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 @@ class OpenSearchVectorStoreIT {
});
}
@Test
void getNativeClientTest() {
getContextRunner().run(context -> {
OpenSearchVectorStore vectorStore = context.getBean("vectorStore", OpenSearchVectorStore.class);
Optional<OpenSearchClient> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {

7
vector-stores/spring-ai-oracle-store/src/main/java/org/springframework/ai/vectorstore/oracle/OracleVectorStore.java
View File

@@ -530,6 +530,13 @@ public class OracleVectorStore extends AbstractObservationVectorStore implements
.similarityMetric(getSimilarityMetric());
}
@Override
public <T> Optional<T> 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();

13
vector-stores/spring-ai-oracle-store/src/test/java/org/springframework/ai/vectorstore/oracle/OracleVectorStoreIT.java
View File

@@ -22,6 +22,7 @@ import java.util.Collections;
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 @@ public class OracleVectorStoreIT {
});
}
@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<JdbcTemplate> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestClient {

7
vector-stores/spring-ai-pgvector-store/src/main/java/org/springframework/ai/vectorstore/pgvector/PgVectorStore.java
View File

@@ -476,6 +476,13 @@ public class PgVectorStore extends AbstractObservationVectorStore implements Ini
return SIMILARITY_TYPE_MAPPING.get(this.distanceType).value();
}
@Override
public <T> Optional<T> 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

9
vector-stores/spring-ai-pgvector-store/src/test/java/org/springframework/ai/vectorstore/pgvector/PgVectorStoreIT.java
View File

@@ -482,6 +482,15 @@ public class PgVectorStoreIT {
});
}
@Test
void getNativeClientTest() {
this.contextRunner.run(context -> {
PgVectorStore vectorStore = context.getBean(PgVectorStore.class);
Optional<JdbcTemplate> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {

7
vector-stores/spring-ai-pinecone-store/src/main/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStore.java
View File

@@ -374,6 +374,13 @@ public class PineconeVectorStore extends AbstractObservationVectorStore {
.fieldName(this.pineconeContentFieldName);
}
@Override
public <T> Optional<T> 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

11
vector-stores/spring-ai-pinecone-store/src/test/java/org/springframework/ai/vectorstore/pinecone/PineconeVectorStoreIT.java
View File

@@ -21,10 +21,12 @@ import java.nio.charset.StandardCharsets;
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 @@ public class PineconeVectorStoreIT {
});
}
@Test
void getNativeClientTest() {
this.contextRunner.run(context -> {
PineconeVectorStore vectorStore = context.getBean(PineconeVectorStore.class);
Optional<PineconeConnection> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
private void cleanupExistingDocuments(VectorStore vectorStore, String query) {
List<Document> existingDocs = searchDocuments(vectorStore, query, DEFAULT_TOP_K);
if (!existingDocs.isEmpty()) {

7
vector-stores/spring-ai-qdrant-store/src/main/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStore.java
View File

@@ -345,6 +345,13 @@ public class QdrantVectorStore extends AbstractObservationVectorStore implements
}
@Override
public <T> Optional<T> 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.

10
vector-stores/spring-ai-qdrant-store/src/test/java/org/springframework/ai/vectorstore/qdrant/QdrantVectorStoreIT.java
View File

@@ -19,6 +19,7 @@ package org.springframework.ai.vectorstore.qdrant;
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 @@ public class QdrantVectorStoreIT {
});
}
@Test
void getNativeClientTest() {
this.contextRunner.run(context -> {
QdrantVectorStore vectorStore = context.getBean(QdrantVectorStore.class);
Optional<QdrantClient> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
public static class TestApplication {

7
vector-stores/spring-ai-redis-store/src/main/java/org/springframework/ai/vectorstore/redis/RedisVectorStore.java
View File

@@ -466,6 +466,13 @@ public class RedisVectorStore extends AbstractObservationVectorStore implements
}
@Override
public <T> Optional<T> 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);
}

10
vector-stores/spring-ai-redis-store/src/test/java/org/springframework/ai/vectorstore/redis/RedisVectorStoreIT.java
View File

@@ -21,6 +21,7 @@ import java.nio.charset.StandardCharsets;
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 @@ class RedisVectorStoreIT {
});
}
@Test
void getNativeClientTest() {
this.contextRunner.run(context -> {
RedisVectorStore vectorStore = context.getBean(RedisVectorStore.class);
Optional<JedisPooled> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {

7
vector-stores/spring-ai-typesense-store/src/main/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStore.java
View File

@@ -378,6 +378,13 @@ public class TypesenseVectorStore extends AbstractObservationVectorStore impleme
.similarityMetric(VectorStoreSimilarityMetric.COSINE.value());
}
@Override
public <T> Optional<T> getNativeClient() {
@SuppressWarnings("unchecked")
T client = (T) this.client;
return Optional.of(client);
}
public static class Builder extends AbstractVectorStoreBuilder<Builder> {
private String collectionName = DEFAULT_COLLECTION_NAME;

10
vector-stores/spring-ai-typesense-store/src/test/java/org/springframework/ai/vectorstore/typesense/TypesenseVectorStoreIT.java
View File

@@ -23,6 +23,7 @@ import java.util.ArrayList;
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 @@ public class TypesenseVectorStoreIT {
});
}
@Test
void getNativeClientTest() {
this.contextRunner.run(context -> {
TypesenseVectorStore vectorStore = context.getBean(TypesenseVectorStore.class);
Optional<TypesenseVectorStore> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration(exclude = { DataSourceAutoConfiguration.class })
public static class TestApplication {

7
vector-stores/spring-ai-weaviate-store/src/main/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStore.java
View File

@@ -443,6 +443,13 @@ public class WeaviateVectorStore extends AbstractObservationVectorStore {
.collectionName(this.weaviateObjectClass);
}
@Override
public <T> Optional<T> getNativeClient() {
@SuppressWarnings("unchecked")
T client = (T) this.weaviateClient;
return Optional.of(client);
}
/**
* Defines the consistency levels for Weaviate operations.
*

10
vector-stores/spring-ai-weaviate-store/src/test/java/org/springframework/ai/vectorstore/weaviate/WeaviateVectorStoreIT.java
View File

@@ -21,6 +21,7 @@ import java.nio.charset.StandardCharsets;
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 @@ public class WeaviateVectorStoreIT {
});
}
@Test
void getNativeClientTest() {
this.contextRunner.run(context -> {
WeaviateVectorStore vectorStore = context.getBean(WeaviateVectorStore.class);
Optional<WeaviateClient> nativeClient = vectorStore.getNativeClient();
assertThat(nativeClient).isPresent();
});
}
@SpringBootConfiguration
@EnableAutoConfiguration
public static class TestApplication {