fix proxy all models test

This commit is contained in:
Ishaan Jaff 2024-08-17 15:54:51 -07:00
parent a856a89d23
commit ec671b491d

View file

@ -278,18 +278,20 @@ def get_remaining_tokens_and_requests_from_request_data(data: Dict) -> Dict[str,
Returns {} when api_key + model rpm/tpm limit is not set
"""
headers = {}
_metadata = data.get("metadata", None) or {}
model_group = get_model_group_from_request_data(data)
# Remaining Requests
remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}"
remaining_requests = _metadata.get(remaining_requests_variable_name, None)
if remaining_requests:
headers[f"x-litellm-key-remaining-requests-{model_group}"] = remaining_requests
# Remaining Tokens
remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}"
remaining_tokens = _metadata.get(remaining_tokens_variable_name, None)
if remaining_tokens:
headers[f"x-litellm-key-remaining-tokens-{model_group}"] = remaining_tokens
return {
f"x-litellm-key-remaining-requests-{model_group}": str(remaining_requests),
f"x-litellm-key-remaining-tokens-{model_group}": str(remaining_tokens),
}
return headers