fix(router): stamp model_group when retrieving a batch

Batch token usage is accounted on the retrieve call, not on create: a provider
only reports token counts once the job finishes, so the usage arrives on
aretrieve_batch and that is the spend log row the tokens land on.

Router.acreate_batch stamps the requested model group into its metadata, but
Router.aretrieve_batch never did. A batch is retrieved by id, so the request
carries no model, and the router fans the lookup out over its deployments -
leaving model_group unset on the one record that carries the tokens.
/global/activity/model groups the spend logs by model_group, so every batch's
tokens were bucketed under an empty group.

Stamp the model group inside the per-deployment retrieve attempt, preferring an
explicitly requested group and otherwise using the model_name of the deployment
that answered, which is unambiguous even when the request named no model. An
existing model_group in the metadata is left untouched, so nothing that already
resolves a group changes.

Scope is limited to aretrieve_batch: acompletion, aresponses and acreate_batch
logging are untouched, and cost/spend attribution by model is unchanged.

Signed-off-by: mynkyu <mynkyu@bucketplace.net>

litellm/router.py

@@ -6137,6 +6137,8 @@ class Router:
         """
         try:
             parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs)
+            requested_model_group: Final = model
+            metadata_variable_name: Final = _get_router_metadata_variable_name(function_name="aretrieve_batch")
             if model is not None:
                 filtered_model_list: (
                     list[DeploymentTypedDict] | list[dict] | dict | None
@@ -6173,6 +6175,13 @@ class Router:
                         kwargs=new_kwargs,
                         function_name="aretrieve_batch",
                     )
+                    ## STAMP THE MODEL GROUP FOR SPEND TRACKING ##
+                    # A batch is retrieved by id, so the request carries no model group of its
+                    # own - only the deployment that answered knows it. Batch token usage lands
+                    # on this retrieve call (the provider reports counts once the job finishes),
+                    # so without this the tokens are logged under an empty model_group.
+                    model_group: Final = requested_model_group or model_name["model_name"]
+                    new_kwargs[metadata_variable_name].setdefault("model_group", model_group)
                     new_kwargs.pop("custom_llm_provider", None)
                     data.pop("custom_llm_provider", None)
                     return await litellm.aretrieve_batch(

tests/test_litellm/test_router_batch_retrieve_model_group.py

@@ -0,0 +1,118 @@
+"""
+model_group attribution on router batch retrieval.
+
+Batch token usage is accounted on the *retrieve* call, not on create: the
+provider only knows the token counts once the job finishes, so
+`LiteLLMBatch.usage` arrives on `aretrieve_batch` and that is the record the
+spend log tokens land on.
+
+`aretrieve_batch` is addressed by batch_id, so the request carries no model,
+and the router fans the lookup out across its deployments. These tests lock
+that the winning deployment's model group is stamped on the emitted
+StandardLoggingPayload, so `/global/activity/model` - which groups the spend
+logs by `model_group` - can attribute those tokens instead of bucketing every
+batch under "".
+"""
+
+import asyncio
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+import litellm
+import litellm.batches.main as bm
+from litellm import Router
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.types.utils import LiteLLMBatch, Usage
+
+MODEL_GROUP = "vertex-gemini-2.5-flash-lite-dev"
+DEPLOYMENT_MODEL = "vertex_ai/gemini-2.5-flash-lite"
+
+
+class _PayloadCollector(CustomLogger):
+    def __init__(self):
+        super().__init__()
+        self.payloads = []
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        self.payloads.append(kwargs.get("standard_logging_object"))
+
+
+@pytest.fixture
+def router():
+    return Router(
+        model_list=[
+            {
+                "model_name": MODEL_GROUP,
+                "litellm_params": {
+                    "model": DEPLOYMENT_MODEL,
+                    "vertex_project": "fake-project",
+                    "vertex_location": "us-central1",
+                    "vertex_credentials": "fake-creds",
+                },
+            }
+        ]
+    )
+
+
+@pytest.fixture
+def collector():
+    logger = _PayloadCollector()
+    previous = litellm.callbacks
+    litellm.callbacks = [logger]
+    try:
+        yield logger
+    finally:
+        litellm.callbacks = previous
+
+
+@pytest.fixture
+def vertex_retrieve():
+    """Mock the vertex provider seam - the only real network boundary."""
+    batch = LiteLLMBatch(
+        id="batch-1",
+        completion_window="24h",
+        created_at=0,
+        endpoint="/v1/chat/completions",
+        input_file_id="file-1",
+        object="batch",
+        status="completed",
+        usage=Usage(prompt_tokens=1000, completion_tokens=200, total_tokens=1200),
+    )
+    seam = MagicMock(name="vertex_ai_batches_instance")
+    seam.retrieve_batch.return_value = batch
+    with patch.object(bm, "vertex_ai_batches_instance", seam):
+        yield seam
+
+
+async def _collected_payload(collector) -> dict:
+    for _ in range(50):  # the success handler runs as a background task
+        payloads = [p for p in collector.payloads if p is not None]
+        if payloads:
+            return payloads[-1]
+        await asyncio.sleep(0.05)
+    raise AssertionError(f"no StandardLoggingPayload was emitted: {collector.payloads}")
+
+
+@pytest.mark.asyncio
+async def test_aretrieve_batch_without_model_stamps_model_group(router, collector, vertex_retrieve):
+    """
+    The proxy retrieves a managed batch by id only - no `model` in the request.
+    The router fans out over its deployments, so the model group is only known
+    from the deployment that answered.
+    """
+    response = await router.aretrieve_batch(batch_id="batch-1")
+
+    assert response.usage.total_tokens == 1200
+    payload = await _collected_payload(collector)
+    assert payload["model"] == DEPLOYMENT_MODEL
+    assert payload["model_group"] == MODEL_GROUP
+
+
+@pytest.mark.asyncio
+async def test_aretrieve_batch_with_model_stamps_requested_model_group(router, collector, vertex_retrieve):
+    """An explicitly requested model group is what gets logged."""
+    await router.aretrieve_batch(model=MODEL_GROUP, batch_id="batch-1")
+
+    payload = await _collected_payload(collector)
+    assert payload["model_group"] == MODEL_GROUP