First draft version for vectorize usage response by maheshrajamani · Pull Request #1865 · stargate/data-api

maheshrajamani · 2025-02-11T17:25:00Z

What this PR does:
Vectorize usage header

Checklist

Changes manually tested
Automated Tests added/updated
Documentation added/updated
CLA Signed: DataStax CLA

maheshrajamani · 2025-02-11T17:25:43Z

src/main/java/io/stargate/sgv2/jsonapi/api/model/command/CommandResult.java

+   *
+   * @return
+   */
+  public RestResponse<CommandResult> toRestResponse(String vectorizeHeader) {


Added method that accepts vectorize header to be returned.

-1 do not want to code how the API returns this header in this PR, we are working on changes for how we do vecorize for the tables etc.

so lets make this PR just about what the embedding provider is doing

maheshrajamani · 2025-02-11T17:26:30Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/DataVectorizer.java

                    EmbeddingProvider.EmbeddingRequestType.INDEX)
-                .map(res -> res.embeddings());
+                .map(
+                    res -> {


Setting value from the aggregated vectorize usage info.

for now just want it on the response object from calling the embedding provider

WIll remove the CDI related code.

maheshrajamani · 2025-02-11T17:27:27Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/EmbeddingProvider.java

+      int batchId, List<float[]> embeddings, VectorizeUsageInfo vectorizeUsageInfo) {
    public static Response of(int batchId, List<float[]> embeddings) {
-      return new Response(batchId, embeddings);
+      return new Response(batchId, embeddings, new VectorizeUsageInfo(0, 0, 0, "", ""));


Added default VectorizeUsageInfo until all embedding provider class get refactored.

this is a record, so it has a constructor that will not enforce that VectorizeUsageInfo is non-null.

either: validate this in the canonical constructor or turn this into a class and validate in a ctor that takes all of the params

we should remove all the of() functions from the record, records should be simple structures with properties if they need a static factory they are prob not the right usage

Adapted the existing code to avoid compilation error in all the embedding providers. Will remove the of method for this record.

maheshrajamani · 2025-02-11T17:28:24Z

...main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/MeteredEmbeddingProvider.java

+              int totalToken = 0;
+              String provider = "";
+              String modelName = "";
              for (Response vectorizedBatch : vectorizedBatches) {


Aggregating the vectorize size from different micro batches.

maheshrajamani · 2025-02-11T17:29:00Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/NetworkUsageInterceptor.java

+import java.io.InputStream;
+import java.util.logging.Logger;
+
+public class NetworkUsageInterceptor implements ClientRequestFilter, ClientResponseFilter {


Interceptor to capture the byte size.

maheshrajamani · 2025-02-11T17:29:34Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/OpenAIEmbeddingProvider.java


-    Uni<EmbeddingResponse> response =
+    // ✅ Create an instance of NetworkUsageInfo and pass it to request properties
+    Uni<jakarta.ws.rs.core.Response> response =


Accept non serialized output as service response

maheshrajamani · 2025-02-11T17:29:49Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/OpenAIEmbeddingProvider.java

-              if (resp.data() == null) {
-                return Response.of(batchId, Collections.emptyList());
+            res -> {
+              EmbeddingResponse embeddingResponse = res.readEntity(EmbeddingResponse.class);


Serialize as EmbeddingResponse

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/VectorizeUsageInfo.java

maheshrajamani · 2025-02-11T17:30:55Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/NetworkUsageInterceptor.java

+          new ByteArrayInputStream(byteArrayOutputStream.toByteArray()));
+    }
+    LOGGER.info("Received Bytes: " + receivedBytes);
+    responseContext.getHeaders().add("sent-bytes", String.valueOf(sentBytes));


Setting to response header to parse in the embedding provider client class.

amorton

reasonable start, take out the Data API returning just make it the embedding provider for now because we are changing things, and some improvments on the code

amorton · 2025-02-11T18:41:39Z

src/main/java/io/stargate/sgv2/jsonapi/api/model/command/CommandResult.java

+   *
+   * @return
+   */
+  public RestResponse<CommandResult> toRestResponse(String vectorizeHeader) {


-1 do not want to code how the API returns this header in this PR, we are working on changes for how we do vecorize for the tables etc.

so lets make this PR just about what the embedding provider is doing

amorton · 2025-02-11T18:43:06Z

src/main/java/io/stargate/sgv2/jsonapi/api/v1/CollectionResource.java

+              String vectorize = null;
+              try {
+
+                if (!Strings.isNullOrEmpty(vectorizeUsageBean.getModel())) {


-1 - just the embedding provider for now

amorton · 2025-02-11T18:44:47Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/DataVectorizer.java

                    EmbeddingProvider.EmbeddingRequestType.INDEX)
-                .map(res -> res.embeddings());
+                .map(
+                    res -> {


for now just want it on the response object from calling the embedding provider

amorton · 2025-02-11T18:47:57Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/EmbeddingProvider.java

+      return new Response(batchId, embeddings, new VectorizeUsageInfo(0, 0, 0, "", ""));
+    }
+
+    public static Response of(


this of() function has the same signature as the record constructor , why is it needed ?

amorton · 2025-02-11T18:51:04Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/EmbeddingProvider.java

+      int batchId, List<float[]> embeddings, VectorizeUsageInfo vectorizeUsageInfo) {
    public static Response of(int batchId, List<float[]> embeddings) {
-      return new Response(batchId, embeddings);
+      return new Response(batchId, embeddings, new VectorizeUsageInfo(0, 0, 0, "", ""));


this is a record, so it has a constructor that will not enforce that VectorizeUsageInfo is non-null.

either: validate this in the canonical constructor or turn this into a class and validate in a ctor that takes all of the params

we should remove all the of() functions from the record, records should be simple structures with properties if they need a static factory they are prob not the right usage

amorton · 2025-02-11T21:07:24Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/VectorizeUsageInfo.java

+public class VectorizeUsageInfo {
+  private int requestSize;
+  private int responseSize;
+  private int totalTokens;


can we also have the input and output tokens ?

Not sure if there is split like that in all the embedding api response. Will check.

Verified Not all provide these, Example cohere response is like `
"billed_units": {

2069 | "input_tokens": 2
2070 | },

`
Huggingface doesn't return any tokens size.

amorton · 2025-02-11T21:08:39Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/VectorizeUsageInfo.java

+  private String model;
+
+  public VectorizeUsageInfo(
+      int requestSize, int responseSize, int totalTokens, String provider, String model) {


do we have an enum for the provider and/or model ?

We don't have an enum as it's configured on the provider configuration yaml.

amorton · 2025-02-11T21:11:33Z

...main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/MeteredEmbeddingProvider.java

+                modelName = vectorizedBatch.vectorizeUsageInfo().getModel();
              }
-              return Response.of(1, result);
+              return Response.of(


see other comments, do not need a static factory on a record

amorton · 2025-02-11T21:14:26Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/NetworkUsageInterceptor.java

+  }
+
+  @Override
+  public void filter(ClientRequestContext requestContext, ClientResponseContext responseContext)


if we do not change the way of counting, and stick with this for an initial release, why not count both the request and the response in this function ?

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/NetworkUsageInterceptor.java

maheshrajamani

I will make the change as suggested and let you know.

maheshrajamani · 2025-02-12T14:42:32Z

...main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/MeteredEmbeddingProvider.java

+                sentBytes += vectorizedBatch.vectorizeUsageInfo().getRequestSize();
+                receivedBytes += vectorizedBatch.vectorizeUsageInfo().getResponseSize();
+                totalToken += vectorizedBatch.vectorizeUsageInfo().getTotalTokens();
+                provider = vectorizedBatch.vectorizeUsageInfo().getProvider();


Currently it is done considering there is only one vectorize per collection/table. Not sure of the code refactoring to support table's multi vectorize.

maheshrajamani · 2025-02-12T14:43:48Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/DataVectorizer.java

                    EmbeddingProvider.EmbeddingRequestType.INDEX)
-                .map(res -> res.embeddings());
+                .map(
+                    res -> {


WIll remove the CDI related code.

maheshrajamani · 2025-02-12T14:45:12Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/EmbeddingProvider.java

+      int batchId, List<float[]> embeddings, VectorizeUsageInfo vectorizeUsageInfo) {
    public static Response of(int batchId, List<float[]> embeddings) {
-      return new Response(batchId, embeddings);
+      return new Response(batchId, embeddings, new VectorizeUsageInfo(0, 0, 0, "", ""));


Adapted the existing code to avoid compilation error in all the embedding providers. Will remove the of method for this record.

maheshrajamani · 2025-02-12T14:51:01Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/VectorizeUsageInfo.java

+public class VectorizeUsageInfo {
+  private int requestSize;
+  private int responseSize;
+  private int totalTokens;


Not sure if there is split like that in all the embedding api response. Will check.

maheshrajamani · 2025-02-12T14:52:13Z

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/VectorizeUsageInfo.java

+  private String model;
+
+  public VectorizeUsageInfo(
+      int requestSize, int responseSize, int totalTokens, String provider, String model) {


We don't have an enum as it's configured on the provider configuration yaml.

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/NetworkUsageInterceptor.java

First draft version for vectorize usage response

5c53563

maheshrajamani requested a review from a team as a code owner February 11, 2025 17:25

maheshrajamani commented Feb 11, 2025

View reviewed changes

Formatted the classes

5efe8da

amorton requested changes Feb 11, 2025

View reviewed changes

maheshrajamani commented Feb 12, 2025

View reviewed changes

maheshrajamani added 2 commits February 12, 2025 10:23

Refactored the NetworkUsageInterceptor as per suggestions

7113bff

Refactored the NetworkUsageInterceptor as per suggestions

96c40f2

tatu-at-datastax reviewed Feb 12, 2025

View reviewed changes

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/NetworkUsageInterceptor.java Outdated Show resolved Hide resolved

tatu-at-datastax reviewed Feb 12, 2025

View reviewed changes

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/NetworkUsageInterceptor.java Outdated Show resolved Hide resolved

tatu-at-datastax reviewed Feb 12, 2025

View reviewed changes

src/main/java/io/stargate/sgv2/jsonapi/service/embedding/operation/NetworkUsageInterceptor.java Outdated Show resolved Hide resolved

maheshrajamani and others added 8 commits February 12, 2025 11:28

Fixed the byte size calculation and also fixed the Logger to use slf4j

b0d6b34

Added the bytes usage and token usage to places where available.

1ad808e

Changes to remove VectorizeUsageBean

16c8083

Refactored vectorize usage from the EGW client response.

741440a

Fixed compile error

42c4b6e

Reverted the unnecessary change

b4eef8d

Fixed typo

88a64f7

Merge branch 'main' into vectorize-usage-tracker

190785f

amorton mentioned this pull request Jun 5, 2025

Usage metering for embedding and reranking #2008

Merged

4 tasks

Conversation

maheshrajamani commented Feb 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

amorton left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Verified Not all provide these, Example cohere response is like ` "billed_units": {

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

maheshrajamani left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

maheshrajamani commented Feb 11, 2025 •

edited

Loading

Verified Not all provide these, Example cohere response is like `
"billed_units": {