scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-16 23:41:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

refactor(vertex): move batch usage modality parsing under llms

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
kerry 2026-09-15 02:03:32 +00:00
parent ca7364fb05
commit a5fc880c90
2 changed files with 34 additions and 32 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm/batches/batch_utils.py
View file

@ -9,8 +9,9 @@ import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS
from litellm.litellm_core_utils.llm_cost_calc.utils import parse_prompt_tokens_details
from litellm.llms.vertex_ai.batches.transformation import vertex_prompt_tokens_details
from litellm.types.llms.openai import Batch
from litellm.types.utils import ModelInfo, PromptTokensDetailsWrapper, Usage
from litellm.types.utils import ModelInfo, Usage
from litellm.utils import token_counter
@ -310,35 +311,6 @@ def _aggregate_batch_cost_usage_models(
)
def _vertex_prompt_tokens_details(
usage_metadata: Mapping[str, object],
) -> PromptTokensDetailsWrapper | None:
raw_details: Final = usage_metadata.get("promptTokensDetails")
if not isinstance(raw_details, list):
return None
def _normalize(detail: object) -> tuple[str, int] | None:
if not isinstance(detail, Mapping):
return None
modality: Final = detail.get("modality")
token_count: Final = detail.get("tokenCount")
if not isinstance(modality, str) or not isinstance(token_count, int):
return None
return modality.upper(), token_count
parsed_details: Final = tuple(_normalize(detail) for detail in raw_details)
normalized: Final = tuple(detail for detail in parsed_details if detail is not None)
if len(normalized) != len(parsed_details):
return None
return PromptTokensDetailsWrapper(
text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")),
audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"),
image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"),
video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"),
)
def calculate_vertex_ai_batch_cost_and_usage(
vertex_ai_batch_responses: list[dict],
model_name: str | None = None,
@ -385,7 +357,7 @@ def calculate_vertex_ai_batch_cost_and_usage(
prompt_tokens=_prompt,
completion_tokens=_completion,
total_tokens=_total,
prompt_tokens_details=_vertex_prompt_tokens_details(usage_metadata),
prompt_tokens_details=vertex_prompt_tokens_details(usage_metadata),
)
try:

32
litellm/llms/vertex_ai/batches/transformation.py
View file

@ -1,3 +1,4 @@
from collections.abc import Mapping
from typing import Any, Final
from urllib.parse import unquote
@ -8,7 +9,36 @@ from litellm.llms.vertex_ai.common_utils import (
)
from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest
from litellm.types.llms.vertex_ai import *
from litellm.types.utils import LiteLLMBatch
from litellm.types.utils import LiteLLMBatch, PromptTokensDetailsWrapper
def vertex_prompt_tokens_details(
usage_metadata: Mapping[str, object],
) -> PromptTokensDetailsWrapper | None:
raw_details: Final = usage_metadata.get("promptTokensDetails")
if not isinstance(raw_details, list):
return None
def _normalize(detail: object) -> tuple[str, int] | None:
if not isinstance(detail, Mapping):
return None
modality: Final = detail.get("modality")
token_count: Final = detail.get("tokenCount")
if not isinstance(modality, str) or not isinstance(token_count, int):
return None
return modality.upper(), token_count
parsed_details: Final = tuple(_normalize(detail) for detail in raw_details)
normalized: Final = tuple(detail for detail in parsed_details if detail is not None)
if len(normalized) != len(parsed_details):
return None
return PromptTokensDetailsWrapper(
text_tokens=sum(token_count for modality, token_count in normalized if modality in ("TEXT", "DOCUMENT")),
audio_tokens=sum(token_count for modality, token_count in normalized if modality == "AUDIO"),
image_tokens=sum(token_count for modality, token_count in normalized if modality == "IMAGE"),
video_tokens=sum(token_count for modality, token_count in normalized if modality == "VIDEO"),
)
class VertexAIBatchTransformation: