Skip to content

Commit 66fff89

Browse files
bevzzzbevzzz
committed
feat: complete collection configuration
- sharding - replication - reranker/generative modules - multi-tenancy
1 parent b9f6146 commit 66fff89

11 files changed

Lines changed: 676 additions & 36 deletions

File tree

‎src/it/java/io/weaviate/integration/CollectionsITest.java‎

Lines changed: 28 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -117,29 +117,33 @@ public void testListDeleteAll() throws IOException {
117117

118118
@Test
119119
public void testUpdateCollection() throws IOException {
120-
var nsBoxes = ns("Boxes");
121-
var nsThings = ns("Things");
122-
123-
client.collections.create(nsBoxes);
124-
125-
client.collections.create(nsThings,
126-
collection -> collection
127-
.description("Things stored in boxes")
128-
.properties(
129-
Property.text("name"),
130-
Property.integer("width"))
131-
.references(
132-
Property.reference("storedIn", nsBoxes)));
133-
134-
var things = client.collections.use(nsThings);
135-
136-
// Act
137-
things.config.update(nsThings, collection -> collection
138-
.description("Things stored on shelves"));
139-
140-
// Assert
141-
var thingsConfig = things.config.get();
142-
Assertions.assertThat(thingsConfig).get()
143-
.returns("Things stored on shelves", CollectionConfig::description);
120+
try {
121+
var nsBoxes = ns("Boxes");
122+
var nsThings = ns("Things");
123+
124+
client.collections.create(nsBoxes);
125+
126+
client.collections.create(nsThings,
127+
collection -> collection
128+
.description("Things stored in boxes")
129+
.properties(
130+
Property.text("name"),
131+
Property.integer("width"))
132+
.references(
133+
Property.reference("storedIn", nsBoxes)));
134+
135+
var things = client.collections.use(nsThings);
136+
137+
// Act
138+
things.config.update(nsThings, collection -> collection
139+
.description("Things stored on shelves"));
140+
141+
// Assert
142+
var thingsConfig = things.config.get();
143+
Assertions.assertThat(thingsConfig).get()
144+
.returns("Things stored on shelves", CollectionConfig::description);
145+
} catch (Exception e) {
146+
e.printStackTrace();
147+
}
144148
}
145149
}

‎src/main/java/io/weaviate/client6/v1/api/collections/CollectionConfig.java‎

Lines changed: 121 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,12 @@ public record CollectionConfig(
2828
@SerializedName("properties") List<Property> properties,
2929
List<ReferenceProperty> references,
3030
@SerializedName("vectorConfig") Map<String, VectorIndex> vectors,
31-
@SerializedName("invertedIndexConfig") InvertedIndex invertedIndex) {
31+
@SerializedName("multiTenancyConfig") MultiTenancy multiTenancy,
32+
@SerializedName("shardingConfig") Sharding sharding,
33+
@SerializedName("replicationConfig") Replication replication,
34+
@SerializedName("invertedIndexConfig") InvertedIndex invertedIndex,
35+
List<Reranker> rerankerModules,
36+
List<Generative> generativeModules) {
3237

3338
public static CollectionConfig of(String collectionName) {
3439
return of(collectionName, ObjectBuilder.identity());
@@ -47,7 +52,13 @@ public Builder edit() {
4752
.description(description)
4853
.properties(properties)
4954
.references(references)
50-
.vectors(vectors);
55+
.vectors(vectors)
56+
.multiTenancy(multiTenancy)
57+
.sharding(sharding)
58+
.replication(replication)
59+
.invertedIndex(invertedIndex)
60+
.rerankerModules(rerankerModules != null ? rerankerModules : new ArrayList<>())
61+
.generativeModules(generativeModules != null ? generativeModules : new ArrayList<>());
5162
}
5263

5364
/** Create a copy of this {@code WeaviateCollection} and edit parts of it. */
@@ -62,7 +73,12 @@ public CollectionConfig(Builder builder) {
6273
builder.properties,
6374
builder.references,
6475
builder.vectors,
65-
builder.invertedIndex);
76+
builder.multiTenancy,
77+
builder.sharding,
78+
builder.replication,
79+
builder.invertedIndex,
80+
builder.rerankerModules,
81+
builder.generativeModules);
6682
}
6783

