Skip to content

Commit bd2deb7

Browse files
bevzzzbevzzz
committed
feat: read provider metadata from generative response
1 parent c8b0770 commit bd2deb7

17 files changed

Lines changed: 271 additions & 42 deletions

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

Lines changed: 110 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,22 @@
11
package io.weaviate.client6.v1.api.collections.generate;
22

3+
import java.util.ArrayList;
34
import java.util.List;
45

6+
import io.weaviate.client6.v1.api.collections.generative.AnthropicGenerative;
7+
import io.weaviate.client6.v1.api.collections.generative.AnyscaleGenerative;
8+
import io.weaviate.client6.v1.api.collections.generative.AwsGenerative;
9+
import io.weaviate.client6.v1.api.collections.generative.CohereGenerative;
10+
import io.weaviate.client6.v1.api.collections.generative.DatabricksGenerative;
511
import io.weaviate.client6.v1.api.collections.generative.DummyGenerative;
12+
import io.weaviate.client6.v1.api.collections.generative.FriendliaiGenerative;
13+
import io.weaviate.client6.v1.api.collections.generative.GoogleGenerative;
14+
import io.weaviate.client6.v1.api.collections.generative.MistralGenerative;
15+
import io.weaviate.client6.v1.api.collections.generative.NvidiaGenerative;
16+
import io.weaviate.client6.v1.api.collections.generative.OllamaGenerative;
17+
import io.weaviate.client6.v1.api.collections.generative.OpenAiGenerative;
18+
import io.weaviate.client6.v1.api.collections.generative.ProviderMetadata;
19+
import io.weaviate.client6.v1.api.collections.generative.XaiGenerative;
620
import io.weaviate.client6.v1.api.collections.query.QueryResponse;
721
import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoGenerative;
822
import io.weaviate.client6.v1.internal.grpc.protocol.WeaviateProtoSearchGet;
@@ -49,6 +63,102 @@ static TaskOutput unmarshalTaskOutput(List<WeaviateProtoGenerative.GenerativeRep
4963
ProviderMetadata providerMetadata = null;
5064
if (metadata.hasDummy()) {
5165
providerMetadata = new DummyGenerative.Metadata();
66+
} else if (metadata.hasAws()) {
67+
providerMetadata = new AwsGenerative.Metadata();
68+
} else if (metadata.hasAnthropic()) {
69+
var anthropic = metadata.getAnthropic();
70+
var usage = anthropic.getUsage();
71+
providerMetadata = new AnthropicGenerative.Metadata(new AnthropicGenerative.Metadata.Usage(
72+
usage.getInputTokens(),
73+
usage.getOutputTokens()));
74+
} else if (metadata.hasAnyscale()) {
75+
providerMetadata = new AnyscaleGenerative.Metadata();
76+
} else if (metadata.hasCohere()) {
77+
var cohere = metadata.getCohere();
78+
var apiVersion = cohere.getApiVersion();
79+
var billedUnits = cohere.getBilledUnits();
80+
var tokens = cohere.getTokens();
81+
providerMetadata = new CohereGenerative.Metadata(
82+
new CohereGenerative.Metadata.ApiVersion(
83+
apiVersion.hasVersion() ? apiVersion.getVersion() : null,
84+
apiVersion.hasIsDeprecated() ? apiVersion.getIsDeprecated() : null,
85+
apiVersion.hasIsExperimental() ? apiVersion.getIsExperimental() : null),
86+
new CohereGenerative.Metadata.BilledUnits(billedUnits.getInputTokens(),
87+
billedUnits.hasOutputTokens() ? billedUnits.getOutputTokens() : null,
88+
billedUnits.hasSearchUnits() ? billedUnits.getSearchUnits() : null,
89+
billedUnits.hasClassifications() ? billedUnits.getClassifications() : null),
90+
new CohereGenerative.Metadata.Tokens(
91+
tokens.hasInputTokens() ? tokens.getInputTokens() : null,
92+
tokens.hasOutputTokens() ? tokens.getOutputTokens() : null),
93+
new ArrayList<>(cohere.getWarnings().getValuesList()));
94+
} else if (metadata.hasDatabricks()) {
95+
var databricks = metadata.getDatabricks();
96+
var usage = databricks.getUsage();
97+
providerMetadata = new DatabricksGenerative.Metadata(new ProviderMetadata.Usage(
98+
usage.hasPromptTokens() ? usage.getPromptTokens() : null,
99+
usage.hasCompletionTokens() ? usage.getCompletionTokens() : null,
100+
usage.hasTotalTokens() ? usage.getTotalTokens() : null));
101+
} else if (metadata.hasFriendliai()) {
102+
var friendliai = metadata.getFriendliai();
103+
var usage = friendliai.getUsage();
104+
providerMetadata = new FriendliaiGenerative.Metadata(new ProviderMetadata.Usage(
105+
usage.hasPromptTokens() ? usage.getPromptTokens() : null,
106+
usage.hasCompletionTokens() ? usage.getCompletionTokens() : null,
107+
usage.hasTotalTokens() ? usage.getTotalTokens() : null));
108+
} else if (metadata.hasMistral()) {
109+
var mistral = metadata.getMistral();
110+
var usage = mistral.getUsage();
111+
providerMetadata = new MistralGenerative.Metadata(new ProviderMetadata.Usage(
112+
usage.hasPromptTokens() ? usage.getPromptTokens() : null,
113+
usage.hasCompletionTokens() ? usage.getCompletionTokens() : null,
114+
usage.hasTotalTokens() ? usage.getTotalTokens() : null));
115+
} else if (metadata.hasNvidia()) {
116+
var nvidia = metadata.getNvidia();
117+
var usage = nvidia.getUsage();
118+
providerMetadata = new NvidiaGenerative.Metadata(new ProviderMetadata.Usage(
119+
usage.hasPromptTokens() ? usage.getPromptTokens() : null,
120+
usage.hasCompletionTokens() ? usage.getCompletionTokens() : null,
121+
usage.hasTotalTokens() ? usage.getTotalTokens() : null));
122+
} else if (metadata.hasOllama()) {
123+
providerMetadata = new OllamaGenerative.Metadata();
124+
} else if (metadata.hasOpenai()) {
125+
var openai = metadata.getOpenai();
126+
var usage = openai.getUsage();
127+
providerMetadata = new OpenAiGenerative.Metadata(new ProviderMetadata.Usage(
128+
usage.hasPromptTokens() ? usage.getPromptTokens() : null,
129+
usage.hasCompletionTokens() ? usage.getCompletionTokens() : null,
130+
usage.hasTotalTokens() ? usage.getTotalTokens() : null));
131+
} else if (metadata.hasGoogle()) {
132+
var google = metadata.getGoogle();
133+
var tokens = google.getMetadata().getTokenMetadata();
134+
var usage = google.getUsageMetadata();
135+
providerMetadata = new GoogleGenerative.Metadata(
136+
new GoogleGenerative.Metadata.TokenMetadata(
137+
new GoogleGenerative.Metadata.TokenCount(
138+
tokens.getInputTokenCount().hasTotalBillableCharacters()
139+
? tokens.getInputTokenCount().getTotalBillableCharacters()
140+
: null,
141+
tokens.getInputTokenCount().hasTotalTokens()
142+
? tokens.getInputTokenCount().getTotalTokens()
143+
: null),
144+
new GoogleGenerative.Metadata.TokenCount(
145+
tokens.getOutputTokenCount().hasTotalBillableCharacters()
146+
? tokens.getOutputTokenCount().getTotalBillableCharacters()
147+
: null,
148+
tokens.getOutputTokenCount().hasTotalTokens()
149+
? tokens.getOutputTokenCount().getTotalTokens()
150+
: null)),
151+
new GoogleGenerative.Metadata.Usage(
152+
usage.hasPromptTokenCount() ? usage.getPromptTokenCount() : null,
153+
usage.hasCandidatesTokenCount() ? usage.getCandidatesTokenCount() : null,
154+
usage.hasTotalTokenCount() ? usage.getTotalTokenCount() : null));
155+
} else if (metadata.hasXai()) {
156+
var xai = metadata.getXai();
157+
var usage = xai.getUsage();
158+
providerMetadata = new XaiGenerative.Metadata(new ProviderMetadata.Usage(
159+
usage.hasPromptTokens() ? usage.getPromptTokens() : null,
160+
usage.hasCompletionTokens() ? usage.getCompletionTokens() : null,
161+
usage.hasTotalTokens() ? usage.getTotalTokens() : null));
52162
}
53163

54164
GenerativeDebug debug = null;

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

Lines changed: 0 additions & 7 deletions
This file was deleted.

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,7 @@
11
package io.weaviate.client6.v1.api.collections.generate;
22

3+
import io.weaviate.client6.v1.api.collections.generative.ProviderMetadata;
4+
35
public record TaskOutput(
46
String text,
57
ProviderMetadata metadata,

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -91,4 +91,15 @@ public AnthropicGenerative build() {
9191
return new AnthropicGenerative(this);
9292
}
9393
}
94+
95+
public static record Metadata(Usage usage) implements ProviderMetadata {
96+
97+
@Override
98+
public Generative.Kind _kind() {
99+
return Generative.Kind.ANTHROPIC;
100+
}
101+
102+
public static record Usage(Long inputTokens, Long outputTokens) {
103+
}
104+
}
94105
}

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

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import com.google.gson.annotations.SerializedName;
66

