feat: add opt-in routing, quota, and observability for LiteLLM Sample by pweiber · Pull Request #355 · GoogleCloudPlatform/service-extensions

pweiber · 2026-07-30T19:51:00Z

Adds 3 optional capabilities to the LiteLLM callout sample. Each one is inert unless configured, so with nothing set the callout behaves exactly as before.

Model routing: A route extension rewrites the x-model-id header before the URL map evaluates routing, so a request can be steered by model alias, by a tier tag, by a weighted A/B split, or by simple-shuffle across a model group. The rewritten header also becomes authoritative over the request body, so a routing decision changes the model that answers, not just the backend it
reaches. Enabled by the router_settings section of the gateway config.

Token management: Virtual keys stored in Memorystore for Redis, enforced in two phases: checked before the upstream call and recorded from the real token usage on the response. Supports token budgets per key and per model, requests
per minute and tokens per minute limits, model allowlists, key expiry, and warn-only soft budgets, returning 401, 402, 403 and 429 as appropriate.
Successful responses carry the LiteLLM usage headers. Keys are defined in Terraform and seeded by a Cloud Run Job, or by a seeder script for ad hoc changes. Enabled by REDIS_HOST.

Observability: One OpenTelemetry span per request with GenAI semantic-convention attributes: provider, model, token usage, cost in USD from LiteLLM's bundled price map, a call id that is also returned as a response header, and a truncated hash of the virtual key. Exported over OTLP to any collector, with per-request opt-out. Enabled by OTEL_EXPORTER_OTLP_ENDPOINT.

Regional LB: A regional Terraform configuration that deploys the gateway with both extensions, Memorystore, and the seed-keys job.

Docs: a WALKTHROUGH.md that steps through each capability with the exact curl and log commands and the expected output for the new scenarios.

Signed-off-by: pweiber <periclesweiber@ciandt.com>

…e to include the improvements Signed-off-by: pweiber <periclesweiber@ciandt.com>

leonm1

Some initial large scoped feedback, I will continue reading through and give more targeted comments as well.

leonm1 · 2026-08-03T14:51:03Z

