Support for token quota with failover

[yanavlasov]

Description:

A token quota is defined on the AIGatewayRouteRuleBackendRef. Token quota will include maximum number of tokens, a conversion function from consumed tokens to consumed quota. Initially conversion function will be a simple multiplication factor and later on can be a CEL expression. quota is specified similarly to rate limit as the number of tokens per time interval.

If a quota is exceeded request is automatically failed over to the next backend ref in the AIGatewayRouteSpec. If all backend refs have exceeded their quota an 429 error response is served.

[mathetake]

I like this one. Will this quota be configured per backend/backendRef instead of per downstream user (attribute) right? One thing i feel tricky to implement is that the "rejection" by nature will happen at cluster filters but at that point ext_proc cannot force a retry since the upstream filters' local reply are not retry-able... i hope there is some workaround

[aabchoo]

I think we should configure per backendRef as quota calculation can differ per model with the same provider.

Would be great if we can add attributes to distinguish usage for quota similar to rate limit distinct header type. Ie goal is each user consumes quota based on their usage

[johnugeorge]

Thanks @yanavlasov for this

Yes. quota calculation is different for each provider and per backendref config makes sense

Can we reuse our cost concept for quota ? It is flexible with CEL, client selectors from headers etc

eg: quota definition to be similar to our cost definition (Can we reuse in

ai-gateway/examples/token_ratelimit/token_ratelimit.yaml

Lines 50 to 63 in a2ba2c9

	- metadataKey: llm_input_token
	type: InputToken
	- metadataKey: llm_cached_input_token
	type: CachedInputToken
	- metadataKey: llm_output_token
	type: OutputToken
	- metadataKey: llm_total_token
	type: TotalToken
	# This configures the token limit based on the CEL expression.
	# For a demonstration purpose, the CEL expression returns 100000000 only when the input token is 3.
	# Otherwise, it returns 0 (no token usage).
	- metadataKey: llm_cel_calculated_token
	type: CEL
	cel: "input_tokens == uint(3) ? 100000000 : 0"

)

And quota usage to similar to our current cost consumption

ai-gateway/examples/token_ratelimit/token_ratelimit.yaml

Lines 106 to 206 in a2ba2c9

	- clientSelectors:
	- headers:
	# Have the rate limit budget be per unique "x-user-id" header value.
	- name: x-user-id
	type: Distinct
	limit:
	# Configures the number of "tokens" allowed per hour, per user.
	requests: 10
	unit: Hour
	cost:
	request:
	from: Number
	# Setting the request cost to zero allows to only check the rate limit budget,
	# and not consume the budget on the request path.
	number: 0
	response:
	from: Metadata
	metadata:
	# This is the fixed namespace for the metadata used by AI Gateway.
	namespace: io.envoy.ai_gateway
	# Limit on the input token.
	key: llm_input_token
	# Repeat the same configuration for a different token type.
	# This configures the output token limit, and it has a different budget than others,
	# so it will be rate limited separately.
	- clientSelectors:
	- headers:
	- name: x-user-id
	type: Distinct
	limit:
	requests: 100
	unit: Hour
	cost:
	request:
	from: Number
	number: 0
	response:
	from: Metadata
	metadata:
	namespace: io.envoy.ai_gateway
	key: llm_output_token
	# Repeat the same configuration for a different token type.
	# This configures the total token limit, and it has a different budget than others,
	# so it will be rate limited separately.
	- clientSelectors:
	- headers:
	- name: x-user-id
	type: Distinct
	limit:
	requests: 1000
	unit: Hour
	cost:
	request:
	from: Number
	number: 0
	response:
	from: Metadata
	metadata:
	namespace: io.envoy.ai_gateway
	key: llm_total_token
	# Repeat the same configuration for a different token type.
	# This configures the cached input token limit, and it has a different budget than others,
	# so it will be rate limited separately.
	- clientSelectors:
	- headers:
	- name: x-user-id
	type: Distinct
	limit:
	requests: 100
	unit: Hour
	cost:
	request:
	from: Number
	number: 0
	response:
	from: Metadata
	metadata:
	namespace: io.envoy.ai_gateway
	key: llm_cached_input_token
	# Repeat the same configuration for a different token type.
	# This configures the token limit based on the CEL expression.
	- clientSelectors:
	- headers:
	- name: x-user-id
	type: Distinct
	limit:
	requests: 1000
	unit: Hour
	cost:
	request:
	from: Number
	number: 0
	response:
	from: Metadata
	metadata:
	namespace: io.envoy.ai_gateway
	key: llm_cel_calculated_token

[yuzisun]

One important thing to note is that we should differentiate token rate limit from token quota failover, the former is between downstream client and ai gateway to answer how much user quota is left, the other is between gateway and upstream model provider to make the route decision based on how much remaining quota is left from provider.