scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-26 01:12:21 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 99

style: trim comments on batch retrieve model group stamp

Browse source
This commit is contained in:
mateo-berri 2026-09-05 20:45:47 -07:00
parent df6990c712
commit 128cb114bd
2 changed files with 2 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm/router.py
View file

@ -6175,11 +6175,8 @@ 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.
# A batch is retrieved by id, so only the deployment that answered knows the
# group, and batch token usage is logged on this retrieve call.
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)

16
tests/test_litellm/test_router.py
View file

@ -958,22 +958,6 @@ async def test_arouter_aretrieve_batch():
assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base"
# ---------------------------------------------------------------------------
# Batch retrieval has to attribute its tokens to a model group.
#
# 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 is read
# off the completed batch's output file during retrieve logging, and that is the
# spend log row the tokens land on. A batch is retrieved by id, so the request
# carries no model and the router fans the lookup out across its deployments -
# the group of the deployment that answered is the only one there is to stamp.
# Leaving it unset files every batch's tokens under an empty model_group, which
# is what /global/activity/model groups the spend logs by.
#
# The provider is faked at the HTTP boundary, so the retrieve call and the usage
# accounting that reads the output file both run for real.
# ---------------------------------------------------------------------------
_BATCH_GROUP = "gemini-batch-group"
_BATCH_DEPLOYMENT_MODEL = "openai/gpt-4o-mini"
_BATCH_API_BASE = "http://localhost:4001/v1"