feat: split billing event types by embedding and reranking - #2583
Merged
Merged
Conversation
Billing events only had the model name, so pricing could not tell a rerank call from an embedding call. Resolve the event_type from ModelType, e.g. internal_reranking_total_tokens and internal_embedding_total_tokens, for every internal/external metric. MODEL_TYPE_UNSPECIFIED has no event types and is logged as an error.
Contributor
📈 Unit Test Coverage Delta vs Main Branch
|
Contributor
Unit Test Coverage Report
|
Contributor
📈 Integration Test Coverage Delta vs Main Branch (dse69-it)
|
Contributor
Integration Test Coverage Report (dse69-it)
|
Contributor
📉 Integration Test Coverage Delta vs Main Branch (hcd-it)
|
Contributor
Integration Test Coverage Report (hcd-it)
|
amorton
requested changes
Sep 16, 2026
| 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 Optional<BillingEventType> of( |
Contributor
There was a problem hiding this comment.
if ModelType#MODEL_TYPE_UNSPECIFIED is passed throw an IllegalArgument error, that value should never have been added it makes no sense.
change to not return Optional, all calls should return a value for any combination of the arg, if not it is an error. Otherwise we will not know how to record the billing.
if possible revert back to using a modern switch{} clause, because this makes the compiler check that all possible values of the enum are covered.
Contributor
Author
There was a problem hiding this comment.
Good call about the compiler!! i am used to the world without it!
Contributor
Author
There was a problem hiding this comment.
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;
};… model type - BillingEventType.of() returns a value (no Optional) via exhaustive switch - throw IllegalArgumentException for ModelType.MODEL_TYPE_UNSPECIFIED - add billingEventName() to ModelType and Metric, event name built from parts - DefaultBilling no longer skips missing event types - remove unneeded test helper
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this PR does:
event_typenow comes from the call'sModelType, e.g.internal_reranking_total_tokens/internal_embedding_total_tokens. 12 event types instead of 6, one per internal/external, embedding/reranking and metric.MODEL_TYPE_UNSPECIFIEDhas no event types, logs an error and emits nothing.Tests: new
BillingEventTypeTestpins the 12 names,DefaultBillingTestcovers both model types and unspecified,BillingS3UploadIntegrationTestexpects only embedding events.Which issue(s) this PR fixes:
None, follow up to #2522
Checklist