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

GH-2954: Opensearch vector store transport improvements

Browse Source 
Fixes: #2954

* fixed authorization to multiple hosts
* ability to provide SSL certs via `SslBundles`
* ability to set connectTimeout/readTimeout

Signed-off-by: Linar Abzaltdinov <abzaltdinov@gmail.com>
This commit is contained in:
Linar Abzaltdinov
2025-05-06 11:55:53 +03:00
committed by Soby Chacko
parent fa5fb53d4c
commit f30ab261b4
4 changed files with 132 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

84
auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-opensearch/src/main/java/org/springframework/ai/vectorstore/opensearch/autoconfigure/OpenSearchVectorStoreAutoConfiguration.java
View File

@@ -17,13 +17,19 @@
package org.springframework.ai.vectorstore.opensearch.autoconfigure;
import java.net.URISyntaxException;
import java.time.Duration;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
import io.micrometer.observation.ObservationRegistry;
import org.apache.hc.client5.http.auth.AuthScope;
import org.apache.hc.client5.http.auth.UsernamePasswordCredentials;
import org.apache.hc.client5.http.config.RequestConfig;
import org.apache.hc.client5.http.impl.auth.BasicCredentialsProvider;
import org.apache.hc.client5.http.impl.nio.PoolingAsyncClientConnectionManagerBuilder;
import org.apache.hc.client5.http.nio.AsyncClientConnectionManager;
import org.apache.hc.client5.http.ssl.ClientTlsStrategyBuilder;
import org.apache.hc.core5.http.HttpHost;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.client.transport.OpenSearchTransport;
@@ -33,7 +39,6 @@ import org.opensearch.client.transport.httpclient5.ApacheHttpClient5TransportBui
import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
import software.amazon.awssdk.http.SdkHttpClient;
import software.amazon.awssdk.http.apache.ApacheHttpClient;
import software.amazon.awssdk.regions.Region;
@@ -50,6 +55,7 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.boot.ssl.SslBundles;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.util.StringUtils;
@@ -99,26 +105,57 @@ public class OpenSearchVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
OpenSearchClient openSearchClient(OpenSearchConnectionDetails connectionDetails) {
HttpHost[] httpHosts = connectionDetails.getUris()
.stream()
.map(s -> createHttpHost(s))
.toArray(HttpHost[]::new);
ApacheHttpClient5TransportBuilder transportBuilder = ApacheHttpClient5TransportBuilder.builder(httpHosts);
Optional.ofNullable(connectionDetails.getUsername())
.map(username -> createBasicCredentialsProvider(httpHosts[0], username,
connectionDetails.getPassword()))
.ifPresent(basicCredentialsProvider -> transportBuilder
.setHttpClientConfigCallback(httpAsyncClientBuilder -> httpAsyncClientBuilder
.setDefaultCredentialsProvider(basicCredentialsProvider)));
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties, Optional<SslBundles> sslBundles) {
HttpHost[] httpHosts = properties.getUris().stream().map(this::createHttpHost).toArray(HttpHost[]::new);
Optional<BasicCredentialsProvider> basicCredentialsProvider = Optional.ofNullable(properties.getUsername())
.map(username -> createBasicCredentialsProvider(httpHosts, username, properties.getPassword()));
var transportBuilder = ApacheHttpClient5TransportBuilder.builder(httpHosts);
transportBuilder.setHttpClientConfigCallback(httpClientBuilder -> {
basicCredentialsProvider.ifPresent(httpClientBuilder::setDefaultCredentialsProvider);
httpClientBuilder.setConnectionManager(createConnectionManager(properties, sslBundles));
httpClientBuilder.setDefaultRequestConfig(createRequestConfig(properties));
return httpClientBuilder;
});
return new OpenSearchClient(transportBuilder.build());
}
private BasicCredentialsProvider createBasicCredentialsProvider(HttpHost httpHost, String username,
private AsyncClientConnectionManager createConnectionManager(OpenSearchVectorStoreProperties properties,
Optional<SslBundles> sslBundles) {
var connectionManagerBuilder = PoolingAsyncClientConnectionManagerBuilder.create();
if (sslBundles.isPresent()) {
Optional.ofNullable(properties.getSslBundle())
.map(bundle -> sslBundles.get().getBundle(bundle))
.map(bundle -> ClientTlsStrategyBuilder.create()
.setSslContext(bundle.createSslContext())
.setTlsVersions(bundle.getOptions().getEnabledProtocols())
.build())
.ifPresent(connectionManagerBuilder::setTlsStrategy);
}
return connectionManagerBuilder.build();
}
private RequestConfig createRequestConfig(OpenSearchVectorStoreProperties properties) {
var requestConfigBuilder = RequestConfig.custom();
Optional.ofNullable(properties.getConnectionTimeout())
.map(Duration::toMillis)
.ifPresent(timeoutMillis -> requestConfigBuilder.setConnectionRequestTimeout(timeoutMillis,
TimeUnit.MILLISECONDS));
Optional.ofNullable(properties.getReadTimeout())
.map(Duration::toMillis)
.ifPresent(
timeoutMillis -> requestConfigBuilder.setResponseTimeout(timeoutMillis, TimeUnit.MILLISECONDS));
return requestConfigBuilder.build();
}
private BasicCredentialsProvider createBasicCredentialsProvider(HttpHost[] httpHosts, String username,
String password) {
BasicCredentialsProvider basicCredentialsProvider = new BasicCredentialsProvider();
basicCredentialsProvider.setCredentials(new AuthScope(httpHost),
new UsernamePasswordCredentials(username, password.toCharArray()));
for (HttpHost httpHost : httpHosts) {
basicCredentialsProvider.setCredentials(new AuthScope(httpHost),
new UsernamePasswordCredentials(username, password.toCharArray()));
}
return basicCredentialsProvider;
}
@@ -159,12 +196,21 @@ public class OpenSearchVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties,
OpenSearchClient openSearchClient(OpenSearchVectorStoreProperties properties, Optional<SslBundles> sslBundles,
AwsOpenSearchConnectionDetails connectionDetails, AwsSdk2TransportOptions options) {
Region region = Region.of(connectionDetails.getRegion());
SdkHttpClient httpClient = ApacheHttpClient.builder().build();
OpenSearchTransport transport = new AwsSdk2Transport(httpClient,
var httpClientBuilder = ApacheHttpClient.builder();
Optional.ofNullable(properties.getConnectionTimeout()).ifPresent(httpClientBuilder::connectionTimeout);
Optional.ofNullable(properties.getReadTimeout()).ifPresent(httpClientBuilder::socketTimeout);
if (sslBundles.isPresent()) {
Optional.ofNullable(properties.getSslBundle())
.map(bundle -> sslBundles.get().getBundle(bundle))
.ifPresent(bundle -> httpClientBuilder
.tlsKeyManagersProvider(() -> bundle.getManagers().getKeyManagers())
.tlsTrustManagersProvider(() -> bundle.getManagers().getTrustManagers()));
}
OpenSearchTransport transport = new AwsSdk2Transport(httpClientBuilder.build(),
connectionDetails.getHost(properties.getAws().getDomainName()),
properties.getAws().getServiceName(), region, options);
return new OpenSearchClient(transport);

39
auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-opensearch/src/main/java/org/springframework/ai/vectorstore/opensearch/autoconfigure/OpenSearchVectorStoreProperties.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.
@@ -16,6 +16,7 @@
package org.springframework.ai.vectorstore.opensearch.autoconfigure;
import java.time.Duration;
import java.util.List;
import org.springframework.ai.vectorstore.properties.CommonVectorStoreProperties;
@@ -39,6 +40,18 @@ public class OpenSearchVectorStoreProperties extends CommonVectorStoreProperties
private String mappingJson;
/**
* SSL Bundle name ({@link org.springframework.boot.ssl.SslBundles}).
*/
private String sslBundle;
/**
*
*/
private Duration connectionTimeout;
private Duration readTimeout;
private Aws aws = new Aws();
public List<String> getUris() {
@@ -81,6 +94,30 @@ public class OpenSearchVectorStoreProperties extends CommonVectorStoreProperties
this.mappingJson = mappingJson;
}
public String getSslBundle() {
return sslBundle;
}
public void setSslBundle(String sslBundle) {
this.sslBundle = sslBundle;
}
public Duration getConnectionTimeout() {
return connectionTimeout;
}
public void setConnectionTimeout(Duration connectionTimeout) {
this.connectionTimeout = connectionTimeout;
}
public Duration getReadTimeout() {
return readTimeout;
}
public void setReadTimeout(Duration readTimeout) {
this.readTimeout = readTimeout;
}
public Aws getAws() {
return this.aws;
}

15
auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-opensearch/src/test/java/org/springframework/ai/vectorstore/opensearch/autoconfigure/AwsOpenSearchVectorStoreAutoConfigurationIT.java
View File

@@ -27,6 +27,7 @@ import net.minidev.json.JSONArray;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.testcontainers.containers.localstack.LocalStackContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -37,8 +38,11 @@ import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.retry.autoconfigure.SpringAiRetryAutoConfiguration;
import org.springframework.ai.transformers.TransformersEmbeddingModel;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.opensearch.OpenSearchVectorStore;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.autoconfigure.ssl.SslAutoConfiguration;
import org.springframework.boot.ssl.SslBundles;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
@@ -138,6 +142,17 @@ class AwsOpenSearchVectorStoreAutoConfigurationIT {
});
}
@Test
public void autoConfigurationWithSslBundles() {
this.contextRunner.withConfiguration(AutoConfigurations.of(SslAutoConfiguration.class)).run(context -> {
assertThat(context.getBeansOfType(SslBundles.class)).isNotEmpty();
assertThat(context.getBeansOfType(OpenSearchClient.class)).isNotEmpty();
assertThat(context.getBeansOfType(OpenSearchVectorStoreProperties.class)).isNotEmpty();
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
assertThat(context.getBean(VectorStore.class)).isInstanceOf(OpenSearchVectorStore.class);
});
}
private String getText(String uri) {
var resource = new DefaultResourceLoader().getResource(uri);
try {

14
auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-opensearch/src/test/java/org/springframework/ai/vectorstore/opensearch/autoconfigure/OpenSearchVectorStoreAutoConfigurationIT.java
View File

@@ -24,6 +24,7 @@ import java.util.Map;
import io.micrometer.observation.tck.TestObservationRegistry;
import org.awaitility.Awaitility;
import org.junit.jupiter.api.Test;
import org.opensearch.client.opensearch.OpenSearchClient;
import org.opensearch.testcontainers.OpensearchContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
@@ -42,6 +43,8 @@ import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.observation.VectorStoreObservationContext;
import org.springframework.ai.vectorstore.opensearch.OpenSearchVectorStore;
import org.springframework.boot.autoconfigure.AutoConfigurations;
import org.springframework.boot.autoconfigure.ssl.SslAutoConfiguration;
import org.springframework.boot.ssl.SslBundles;
import org.springframework.boot.test.context.FilteredClassLoader;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
@@ -171,6 +174,17 @@ class OpenSearchVectorStoreAutoConfigurationIT {
});
}
@Test
public void autoConfigurationWithSslBundles() {
this.contextRunner.withConfiguration(AutoConfigurations.of(SslAutoConfiguration.class)).run(context -> {
assertThat(context.getBeansOfType(SslBundles.class)).isNotEmpty();
assertThat(context.getBeansOfType(OpenSearchClient.class)).isNotEmpty();
assertThat(context.getBeansOfType(OpenSearchVectorStoreProperties.class)).isNotEmpty();
assertThat(context.getBeansOfType(VectorStore.class)).isNotEmpty();
assertThat(context.getBean(VectorStore.class)).isInstanceOf(OpenSearchVectorStore.class);
});
}
private String getText(String uri) {
var resource = new DefaultResourceLoader().getResource(uri);
try {