feat: split billing event types by embedding and reranking

src/main/java/io/stargate/sgv2/jsonapi/config/BillingConfig.java

@@ -59,9 +59,9 @@ public interface BillingConfig {
    *   jsonapi:
    *     billing:
    *       enabled-event-types:
-   *         - INTERNAL_MODEL_TOTAL_TOKENS
-   *         - EXTERNAL_MODEL_TOTAL_TOKENS
-   *         - INTERNAL_MODEL_EGRESS_BYTES
+   *         - INTERNAL_RERANKING_TOTAL_TOKENS
+   *         - EXTERNAL_RERANKING_TOTAL_TOKENS
+   *         - INTERNAL_RERANKING_EGRESS_BYTES
    * </pre>
    */
   Optional<Set<BillingEventType>> enabledEventTypes();

src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEvent.java

@@ -22,7 +22,7 @@
  *   "id": "8c0e9b8a-1d3a-4f6b-9c0d-1234567890ab",
  *   "timestamp": "2026-05-20T14:23:11.482Z",
  *   "product": "serverless",
- *   "event_type": "internal_model_total_tokens",
+ *   "event_type": "internal_reranking_total_tokens",
  *   "properties": {
  *     "usage": 7,
  *     "region": "us-west-2",
@@ -37,7 +37,7 @@
  * @param id Unique random-based (UUID v4) identifier for this event.
  * @param timestamp ISO 8601 timestamp of when the event was created.
  * @param product Product identifier, e.g. {@code "serverless"}.
- * @param eventType One of the six {@link BillingEventType} values.
+ * @param eventType One of the {@link BillingEventType} values.
  * @param properties Usage details including the billable amount, region, resource and model
  *     identifiers.
  */

src/main/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventType.java

@@ -1,6 +1,7 @@
 package io.stargate.sgv2.jsonapi.service.billing;
 
 import com.fasterxml.jackson.annotation.JsonValue;
+import io.stargate.sgv2.jsonapi.service.provider.ModelType;
 import java.util.EnumSet;
 import java.util.Set;
 
@@ -10,66 +11,122 @@
  * <p>Each event represents a single billable metric of a model call. The {@code internal_*}
  * variants are emitted for providers configured in {@link
  * io.stargate.sgv2.jsonapi.config.BillingConfig#internalModelProviders()}; everything else is
- * {@code external_*}.
+ * {@code external_*}. Each has an {@code embedding} and a {@code reranking} variant, from the
+ * {@link ModelType} of the call.
  *
  * <ul>
- *   <li>{@link #INTERNAL_MODEL_TOTAL_TOKENS} / {@link #EXTERNAL_MODEL_TOTAL_TOKENS} — total tokens
- *       reported by the model.
- *   <li>{@link #INTERNAL_MODEL_EGRESS_BYTES} / {@link #EXTERNAL_MODEL_EGRESS_BYTES} — bytes sent
- *       from the data plane to the model (request payload).
- *   <li>{@link #INTERNAL_MODEL_INGRESS_BYTES} / {@link #EXTERNAL_MODEL_INGRESS_BYTES} — bytes
- *       received from the model back to the data plane (response payload).
+ *   <li>{@code *_TOTAL_TOKENS} — total tokens reported by the model.
+ *   <li>{@code *_EGRESS_BYTES} — bytes sent from the data plane to the model (request payload).
+ *   <li>{@code *_INGRESS_BYTES} — bytes received from the model back to the data plane (response
+ *       payload).
  * </ul>
  */
 public enum BillingEventType {
-  INTERNAL_MODEL_TOTAL_TOKENS("internal_model_total_tokens", true, Metric.TOTAL_TOKENS),
-  EXTERNAL_MODEL_TOTAL_TOKENS("external_model_total_tokens", false, Metric.TOTAL_TOKENS),
-  INTERNAL_MODEL_EGRESS_BYTES("internal_model_egress_bytes", true, Metric.EGRESS_BYTES),
-  EXTERNAL_MODEL_EGRESS_BYTES("external_model_egress_bytes", false, Metric.EGRESS_BYTES),
-  INTERNAL_MODEL_INGRESS_BYTES("internal_model_ingress_bytes", true, Metric.INGRESS_BYTES),
-  EXTERNAL_MODEL_INGRESS_BYTES("external_model_ingress_bytes", false, Metric.INGRESS_BYTES);
+  INTERNAL_EMBEDDING_TOTAL_TOKENS(true, ModelType.EMBEDDING, Metric.TOTAL_TOKENS),
+  INTERNAL_RERANKING_TOTAL_TOKENS(true, ModelType.RERANKING, Metric.TOTAL_TOKENS),
+  EXTERNAL_EMBEDDING_TOTAL_TOKENS(false, ModelType.EMBEDDING, Metric.TOTAL_TOKENS),
+  EXTERNAL_RERANKING_TOTAL_TOKENS(false, ModelType.RERANKING, Metric.TOTAL_TOKENS),
+
+  INTERNAL_EMBEDDING_EGRESS_BYTES(true, ModelType.EMBEDDING, Metric.EGRESS_BYTES),
+  INTERNAL_RERANKING_EGRESS_BYTES(true, ModelType.RERANKING, Metric.EGRESS_BYTES),
+  EXTERNAL_EMBEDDING_EGRESS_BYTES(false, ModelType.EMBEDDING, Metric.EGRESS_BYTES),
+  EXTERNAL_RERANKING_EGRESS_BYTES(false, ModelType.RERANKING, Metric.EGRESS_BYTES),
+
+  INTERNAL_EMBEDDING_INGRESS_BYTES(true, ModelType.EMBEDDING, Metric.INGRESS_BYTES),
+  INTERNAL_RERANKING_INGRESS_BYTES(true, ModelType.RERANKING, Metric.INGRESS_BYTES),
+  EXTERNAL_EMBEDDING_INGRESS_BYTES(false, ModelType.EMBEDDING, Metric.INGRESS_BYTES),
+  EXTERNAL_RERANKING_INGRESS_BYTES(false, ModelType.RERANKING, Metric.INGRESS_BYTES);
 
   /** The billable metric a {@link BillingEventType} measures. */
   public enum Metric {
-    TOTAL_TOKENS,
-    EGRESS_BYTES,
-    INGRESS_BYTES
+    TOTAL_TOKENS("total_tokens"),
+    EGRESS_BYTES("egress_bytes"),
+    INGRESS_BYTES("ingress_bytes");
+
+    private final String billingEventName;
+
+    Metric(String billingEventName) {
+      this.billingEventName = billingEventName;
+    }
+
+    /** Name of the metric used in the billing event name. */
+    public String billingEventName() {
+      return billingEventName;
+    }
   }
 
   public static final Set<BillingEventType> ALL = Set.copyOf(EnumSet.allOf(BillingEventType.class));
 
+  private static final String INTERNAL = "internal";
+  private static final String EXTERNAL = "external";
+
   private final String eventName;
   private final boolean internal;
+  private final ModelType modelType;
   private final Metric metric;
 
-  BillingEventType(String eventName, boolean internal, Metric metric) {
-    // Event names are emitted lower-case in the JSON billing event payload
-    this.eventName = eventName.toLowerCase();
+  BillingEventType(boolean internal, ModelType modelType, Metric metric) {
+    this.eventName = eventName(internal, modelType, metric);
     this.internal = internal;
+    this.modelType = modelType;
     this.metric = metric;
   }
 
+  /** Builds the event name e.g. {@code internal_embedding_total_tokens} */
+  private static String eventName(boolean internal, ModelType modelType, Metric metric) {
+    return String.join(
+        "_",
+        internal ? INTERNAL : EXTERNAL,
+        modelType.billingEventName(),
+        metric.billingEventName());
+  }
+
   /** Lower-case event_type string used in the JSON billing event. */
   @JsonValue
   public String eventName() {
     return eventName;
   }
 
+  public ModelType modelType() {
+    return modelType;
+  }
+
   public Metric metric() {
     return metric;
   }
 
   /**
-   * Resolves the event type for a given metric and provider classification.
+   * Resolves the event type for a given model type, metric and provider classification.
    *
+   * @param modelType the type of model that was called
    * @param metric which billable metric we are emitting
    * @param internal {@code true} if the model provider is configured as internal
+   * @throws IllegalArgumentException if the modelType is {@link ModelType#MODEL_TYPE_UNSPECIFIED}
    */
-  public static BillingEventType of(Metric metric, boolean internal) {
-    return switch (metric) {
-      case TOTAL_TOKENS -> internal ? INTERNAL_MODEL_TOTAL_TOKENS : EXTERNAL_MODEL_TOTAL_TOKENS;
-      case EGRESS_BYTES -> internal ? INTERNAL_MODEL_EGRESS_BYTES : EXTERNAL_MODEL_EGRESS_BYTES;
-      case INGRESS_BYTES -> internal ? INTERNAL_MODEL_INGRESS_BYTES : EXTERNAL_MODEL_INGRESS_BYTES;
+  public static BillingEventType of(ModelType modelType, Metric metric, boolean internal) {
+    return switch (modelType) {
+      case MODEL_TYPE_UNSPECIFIED ->
+          throw new IllegalArgumentException(
+              "BillingEventType.of() - modelType must be specified, modelType=%s, metric=%s"
+                  .formatted(modelType, metric));
+      case EMBEDDING ->
+          switch (metric) {
+            case TOTAL_TOKENS ->
+                internal ? INTERNAL_EMBEDDING_TOTAL_TOKENS : EXTERNAL_EMBEDDING_TOTAL_TOKENS;
+            case EGRESS_BYTES ->
+                internal ? INTERNAL_EMBEDDING_EGRESS_BYTES : EXTERNAL_EMBEDDING_EGRESS_BYTES;
+            case INGRESS_BYTES ->
+                internal ? INTERNAL_EMBEDDING_INGRESS_BYTES : EXTERNAL_EMBEDDING_INGRESS_BYTES;
+          };
+      case RERANKING ->
+          switch (metric) {
+            case TOTAL_TOKENS ->
+                internal ? INTERNAL_RERANKING_TOTAL_TOKENS : EXTERNAL_RERANKING_TOTAL_TOKENS;
+            case EGRESS_BYTES ->
+                internal ? INTERNAL_RERANKING_EGRESS_BYTES : EXTERNAL_RERANKING_EGRESS_BYTES;
+            case INGRESS_BYTES ->
+                internal ? INTERNAL_RERANKING_INGRESS_BYTES : EXTERNAL_RERANKING_INGRESS_BYTES;
+          };
     };
   }
 }

src/main/java/io/stargate/sgv2/jsonapi/service/billing/DefaultBilling.java

@@ -41,7 +41,7 @@ public class DefaultBilling implements Billing {
 
   private static final ObjectWriter OBJECT_WRITER = new ObjectMapper().writer();
 
-  // There are three metrics, but 6 Event Types, because 3 each for internal or external provider.
+  // There are three metrics, each has internal / external and embedding / reranking event types.
   // this is the mapping of the metric and where we get it from
   private static final List<Pair<BillingEventType.Metric, Function<ModelUsage, Integer>>>
       METRICS_PER_USAGE =
@@ -90,10 +90,14 @@ public void emitEvent(ModelUsage modelUsage) {
   /**
    * Builds the list of billing events for one {@link ModelUsage}: one event per billable metric
    * (total tokens, egress bytes, ingress bytes), with the {@code internal_*} or {@code external_*}
-   * variant chosen based on {@link BillingConfig#internalModelProviders()}.
+   * variant chosen based on {@link BillingConfig#internalModelProviders()}, and the {@code
+   * embedding} or {@code reranking} variant from {@link ModelUsage#modelType()}.
    *
    * <p>All events from a single {@code ModelUsage} share one timestamp so they can be correlated in
    * the billing logs.
+   *
+   * @throws IllegalArgumentException if the model type is {@link
+   *     io.stargate.sgv2.jsonapi.service.provider.ModelType#MODEL_TYPE_UNSPECIFIED}
    */
   @VisibleForTesting
   List<BillingEvent> buildEvents(ModelUsage modelUsage) {
@@ -109,7 +113,7 @@ List<BillingEvent> buildEvents(ModelUsage modelUsage) {
     for (Pair<BillingEventType.Metric, Function<ModelUsage, Integer>> pair : METRICS_PER_USAGE) {
 
       var metric = pair.getLeft();
-      var eventType = BillingEventType.of(metric, internal);
+      var eventType = BillingEventType.of(modelUsage.modelType(), metric, internal);
       var supplier = pair.getRight();
 
       if (!enabledEventTypes.contains(eventType)) {

src/main/java/io/stargate/sgv2/jsonapi/service/provider/ModelType.java

@@ -24,6 +24,21 @@ public static Optional<ModelType> fromEmbeddingGateway(
     };
   }
 
+  /**
+   * Name of the model type used in the billing event name.
+   *
+   * @throws IllegalArgumentException for {@link #MODEL_TYPE_UNSPECIFIED}, it cannot be billed
+   */
+  public String billingEventName() {
+    return switch (this) {
+      case MODEL_TYPE_UNSPECIFIED ->
+          throw new IllegalArgumentException(
+              "ModelType.billingEventName() - MODEL_TYPE_UNSPECIFIED has no billing event name");
+      case EMBEDDING -> "embedding";
+      case RERANKING -> "reranking";
+    };
+  }
+
   public EmbeddingGateway.ModelUsage.ModelType toEmbeddingGateway() {
     return switch (this) {
       case MODEL_TYPE_UNSPECIFIED -> EmbeddingGateway.ModelUsage.ModelType.MODEL_TYPE_UNSPECIFIED;

src/test/java/io/stargate/sgv2/jsonapi/api/v1/BillingS3UploadIntegrationTest.java

@@ -12,6 +12,7 @@
 import io.quarkus.test.junit.QuarkusIntegrationTest;
 import io.stargate.sgv2.jsonapi.TestConstants;
 import io.stargate.sgv2.jsonapi.service.billing.BillingEventType;
+import io.stargate.sgv2.jsonapi.service.provider.ModelType;
 import io.stargate.sgv2.jsonapi.testresource.DseTestResource;
 import io.stargate.sgv2.jsonapi.testresource.S3MockTestResource;
 import java.net.URI;
@@ -56,8 +57,13 @@ public class BillingS3UploadIntegrationTest extends AbstractCollectionIntegratio
   private static final int DOCUMENT_COUNT = 10;
   private static final Pattern KEY_PATTERN =
       Pattern.compile("data-api/\\d{4}/\\d{2}/\\d{2}/\\d{2}/\\d{2}/[0-9a-f-]{36}\\.jsonl");
+  // inserts are vectorized, so only embedding events
   private static final Set<String> EVENT_TYPES =
-      new HashSet<>(BillingEventType.ALL.stream().map(BillingEventType::eventName).toList());
+      new HashSet<>(
+          BillingEventType.ALL.stream()
+              .filter(type -> type.modelType() == ModelType.EMBEDDING)
+              .map(BillingEventType::eventName)
+              .toList());
 
   @BeforeAll
   public void setup() {

src/test/java/io/stargate/sgv2/jsonapi/service/billing/BillingEventTypeTest.java

@@ -0,0 +1,53 @@
+package io.stargate.sgv2.jsonapi.service.billing;
+
+import static io.stargate.sgv2.jsonapi.service.billing.BillingEventType.Metric.EGRESS_BYTES;
+import static io.stargate.sgv2.jsonapi.service.billing.BillingEventType.Metric.INGRESS_BYTES;
+import static io.stargate.sgv2.jsonapi.service.billing.BillingEventType.Metric.TOTAL_TOKENS;
+import static io.stargate.sgv2.jsonapi.service.provider.ModelType.EMBEDDING;
+import static io.stargate.sgv2.jsonapi.service.provider.ModelType.RERANKING;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+import io.stargate.sgv2.jsonapi.service.provider.ModelType;
+import java.util.stream.Stream;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.Arguments;
+import org.junit.jupiter.params.provider.EnumSource;
+import org.junit.jupiter.params.provider.MethodSource;
+
+/** Tests for {@link BillingEventType} */
+public class BillingEventTypeTest {
+
+  @ParameterizedTest(name = "{3}")
+  @MethodSource("eventTypes")
+  public void resolvesEventName(
+      ModelType modelType, BillingEventType.Metric metric, boolean internal, String eventName) {
+    // downstream pricing matches on these names, do not change them
+    assertThat(BillingEventType.of(modelType, metric, internal).eventName()).isEqualTo(eventName);
+  }
+
+  private static Stream<Arguments> eventTypes() {
+    return Stream.of(
+        Arguments.of(EMBEDDING, TOTAL_TOKENS, true, "internal_embedding_total_tokens"),
+        Arguments.of(RERANKING, TOTAL_TOKENS, true, "internal_reranking_total_tokens"),
+        Arguments.of(EMBEDDING, TOTAL_TOKENS, false, "external_embedding_total_tokens"),
+        Arguments.of(RERANKING, TOTAL_TOKENS, false, "external_reranking_total_tokens"),
+        Arguments.of(EMBEDDING, EGRESS_BYTES, true, "internal_embedding_egress_bytes"),
+        Arguments.of(RERANKING, EGRESS_BYTES, true, "internal_reranking_egress_bytes"),
+        Arguments.of(EMBEDDING, EGRESS_BYTES, false, "external_embedding_egress_bytes"),
+        Arguments.of(RERANKING, EGRESS_BYTES, false, "external_reranking_egress_bytes"),
+        Arguments.of(EMBEDDING, INGRESS_BYTES, true, "internal_embedding_ingress_bytes"),
+        Arguments.of(RERANKING, INGRESS_BYTES, true, "internal_reranking_ingress_bytes"),
+        Arguments.of(EMBEDDING, INGRESS_BYTES, false, "external_embedding_ingress_bytes"),
+        Arguments.of(RERANKING, INGRESS_BYTES, false, "external_reranking_ingress_bytes"));
+  }
+
+  @ParameterizedTest(name = "{0}")
+  @EnumSource(BillingEventType.Metric.class)
+  public void unspecifiedModelTypeThrows(BillingEventType.Metric metric) {
+    assertThatThrownBy(() -> BillingEventType.of(ModelType.MODEL_TYPE_UNSPECIFIED, metric, true))
+        .isInstanceOf(IllegalArgumentException.class);
+    assertThatThrownBy(() -> BillingEventType.of(ModelType.MODEL_TYPE_UNSPECIFIED, metric, false))
+        .isInstanceOf(IllegalArgumentException.class);
+  }
+}