Merge pull request #22223 from emerzon/feat/vertex-gemini-3-1-flash-image-preview-pricing

feat(vertex): add gemini-3.1-flash-image-preview to model DB
This commit is contained in:
Sameer Kankute
2026-02-27 18:10:04 +05:30
committed by GitHub
6 changed files with 334 additions and 12 deletions
@@ -8,9 +8,11 @@ from litellm._logging import verbose_logger
from litellm.types.utils import (
CacheCreationTokenDetails,
CallTypes,
CompletionTokensDetailsWrapper,
ImageResponse,
ModelInfo,
PassthroughCallTypes,
PromptTokensDetailsWrapper,
ServiceTier,
Usage,
)
@@ -767,6 +769,64 @@ def generic_cost_per_token( # noqa: PLR0915
return prompt_cost, completion_cost
def calculate_image_response_cost_from_usage(
model: str,
image_response: ImageResponse,
custom_llm_provider: str,
) -> Optional[float]:
"""
Calculate image generation cost from usage metadata when available.
Returns:
Optional[float]: total cost from token usage, or None when usage metadata
is missing/incomplete and caller should fall back to flat per-image pricing.
"""
usage = image_response.usage
if usage is None:
return None
prompt_tokens = usage.input_tokens
completion_tokens = usage.output_tokens
total_tokens = usage.total_tokens
if prompt_tokens is None or completion_tokens is None or total_tokens is None:
return None
# ImageResponse may carry a default zeroed usage object even when provider
# usage metadata is absent. Treat this as missing usage and fall back.
if prompt_tokens == 0 and completion_tokens == 0 and total_tokens == 0:
return None
input_tokens_details = getattr(usage, "input_tokens_details", None)
prompt_tokens_details: Optional[PromptTokensDetailsWrapper] = None
if input_tokens_details is not None:
prompt_tokens_details = PromptTokensDetailsWrapper(
text_tokens=getattr(input_tokens_details, "text_tokens", None),
image_tokens=getattr(input_tokens_details, "image_tokens", None),
cached_tokens=0,
)
normalized_usage = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=total_tokens,
prompt_tokens_details=prompt_tokens_details,
completion_tokens_details=CompletionTokensDetailsWrapper(
text_tokens=0,
image_tokens=completion_tokens,
reasoning_tokens=0,
audio_tokens=0,
),
)
prompt_cost, completion_cost = generic_cost_per_token(
model=model,
usage=normalized_usage,
custom_llm_provider=custom_llm_provider,
)
return prompt_cost + completion_cost
class CostCalculatorUtils:
@staticmethod
def _call_type_has_image_response(call_type: str) -> bool:
@@ -5,6 +5,9 @@ Google AI Image Generation Cost Calculator
from typing import Any
import litellm
from litellm.litellm_core_utils.llm_cost_calc.utils import (
calculate_image_response_cost_from_usage,
)
from litellm.types.utils import ImageResponse
@@ -13,13 +16,22 @@ def cost_calculator(
image_response: Any,
) -> float:
"""
Vertex AI Image Generation Cost Calculator
Google AI Image Generation Cost Calculator
"""
_model_info = litellm.get_model_info(
model=model,
custom_llm_provider="gemini",
)
if isinstance(image_response, ImageResponse):
token_based_cost = calculate_image_response_cost_from_usage(
model=model,
image_response=image_response,
custom_llm_provider="gemini",
)
if token_based_cost is not None:
return token_based_cost
output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
num_images: int = 0
if isinstance(image_response, ImageResponse):
@@ -3,6 +3,9 @@ Vertex AI Image Generation Cost Calculator
"""
import litellm
from litellm.litellm_core_utils.llm_cost_calc.utils import (
calculate_image_response_cost_from_usage,
)
from litellm.types.utils import ImageResponse
@@ -18,6 +21,14 @@ def cost_calculator(
custom_llm_provider="vertex_ai",
)
token_based_cost = calculate_image_response_cost_from_usage(
model=model,
image_response=image_response,
custom_llm_provider="vertex_ai",
)
if token_based_cost is not None:
return token_based_cost
output_cost_per_image: float = _model_info.get("output_cost_per_image") or 0.0
num_images: int = 0
if image_response.data:
@@ -14194,6 +14194,38 @@
"supports_vision": true,
"supports_web_search": true
},
"gemini-3.1-flash-image-preview": {
"input_cost_per_image": 0.00056,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 3e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions",
"/v1/batch"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
"supports_web_search": true
},
"deep-research-pro-preview-12-2025": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
@@ -31545,6 +31577,19 @@
"output_cost_per_token_batches": 6e-06,
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
},
"vertex_ai/gemini-3.1-flash-image-preview": {
"input_cost_per_image": 0.00056,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 3e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
},
"vertex_ai/deep-research-pro-preview-12-2025": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
+45
View File
@@ -14194,6 +14194,38 @@
"supports_vision": true,
"supports_web_search": true
},
"gemini-3.1-flash-image-preview": {
"input_cost_per_image": 0.00056,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 3e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models",
"supported_endpoints": [
"/v1/chat/completions",
"/v1/completions",
"/v1/batch"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"text",
"image"
],
"supports_function_calling": false,
"supports_prompt_caching": true,
"supports_response_schema": true,
"supports_system_messages": true,
"supports_vision": true,
"supports_web_search": true
},
"deep-research-pro-preview-12-2025": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
@@ -31545,6 +31577,19 @@
"output_cost_per_token_batches": 6e-06,
"source": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-pro-image"
},
"vertex_ai/gemini-3.1-flash-image-preview": {
"input_cost_per_image": 0.00056,
"input_cost_per_token": 5e-07,
"litellm_provider": "vertex_ai-language-models",
"max_input_tokens": 65536,
"max_output_tokens": 32768,
"max_tokens": 32768,
"mode": "image_generation",
"output_cost_per_image": 0.0672,
"output_cost_per_image_token": 6e-05,
"output_cost_per_token": 3e-06,
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing#gemini-models"
},
"vertex_ai/deep-research-pro-preview-12-2025": {
"input_cost_per_image": 0.0011,
"input_cost_per_token": 2e-06,
@@ -9,9 +9,19 @@ import litellm
from litellm.litellm_core_utils.llm_cost_calc.tool_call_cost_tracking import (
StandardBuiltInToolCostTracking,
)
from litellm.llms.gemini.image_generation.cost_calculator import (
cost_calculator as gemini_image_generation_cost_calculator,
)
from litellm.llms.vertex_ai.image_generation.cost_calculator import (
cost_calculator as vertex_image_generation_cost_calculator,
)
from litellm.types.llms.openai import FileSearchTool, WebSearchOptions
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
ImageObject,
ImageResponse,
ImageUsage,
ImageUsageInputTokensDetails,
ModelInfo,
ModelResponse,
PromptTokensDetailsWrapper,
@@ -766,7 +776,14 @@ def test_service_tier_fallback_pricing():
assert abs(std_cost[1] - expected_standard_completion) < 1e-10, f"Standard completion cost mismatch: {std_cost[1]} vs {expected_standard_completion}"
def test_gemini_image_generation_cost_with_zero_text_tokens():
@pytest.mark.parametrize(
"model",
[
"gemini-3-pro-image-preview",
"gemini-3.1-flash-image-preview",
],
)
def test_gemini_image_generation_cost_with_zero_text_tokens(model: str):
"""
Test that image_tokens are correctly costed when text_tokens=0.
@@ -779,7 +796,6 @@ def test_gemini_image_generation_cost_with_zero_text_tokens():
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gemini-3-pro-image-preview"
custom_llm_provider = "vertex_ai"
# Usage from the issue: text_tokens=0, image_tokens=1120, reasoning_tokens=225
@@ -809,9 +825,9 @@ def test_gemini_image_generation_cost_with_zero_text_tokens():
# Expected costs:
# - text_tokens: 0 * output_cost_per_token = 0
# - image_tokens: 1120 * output_cost_per_image_token = 1120 * 1.2e-04 = 0.1344
# - reasoning_tokens: 225 * output_cost_per_token = 225 * 1.2e-05 = 0.0027
# Total completion: ~0.1371
# - image_tokens: 1120 * output_cost_per_image_token
# - reasoning_tokens: 225 * output_cost_per_token
# Total completion should include both image + reasoning costs.
output_cost_per_image_token = model_cost_map.get("output_cost_per_image_token", 0)
output_cost_per_token = model_cost_map.get("output_cost_per_token", 0)
@@ -820,18 +836,151 @@ def test_gemini_image_generation_cost_with_zero_text_tokens():
expected_reasoning_cost = 225 * output_cost_per_token # reasoning uses base token cost
expected_completion_cost = expected_image_cost + expected_reasoning_cost
# The bug was: all 1345 tokens were treated as text = 1345 * 1.2e-05 = 0.01614
# Fixed: image_tokens use image pricing = ~0.137
assert completion_cost > 0.10, (
f"Completion cost should be > $0.10 (image tokens are expensive), got ${completion_cost:.6f}. "
f"Bug: tokens may be incorrectly treated as text tokens."
# The bug was: all completion tokens were treated as text tokens only.
bugged_text_only_cost = 1345 * output_cost_per_token
assert completion_cost > bugged_text_only_cost * 2, (
f"Completion cost should be significantly larger than text-only bugged path. "
f"Expected > {bugged_text_only_cost * 2:.6f}, got {completion_cost:.6f}"
)
assert round(completion_cost, 4) == round(expected_completion_cost, 4), (
f"Expected completion cost ${expected_completion_cost:.6f}, got ${completion_cost:.6f}"
)
def test_vertex_image_generation_cost_prefers_token_usage_metadata():
"""
When usage metadata exists on image responses, Vertex image generation cost
should be calculated from token pricing, not flat output_cost_per_image.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gemini-3.1-flash-image-preview"
model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")
input_text_tokens = 50
input_image_tokens = 1120
output_image_tokens = 1120
prompt_tokens = input_text_tokens + input_image_tokens
image_response = ImageResponse(
data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")],
usage=ImageUsage(
input_tokens=prompt_tokens,
input_tokens_details=ImageUsageInputTokensDetails(
text_tokens=input_text_tokens,
image_tokens=input_image_tokens,
),
output_tokens=output_image_tokens,
total_tokens=prompt_tokens + output_image_tokens,
),
)
cost = vertex_image_generation_cost_calculator(
model=model,
image_response=image_response,
)
expected_prompt_cost = prompt_tokens * model_info["input_cost_per_token"]
expected_completion_cost = output_image_tokens * model_info["output_cost_per_image_token"]
expected_total_cost = expected_prompt_cost + expected_completion_cost
assert round(cost, 10) == round(expected_total_cost, 10)
# Ensure this is not falling back to flat per-image pricing.
assert cost != len(image_response.data) * model_info["output_cost_per_image"]
def test_vertex_image_generation_cost_falls_back_to_flat_image_pricing():
"""
Without usage metadata, Vertex image generation cost should fall back to
output_cost_per_image * number_of_images.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gemini-3.1-flash-image-preview"
model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")
image_response = ImageResponse(
data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]
)
cost = vertex_image_generation_cost_calculator(
model=model,
image_response=image_response,
)
expected_cost = len(image_response.data) * model_info["output_cost_per_image"]
assert round(cost, 10) == round(expected_cost, 10)
def test_gemini_image_generation_cost_prefers_token_usage_metadata():
"""
When usage metadata exists on image responses, Gemini image generation cost
should be calculated from token pricing, not flat output_cost_per_image.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gemini/gemini-3-pro-image-preview"
model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini")
input_text_tokens = 20
input_image_tokens = 1120
output_image_tokens = 1120
prompt_tokens = input_text_tokens + input_image_tokens
image_response = ImageResponse(
data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")],
usage=ImageUsage(
input_tokens=prompt_tokens,
input_tokens_details=ImageUsageInputTokensDetails(
text_tokens=input_text_tokens,
image_tokens=input_image_tokens,
),
output_tokens=output_image_tokens,
total_tokens=prompt_tokens + output_image_tokens,
),
)
cost = gemini_image_generation_cost_calculator(
model=model,
image_response=image_response,
)
expected_prompt_cost = prompt_tokens * model_info["input_cost_per_token"]
expected_completion_cost = output_image_tokens * model_info["output_cost_per_image_token"]
expected_total_cost = expected_prompt_cost + expected_completion_cost
assert round(cost, 10) == round(expected_total_cost, 10)
# Ensure this is not falling back to flat per-image pricing.
assert cost != len(image_response.data) * model_info["output_cost_per_image"]
def test_gemini_image_generation_cost_falls_back_to_flat_image_pricing():
"""
Without usage metadata, Gemini image generation cost should fall back to
output_cost_per_image * number_of_images.
"""
os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
litellm.model_cost = litellm.get_model_cost_map(url="")
model = "gemini/gemini-3-pro-image-preview"
model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini")
image_response = ImageResponse(
data=[ImageObject(b64_json="img1"), ImageObject(b64_json="img2")]
)
cost = gemini_image_generation_cost_calculator(
model=model,
image_response=image_response,
)
expected_cost = len(image_response.data) * model_info["output_cost_per_image"]
assert round(cost, 10) == round(expected_cost, 10)
def test_bedrock_anthropic_prompt_caching():
"""Test Bedrock Anthropic models with prompt caching return correct costs."""
model = "us.anthropic.claude-sonnet-4-5-20250929-v1:0"