Chroma Vector Store: replace RestTemplate by RestClient

- replace RestTemplate by  RestClient but default to SimpleClientHttpRequestFactory as Chroma seems to have issues with HTTP2.
 - update the documentation.
 - update the ghcr.io/chroma-core/chroma version to 0.5.0. Fix the withBasicAuthCredentials.

Co-authored-by: Christian Tzolov <ctzolov@vmware.com>
This commit is contained in:
Eddú Meléndez
2024-05-11 21:24:16 +02:00
committed by Christian Tzolov
parent 46e47849ce
commit 704dec20d9
8 changed files with 155 additions and 127 deletions

View File

@@ -205,19 +205,20 @@ TIP: Refer to the xref:getting-started.adoc#dependency-management[Dependency Man
=== Sample Code
Create a `RestTemplate` instance with proper ChromaDB authorization configurations and Use it to create a `ChromaApi` instance:
Create a `RestClient.Builder` instance with proper ChromaDB authorization configurations and Use it to create a `ChromaApi` instance:
[source,java]
----
@Bean
public RestTemplate restTemplate() {
return new RestTemplate();
public RestClient.Builder builder() {
return RestClient.builder().requestFactory(new SimpleClientHttpRequestFactory());
}
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
public ChromaApi chromaApi(RestClient.Builder restClientBuilder) {
String chromaUrl = "http://localhost:8000";
ChromaApi chromaApi = new ChromaApi(chromaUrl, restTemplate);
ChromaApi chromaApi = new ChromaApi(chromaUrl, restClientBuilder);
return chromaApi;
}
----

View File

@@ -15,8 +15,6 @@
*/
package org.springframework.ai.autoconfigure.vectorstore.chroma;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.ChromaVectorStore;
@@ -25,15 +23,18 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
import org.springframework.boot.context.properties.EnableConfigurationProperties;
import org.springframework.context.annotation.Bean;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.util.StringUtils;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.client.RestClient;
import com.fasterxml.jackson.databind.ObjectMapper;
/**
* @author Christian Tzolov
* @author Eddú Meléndez
*/
@AutoConfiguration
@ConditionalOnClass({ EmbeddingModel.class, RestTemplate.class, ChromaVectorStore.class, ObjectMapper.class })
@ConditionalOnClass({ EmbeddingModel.class, RestClient.class, ChromaVectorStore.class, ObjectMapper.class })
@EnableConfigurationProperties({ ChromaApiProperties.class, ChromaVectorStoreProperties.class })
public class ChromaVectorStoreAutoConfiguration {
@@ -45,18 +46,18 @@ public class ChromaVectorStoreAutoConfiguration {
@Bean
@ConditionalOnMissingBean
public RestTemplate restTemplate() {
return new RestTemplate();
public RestClient.Builder builder() {
return RestClient.builder().requestFactory(new SimpleClientHttpRequestFactory());
}
@Bean
@ConditionalOnMissingBean
public ChromaApi chromaApi(ChromaApiProperties apiProperties, RestTemplate restTemplate,
public ChromaApi chromaApi(ChromaApiProperties apiProperties, RestClient.Builder restClientBuilder,
ChromaConnectionDetails connectionDetails) {
String chromaUrl = String.format("%s:%s", connectionDetails.getHost(), connectionDetails.getPort());
var chromaApi = new ChromaApi(chromaUrl, restTemplate, new ObjectMapper());
var chromaApi = new ChromaApi(chromaUrl, restClientBuilder, new ObjectMapper());
if (StringUtils.hasText(apiProperties.getKeyToken())) {
chromaApi.withKeyToken(apiProperties.getKeyToken());

View File

@@ -43,7 +43,7 @@ import static org.assertj.core.api.Assertions.assertThat;
public class ChromaVectorStoreAutoConfigurationIT {
@Container
static ChromaDBContainer chroma = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.15");
static ChromaDBContainer chroma = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.5.0");
private final ApplicationContextRunner contextRunner = new ApplicationContextRunner()
.withConfiguration(AutoConfigurations.of(ChromaVectorStoreAutoConfiguration.class))

View File

@@ -19,49 +19,55 @@ import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.function.Consumer;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import org.springframework.ai.chroma.ChromaApi.QueryRequest.Include;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.HttpHeaders;
import org.springframework.http.MediaType;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.http.client.support.BasicAuthenticationInterceptor;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.client.HttpServerErrorException;
import org.springframework.web.client.RestClient;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.chroma.ChromaApi.QueryRequest.Include;
import org.springframework.http.HttpEntity;
import org.springframework.http.HttpHeaders;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.client.support.BasicAuthenticationInterceptor;
import org.springframework.util.CollectionUtils;
import org.springframework.util.StringUtils;
import org.springframework.web.client.HttpServerErrorException;
import org.springframework.web.client.RestTemplate;
/**
* Single-class Chroma API implementation based on the (unofficial) Chroma REST API.
*
* @author Christian Tzolov
* @author Eddú Meléndez
*/
public class ChromaApi {
// Regular expression pattern that looks for a message inside the ValueError(...).
private static Pattern VALUE_ERROR_PATTERN = Pattern.compile("ValueError\\('([^']*)'\\)");
private final String baseUrl;
private final RestTemplate restTemplate;
private RestClient restClient;
private final ObjectMapper objectMapper;
private String keyToken;
public ChromaApi(String baseUrl, RestTemplate restTemplate) {
this(baseUrl, restTemplate, new ObjectMapper());
public ChromaApi(String baseUrl) {
this(baseUrl, RestClient.builder().requestFactory(new SimpleClientHttpRequestFactory()), new ObjectMapper());
}
public ChromaApi(String baseUrl, RestTemplate restTemplate, ObjectMapper objectMapper) {
this.baseUrl = baseUrl;
this.restTemplate = restTemplate;
public ChromaApi(String baseUrl, RestClient.Builder restClientBuilder) {
this(baseUrl, restClientBuilder, new ObjectMapper());
}
public ChromaApi(String baseUrl, RestClient.Builder restClientBuilder, ObjectMapper objectMapper) {
Consumer<HttpHeaders> defaultHeaders = headers -> {
headers.setContentType(MediaType.APPLICATION_JSON);
};
this.restClient = restClientBuilder.baseUrl(baseUrl).defaultHeaders(defaultHeaders).build();
this.objectMapper = objectMapper;
}
@@ -82,7 +88,9 @@ public class ChromaApi {
* @param password Credentials password.
*/
public ChromaApi withBasicAuthCredentials(String username, String password) {
this.restTemplate.getInterceptors().add(new BasicAuthenticationInterceptor(username, password));
this.restClient = this.restClient.mutate()
.requestInterceptor(new BasicAuthenticationInterceptor(username, password))
.build();
return this;
}
@@ -265,9 +273,23 @@ public class ChromaApi {
public Collection createCollection(CreateCollectionRequest createCollectionRequest) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections", HttpMethod.POST,
this.getHttpEntityFor(createCollectionRequest), Collection.class)
return this.restClient.post()
.uri("/api/v1/collections")
.headers(this::httpHeaders)
.body(createCollectionRequest)
.retrieve()
.toEntity(Collection.class)
.getBody();
}
public Map<String, Object> createCollection2(CreateCollectionRequest createCollectionRequest) {
return this.restClient.post()
.uri("/api/v1/collections")
.headers(this::httpHeaders)
.body(createCollectionRequest)
.retrieve()
.toEntity(Map.class)
.getBody();
}
@@ -278,16 +300,21 @@ public class ChromaApi {
*/
public void deleteCollection(String collectionName) {
this.restTemplate.exchange(this.baseUrl + "/api/v1/collections/{collection_name}", HttpMethod.DELETE,
new HttpEntity<>(httpHeaders()), Void.class, collectionName);
this.restClient.delete()
.uri("/api/v1/collections/{collection_name}", collectionName)
.headers(this::httpHeaders)
.retrieve()
.toBodilessEntity();
}
public Collection getCollection(String collectionName) {
try {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_name}", HttpMethod.GET,
new HttpEntity<>(httpHeaders()), Collection.class, collectionName)
return this.restClient.get()
.uri("/api/v1/collections/{collection_name}", collectionName)
.headers(this::httpHeaders)
.retrieve()
.toEntity(Collection.class)
.getBody();
}
catch (HttpServerErrorException e) {
@@ -305,9 +332,11 @@ public class ChromaApi {
public List<Collection> listCollections() {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections", HttpMethod.GET, new HttpEntity<>(httpHeaders()),
CollectionList.class)
return this.restClient.get()
.uri("/api/v1/collections")
.headers(this::httpHeaders)
.retrieve()
.toEntity(CollectionList.class)
.getBody();
}
@@ -317,41 +346,55 @@ public class ChromaApi {
public void upsertEmbeddings(String collectionId, AddEmbeddingsRequest embedding) {
this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_id}/upsert", HttpMethod.POST,
this.getHttpEntityFor(embedding), Boolean.class, collectionId)
.getBody();
this.restClient.post()
.uri("/api/v1/collections/{collection_id}/upsert", collectionId)
.headers(this::httpHeaders)
.body(embedding)
.retrieve()
.toBodilessEntity();
}
public List<String> deleteEmbeddings(String collectionId, DeleteEmbeddingsRequest deleteRequest) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_id}/delete", HttpMethod.POST,
this.getHttpEntityFor(deleteRequest), List.class, collectionId)
return this.restClient.post()
.uri("/api/v1/collections/{collection_id}/delete", collectionId)
.headers(this::httpHeaders)
.body(deleteRequest)
.retrieve()
.toEntity(new ParameterizedTypeReference<List<String>>() {
})
.getBody();
}
public Long countEmbeddings(String collectionId) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_id}/count", HttpMethod.GET,
new HttpEntity<>(httpHeaders()), Long.class, collectionId)
return this.restClient.get()
.uri("/api/v1/collections/{collection_id}/count", collectionId)
.headers(this::httpHeaders)
.retrieve()
.toEntity(Long.class)
.getBody();
}
public QueryResponse queryCollection(String collectionId, QueryRequest queryRequest) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_id}/query", HttpMethod.POST,
this.getHttpEntityFor(queryRequest), QueryResponse.class, collectionId)
return this.restClient.post()
.uri("/api/v1/collections/{collection_id}/query", collectionId)
.headers(this::httpHeaders)
.body(queryRequest)
.retrieve()
.toEntity(QueryResponse.class)
.getBody();
}
public GetEmbeddingResponse getEmbeddings(String collectionId, GetEmbeddingsRequest getEmbeddingsRequest) {
return this.restTemplate
.exchange(this.baseUrl + "/api/v1/collections/{collection_id}/get", HttpMethod.POST,
this.getHttpEntityFor(getEmbeddingsRequest), GetEmbeddingResponse.class, collectionId)
return this.restClient.post()
.uri("/api/v1/collections/{collection_id}/get", collectionId)
.headers(this::httpHeaders)
.body(getEmbeddingsRequest)
.retrieve()
.toEntity(GetEmbeddingResponse.class)
.getBody();
}
@@ -365,17 +408,10 @@ public class ChromaApi {
}
}
private <T> HttpEntity<T> getHttpEntityFor(T body) {
return new HttpEntity<>(body, httpHeaders());
}
private HttpHeaders httpHeaders() {
HttpHeaders headers = new HttpHeaders();
headers.setContentType(MediaType.APPLICATION_JSON);
private void httpHeaders(HttpHeaders headers) {
if (StringUtils.hasText(this.keyToken)) {
headers.setBearerAuth(this.keyToken);
}
return headers;
}
private String getValueErrorMessage(String logString) {

View File

@@ -15,26 +15,24 @@
*/
package org.springframework.ai.chroma;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.springframework.ai.chroma.ChromaApi.AddEmbeddingsRequest;
import org.springframework.ai.chroma.ChromaApi.Collection;
import org.springframework.ai.chroma.ChromaApi.GetEmbeddingsRequest;
import org.springframework.ai.chroma.ChromaApi.QueryRequest;
import org.springframework.web.client.RestTemplate;
import static org.assertj.core.api.Assertions.assertThat;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.context.annotation.Bean;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
/**
* @author Christian Tzolov
@@ -45,7 +43,7 @@ import static org.assertj.core.api.Assertions.assertThat;
public class ChromaApiIT {
@Container
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.22");
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.12");
@Autowired
ChromaApi chroma;
@@ -179,13 +177,8 @@ public class ChromaApiIT {
public static class Config {
@Bean
public RestTemplate restTemplate() {
return new RestTemplate();
}
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
return new ChromaApi(chromaContainer.getEndpoint(), restTemplate);
public ChromaApi chromaApi() {
return new ChromaApi(chromaContainer.getEndpoint());
}
}

View File

@@ -15,28 +15,27 @@
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.web.client.RestTemplate;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.web.client.RestClient;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.testcontainers.utility.MountableFile;
import static org.assertj.core.api.Assertions.assertThat;
/**
* ChromaDB with Basic Authentication:
* https://docs.trychroma.com/usage-guide#basic-authentication
@@ -55,7 +54,7 @@ public class BasicAuthChromaWhereIT {
* https://docs.trychroma.com/usage-guide#basic-authentication
*/
@Container
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.22")
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.5.0")
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS_FILE", "/chroma/server.htpasswd")
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER",
"chromadb.auth.providers.HtpasswdFileServerAuthCredentialsProvider")
@@ -96,14 +95,13 @@ public class BasicAuthChromaWhereIT {
public static class TestApplication {
@Bean
public RestTemplate restTemplate() {
return new RestTemplate();
public RestClient.Builder builder() {
return RestClient.builder().requestFactory(new SimpleClientHttpRequestFactory());
}
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
return new ChromaApi(chromaContainer.getEndpoint(), restTemplate).withBasicAuthCredentials("admin",
"password");
public ChromaApi chromaApi(RestClient.Builder builder) {
return new ChromaApi(chromaContainer.getEndpoint(), builder).withBasicAuthCredentials("admin", "password");
}
@Bean

View File

@@ -15,16 +15,14 @@
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import org.junit.jupiter.api.Test;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
@@ -33,9 +31,11 @@ import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.web.client.RestTemplate;
import static org.assertj.core.api.Assertions.assertThat;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.web.client.RestClient;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
/**
* @author Christian Tzolov
@@ -45,7 +45,7 @@ import static org.assertj.core.api.Assertions.assertThat;
public class ChromaVectorStoreIT {
@Container
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.22");
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.5.0");
List<Document> documents = List.of(
new Document("Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!! Spring AI rocks!!",
@@ -202,13 +202,13 @@ public class ChromaVectorStoreIT {
public static class TestApplication {
@Bean
public RestTemplate restTemplate() {
return new RestTemplate();
public RestClient.Builder builder() {
return RestClient.builder().requestFactory(new SimpleClientHttpRequestFactory());
}
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
return new ChromaApi(chromaContainer.getEndpoint(), restTemplate);
public ChromaApi chromaApi(RestClient.Builder builder) {
return new ChromaApi(chromaContainer.getEndpoint(), builder);
}
@Bean

View File

@@ -15,26 +15,25 @@
*/
package org.springframework.ai.vectorstore;
import static org.assertj.core.api.Assertions.assertThat;
import java.util.List;
import java.util.Map;
import org.junit.jupiter.api.Test;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
import org.springframework.ai.chroma.ChromaApi;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.ai.openai.OpenAiEmbeddingModel;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.boot.SpringBootConfiguration;
import org.springframework.boot.test.context.runner.ApplicationContextRunner;
import org.springframework.context.annotation.Bean;
import org.springframework.web.client.RestTemplate;
import static org.assertj.core.api.Assertions.assertThat;
import org.springframework.http.client.SimpleClientHttpRequestFactory;
import org.springframework.web.client.RestClient;
import org.testcontainers.chromadb.ChromaDBContainer;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.junit.jupiter.Testcontainers;
/**
* ChromaDB with static API Token Authentication:
@@ -57,7 +56,7 @@ public class TokenSecuredChromaWhereIT {
* https://docs.trychroma.com/usage-guide#static-api-token-authentication
*/
@Container
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.4.22")
static ChromaDBContainer chromaContainer = new ChromaDBContainer("ghcr.io/chroma-core/chroma:0.5.0")
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS", CHROMA_SERVER_AUTH_CREDENTIALS)
.withEnv("CHROMA_SERVER_AUTH_CREDENTIALS_PROVIDER",
"chromadb.auth.token.TokenConfigServerAuthCredentialsProvider")
@@ -127,13 +126,13 @@ public class TokenSecuredChromaWhereIT {
public static class TestApplication {
@Bean
public RestTemplate restTemplate() {
return new RestTemplate();
public RestClient.Builder builder() {
return RestClient.builder().requestFactory(new SimpleClientHttpRequestFactory());
}
@Bean
public ChromaApi chromaApi(RestTemplate restTemplate) {
var chromaApi = new ChromaApi(chromaContainer.getEndpoint(), restTemplate);
public ChromaApi chromaApi(RestClient.Builder builder) {
var chromaApi = new ChromaApi(chromaContainer.getEndpoint(), builder);
chromaApi.withKeyToken(CHROMA_SERVER_AUTH_CREDENTIALS);
return chromaApi;
}