scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-03 02:22:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 80

fix(azure_ai): preserve zero num_output_tokens in MAI usage normalization

Browse source
This commit is contained in:
mateo-berri 2026-06-08 20:35:07 +00:00
parent d9e4eaab50
commit 99a39c919f
No known key found for this signature in database
2 changed files with 18 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
litellm/llms/azure_ai/image_generation/mai_transformation.py
View file

@ -91,11 +91,11 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
output_tokens = normalized_usage.get("output_tokens")
if output_tokens is None:
output_tokens = (
normalized_usage.get("num_output_tokens")
or normalized_usage.get("output_image_tokens")
or 0
)
output_tokens = normalized_usage.get("num_output_tokens")
if output_tokens is None:
output_tokens = normalized_usage.get("output_image_tokens")
if output_tokens is None:
output_tokens = 0
input_tokens = normalized_usage.get("input_tokens")
if input_tokens is None:

13
tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
View file

@ -246,6 +246,19 @@ class TestAzureMAIImageGeneration:
assert image_response.usage.input_tokens == 22
assert image_response.usage.total_tokens == 1046
def test_normalize_mai_usage_preserves_zero_output_tokens(self):
config = AzureFoundryMAIImageGenerationConfig()
normalized = config.normalize_mai_image_usage(
{
"num_output_tokens": 0,
"output_image_tokens": 1024,
"num_input_text_tokens": 22,
}
)
assert normalized["output_tokens"] == 0
assert normalized["input_tokens"] == 22
assert normalized["total_tokens"] == 22
def test_azure_sync_image_generation_uses_mai_response_transform(self):
raw_response = MagicMock(spec=httpx.Response)
raw_response.json.return_value = {