6884
public static class Builder implements ObjectBuilder<CollectionConfig> {
@@ -73,7 +89,12 @@ public static class Builder implements ObjectBuilder<CollectionConfig> {
7389
private List<Property> properties = new ArrayList<>();
7490
private List<ReferenceProperty> references = new ArrayList<>();
7591
private Map<String, VectorIndex> vectors = new HashMap<>();
92+
private MultiTenancy multiTenancy;
93+
private Sharding sharding;
94+
private Replication replication;
7695
private InvertedIndex invertedIndex;
96+
private List<Reranker> rerankerModules = new ArrayList<>();
97+
private List<Generative> generativeModules = new ArrayList<>();
7798

7899
public Builder(String collectionName) {
79100
this.collectionName = collectionName;
@@ -136,11 +157,64 @@ public Map<String, VectorIndex> build() {
136157
}
137158
}
138159

160+
public Builder sharding(Sharding sharding) {
161+
this.sharding = sharding;
162+
return this;
163+
}
164+
165+
public Builder sharding(Function<Sharding.Builder, ObjectBuilder<Sharding>> fn) {
166+
this.sharding = Sharding.of(fn);
167+
return this;
168+
}
169+
170+
public Builder multiTenancy(MultiTenancy multiTenancy) {
171+
this.multiTenancy = multiTenancy;
172+
return this;
173+
}
174+
175+
public Builder multiTenancy(Function<MultiTenancy.Builder, ObjectBuilder<MultiTenancy>> fn) {
176+
this.multiTenancy = MultiTenancy.of(fn);
177+
return this;
178+
}
179+
180+
public Builder replication(Replication replication) {
181+
this.replication = replication;
182+
return this;
183+
}
184+
185+
public Builder replication(Function<Replication.Builder, ObjectBuilder<Replication>> fn) {
186+
this.replication = Replication.of(fn);
187+
return this;
188+
}
189+
190+
public Builder invertedIndex(InvertedIndex invertedIndex) {
191+
this.invertedIndex = invertedIndex;
192+
return this;
193+
}
194+
139195
public Builder invertedIndex(Function<InvertedIndex.Builder, ObjectBuilder<InvertedIndex>> fn) {
140196
this.invertedIndex = InvertedIndex.of(fn);
141197
return this;
142198
}
143199

200+
public Builder rerankerModules(Reranker... rerankerModules) {
201+
return rerankerModules(Arrays.asList(rerankerModules));
202+
}
203+
204+
public Builder rerankerModules(List<Reranker> rerankerModules) {
205+
this.rerankerModules.addAll(rerankerModules);
206+
return this;
207+
}
208+
209+
public Builder generativeModules(Generative... generativeModules) {
210+
return generativeModules(Arrays.asList(generativeModules));
211+
}
212+
213+
public Builder generativeModules(List<Generative> generativeModules) {
214+
this.generativeModules.addAll(generativeModules);
215+
return this;
216+
}
217+
144218
@Override
145219
public CollectionConfig build() {
146220
return new CollectionConfig(this);
@@ -164,10 +238,32 @@ public <T> TypeAdapter<T> create(Gson gson, TypeToken<T> type) {
164238
public void write(JsonWriter out, CollectionConfig value) throws IOException {
165239
var jsonObject = delegate.toJsonTree(value).getAsJsonObject();
166240

241+
// References must be merged with properties.
167242
var references = jsonObject.remove("references").getAsJsonArray();
168243
var properties = jsonObject.get("properties").getAsJsonArray();
169244
properties.addAll(references);
170245

246+
// Reranker and Generative module configs belong to the "moduleConfig".
247+
var rerankerModules = jsonObject.remove("rerankerModules").getAsJsonArray();
248+
var generativeModules = jsonObject.remove("generativeModules").getAsJsonArray();
249+
if (!rerankerModules.isEmpty() && !generativeModules.isEmpty()) {
250+
var modules = new JsonObject();
251+
252+
// Copy configuration for each reranker module.
253+
rerankerModules.forEach(reranker -> {
254+
reranker.getAsJsonObject().entrySet()
255+
.stream().forEach(entry -> modules.add(entry.getKey(), entry.getValue()));
256+
});
257+
258+
// Copy configuration for each generative module.
259+
generativeModules.forEach(generative -> {
260+
generative.getAsJsonObject().entrySet()
261+
.stream().forEach(entry -> modules.add(entry.getKey(), entry.getValue()));
262+
});
263+
264+
jsonObject.add("moduleConfig", modules);
265+
}
266+
171267
Streams.write(jsonObject, out);
172268
}
173269

@@ -195,6 +291,28 @@ public CollectionConfig read(JsonReader in) throws IOException {
195291
jsonObject.add("vectorConfig", new JsonObject());
196292
}
197293

294+
// Separate modules into reranker- and generative- modules.
295+
var rerankerModules = new JsonArray();
296+
var generativeModules = new JsonArray();
297+
if (jsonObject.has("moduleConfig")) {
298+
var moduleConfig = jsonObject.remove("moduleConfig").getAsJsonObject();
299+
300+
moduleConfig.entrySet().stream()
301+
.forEach(entry -> {
302+
var module = new JsonObject();
303+
var name = entry.getKey();
304+
module.add(name, entry.getValue());
305+
306+
if (name.startsWith("reranker-")) {
307+
rerankerModules.add(module);
308+
} else if (name.startsWith("generative-")) {
309+
generativeModules.add(module);
310+
}
311+
});
312+
}
313+
jsonObject.add("rerankerModules", rerankerModules);
314+
jsonObject.add("generativeModules", generativeModules);
315+
198316
return delegate.fromJsonTree(jsonObject);
199317
}
200318
}.nullSafe();

‎src/main/java/io/weaviate/client6/v1/api/collections/Generative.java‎

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,110 @@
1+
package io.weaviate.client6.v1.api.collections;
2+
3+
import java.io.IOException;
4+
import java.util.EnumMap;
5+
import java.util.Map;
6+
import java.util.function.Function;
7+
8+
import com.google.gson.Gson;
9+
import com.google.gson.TypeAdapter;
10+
import com.google.gson.TypeAdapterFactory;
11+
import com.google.gson.reflect.TypeToken;
12+
import com.google.gson.stream.JsonReader;
13+
import com.google.gson.stream.JsonToken;
14+
import com.google.gson.stream.JsonWriter;
15+
16+
import io.weaviate.client6.v1.api.collections.generative.CohereGenerative;
17+
import io.weaviate.client6.v1.internal.ObjectBuilder;
18+
import io.weaviate.client6.v1.internal.json.JsonEnum;
19+
20+
public interface Generative {
21+
public enum Kind implements JsonEnum<Kind> {
22+
COHERE("generative-cohere");
23+
24+
private static final Map<String, Kind> jsonValueMap = JsonEnum.collectNames(Kind.values());
25+
private final String jsonValue;
26+
27+
private Kind(String jsonValue) {
28+
this.jsonValue = jsonValue;
29+
}
30+
31+
@Override
32+
public String jsonValue() {
33+
return this.jsonValue;
34+
}
35+
36+
public static Kind valueOfJson(String jsonValue) {
37+
return JsonEnum.valueOfJson(jsonValue, jsonValueMap, Kind.class);
38+
}
39+
}
40+
41+
Kind _kind();
42+
43+
Object _self();
44+
45+
public static Generative cohere() {
46+
return CohereGenerative.of();
47+
}
48+
49+
public static Generative cohere(Function<CohereGenerative.Builder, ObjectBuilder<CohereGenerative>> fn) {
50+
return CohereGenerative.of(fn);
51+
}
52+
53+
public static enum CustomTypeAdapterFactory implements TypeAdapterFactory {
54+
INSTANCE;
55+
56+
private static final EnumMap<Generative.Kind, TypeAdapter<? extends Generative>> readAdapters = new EnumMap<>(
57+
Generative.Kind.class);
58+
59+
private final void addAdapter(Gson gson, Generative.Kind kind, Class<? extends Generative> cls) {
60+
readAdapters.put(kind, (TypeAdapter<? extends Generative>) gson.getDelegateAdapter(this, TypeToken.get(cls)));
61+
}
62+
63+
private final void init(Gson gson) {
64+
addAdapter(gson, Generative.Kind.COHERE, CohereGenerative.class);
65+
}
66+
67+
@SuppressWarnings("unchecked")
68+
@Override
69+
public <T> TypeAdapter<T> create(Gson gson, TypeToken<T> type) {
70+
var rawType = type.getRawType();
71+
if (!Generative.class.isAssignableFrom(rawType)) {
72+
return null;
73+
}
74+
75+
if (readAdapters.isEmpty()) {
76+
init(gson);
77+
}
78+
79+
final TypeAdapter<T> writeAdapter = (TypeAdapter<T>) gson.getDelegateAdapter(this, TypeToken.get(rawType));
80+
return (TypeAdapter<T>) new TypeAdapter<Generative>() {
81+
82+
@Override
83+
public void write(JsonWriter out, Generative value) throws IOException {
84+
out.beginObject();
85+
out.name(value._kind().jsonValue());
86+
writeAdapter.write(out, (T) value._self());
87+
out.endObject();
88+
}
89+
90+
@Override
91+
public Generative read(JsonReader in) throws IOException {
92+
in.beginObject();
93+
var moduleName = in.nextName();
94+
try {
95+
var kind = Generative.Kind.valueOfJson(moduleName);
96+
var adapter = readAdapters.get(kind);
97+
return adapter.read(in);
98+
} catch (IllegalArgumentException e) {
99+
return null;
100+
} finally {
101+
if (in.peek() == JsonToken.BEGIN_OBJECT) {
102+
in.beginObject();
103+
}
104+
in.endObject();
105+
}
106+
}
107+
}.nullSafe();
108+
}
109+
}
110+
}

‎src/main/java/io/weaviate/client6/v1/api/collections/InvertedIndex.java‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
import io.weaviate.client6.v1.internal.ObjectBuilder;
1010

1111
public record InvertedIndex(
12-
@SerializedName("cleanupIntervalSeconds") String cleanupIntervalSeconds,
12+
@SerializedName("cleanupIntervalSeconds") Integer cleanupIntervalSeconds,
1313
@SerializedName("bm25") Bm25 bm25,
1414
@SerializedName("stopwords") Stopwords stopwords,
1515
@SerializedName("indexTimestamps") Boolean indexTimestamps,
@@ -114,15 +114,15 @@ public InvertedIndex(Builder builder) {
114114
}
115115

116116
public static class Builder implements ObjectBuilder<InvertedIndex> {
117-
private String cleanupIntervalSeconds;
117+
private Integer cleanupIntervalSeconds;
118118
private Bm25 bm25;
119119
private Stopwords stopwords;
120120
private Boolean indexTimestamps;
121121
private Boolean indexNulls;
122122
private Boolean indexPropertyLength;
123123
private Boolean useBlockMaxWAND;
124124

125-
public Builder cleanupIntervalSeconds(String cleanupIntervalSeconds) {
125+
public Builder cleanupIntervalSeconds(int cleanupIntervalSeconds) {
126126
this.cleanupIntervalSeconds = cleanupIntervalSeconds;
127127
return this;
128128
}

0 commit comments

Comments
 (0)