77
import io.weaviate.client6.v1.api.collections.Generative;
8+
import io.weaviate.client6.v1.api.collections.generate.ProviderMetadata;
89
import io.weaviate.client6.v1.internal.ObjectBuilder;
910

1011
public record AnyscaleGenerative(
@@ -66,4 +67,12 @@ public AnyscaleGenerative build() {
6667
return new AnyscaleGenerative(this);
6768
}
6869
}
70+
71+
public static record Metadata() implements ProviderMetadata {
72+
73+
@Override
74+
public Generative.Kind _kind() {
75+
return Generative.Kind.ANYSCALE;
76+
}
77+
}
6978
}

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

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
import com.google.gson.annotations.SerializedName;
66

77
import io.weaviate.client6.v1.api.collections.Generative;
8+
import io.weaviate.client6.v1.api.collections.generate.ProviderMetadata;
89
import io.weaviate.client6.v1.internal.ObjectBuilder;
910

1011
public record AwsGenerative(
@@ -14,7 +15,7 @@ public record AwsGenerative(
1415
@SerializedName("model") String model) implements Generative {
1516

1617
@Override
17-
public Kind _kind() {
18+
public Generative.Kind _kind() {
1819
return Generative.Kind.AWS;
1920
}
2021

@@ -68,4 +69,12 @@ public AwsGenerative build() {
6869
return new AwsGenerative(this);
6970
}
7071
}
72+
73+
public static record Metadata() implements ProviderMetadata {
74+
75+
@Override
76+
public Generative.Kind _kind() {
77+
return Generative.Kind.AWS;
78+
}
79+
}
7180
}

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

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -108,4 +108,23 @@ public CohereGenerative build() {
108108
return new CohereGenerative(this);
109109
}
110110
}
111+
112+
public static record Metadata(ApiVersion apiVersion, BilledUnits billedUnits, Tokens tokens, List<String> warnings)
113+
implements ProviderMetadata {
114+
115+
@Override
116+
public Generative.Kind _kind() {
117+
return Generative.Kind.COHERE;
118+
}
119+
120+
public static record ApiVersion(String version, Boolean deprecated, Boolean experimental) {
121+
}
122+
123+
public static record BilledUnits(Double inputTokens, Double outputTokens, Double searchUnits,
124+
Double classifications) {
125+
}
126+
127+
public static record Tokens(Double inputTokens, Double outputTokens) {
128+
}
129+
}
111130
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,4 +85,12 @@ public DatabricksGenerative build() {
8585
return new DatabricksGenerative(this);
8686
}
8787
}
88+
89+
public static record Metadata(ProviderMetadata.Usage usage) implements ProviderMetadata {
90+
91+
@Override
92+
public Generative.Kind _kind() {
93+
return Generative.Kind.DATABRICKS;
94+
}
95+
}
8896
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,4 +77,12 @@ public FriendliaiGenerative build() {
7777
return new FriendliaiGenerative(this);
7878
}
7979
}
80+
81+
public static record Metadata(ProviderMetadata.Usage usage) implements ProviderMetadata {
82+
83+
@Override
84+
public Generative.Kind _kind() {
85+
return Generative.Kind.FRIENDLIAI;
86+
}
87+
}
8088
}

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,4 +103,21 @@ public GoogleGenerative build() {
103103
return new GoogleGenerative(this);
104104
}
105105
}
106+
107+
public static record Metadata(TokenMetadata tokens, Usage usage) implements ProviderMetadata {
108+
109+
@Override
110+
public Generative.Kind _kind() {
111+
return Generative.Kind.GOOGLE;
112+
}
113+
114+
public static record TokenCount(Long totalBillableCharacters, Long totalTokens) {
115+
}
116+
117+
public static record TokenMetadata(TokenCount inputTokens, TokenCount outputTokens) {
118+
}
119+
120+
public static record Usage(Long promptTokenCount, Long candidatesTokenCount, Long totalTokenCount) {
121+
}
122+
}
106123
}

0 commit comments

Comments
 (0)