Skip to content

Commit 478b7a6

Browse files
bevzzzbevzzz
committed
feat(rag): add dynamic providers for Anthropic/Anyscale/Aws/Cohere
1 parent f18ddf2 commit 478b7a6

5 files changed

Lines changed: 591 additions & 7 deletions

File tree

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

Lines changed: 38 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ void appendTo(WeaviateProtoGenerative.GenerativeSearch.Builder req) {
5858
}
5959
}
6060

61-
public record Single(String prompt, boolean debug) {
61+
public record Single(String prompt, boolean debug, List<DynamicProvider> providers) {
6262
public static Single of(String prompt) {
6363
return of(prompt, ObjectBuilder.identity());
6464
}
@@ -68,11 +68,12 @@ public static Single of(String prompt, Function<Builder, ObjectBuilder<Single>>
6868
}
6969

7070
public Single(Builder builder) {
71-
this(builder.prompt, builder.debug);
71+
this(builder.prompt, builder.debug, builder.providers);
7272
}
7373

7474
public static class Builder implements ObjectBuilder<Single> {
7575
private final String prompt;
76+
private final List<DynamicProvider> providers = new ArrayList<>();
7677
private boolean debug = false;
7778

7879
public Builder(String prompt) {
@@ -84,21 +85,36 @@ public Builder debug(boolean enable) {
8485
return this;
8586
}
8687

88+
public Builder generativeProvider(DynamicProvider provider) {
89+
providers.clear(); // Protobuf allows `repeated` but the server expects there to be 1.
90+
providers.add(provider);
91+
return this;
92+
}
93+
8794
@Override
8895
public Single build() {
8996
return new Single(this);
9097
}
9198
}
9299

93100
public void appendTo(WeaviateProtoGenerative.GenerativeSearch.Builder req) {
101+
var ragProviders = providers.stream()
102+
.map(provider -> {
103+
var proto = WeaviateProtoGenerative.GenerativeProvider.newBuilder();
104+
provider.appendTo(proto);
105+
return proto.build();
106+
})
107+
.toList();
108+
94109
req.setSingle(
95110
WeaviateProtoGenerative.GenerativeSearch.Single.newBuilder()
96111
.setPrompt(prompt)
97-
.setDebug(debug));
112+
.setDebug(debug)
113+
.addAllQueries(ragProviders));
98114
}
99115
}
100116

101-
public record Grouped(String prompt, boolean debug, List<String> properties) {
117+
public record Grouped(String prompt, boolean debug, List<String> properties, List<DynamicProvider> providers) {
102118
public static Grouped of(String prompt) {
103119
return of(prompt, ObjectBuilder.identity());
104120
}
@@ -108,11 +124,12 @@ public static Grouped of(String prompt, Function<Builder, ObjectBuilder<Grouped>
108124
}
109125

110126
public Grouped(Builder builder) {
111-
this(builder.prompt, builder.debug, builder.properties);
127+
this(builder.prompt, builder.debug, builder.properties, builder.providers);
112128
}
113129

114130
public static class Builder implements ObjectBuilder<Grouped> {
115131
private final String prompt;
132+
private final List<DynamicProvider> providers = new ArrayList<>();
116133
private final List<String> properties = new ArrayList<>();
117134
private boolean debug = false;
118135

@@ -129,6 +146,12 @@ public Builder properties(List<String> properties) {
129146
return this;
130147
}
131148

149+
public Builder generativeProvider(DynamicProvider provider) {
150+
providers.clear(); // Protobuf allows `repeated` but the server expects there to be 1.
151+
providers.add(provider);
152+
return this;
153+
}
154+
132155
public Builder debug(boolean enable) {
133156
this.debug = enable;
134157
return this;
@@ -151,6 +174,16 @@ public void appendTo(WeaviateProtoGenerative.GenerativeSearch.Builder req) {
151174
.addAllValues(properties));
152175

153176
}
177+
178+
var ragProviders = providers.stream()
179+
.map(provider -> {
180+
var proto = WeaviateProtoGenerative.GenerativeProvider.newBuilder();
181+
provider.appendTo(proto);
182+
return proto.build();
183+
})
184+
.toList();
185+
grouped.addAllQueries(ragProviders);
186+
154187
req.setGrouped(grouped);
155188
}
156189
}

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

Lines changed: 178 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,13 +8,17 @@
88
import com.google.gson.annotations.SerializedName;
99

1010
import io.weaviate.client6.v1.api.collections.Generative;
11+
import io.weaviate.client6.v1.api.collections.generate.DynamicProvider;
1112
import io.weaviate.client6.v1.internal.ObjectBuilder;
13+
import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoBase;
14+
import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoGenerative;
1215

1316
public record AnthropicGenerative(
1417
@SerializedName("model") String model,
1518
@SerializedName("maxTokens") Integer maxTokens,
1619
@SerializedName("temperature") Float temperature,
1720
@SerializedName("topK") Integer topK,
21+
@SerializedName("topP") Float topP,
1822
@SerializedName("stopSequences") List<String> stopSequences) implements Generative {
1923

2024
@Override
@@ -41,21 +45,30 @@ public AnthropicGenerative(Builder builder) {
4145
builder.maxTokens,
4246
builder.temperature,
4347
builder.topK,
48+
builder.topP,
4449
builder.stopSequences);
4550
}
4651

4752
public static class Builder implements ObjectBuilder<AnthropicGenerative> {
4853
private Integer topK;
54+
private Float topP;
4955
private String model;
5056
private Integer maxTokens;
5157
private Float temperature;
52-
private List<String> stopSequences = new ArrayList<>();
58+
private final List<String> stopSequences = new ArrayList<>();
5359

60+
/** Top K value for sampling. */
5461
public Builder topK(int topK) {
5562
this.topK = topK;
5663
return this;
5764
}
5865

66+
/** Top P value for nucleus sampling. */
67+
public Builder topP(float topP) {
68+
this.topP = topP;
69+
return this;
70+
}
71+
5972
/** Select generative model. */
6073
public Builder model(String model) {
6174
this.model = model;
@@ -68,12 +81,18 @@ public Builder maxTokens(int maxTokens) {
6881
return this;
6982
}
7083

84+
/**
85+
* Set tokens which should signal the model to stop generating further output.
86+
*/
7187
public Builder stopSequences(String... stopSequences) {
7288
return stopSequences(Arrays.asList(stopSequences));
7389
}
7490

91+
/**
92+
* Set tokens which should signal the model to stop generating further output.
93+
*/
7594
public Builder stopSequences(List<String> stopSequences) {
76-
this.stopSequences = stopSequences;
95+
this.stopSequences.addAll(stopSequences);
7796
return this;
7897
}
7998

@@ -102,4 +121,161 @@ public Generative.Kind _kind() {
102121
public static record Usage(Long inputTokens, Long outputTokens) {
103122
}
104123
}
124+
125+
public static record Provider(
126+
String baseUrl,
127+
Integer maxTokens,
128+
String model,
129+
Float temperature,
130+
Integer topK,
131+
Float topP,
132+
List<String> stopSequences,
133+
List<String> images,
134+
List<String> imageProperties) implements DynamicProvider {
135+
136+
public static Provider of(
137+
Function<AnthropicGenerative.Provider.Builder, ObjectBuilder<AnthropicGenerative.Provider>> fn) {
138+
return fn.apply(new Builder()).build();
139+
}
140+
141+
@Override
142+
public void appendTo(
143+
io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoGenerative.GenerativeProvider.Builder req) {
144+
var provider = WeaviateProtoGenerative.GenerativeAnthropic.newBuilder();
145+
if (baseUrl != null) {
146+
provider.setBaseUrl(baseUrl);
147+
}
148+
if (maxTokens != null) {
149+
provider.setMaxTokens(maxTokens);
150+
}
151+
if (model != null) {
152+
provider.setModel(model);
153+
}
154+
if (temperature != null) {
155+
provider.setTemperature(temperature);
156+
}
157+
if (topK != null) {
158+
provider.setTopK(topK);
159+
}
160+
if (topP != null) {
161+
provider.setTopP(topP);
162+
}
163+
164+
if (stopSequences != null) {
165+
provider.setStopSequences(WeaviateProtoBase.TextArray.newBuilder()
166+
.addAllValues(stopSequences));
167+
}
168+
if (images != null) {
169+
provider.setImages(WeaviateProtoBase.TextArray.newBuilder()
170+
.addAllValues(images));
171+
}
172+
if (imageProperties != null) {
173+
provider.setImageProperties(WeaviateProtoBase.TextArray.newBuilder()
174+
.addAllValues(imageProperties));
175+
}
176+
req.setAnthropic(provider);
177+
}
178+
179+
public Provider(Builder builder) {
180+
this(
181+
builder.baseUrl,
182+
builder.maxTokens,
183+
builder.model,
184+
builder.temperature,
185+
builder.topK,
186+
builder.topP,
187+
builder.stopSequences,
188+
builder.images,
189+
builder.imageProperties);
190+
}
191+
192+
public static class Builder implements ObjectBuilder<AnthropicGenerative.Provider> {
193+
private String baseUrl;
194+
private Integer topK;
195+
private Float topP;
196+
private String model;
197+
private Integer maxTokens;
198+
private Float temperature;
199+
private final List<String> stopSequences = new ArrayList<>();
200+
private final List<String> images = new ArrayList<>();
201+
private final List<String> imageProperties = new ArrayList<>();
202+
203+
/** Base URL of the generative provider. */
204+
public Builder baseUrl(String baseUrl) {
205+
this.baseUrl = baseUrl;
206+
return this;
207+
}
208+
209+
/** Top K value for sampling. */
210+
public Builder topK(int topK) {
211+
this.topK = topK;
212+
return this;
213+
}
214+
215+
/** Top P value for nucleus sampling. */
216+
public Builder topP(float topP) {
217+
this.topP = topP;
218+
return this;
219+
}
220+
221+
/** Select generative model. */
222+
public Builder model(String model) {
223+
this.model = model;
224+
return this;
225+
}
226+
227+
/** Limit the number of tokens to generate in the response. */
228+
public Builder maxTokens(int maxTokens) {
229+
this.maxTokens = maxTokens;
230+
return this;
231+
}
232+
233+
/**
234+
* Set tokens which should signal the model to stop generating further output.
235+
*/
236+
public Builder stopSequences(String... stopSequences) {
237+
return stopSequences(Arrays.asList(stopSequences));
238+
}
239+
240+
/**
241+
* Set tokens which should signal the model to stop generating further output.
242+
*/
243+
public Builder stopSequences(List<String> stopSequences) {
244+
this.stopSequences.addAll(stopSequences);
245+
return this;
246+
}
247+
248+
public Builder images(String... images) {
249+
return images(Arrays.asList(images));
250+
}
251+
252+
public Builder images(List<String> images) {
253+
this.images.addAll(images);
254+
return this;
255+
}
256+
257+
public Builder imageProperties(String... imageProperties) {
258+
return imageProperties(Arrays.asList(imageProperties));
259+
}
260+
261+
public Builder imageProperties(List<String> imageProperties) {
262+
this.imageProperties.addAll(imageProperties);
263+
return this;
264+
}
265+
266+
/**
267+
* Control the randomness of the model's output.
268+
* Higher values make output more random.
269+
*/
270+
public Builder temperature(float temperature) {
271+
this.temperature = temperature;
272+
return this;
273+
}
274+
275+
@Override
276+
public AnthropicGenerative.Provider build() {
277+
return new AnthropicGenerative.Provider(this);
278+
}
279+
}
280+
}
105281
}

0 commit comments

Comments
 (0)