11package io .weaviate .containers ;
22
33import java .io .IOException ;
4+ import java .util .Arrays ;
5+ import java .util .HashMap ;
46import java .util .HashSet ;
7+ import java .util .Map ;
58import java .util .Set ;
69
710import org .testcontainers .weaviate .WeaviateContainer ;
1013import io .weaviate .client6 .WeaviateClient ;
1114
1215public class Weaviate extends WeaviateContainer {
13- private static WeaviateClient clientInstance ;
16+ private WeaviateClient clientInstance ;
1417
1518 public static final String VERSION = "1.29.0" ;
1619 public static final String DOCKER_IMAGE = "semitechnologies/weaviate" ;
@@ -42,14 +45,14 @@ public static Weaviate.Builder custom() {
4245
4346 public static class Builder {
4447 private String versionTag ;
45- private Set <String > enableModules ;
48+ private Set <String > enableModules = new HashSet <>() ;
4649 private String defaultVectorizerModule ;
47- private String contextionaryUrl ;
4850 private boolean telemetry ;
4951
52+ private Map <String , String > environment = new HashMap <>();
53+
5054 public Builder () {
5155 this .versionTag = VERSION ;
52- this .enableModules = new HashSet <>();
5356 this .telemetry = false ;
5457 }
5558
@@ -58,43 +61,46 @@ public Builder withVersion(String version) {
5861 return this ;
5962 }
6063
61- public Builder addModule (String module ) {
62- enableModules .add ( module );
64+ public Builder addModules (String ... modules ) {
65+ enableModules .addAll ( Arrays . asList ( modules ) );
6366 return this ;
6467 }
6568
6669 public Builder withDefaultVectorizer (String module ) {
67- addModule (module );
68- defaultVectorizerModule = module ;
70+ addModules (module );
71+ environment . put ( "DEFAULT_VECTORIZER_MODULE" , module ) ;
6972 return this ;
7073 }
7174
7275 public Builder withContextionaryUrl (String url ) {
73- contextionaryUrl = url ;
76+ addModules (Contextionary .MODULE );
77+ environment .put ("CONTEXTIONARY_URL" , url );
78+ return this ;
79+ }
80+
81+ public Builder withImageInference (String url , String module ) {
82+ addModules (module );
83+ environment .put ("IMAGE_INFERENCE_API" , "http://" + url );
7484 return this ;
7585 }
7686
77- public Builder enableTelemetry () {
78- telemetry = true ;
87+ public Builder enableTelemetry (boolean enable ) {
88+ telemetry = enable ;
7989 return this ;
8090 }
8191
8292 public Weaviate build () {
8393 var c = new Weaviate (DOCKER_IMAGE + ":" + versionTag );
8494
8595 if (!enableModules .isEmpty ()) {
96+ c .withEnv ("ENABLE_API_BASED_MODULES" , "'true'" );
8697 c .withEnv ("ENABLE_MODULES" , String .join ("," , enableModules ));
8798 }
88- if (defaultVectorizerModule != null ) {
89- c .withEnv ("DEFAULT_VECTORIZER_MODULE" , defaultVectorizerModule );
90- }
91- if (contextionaryUrl != null ) {
92- c .withEnv ("CONTEXTIONARY_URL" , contextionaryUrl );
93- }
9499 if (!telemetry ) {
95100 c .withEnv ("DISABLE_TELEMETRY" , "true" );
96101 }
97102
103+ environment .forEach ((name , value ) -> c .withEnv (name , value ));
98104 c .withCreateContainerCmdModifier (cmd -> cmd .withHostName ("weaviate" ));
99105 return c ;
100106 }
0 commit comments