From 0ee884d4da5efee554849ddda9eb5ff3b33c746c Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 14 Nov 2024 17:46:42 -0800
Subject: [PATCH] fix imagegeneration output_cost_per_image on model cost map

---
 litellm/cost_calculator.py                    | 11 +++++---
 .../image_generation/cost_calculator.py       | 25 +++++++++++++++++++
 ...odel_prices_and_context_window_backup.json |  6 ++---
 model_prices_and_context_window.json          |  6 ++---
 4 files changed, 39 insertions(+), 9 deletions(-)
 create mode 100644 litellm/llms/vertex_ai_and_google_ai_studio/image_generation/cost_calculator.py

diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index 0aa8a8e36..50bed6fe9 100644
--- a/litellm/cost_calculator.py
+++ b/litellm/cost_calculator.py
@@ -46,6 +46,9 @@ from litellm.llms.OpenAI.cost_calculation import (
 from litellm.llms.OpenAI.cost_calculation import cost_per_token as openai_cost_per_token
 from litellm.llms.OpenAI.cost_calculation import cost_router as openai_cost_router
 from litellm.llms.together_ai.cost_calculator import get_model_params_and_category
+from litellm.llms.vertex_ai_and_google_ai_studio.image_generation.cost_calculator import (
+    cost_calculator as vertex_ai_image_cost_calculator,
+)
 from litellm.types.llms.openai import HttpxBinaryResponseContent
 from litellm.types.rerank import RerankResponse
 from litellm.types.router import SPECIAL_MODEL_INFO_PARAMS
@@ -667,9 +670,11 @@ def completion_cost(  # noqa: PLR0915
         ):
             ### IMAGE GENERATION COST CALCULATION ###
             if custom_llm_provider == "vertex_ai":
-                # https://cloud.google.com/vertex-ai/generative-ai/pricing
-                # Vertex Charges Flat $0.20 per image
-                return 0.020
+                if isinstance(completion_response, ImageResponse):
+                    return vertex_ai_image_cost_calculator(
+                        model=model,
+                        image_response=completion_response,
+                    )
             elif custom_llm_provider == "bedrock":
                 if isinstance(completion_response, ImageResponse):
                     return bedrock_image_cost_calculator(
diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/cost_calculator.py b/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/cost_calculator.py
new file mode 100644
index 000000000..2d7fa37f7
--- /dev/null
+++ b/litellm/llms/vertex_ai_and_google_ai_studio/image_generation/cost_calculator.py
@@ -0,0 +1,25 @@
+"""
+Vertex AI Image Generation Cost Calculator
+"""
+
+from typing import Optional
+
+import litellm
+from litellm.types.utils import ImageResponse
+
+
+def cost_calculator(
+    model: str,
+    image_response: ImageResponse,
+) -> float:
+    """
+    Vertex AI Image Generation Cost Calculator
+    """
+    _model_info = litellm.get_model_info(
+        model=model,
+        custom_llm_provider="vertex_ai",
+    )
+
+    output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
+    num_images: int = len(image_response.data)
+    return output_cost_per_image * num_images
diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json
index 137818f64..cae3bee12 100644
--- a/litellm/model_prices_and_context_window_backup.json
+++ b/litellm/model_prices_and_context_window_backup.json
@@ -2986,19 +2986,19 @@
         "supports_function_calling": true
     },
     "vertex_ai/imagegeneration@006": {
-        "cost_per_image": 0.020,
+        "output_cost_per_image": 0.020,
         "litellm_provider": "vertex_ai-image-models",
         "mode": "image_generation",
         "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
     },
     "vertex_ai/imagen-3.0-generate-001": {
-        "cost_per_image": 0.04,
+        "output_cost_per_image": 0.04,
         "litellm_provider": "vertex_ai-image-models",
         "mode": "image_generation",
         "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
     },
     "vertex_ai/imagen-3.0-fast-generate-001": {
-        "cost_per_image": 0.02,
+        "output_cost_per_image": 0.02,
         "litellm_provider": "vertex_ai-image-models",
         "mode": "image_generation",
         "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json
index 137818f64..cae3bee12 100644
--- a/model_prices_and_context_window.json
+++ b/model_prices_and_context_window.json
@@ -2986,19 +2986,19 @@
         "supports_function_calling": true
     },
     "vertex_ai/imagegeneration@006": {
-        "cost_per_image": 0.020,
+        "output_cost_per_image": 0.020,
         "litellm_provider": "vertex_ai-image-models",
         "mode": "image_generation",
         "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
     },
     "vertex_ai/imagen-3.0-generate-001": {
-        "cost_per_image": 0.04,
+        "output_cost_per_image": 0.04,
         "litellm_provider": "vertex_ai-image-models",
         "mode": "image_generation",
         "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
     },
     "vertex_ai/imagen-3.0-fast-generate-001": {
-        "cost_per_image": 0.02,
+        "output_cost_per_image": 0.02,
         "litellm_provider": "vertex_ai-image-models",
         "mode": "image_generation",
         "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"