fix img gen cost

ishaan-jaff · ishaan-jaff · commit c8cedbed202a · 2025-01-12T16:31:04.000-08:00
diff --git a/litellm/main.py b/litellm/main.py
@@ -4434,6 +4434,8 @@ def image_generation(  # noqa: PLR0915
             },
             custom_llm_provider=custom_llm_provider,
         )
+        if "custom_llm_provider" not in logging.model_call_details:
+            logging.model_call_details["custom_llm_provider"] = custom_llm_provider
         if mock_response is not None:
             return mock_image_generation(model=model, mock_response=mock_response)
 

Original file line number	Diff line number	Diff line change
`@@ -4434,6 +4434,8 @@ def image_generation( # noqa: PLR0915`
`4434`	`4434`	`},`
`4435`	`4435`	`custom_llm_provider=custom_llm_provider,`
`4436`	`4436`	`)`
	`4437`	`+ if "custom_llm_provider" not in logging.model_call_details:`
	`4438`	`+ logging.model_call_details["custom_llm_provider"] = custom_llm_provider`
`4437`	`4439`	`if mock_response is not None:`
`4438`	`4440`	`return mock_image_generation(model=model, mock_response=mock_response)`
`4439`	`4441`