+    """Request-leg pre-check: unknown/expired key, model allowlist, RPM, TPM,
+    accumulated token budget, and per-model budget (model_max_budget).
+
+    Quota is opt-in per request: a request that presents no virtual key (no


Quota can't be opt-in per request: if you exceed your quota, you can just remove the API key. Please make it so that quota is either enabled for all or none requests, or provide an opt-in configuration option (i.e. defaulted to off) which allows unauthenticated requests to bypass quota.

Done, let me know if that's fine now

leonm1 · 2026-08-03T14:57:31Z

This optional_params is passed to LiteLLM's get_complete_url() function directly, and req_map is parsed from the request body. This allows adversaries to smuggle parameters to get_complete_url and override the behavior of endpoint picking.

Validate all the parameters that get used, and only pass explicitly validated parameters that we plan to support.

Done, let me know if that's fine now

leonm1 · 2026-08-03T15:00:36Z

+from datetime import datetime, timezone
+from typing import Any, NamedTuple, Optional
+
+import redis


This callout server could be fairly high concurrency. We should be using redis.asyncio.Redis rather than sync redis.

Thanks for all the review, for the async redis I believe we would need an await point in the callout, and our current sdk is still sync. I created a draft to move it to asyncio in #315 while we were load testing it some months ago. I think this affects limits.aio too. What do you think?

I will review #315 in parallel, it looks fairly small in comparison. We can implement the handler here with asyncio and, worst case, run the async handlers from a sync method with asyncio.run(). It'll probably be slower in the short term (waiting for async from sync provides no performance benefit afaict), but once we can land #315 we can make it async native, which will greatly increase performance.

Done, I tried asyncio.run() per call and it does not survive a second request. It closes its loop on return while the Redis pool outlives it, so the next call gets a connection bound to a dead loop:
RuntimeError: Event loop is closed
then got
Future attached to a different loop
Running the coroutines on one long-lived loop instead fixed it with no other change. That loop is ~20 lines and I believe it can be deleted the moment #315 lands.

leonm1 · 2026-08-03T16:58:53Z

+    if not _enabled or not api_key:
+        return {}
+    try:
+        cfg = _client.hgetall(f"key:{api_key}")


This will fetch the config from redis three times per request: once in check(), once in usage_headers(), and once in record() (that'll be for the request headers, response headers, and response body iiuc).

That amplifies the incoming traffic into significantly more per request. We should use a provider with short TTL (1-5s) cache for the config (best) to share it across requests, or at the very least fetch it only once per request and store it in the state of the grpc stream handler (okay).

Done, let me know if that's fine now

leonm1 · 2026-08-03T17:10:02Z

+        if allowed_models and not _model_allowed(model, allowed_models):
+            return QuotaDecision(False, 403, "model not allowed for key")
+        rpm_limit = int(cfg.get("rpm_limit", 0) or 0)
+        if rpm_limit:


Rather than implememting limits ourselves, we should use a library that implements all of this and is well tested. limits.aio has a redis library that will:

Use asyncio for communication with redis, with a single pipelined request rather than multiple round trips

Use atomic redis-side accounting for the limits, which is not subject to concurrency races when querying and updating concurrently

Has unit and integration tests asserting its correctness

Adopted for rpm_limit: limits.aio with MovingWindowRateLimiter over the same Redis. But for tpm_limit did not fit and stays a counter. limits declines to record a hit whose cost exceeds the limit, and a single response routinely costs more than a whole minute's token allowance, so the window never filled and the limit never fired: hit(cost=30) against a limit of 10 returns False and leaves remaining: 10. From the test I did, tpm went [200, 200, 200] on the library and [200, 429, 429] on the counter. Maybe after the #315 we can try again to see if it works properly or if it was a configuration issue from my side

leonm1 · 2026-08-03T17:12:32Z

@@ -91,6 +98,25 @@
    "host", ":authority", ":path", "content-length", "content-type",


Also authorization header?

I believe it could not go in that set. _MANAGED_HEADERS (:107) is a skip-list over LiteLLM's output headers, since the callout owns :path, :authority and the content headers itself, adding authorization there would skip LiteLLM's own Authorization: Bearer <ADC token> and send every Vertex request out unauthenticated. It is handled at :525 instead: when LiteLLM emits an authorization the rewrite overwrites the caller's value (OVERWRITE_IF_EXISTS_OR_ADD), and when the provider authenticates some other way and emits none (Anthropic's x-api-key), authorization is added to remove_headers. Either way the caller's virtual key stops at the gateway.

leonm1 · 2026-08-03T17:19:28Z

+    section = _config.get("router_settings")
+    if not section:
+        return None
+    tag_rules: dict[tuple[str, str], str] = {}


The second parameter to .get() is a default, but it is not returned in case the yaml section is explicitly set but null. It's safer to explicitly fall back with a falsiness check:

# If section.get("tag_rules") returns None, `None or []` evaluates to `[]` for r in (section.get("tag_rules") or []):

Done, let me know if that's fine now

leonm1 · 2026-08-03T17:21:30Z

+    total = sum(weight for _, weight in members)
+    if total <= 0:
+        return {}
+    key = headers.get(hash_header) or model_id


Falling back to the model_id will break split routing. All requests will go to the same backend in the failure case. If there is no header to use for determinism, we should fail safe and use a generated random value for the request.

Done, let me know if that's fine now

leonm1 · 2026-08-03T18:09:33Z

+
+def quota_settings() -> QuotaSettings:
+    section = _config.get("general_settings") or {}
+    return QuotaSettings(fail_open=bool(section.get("quota_fail_open", True)))


LLMs are expensive, we should probably fail closed

Done, let me know if that's fine now

leonm1 · 2026-08-03T18:13:37Z

+        return
+    span.set_attribute("gen_ai.operation.name", operation)
+    span.set_attribute("gen_ai.system", provider or "")
+    span.set_attribute("gen_ai.request.model", model or "")


Should we also log original model from x-request-id?

Added as litellm.request.original_model, under litellm.

leonm1 · 2026-08-03T18:14:55Z

+    if cost is not None:
+        span.set_attribute("gen_ai.cost.total_cost", cost)
+    if call_id:
+        span.set_attribute("litellm.call_id", call_id)


Should we also log the strategy that caused the model to be chosen?

Added as litellm.routing_strategy

leonm1 · 2026-08-03T18:16:53Z

+
+### Telemetry (OpenTelemetry)
+
+Each LLM request emits one span (`llm.request`) with GenAI


The opentelemetry spec says model requests should be logged with a specific span name:

https://github.com/open-telemetry/semantic-conventions-genai/blob/main/docs/gen-ai/gen-ai-spans.md#spans

{gen_ai.operation.name} {gen_ai.request.model}

Done, let me know if that's fine now

leonm1 · 2026-08-03T18:18:36Z

+headers that do not yet include that request's own usage, since spend is
+recorded after the response body is read.
+
+### Telemetry (OpenTelemetry)


This exposes spans, but we should also expose metrics: https://github.com/open-telemetry/semantic-conventions-genai/blob/main/docs/gen-ai/gen-ai-metrics.md

Done, let me know if that's fine now

hillsp · 2026-08-03T15:24:57Z

+
+# Quota reject HTTP status codes mapped to Envoy StatusCode values.
+_QUOTA_STATUS = {
+    401: StatusCode.Unauthorized,


If I understand correctly, StatusCode.Unauthorized is an enum with the value "401". Is the mapping here providing any value? If you are worried that the int must be a valid value in the proto, you could use something like:

try: StatusCode.Name(http_status_int) http_status_valid = True except ValueError: http_status_valid = False

or (less readably but apparently faster):

http_status_valid = http_status_int in StatusCode.DESCRIPTOR.values_by_number

No value, the enum values are the HTTP codes, so it was an identity map. Dropped, and the status is passed straight through, with a values_by_number guard so an unmapped code degrades to 500 rather than raising inside the proto.

hillsp · 2026-08-03T15:33:30Z

+_config: dict[str, Any] = {}
+
+
+def load() -> None:


Rather than a free function that references a global, I would typically expect callers to use a constructor or factory that returns the config, something like:

config = RoutingSettings.load()

Is there a reason to prefer the module global? If you need to override this for testing, you can mock out the function.

Done, GatewayConfig.load() and .from_dict() are classmethods that return a value, and the module global is gone. Same shape applied to quota, telemetry and routing.

hillsp · 2026-08-03T15:37:33Z

+        tag = r.get("tag") if isinstance(r, dict) else None
+        target = r.get("target") if isinstance(r, dict) else None
+        if not provider_prefix or not tag or not target:
+            logging.warning(


Rather than logging a warning, what do you think about instead raising for any error you find in the config? Since this parser is run at startup, it seems reasonable to fail to start if the config is bad.

In that world, it's probably a good idea to do this parsing as soon as you load the config (and save the result in a [member] variable), rather than waiting until routing_settings is called.

Agreed, the whole file is parsed in from_dict at startup and anything unusable raises

hillsp · 2026-08-03T15:39:26Z

+        return None
+    tag_rules: dict[tuple[str, str], str] = {}
+    for r in section.get("tag_rules", []):
+        provider_prefix = r.get("provider_prefix") if isinstance(r, dict) \


All of this code seems to depend on isinstance(r, dict) being true. I think it would make sense to check that up-front and error handle it specifically, then the code below can unconditionally reference r. I think that would make the rest of the code easier to read. What do you think?

Done, _section() does the check once and every parser entry point goes through it, so a non-mapping section fails with the section name rather than an AttributeError deeper in.

hillsp · 2026-08-03T15:40:55Z

+    for group, members in section.get("weighted_groups", {}).items():
+        parsed: list[tuple[str, int]] = []
+        for m in members:
+            model = m.get("model") if isinstance(m, dict) else None


As above on isinstance for m.

Same fix, a non-mapping member raises with the group name and the offending entry, instead of being assumed to be a mapping.

hillsp · 2026-08-03T18:39:20Z

    ) -> service_pb2.ProcessingResponse | None:
+        state = _state(context)
+        if state.route_mode:
+            header_map = {


I think it would be good to move this route mode handling to its own function. WDYT?

Agreed, it is _route_extension_response() now.

hillsp · 2026-08-03T18:41:59Z

        except json.JSONDecodeError as e:
            logging.warning("Invalid JSON body: %s", e)
-            return callout_tools.header_immediate_response(StatusCode.BadRequest)
+            self._end_span(state, 400)


WDYT about a single function that both updates the telemetry and calls header_immediate_response? This seems especially useful now that we know that StatusCode.BadRequest == 400.

Agreed, it is _reject(state, status) now

hillsp · 2026-08-03T18:49:46Z

        if state.provider in ("vertex_ai", "vertex_ai_beta"):
-            from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
-                ModelResponseIterator,
+            from litellm.llms.vertex_ai.gemini import (


There is a performance penalty to importing a python package inline, rather than once at the top of the file, and this is a hot function. Is that a reason you decided to do it this way?

No good reason, it was a misconfiguration from my IDE, should be fine now

hillsp · 2026-08-03T18:58:55Z

+def init_tracer() -> None:
+    """Initialize the global tracer once. No-op without an OTLP endpoint."""
+    global _tracer, _enabled
+    if _enabled:


Do we need _enabled, or can we just check if _tracer is None?

Gone, same as in quota. enabled() is self._tracer is not None, and a Telemetry with no tracer is a working no-op, so nothing branches on it.

hillsp · 2026-08-03T19:03:04Z

+        resource=Resource.create({"service.name": service_name}))
+    provider.add_span_processor(BatchSpanProcessor(OTLPSpanExporter()))
+    _install_provider(provider)
+    atexit.register(provider.shutdown)


Do you need this? The internet seems to think that TracerProvider will do this for you.

Just checked, the TracerProvider.__init__ defaults shutdown_on_exit=True and registers it already. Removed the duplicate.

Signed-off-by: pweiber <periclesweiber@ciandt.com>

…hers Signed-off-by: pweiber <periclesweiber@ciandt.com>

Signed-off-by: pweiber <periclesweiber@ciandt.com>

hillsp

Thanks for the improvements! I would have resolved a lot of the comment threads but I it turns out that I don't have permission. Please resolve anything that I didn't explicitly comment on.

hillsp · 2026-08-12T20:38:16Z

+    return int(time.time()) // _window_seconds(window)
+
+
+def _parse_list(raw: str) -> list[str]:


Thanks for the change here and in _parse_model_budgets. I notice that you removed the try/catch for JSONDecodeError. Are you confident that we don't need that any more, or should these return empty in that case?

hillsp · 2026-08-12T21:14:26Z

+    return _enabled
+
+
+def _window_seconds(window: str) -> int:


Appreciate you moving this out of the request parsing. If it turns out that we need to consume a LiteLLM keys file, we could easily add that parsing into seed_keys.py in the future, without breaking anything. Thanks!

hillsp · 2026-08-12T21:42:11Z

+    if not _enabled or not api_key:
+        return {}
+    try:
+        cfg = _client.hgetall(f"key:{api_key}")


Thanks for those changes - that should have improved the performance for sure.

Looking at the code you have now, I think it would be a nice cleanup to wrap the API key, model and config into a class, something like this (much simplified):

class QuotaHandler: @classmethod def create(cls, async_bridge, quota, api_key, model): cfg = async_bridge.run(quota.key_config(state.api_key)) return QuotaHandler(async_runner, quota, api_key, model, cfg) def check(self): try: return self.async_bridge.run(self.quota.check(self.api_key, self.model, self.cfg)) except Exception: return self.quota.unavailable() if self.quota.enabled: state.quota_handler = QuotaHandler.create(self._async, self.quota, state.api_key, model) decision = state.quota_handler.check() else: sate.quota_handler = None [...] if state.quota_handler: self.quota_handler.record()

That's a totally optional suggestion - what you have solves the performance problem.

pweiber added 4 commits July 30, 2026 15:40

feat: adding routing strategies, quota, seeding keys and telemetry

b0411b1

Signed-off-by: pweiber <periclesweiber@ciandt.com>

feat: adding tests

09db769

Signed-off-by: pweiber <periclesweiber@ciandt.com>

feat: adding regional terraform and configuration examples

7317ae1

Signed-off-by: pweiber <periclesweiber@ciandt.com>

docs: adding feature walkthrough for each scenario and updating readm…

578c076

…e to include the improvements Signed-off-by: pweiber <periclesweiber@ciandt.com>

pweiber requested a review from a team as a code owner July 30, 2026 19:51

Merge branch 'main' into py-litellm-improv

c4b79f8

martijneken assigned leonm1 Aug 3, 2026

leonm1 requested changes Aug 3, 2026

View reviewed changes

leonm1 reviewed Aug 3, 2026

View reviewed changes

hillsp reviewed Aug 3, 2026

View reviewed changes

pweiber added 11 commits August 7, 2026 17:24

fix: adjusting config examples

f7eac2b

Signed-off-by: pweiber <periclesweiber@ciandt.com>

fix: adjusting gateway config

50704c6

Signed-off-by: pweiber <periclesweiber@ciandt.com>

fix: adjusting quota

5b1456e

Signed-off-by: pweiber <periclesweiber@ciandt.com>

fix: adjusting routing

f2ad616

Signed-off-by: pweiber <periclesweiber@ciandt.com>

fix: adjusting the seed keys

312cad8

Signed-off-by: pweiber <periclesweiber@ciandt.com>

fix: adjusting the callout to use async redis

2eb1e0a

Signed-off-by: pweiber <periclesweiber@ciandt.com>

fix: adjusting telemetry to include metrics

671a7a7

Signed-off-by: pweiber <periclesweiber@ciandt.com>

fix: adjusting requirements files

e77fe3f

Signed-off-by: pweiber <periclesweiber@ciandt.com>

fix: updating tests to consider the review changes

368325c

Signed-off-by: pweiber <periclesweiber@ciandt.com>

fix: updating terraform config to use redis auth, updating paths matc…

0ec46d2

…hers Signed-off-by: pweiber <periclesweiber@ciandt.com>

docs: updating docs to reflect the review changes

b7008a8

Signed-off-by: pweiber <periclesweiber@ciandt.com>

hillsp reviewed Aug 12, 2026

View reviewed changes

		@@ -91,6 +98,25 @@
		"host", ":authority", ":path", "content-length", "content-type",


		### Telemetry (OpenTelemetry)

		Each LLM request emits one span (`llm.request`) with GenAI

		return int(time.time()) // _window_seconds(window)


		def _parse_list(raw: str) -> list[str]:

Uh oh!

Conversation

pweiber commented Jul 30, 2026

Uh oh!

leonm1 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

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!

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!

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