Skip to content

Commit

Permalink
Add ServiceConnection support for
Browse files Browse the repository at this point in the history
 * ChromaDBContainer
 * MilvusContainer
 * QdrantContainer
 * RedisStackContainer
 * WeaviateContainer
 * OllamaContainer

Add docs
  • Loading branch information
eddumelendez authored and markpollack committed Mar 19, 2024
1 parent 666474c commit 8503078
Show file tree
Hide file tree
Showing 37 changed files with 1,504 additions and 17 deletions.
1 change: 1 addition & 0 deletions document-readers/pdf-reader/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -55,6 +55,7 @@
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>

Expand Down
1 change: 1 addition & 0 deletions models/spring-ai-ollama/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>
</dependencies>
Expand Down
2 changes: 2 additions & 0 deletions models/spring-ai-postgresml/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -54,12 +54,14 @@
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>

<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>postgresql</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>

Expand Down
1 change: 1 addition & 0 deletions models/spring-ai-transformers/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,7 @@
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>junit-jupiter</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>

Expand Down
3 changes: 2 additions & 1 deletion pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@
<module>spring-ai-retry</module>
<module>vector-stores/spring-ai-mongodb-atlas-store</module>
<module>spring-ai-spring-boot-starters/spring-ai-starter-mongodb-atlas-store</module>
<module>spring-ai-spring-boot-testcontainers</module>
</modules>

<organization>
Expand Down Expand Up @@ -135,7 +136,7 @@
<qdrant.version>1.7.1</qdrant.version>

<!-- testing dependencies -->
<testcontainers.version>1.19.6</testcontainers.version>
<testcontainers.version>1.19.7</testcontainers.version>

<!-- documentation dependencies -->
<io.spring.maven.antora-version>0.0.4</io.spring.maven.antora-version>
Expand Down
6 changes: 6 additions & 0 deletions spring-ai-bom/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -313,6 +313,12 @@
<artifactId>spring-ai-mongodb-atlas-store-spring-boot-starter</artifactId>
<version>${project.version}</version>
</dependency>

<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-spring-boot-testcontainers</artifactId>
<version>${project.version}</version>
</dependency>
</dependencies>
</dependencyManagement>

Expand Down
1 change: 1 addition & 0 deletions spring-ai-docs/src/main/antora/modules/ROOT/nav.adoc
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,7 @@
** xref:api/etl-pipeline.adoc[]
** xref:api/testing.adoc[]
** xref:api/generic-model.adoc[]
* xref:api/testcontainers.adoc[Testcontainers]
* xref:contribution-guidelines.adoc[Contribution Guidelines]
* Appendices
** xref:upgrade-notes.adoc[]
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
[[testcontainers]]
= Testcontainers

== Service Connections

The following service connection factories are provided in the spring-ai-spring-boot-testcontainers jar:

[cols="|,|"]
|====
| Connection Details | Matched on
| `ChromaConnectionDetails`
| Containers of type `ChromaDBContainer`

| `MilvusServiceClientConnectionDetails`
| Containers of type `MilvusContainer`

| `OllamaConnectionDetails`
| Containers of type `OllamaContainer`

| `QdrantConnectionDetails`
| Containers of type `QdrantContainer`

| `RedisConnectionDetails`
| Containers of type `RedisStackContainer`

| `WeaviateConnectionDetails`
| Containers of type `WeaviateContainer`
|====
3 changes: 2 additions & 1 deletion spring-ai-spring-boot-autoconfigure/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -273,6 +273,7 @@
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>postgresql</artifactId>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>

