fix imagegeneration output_cost_per_image on model cost map (#6752)

This commit is contained in:
Ishaan Jaff 2024-11-14 20:37:21 -08:00 committed by GitHub
parent 0585fd56d6
commit c03351328f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 39 additions and 9 deletions

View file

@ -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_per_token as openai_cost_per_token
from litellm.llms.OpenAI.cost_calculation import cost_router as openai_cost_router 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.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.llms.openai import HttpxBinaryResponseContent
from litellm.types.rerank import RerankResponse from litellm.types.rerank import RerankResponse
from litellm.types.router import SPECIAL_MODEL_INFO_PARAMS from litellm.types.router import SPECIAL_MODEL_INFO_PARAMS
@ -667,9 +670,11 @@ def completion_cost( # noqa: PLR0915
): ):
### IMAGE GENERATION COST CALCULATION ### ### IMAGE GENERATION COST CALCULATION ###
if custom_llm_provider == "vertex_ai": if custom_llm_provider == "vertex_ai":
# https://cloud.google.com/vertex-ai/generative-ai/pricing if isinstance(completion_response, ImageResponse):
# Vertex Charges Flat $0.20 per image return vertex_ai_image_cost_calculator(
return 0.020 model=model,
image_response=completion_response,
)
elif custom_llm_provider == "bedrock": elif custom_llm_provider == "bedrock":
if isinstance(completion_response, ImageResponse): if isinstance(completion_response, ImageResponse):
return bedrock_image_cost_calculator( return bedrock_image_cost_calculator(

View file

@ -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

View file

@ -2986,19 +2986,19 @@
"supports_function_calling": true "supports_function_calling": true
}, },
"vertex_ai/imagegeneration@006": { "vertex_ai/imagegeneration@006": {
"cost_per_image": 0.020, "output_cost_per_image": 0.020,
"litellm_provider": "vertex_ai-image-models", "litellm_provider": "vertex_ai-image-models",
"mode": "image_generation", "mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
}, },
"vertex_ai/imagen-3.0-generate-001": { "vertex_ai/imagen-3.0-generate-001": {
"cost_per_image": 0.04, "output_cost_per_image": 0.04,
"litellm_provider": "vertex_ai-image-models", "litellm_provider": "vertex_ai-image-models",
"mode": "image_generation", "mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
}, },
"vertex_ai/imagen-3.0-fast-generate-001": { "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", "litellm_provider": "vertex_ai-image-models",
"mode": "image_generation", "mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"

View file

@ -2986,19 +2986,19 @@
"supports_function_calling": true "supports_function_calling": true
}, },
"vertex_ai/imagegeneration@006": { "vertex_ai/imagegeneration@006": {
"cost_per_image": 0.020, "output_cost_per_image": 0.020,
"litellm_provider": "vertex_ai-image-models", "litellm_provider": "vertex_ai-image-models",
"mode": "image_generation", "mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
}, },
"vertex_ai/imagen-3.0-generate-001": { "vertex_ai/imagen-3.0-generate-001": {
"cost_per_image": 0.04, "output_cost_per_image": 0.04,
"litellm_provider": "vertex_ai-image-models", "litellm_provider": "vertex_ai-image-models",
"mode": "image_generation", "mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"
}, },
"vertex_ai/imagen-3.0-fast-generate-001": { "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", "litellm_provider": "vertex_ai-image-models",
"mode": "image_generation", "mode": "image_generation",
"source": "https://cloud.google.com/vertex-ai/generative-ai/pricing" "source": "https://cloud.google.com/vertex-ai/generative-ai/pricing"