Skip to content

Commit fd4beef

Browse files
Hazel-DatastaxHazel-Datastax
committed
Revert "Usage metering for embedding and reranking (#2008)"
This reverts commit 48857a2.
1 parent 346d6de commit fd4beef

72 files changed

Lines changed: 2707 additions & 4110 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎src/main/java/io/stargate/sgv2/jsonapi/api/model/command/impl/CreateCollectionCommand.java‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import io.stargate.sgv2.jsonapi.api.model.command.CollectionOnlyCommand;
66
import io.stargate.sgv2.jsonapi.api.model.command.CommandName;
77
import io.stargate.sgv2.jsonapi.config.constants.DocumentConstants;
8-
import io.stargate.sgv2.jsonapi.config.constants.ServiceDescConstants;
8+
import io.stargate.sgv2.jsonapi.config.constants.RerankingConstants;
99
import io.stargate.sgv2.jsonapi.exception.ErrorCodeV1;
1010
import io.stargate.sgv2.jsonapi.service.schema.collections.DocumentPath;
1111
import io.stargate.sgv2.jsonapi.service.schema.naming.NamingRules;
@@ -276,28 +276,28 @@ public record RerankServiceDesc(
276276
description = "Registered reranking service provider",
277277
type = SchemaType.STRING,
278278
implementation = String.class)
279-
@JsonProperty(ServiceDescConstants.PROVIDER)
279+
@JsonProperty(RerankingConstants.RerankingService.PROVIDER)
280280
String provider,
281281
@Schema(
282282
description = "Registered reranking service model",
283283
type = SchemaType.STRING,
284284
implementation = String.class)
285-
@JsonProperty(ServiceDescConstants.MODEL_NAME)
285+
@JsonProperty(RerankingConstants.RerankingService.MODEL_NAME)
286286
String modelName,
287287
@Valid
288288
@Nullable
289289
@Schema(
290290
description = "Authentication config for chosen reranking service",
291291
type = SchemaType.OBJECT)
292-
@JsonProperty(ServiceDescConstants.AUTHENTICATION)
292+
@JsonProperty(RerankingConstants.RerankingService.AUTHENTICATION)
293293
@JsonInclude(JsonInclude.Include.NON_NULL)
294294
Map<String, String> authentication,
295295
@Nullable
296296
@Schema(
297297
description =
298298
"Optional parameters that match the messageTemplate provided for the reranking provider",
299299
type = SchemaType.OBJECT)
300-
@JsonProperty(ServiceDescConstants.PARAMETERS)
300+
@JsonProperty(RerankingConstants.RerankingService.PARAMETERS)
301301
@JsonInclude(JsonInclude.Include.NON_NULL)
302302
Map<String, Object> parameters) {
303303

‎src/main/java/io/stargate/sgv2/jsonapi/api/model/command/impl/VectorizeConfig.java‎

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,7 @@
44
import com.fasterxml.jackson.annotation.JsonProperty;
55
import io.stargate.sgv2.jsonapi.config.constants.VectorConstants;
66
import io.stargate.sgv2.jsonapi.exception.ErrorCodeV1;
7-
import io.stargate.sgv2.jsonapi.service.embedding.operation.HuggingFaceDedicatedEmbeddingProvider;
8-
import io.stargate.sgv2.jsonapi.service.provider.ModelProvider;
7+
import io.stargate.sgv2.jsonapi.service.embedding.configuration.ProviderConstants;
98
import jakarta.validation.Valid;
109
import jakarta.validation.constraints.*;
1110
import java.util.*;
@@ -49,30 +48,24 @@ public VectorizeConfig(
4948
String modelName,
5049
Map<String, String> authentication,
5150
Map<String, Object> parameters) {
52-
5351
if (provider == null) {
5452
throw ErrorCodeV1.INVALID_CREATE_COLLECTION_OPTIONS.toApiException(
5553
"'provider' in required property for 'vector.service' Object value");
5654
}
57-
5855
this.provider = provider;
59-
6056
// HuggingfaceDedicated does not need user to specify model explicitly
6157
// If user specifies modelName other than endpoint-defined-model, will error out
6258
// By default, huggingfaceDedicated provider use endpoint-defined-model as placeholder
63-
if (ModelProvider.HUGGINGFACE_DEDICATED.apiName().equals(provider)) {
59+
if (ProviderConstants.HUGGINGFACE_DEDICATED.equals(provider)) {
6460
if (modelName == null) {
65-
modelName =
66-
HuggingFaceDedicatedEmbeddingProvider.HUGGINGFACE_DEDICATED_ENDPOINT_DEFINED_MODEL;
67-
} else if (!modelName.equals(
68-
HuggingFaceDedicatedEmbeddingProvider.HUGGINGFACE_DEDICATED_ENDPOINT_DEFINED_MODEL)) {
61+
modelName = ProviderConstants.HUGGINGFACE_DEDICATED_DEFINED_MODEL;
62+
} else if (!modelName.equals(ProviderConstants.HUGGINGFACE_DEDICATED_DEFINED_MODEL)) {
6963
throw ErrorCodeV1.INVALID_CREATE_COLLECTION_OPTIONS.toApiException(
7064
"'modelName' is not needed for embedding provider %s explicitly, only '%s' is accepted",
71-
ModelProvider.HUGGINGFACE_DEDICATED,
72-
HuggingFaceDedicatedEmbeddingProvider.HUGGINGFACE_DEDICATED_ENDPOINT_DEFINED_MODEL);
65+
ProviderConstants.HUGGINGFACE_DEDICATED,
66+
ProviderConstants.HUGGINGFACE_DEDICATED_DEFINED_MODEL);
7367
}
7468
}
75-
7669
this.modelName = modelName;
7770
if (authentication != null && !authentication.isEmpty()) {
7871
Map<String, String> updatedAuth = new HashMap<>();

‎src/main/java/io/stargate/sgv2/jsonapi/api/request/EmbeddingCredentials.java‎

Lines changed: 1 addition & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,9 @@
66
* EmbeddingCredentials is a record that holds the embedding provider credentials for the embedding
77
* service passed as header.
88
*
9-
* <p>Includes the tenantID, so we can fully identify the usage when creating the {@link
10-
* io.stargate.sgv2.jsonapi.service.provider.ModelUsage}
11-
*
12-
* @param tenantId - Tenant Id that called the API.
139
* @param apiKey - API token for the embedding service
1410
* @param accessId - Access Id used for AWS Bedrock embedding service
1511
* @param secretId - Secret Id used for AWS Bedrock embedding service
1612
*/
1713
public record EmbeddingCredentials(
18-
String tenantId,
19-
Optional<String> apiKey,
20-
Optional<String> accessId,
21-
Optional<String> secretId) {}
14+
Optional<String> apiKey, Optional<String> accessId, Optional<String> secretId) {}

‎src/main/java/io/stargate/sgv2/jsonapi/api/request/EmbeddingCredentialsSupplier.java‎

Lines changed: 1 addition & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -79,14 +79,10 @@ public EmbeddingCredentials create(
7979
&& collectionSupportsNoneAuth) {
8080
var authToken = requestContext.getHttpHeaders().getHeader(this.authTokenHeaderName);
8181
return new EmbeddingCredentials(
82-
requestContext.getTenantId().orElse(""),
83-
Optional.ofNullable(authToken),
84-
Optional.empty(),
85-
Optional.empty());
82+
Optional.ofNullable(authToken), Optional.empty(), Optional.empty());
8683
}
8784

8885
return new EmbeddingCredentials(
89-
requestContext.getTenantId().orElse(""),
9086
Optional.ofNullable(embeddingApi),
9187
Optional.ofNullable(accessId),
9288
Optional.ofNullable(secretId));

‎src/main/java/io/stargate/sgv2/jsonapi/api/request/RequestContext.java‎

Lines changed: 15 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,6 @@
22

33
import com.fasterxml.uuid.Generators;
44
import com.fasterxml.uuid.NoArgGenerator;
5-
import com.google.common.annotations.VisibleForTesting;
65
import io.stargate.sgv2.jsonapi.api.request.tenant.DataApiTenantResolver;
76
import io.stargate.sgv2.jsonapi.api.request.token.DataApiTokenResolver;
87
import io.stargate.sgv2.jsonapi.config.constants.HttpConstants;
@@ -36,25 +35,20 @@ public class RequestContext {
3635

3736
private final String userAgent;
3837

39-
/** FOR TESTING ONLY - so we can bypass pulling things the headers, still messy, getting better */
40-
@VisibleForTesting
41-
public RequestContext(
42-
Optional<String> tenantId,
43-
Optional<String> cassandraToken,
44-
RerankingCredentials rerankingCredentials,
45-
String userAgent) {
38+
/**
39+
* Constructor that will be useful in the offline library mode, where only the tenant will be set
40+
* and accessed.
41+
*
42+
* @param tenantId Tenant Id
43+
*/
44+
public RequestContext(Optional<String> tenantId) {
4645
this.tenantId = tenantId;
47-
this.cassandraToken = cassandraToken;
48-
embeddingCredentialsSupplier =
49-
new EmbeddingCredentialsSupplier(
50-
HttpConstants.AUTHENTICATION_TOKEN_HEADER_NAME,
51-
HttpConstants.EMBEDDING_AUTHENTICATION_TOKEN_HEADER_NAME,
52-
HttpConstants.EMBEDDING_AUTHENTICATION_ACCESS_ID_HEADER_NAME,
53-
HttpConstants.EMBEDDING_AUTHENTICATION_SECRET_ID_HEADER_NAME);
54-
this.rerankingCredentials = rerankingCredentials;
55-
this.userAgent = userAgent;
56-
this.httpHeaders = new HttpHeaderAccess(io.vertx.core.MultiMap.caseInsensitiveMultiMap());
46+
cassandraToken = Optional.empty();
47+
embeddingCredentialsSupplier = null;
48+
rerankingCredentials = null;
49+
httpHeaders = null;
5750
requestId = generateRequestId();
51+
userAgent = null;
5852
}
5953

6054
@Inject
@@ -83,14 +77,11 @@ public RequestContext(
8377
HeaderBasedRerankingKeyResolver.resolveRerankingKey(routingContext);
8478
rerankingCredentials =
8579
rerankingApiKeyFromHeader
86-
.map(apiKey -> new RerankingCredentials(this.tenantId.orElse(""), Optional.of(apiKey)))
80+
.map(apiKey -> new RerankingCredentials(Optional.of(apiKey)))
8781
.orElse(
8882
this.cassandraToken
89-
.map(
90-
cassandraToken ->
91-
new RerankingCredentials(
92-
this.tenantId.orElse(""), Optional.of(cassandraToken)))
93-
.orElse(new RerankingCredentials(this.tenantId.orElse(""), Optional.empty())));
83+
.map(cassandraToken -> new RerankingCredentials(Optional.of(cassandraToken)))
84+
.orElse(new RerankingCredentials(Optional.empty())));
9485
}
9586

9687
private static String generateRequestId() {

‎src/main/java/io/stargate/sgv2/jsonapi/api/request/RerankingCredentials.java‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,5 @@
77
* resolved from the request header 'reranking-api-key', if it is not present, then we will use the
88
* cassandra token as the reranking api key. Note, both cassandra token and reranking-api-key could
99
* be absent in Data API request, although it is invalid for authentication.
10-
*
11-
* <p>Includes the tenantId, so we can fully identify the usage when creating the {@link
12-
* io.stargate.sgv2.jsonapi.service.provider.ModelUsage}
1310
*/
14-
public record RerankingCredentials(String tenantId, Optional<String> apiKey) {}
11+
public record RerankingCredentials(Optional<String> apiKey) {}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -255,7 +255,7 @@ public Uni<RestResponse<CommandResult>> postCommand(
255255

256256
if (vectorColDef != null && vectorColDef.vectorizeDefinition() != null) {
257257
embeddingProvider =
258-
embeddingProviderFactory.create(
258+
embeddingProviderFactory.getConfiguration(
259259
requestContext.getTenantId(),
260260
requestContext.getCassandraToken(),
261261
vectorColDef.vectorizeDefinition().provider(),

‎src/main/java/io/stargate/sgv2/jsonapi/config/constants/RerankingConstants.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,4 +6,6 @@ interface CollectionRerankingOptions {
66
String ENABLED = "enabled";
77
String SERVICE = ServiceDescConstants.SERVICE;
88
}
9+
10+
interface RerankingService extends ServiceDescConstants {}
911
}

‎src/main/java/io/stargate/sgv2/jsonapi/config/constants/ServiceDescConstants.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
package io.stargate.sgv2.jsonapi.config.constants;
22

33
/** Common service description constants shared between vector and reranking */
4-
public interface ServiceDescConstants {
4+
interface ServiceDescConstants {
55
String SERVICE = "service";
66
String PROVIDER = "provider";
77
String MODEL_NAME = "modelName";

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

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -52,8 +52,7 @@ public DataVectorizer(
5252
SchemaObject schemaObject) {
5353
this.embeddingProvider = embeddingProvider;
5454
this.nodeFactory = nodeFactory;
55-
this.embeddingCredentials =
56-
Objects.requireNonNull(embeddingCredentials, "embeddingCredentials must not be null");
55+
this.embeddingCredentials = embeddingCredentials;
5756
this.schemaObject = schemaObject;
5857
}
5958

@@ -176,7 +175,7 @@ public Uni<float[]> vectorize(String vectorizeContent) {
176175
List.of(vectorizeContent),
177176
embeddingCredentials,
178177
EmbeddingProvider.EmbeddingRequestType.INDEX)
179-
.map(EmbeddingProvider.BatchedEmbeddingResponse::embeddings);
178+
.map(EmbeddingProvider.Response::embeddings);
180179
return vectors
181180
.onItem()
182181
.transform(
@@ -304,7 +303,7 @@ private Uni<List<float[]>> vectorizeTexts(
304303

305304
return embeddingProvider
306305
.vectorize(1, textsToVectorize, embeddingCredentials, requestType)
307-
.map(EmbeddingProvider.BatchedEmbeddingResponse::embeddings)
306+
.map(EmbeddingProvider.Response::embeddings)
308307
.onItem()
309308
.transform(
310309
vectorData -> {

0 commit comments

Comments
 (0)