Expand Down Expand Up @@ -306,7 +307,7 @@
<dependency>
<groupId>org.testcontainers</groupId>
<artifactId>qdrant</artifactId>
<version>1.19.6</version>
<version>${testcontainers.version}</version>
<scope>test</scope>
</dependency>

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,6 +31,7 @@
* {@link AutoConfiguration Auto-configuration} for Ollama Chat Client.
*
* @author Christian Tzolov
* @author Eddú Meléndez
* @since 0.8.0
*/
@AutoConfiguration(after = RestClientAutoConfiguration.class)
Expand All @@ -39,10 +40,16 @@
OllamaConnectionProperties.class })
public class OllamaAutoConfiguration {

@Bean
@ConditionalOnMissingBean(OllamaConnectionDetails.class)
public PropertiesOllamaConnectionDetails ollamaConnectionDetails(OllamaConnectionProperties properties) {
return new PropertiesOllamaConnectionDetails(properties);
}

@Bean
@ConditionalOnMissingBean
public OllamaApi ollamaApi(OllamaConnectionProperties properties, RestClient.Builder restClientBuilder) {
return new OllamaApi(properties.getBaseUrl(), restClientBuilder);
public OllamaApi ollamaApi(OllamaConnectionDetails connectionDetails, RestClient.Builder restClientBuilder) {
return new OllamaApi(connectionDetails.getBaseUrl(), restClientBuilder);
}

@Bean
Expand All @@ -65,4 +72,19 @@ public OllamaEmbeddingClient ollamaEmbeddingClient(OllamaApi ollamaApi, OllamaEm
.withDefaultOptions(properties.getOptions());
}

private static class PropertiesOllamaConnectionDetails implements OllamaConnectionDetails {

private final OllamaConnectionProperties properties;

PropertiesOllamaConnectionDetails(OllamaConnectionProperties properties) {
this.properties = properties;
}

@Override
public String getBaseUrl() {
return this.properties.getBaseUrl();
}

}

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
/*
* Copyright 2023 - 2024 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.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.ollama;

import org.springframework.boot.autoconfigure.service.connection.ConnectionDetails;

/**
* @author Eddú Meléndez
*/
public interface OllamaConnectionDetails extends ConnectionDetails {

String getBaseUrl();

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
/*
* Copyright 2023 - 2024 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.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.vectorstore.chroma;

import org.springframework.boot.autoconfigure.service.connection.ConnectionDetails;

/**
* @author Eddú Meléndez
*/
public interface ChromaConnectionDetails extends ConnectionDetails {

String getHost();

int getPort();

}
Original file line number Diff line number Diff line change
Expand Up @@ -30,12 +30,19 @@

/**
* @author Christian Tzolov
* @author Eddú Meléndez
*/
@AutoConfiguration
@ConditionalOnClass({ EmbeddingClient.class, RestTemplate.class, ChromaVectorStore.class, ObjectMapper.class })
@EnableConfigurationProperties({ ChromaApiProperties.class, ChromaVectorStoreProperties.class })
public class ChromaVectorStoreAutoConfiguration {

@Bean
@ConditionalOnMissingBean(ChromaConnectionDetails.class)
PropertiesChromaConnectionDetails chromaConnectionDetails(ChromaApiProperties properties) {
return new PropertiesChromaConnectionDetails(properties);
}

@Bean
@ConditionalOnMissingBean
public RestTemplate restTemplate() {
Expand All @@ -44,9 +51,10 @@ public RestTemplate restTemplate() {

@Bean
@ConditionalOnMissingBean
public ChromaApi chromaApi(ChromaApiProperties apiProperties, RestTemplate restTemplate) {
public ChromaApi chromaApi(ChromaApiProperties apiProperties, RestTemplate restTemplate,
ChromaConnectionDetails connectionDetails) {

String chromaUrl = String.format("%s:%s", apiProperties.getHost(), apiProperties.getPort());
String chromaUrl = String.format("%s:%s", connectionDetails.getHost(), connectionDetails.getPort());

var chromaApi = new ChromaApi(chromaUrl, restTemplate, new ObjectMapper());

Expand All @@ -67,4 +75,24 @@ public ChromaVectorStore vectorStore(EmbeddingClient embeddingClient, ChromaApi
return new ChromaVectorStore(embeddingClient, chromaApi, storeProperties.getCollectionName());
}

private static class PropertiesChromaConnectionDetails implements ChromaConnectionDetails {

private final ChromaApiProperties properties;

PropertiesChromaConnectionDetails(ChromaApiProperties properties) {
this.properties = properties;
}

@Override
public String getHost() {
return this.properties.getHost();
}

@Override
public int getPort() {
return this.properties.getPort();
}

}

}
Original file line number Diff line number Diff line change
@@ -0,0 +1,29 @@
/*
* Copyright 2023 - 2024 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.
* You may obtain a copy of the License at
*
* https://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.springframework.ai.autoconfigure.vectorstore.milvus;

import org.springframework.boot.autoconfigure.service.connection.ConnectionDetails;

/**
* @author Eddú Meléndez
*/
public interface MilvusServiceClientConnectionDetails extends ConnectionDetails {

String getHost();

int getPort();

}
Original file line number Diff line number Diff line change
Expand Up @@ -34,12 +34,20 @@

/**
* @author Christian Tzolov
* @author Eddú Meléndez
*/
@AutoConfiguration
@ConditionalOnClass({ MilvusVectorStore.class, EmbeddingClient.class })
@EnableConfigurationProperties({ MilvusServiceClientProperties.class, MilvusVectorStoreProperties.class })
public class MilvusVectorStoreAutoConfiguration {

@Bean
@ConditionalOnMissingBean(MilvusServiceClientConnectionDetails.class)
PropertiesMilvusServiceClientConnectionDetails milvusServiceClientConnectionDetails(
MilvusServiceClientProperties properties) {
return new PropertiesMilvusServiceClientConnectionDetails(properties);
}

@Bean
@ConditionalOnMissingBean
public MilvusVectorStore vectorStore(MilvusServiceClient milvusClient, EmbeddingClient embeddingClient,
Expand All @@ -60,11 +68,11 @@ public MilvusVectorStore vectorStore(MilvusServiceClient milvusClient, Embedding
@Bean
@ConditionalOnMissingBean
public MilvusServiceClient milvusClient(MilvusVectorStoreProperties serverProperties,
MilvusServiceClientProperties clientProperties) {
MilvusServiceClientProperties clientProperties, MilvusServiceClientConnectionDetails connectionDetails) {

var builder = ConnectParam.newBuilder()
.withHost(clientProperties.getHost())
.withPort(clientProperties.getPort())
.withHost(connectionDetails.getHost())
.withPort(connectionDetails.getPort())
.withDatabaseName(serverProperties.getDatabaseName())
.withConnectTimeout(clientProperties.getConnectTimeoutMs(), TimeUnit.MILLISECONDS)
.withKeepAliveTime(clientProperties.getKeepAliveTimeMs(), TimeUnit.MILLISECONDS)
Expand Down Expand Up @@ -105,4 +113,25 @@ public MilvusServiceClient milvusClient(MilvusVectorStoreProperties serverProper
return new MilvusServiceClient(builder.build());
}

private static class PropertiesMilvusServiceClientConnectionDetails
implements MilvusServiceClientConnectionDetails {

private final MilvusServiceClientProperties properties;

PropertiesMilvusServiceClientConnectionDetails(MilvusServiceClientProperties properties) {
this.properties = properties;
}

@Override
public String getHost() {
return this.properties.getHost();
}

@Override
public int getPort() {
return this.properties.getPort();
}

}

}
Loading

0 comments on commit 8503078

Please sign in to comment.