mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
Revert batch utils with original logic
This commit is contained in:
Sameer Kankute
parent
385f4e8eee
commit
854183e3b9
1 changed files with 154 additions and 260 deletions
|
|
@ -9,7 +9,7 @@ from litellm._logging import verbose_logger
|
|||
from litellm._uuid import uuid
|
||||
from litellm.types.llms.openai import Batch
|
||||
from litellm.types.utils import CallTypes, ModelResponse, Usage
|
||||
from litellm.utils import token_counter, ProviderConfigManager
|
||||
from litellm.utils import token_counter
|
||||
|
||||
|
||||
async def calculate_batch_cost_and_usage(
|
||||
|
|
@ -30,9 +30,7 @@ async def calculate_batch_cost_and_usage(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
)
|
||||
batch_models = _get_batch_models_from_file_content(
|
||||
file_content_dictionary, model_name, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name)
|
||||
|
||||
return batch_cost, batch_usage, batch_models
|
||||
|
||||
|
|
@ -60,136 +58,14 @@ async def _handle_completed_batch(
|
|||
model_name=model_name,
|
||||
)
|
||||
|
||||
batch_models = _get_batch_models_from_file_content(
|
||||
file_content_dictionary, model_name, custom_llm_provider=custom_llm_provider
|
||||
)
|
||||
batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name)
|
||||
|
||||
return batch_cost, batch_usage, batch_models
|
||||
|
||||
|
||||
def transform_raw_provider_response_to_openai(
|
||||
raw_response: dict,
|
||||
custom_llm_provider: str,
|
||||
model: Optional[str] = None,
|
||||
messages: Optional[list] = None,
|
||||
) -> ModelResponse:
|
||||
"""
|
||||
Unified method to transform any raw LLM provider response to OpenAI format.
|
||||
|
||||
Args:
|
||||
raw_response: Raw response dictionary from any provider (Anthropic, OpenAI, etc.)
|
||||
custom_llm_provider: Provider name (e.g., "anthropic", "openai", "vertex_ai")
|
||||
model: Model name (optional, will try to extract from raw_response if not provided)
|
||||
messages: Original messages list (optional, defaults to empty list)
|
||||
|
||||
Returns:
|
||||
ModelResponse: OpenAI-compatible response object
|
||||
"""
|
||||
# Lazy import to avoid circular dependency
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
|
||||
# Extract model from response if not provided
|
||||
if model is None:
|
||||
model = raw_response.get("model", "unknown-model")
|
||||
|
||||
# Default messages if not provided
|
||||
if messages is None:
|
||||
messages = []
|
||||
|
||||
# Get provider config using ProviderConfigManager
|
||||
provider_config = ProviderConfigManager.get_provider_chat_config(
|
||||
model=model,
|
||||
provider=litellm.LlmProviders(custom_llm_provider)
|
||||
)
|
||||
|
||||
if provider_config is None:
|
||||
raise ValueError(f"Could not get config for provider: {custom_llm_provider}")
|
||||
|
||||
# Create a mock httpx.Response from the dict
|
||||
response_text = json.dumps(raw_response)
|
||||
mock_httpx_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=response_text.encode('utf-8'),
|
||||
headers={"content-type": "application/json"}
|
||||
)
|
||||
|
||||
# Create a minimal logging object
|
||||
logging_obj = Logging(
|
||||
model=model,
|
||||
messages=messages,
|
||||
stream=False,
|
||||
call_type=CallTypes.completion.value,
|
||||
start_time=time.time(),
|
||||
litellm_call_id=None,
|
||||
function_id=None,
|
||||
)
|
||||
|
||||
# Create empty ModelResponse to be populated
|
||||
model_response = ModelResponse()
|
||||
|
||||
# Call transform_response on the provider config
|
||||
transformed_response = provider_config.transform_response(
|
||||
model=model,
|
||||
raw_response=mock_httpx_response,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
request_data={},
|
||||
messages=messages,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
encoding=litellm.encoding,
|
||||
api_key=None,
|
||||
json_mode=None,
|
||||
)
|
||||
|
||||
return transformed_response
|
||||
|
||||
|
||||
def _extract_raw_response_from_batch_item(
|
||||
batch_item: dict,
|
||||
custom_llm_provider: str,
|
||||
) -> Optional[dict]:
|
||||
"""
|
||||
Extract the raw provider response from a batch output file item.
|
||||
|
||||
Handles different batch output formats:
|
||||
- Anthropic: {"result": {"type": "succeeded", "message": {...}}}
|
||||
- Vertex AI: {"status": "JOB_STATE_SUCCEEDED", "response": {...}}
|
||||
- OpenAI/Azure: {"response": {"status_code": 200, "body": {...}}}
|
||||
|
||||
Args:
|
||||
batch_item: A single item from the batch output file
|
||||
custom_llm_provider: Provider name (e.g., "anthropic", "openai", "vertex_ai")
|
||||
|
||||
Returns:
|
||||
Raw response dict or None if not successful
|
||||
"""
|
||||
# Anthropic format: {"result": {"type": "succeeded", "message": {...}}}
|
||||
if custom_llm_provider == "anthropic":
|
||||
result = batch_item.get("result", {})
|
||||
if result.get("type") == "succeeded":
|
||||
return result.get("message", {})
|
||||
return None
|
||||
|
||||
# Vertex AI format: {"status": "JOB_STATE_SUCCEEDED", "response": {...}}
|
||||
if custom_llm_provider == "vertex_ai":
|
||||
if batch_item.get("status") == "JOB_STATE_SUCCEEDED":
|
||||
return batch_item.get("response", {})
|
||||
return None
|
||||
|
||||
# OpenAI/Azure format: {"response": {"status_code": 200, "body": {...}}}
|
||||
# Default to OpenAI format for openai, azure, hosted_vllm, etc.
|
||||
response = batch_item.get("response", {})
|
||||
if response.get("status_code") == 200:
|
||||
return response.get("body", {})
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def _get_batch_models_from_file_content(
|
||||
file_content_dictionary: List[dict],
|
||||
model_name: Optional[str] = None,
|
||||
custom_llm_provider: str = "openai",
|
||||
) -> List[str]:
|
||||
"""
|
||||
Get the models from the file content
|
||||
|
|
@ -198,16 +74,11 @@ def _get_batch_models_from_file_content(
|
|||
return [model_name]
|
||||
batch_models = []
|
||||
for _item in file_content_dictionary:
|
||||
if _batch_response_was_successful(_item, custom_llm_provider=custom_llm_provider):
|
||||
# Extract raw response using generalized method
|
||||
raw_response = _extract_raw_response_from_batch_item(
|
||||
batch_item=_item,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
if raw_response:
|
||||
_model = raw_response.get("model")
|
||||
if _model:
|
||||
batch_models.append(_model)
|
||||
if _batch_response_was_successful(_item):
|
||||
_response_body = _get_response_from_batch_job_output_file(_item)
|
||||
_model = _response_body.get("model")
|
||||
if _model:
|
||||
batch_models.append(_model)
|
||||
return batch_models
|
||||
|
||||
|
||||
|
|
@ -219,54 +90,99 @@ def _batch_cost_calculator(
|
|||
"""
|
||||
Calculate the cost of a batch based on the output file id
|
||||
"""
|
||||
total_cost: float = 0.0
|
||||
|
||||
for batch_item in file_content_dictionary:
|
||||
if not _batch_response_was_successful(batch_item, custom_llm_provider=custom_llm_provider):
|
||||
continue
|
||||
|
||||
# Extract raw response from batch item
|
||||
raw_response = _extract_raw_response_from_batch_item(
|
||||
batch_item=batch_item,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
if raw_response is None:
|
||||
continue
|
||||
|
||||
# Extract model from response if not provided
|
||||
actual_model = model_name or raw_response.get("model")
|
||||
if actual_model is None:
|
||||
verbose_logger.warning("Could not determine model for batch item, skipping cost calculation")
|
||||
continue
|
||||
|
||||
try:
|
||||
# Transform to OpenAI format using generalized method
|
||||
openai_format_response = transform_raw_provider_response_to_openai(
|
||||
raw_response=raw_response,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model=actual_model,
|
||||
messages=[], # Messages not needed for cost calculation
|
||||
)
|
||||
|
||||
# Calculate cost using standard OpenAI cost calculation
|
||||
cost = litellm.completion_cost(
|
||||
completion_response=openai_format_response,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
total_cost += cost
|
||||
verbose_logger.debug("item_cost=%s, total_cost=%s", cost, total_cost)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"Error calculating cost for batch item: {e}. Skipping this item."
|
||||
)
|
||||
continue
|
||||
# Handle Vertex AI with specialized method
|
||||
if custom_llm_provider == "vertex_ai" and model_name:
|
||||
batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
|
||||
verbose_logger.debug("vertex_ai_total_cost=%s", batch_cost)
|
||||
return batch_cost
|
||||
|
||||
# For other providers, use the existing logic
|
||||
total_cost = _get_batch_job_cost_from_file_content(
|
||||
file_content_dictionary=file_content_dictionary,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
verbose_logger.debug("total_cost=%s", total_cost)
|
||||
return total_cost
|
||||
|
||||
|
||||
def calculate_vertex_ai_batch_cost_and_usage(
|
||||
vertex_ai_batch_responses: List[dict],
|
||||
model_name: Optional[str] = None,
|
||||
) -> Tuple[float, Usage]:
|
||||
"""
|
||||
Calculate both cost and usage from Vertex AI batch responses
|
||||
"""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
VertexGeminiConfig,
|
||||
)
|
||||
total_cost = 0.0
|
||||
total_tokens = 0
|
||||
prompt_tokens = 0
|
||||
completion_tokens = 0
|
||||
|
||||
for response in vertex_ai_batch_responses:
|
||||
if response.get("status") == "JOB_STATE_SUCCEEDED": # Check if response was successful
|
||||
# Transform Vertex AI response to OpenAI format if needed
|
||||
|
||||
# Create required arguments for the transformation method
|
||||
model_response = ModelResponse()
|
||||
|
||||
# Ensure model_name is not None
|
||||
actual_model_name = model_name or "gemini-2.5-flash"
|
||||
|
||||
# Create a real LiteLLM logging object
|
||||
logging_obj = Logging(
|
||||
model=actual_model_name,
|
||||
messages=[{"role": "user", "content": "batch_request"}],
|
||||
stream=False,
|
||||
call_type=CallTypes.aretrieve_batch,
|
||||
start_time=time.time(),
|
||||
litellm_call_id="batch_" + str(uuid.uuid4()),
|
||||
function_id="batch_processing",
|
||||
litellm_trace_id=str(uuid.uuid4()),
|
||||
kwargs={"optional_params": {}}
|
||||
)
|
||||
|
||||
# Add the optional_params attribute that the Vertex AI transformation expects
|
||||
logging_obj.optional_params = {}
|
||||
raw_response = httpx.Response(200) # Mock response object
|
||||
|
||||
openai_format_response = VertexGeminiConfig()._transform_google_generate_content_to_openai_model_response(
|
||||
completion_response=response["response"],
|
||||
model_response=model_response,
|
||||
model=actual_model_name,
|
||||
logging_obj=logging_obj,
|
||||
raw_response=raw_response,
|
||||
)
|
||||
|
||||
# Calculate cost using existing function
|
||||
cost = litellm.completion_cost(
|
||||
completion_response=openai_format_response,
|
||||
custom_llm_provider="vertex_ai",
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
total_cost += cost
|
||||
|
||||
# Extract usage from the transformed response
|
||||
usage_obj = getattr(openai_format_response, 'usage', None)
|
||||
if usage_obj:
|
||||
usage = usage_obj
|
||||
else:
|
||||
# Fallback: create usage from response dict
|
||||
response_dict = openai_format_response.dict() if hasattr(openai_format_response, 'dict') else {}
|
||||
usage = _get_batch_job_usage_from_response_body(response_dict)
|
||||
|
||||
total_tokens += usage.total_tokens
|
||||
prompt_tokens += usage.prompt_tokens
|
||||
completion_tokens += usage.completion_tokens
|
||||
|
||||
return total_cost, Usage(
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
)
|
||||
|
||||
|
||||
async def _get_batch_output_file_content_as_dictionary(
|
||||
batch: Batch,
|
||||
|
|
@ -307,6 +223,34 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]:
|
|||
raise e
|
||||
|
||||
|
||||
def _get_batch_job_cost_from_file_content(
|
||||
file_content_dictionary: List[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
|
||||
) -> float:
|
||||
"""
|
||||
Get the cost of a batch job from the file content
|
||||
"""
|
||||
try:
|
||||
total_cost: float = 0.0
|
||||
# parse the file content as json
|
||||
verbose_logger.debug(
|
||||
"file_content_dictionary=%s", json.dumps(file_content_dictionary, indent=4)
|
||||
)
|
||||
for _item in file_content_dictionary:
|
||||
if _batch_response_was_successful(_item):
|
||||
_response_body = _get_response_from_batch_job_output_file(_item)
|
||||
total_cost += litellm.completion_cost(
|
||||
completion_response=_response_body,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
verbose_logger.debug("total_cost=%s", total_cost)
|
||||
return total_cost
|
||||
except Exception as e:
|
||||
verbose_logger.error("error in _get_batch_job_cost_from_file_content", e)
|
||||
raise e
|
||||
|
||||
|
||||
def _get_batch_job_total_usage_from_file_content(
|
||||
file_content_dictionary: List[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
|
||||
|
|
@ -315,64 +259,26 @@ def _get_batch_job_total_usage_from_file_content(
|
|||
"""
|
||||
Get the tokens of a batch job from the file content
|
||||
"""
|
||||
from litellm.cost_calculator import BaseTokenUsageProcessor
|
||||
# Handle Vertex AI with specialized method
|
||||
if custom_llm_provider == "vertex_ai" and model_name:
|
||||
_, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
|
||||
return batch_usage
|
||||
|
||||
all_usage: List[Usage] = []
|
||||
|
||||
for batch_item in file_content_dictionary:
|
||||
if not _batch_response_was_successful(batch_item, custom_llm_provider=custom_llm_provider):
|
||||
continue
|
||||
|
||||
# Extract raw response from batch item
|
||||
raw_response = _extract_raw_response_from_batch_item(
|
||||
batch_item=batch_item,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
)
|
||||
|
||||
if raw_response is None:
|
||||
continue
|
||||
|
||||
# Extract model from response if not provided
|
||||
actual_model = model_name or raw_response.get("model")
|
||||
if actual_model is None:
|
||||
verbose_logger.warning("Could not determine model for batch item, skipping usage calculation")
|
||||
continue
|
||||
|
||||
try:
|
||||
# Transform to OpenAI format using generalized method
|
||||
openai_format_response = transform_raw_provider_response_to_openai(
|
||||
raw_response=raw_response,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model=actual_model,
|
||||
messages=[], # Messages not needed for usage extraction
|
||||
)
|
||||
|
||||
# Extract usage from transformed response
|
||||
usage_obj = getattr(openai_format_response, 'usage', None)
|
||||
if usage_obj:
|
||||
all_usage.append(usage_obj)
|
||||
else:
|
||||
# Fallback: try to extract from response dict
|
||||
response_dict = openai_format_response.dict() if hasattr(openai_format_response, 'dict') else {}
|
||||
usage = _get_batch_job_usage_from_response_body(response_dict)
|
||||
if usage and usage.total_tokens > 0:
|
||||
all_usage.append(usage)
|
||||
except Exception as e:
|
||||
verbose_logger.warning(
|
||||
f"Error extracting usage for batch item: {e}. Skipping this item."
|
||||
)
|
||||
continue
|
||||
|
||||
# Combine all usage objects
|
||||
if all_usage:
|
||||
combined_usage = BaseTokenUsageProcessor.combine_usage_objects(all_usage)
|
||||
return combined_usage
|
||||
|
||||
# Return empty usage if no valid responses
|
||||
# For other providers, use the existing logic
|
||||
total_tokens: int = 0
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
for _item in file_content_dictionary:
|
||||
if _batch_response_was_successful(_item):
|
||||
_response_body = _get_response_from_batch_job_output_file(_item)
|
||||
usage: Usage = _get_batch_job_usage_from_response_body(_response_body)
|
||||
total_tokens += usage.total_tokens
|
||||
prompt_tokens += usage.prompt_tokens
|
||||
completion_tokens += usage.completion_tokens
|
||||
return Usage(
|
||||
total_tokens=0,
|
||||
prompt_tokens=0,
|
||||
completion_tokens=0,
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
)
|
||||
|
||||
def _get_batch_job_input_file_usage(
|
||||
|
|
@ -412,30 +318,18 @@ def _get_batch_job_usage_from_response_body(response_body: dict) -> Usage:
|
|||
return usage
|
||||
|
||||
|
||||
def _batch_response_was_successful(
|
||||
batch_job_output_file: dict,
|
||||
custom_llm_provider: str = "openai",
|
||||
) -> bool:
|
||||
def _get_response_from_batch_job_output_file(batch_job_output_file: dict) -> Any:
|
||||
"""
|
||||
Check if the batch job response was successful.
|
||||
|
||||
Args:
|
||||
batch_job_output_file: A single item from the batch output file
|
||||
custom_llm_provider: Provider name (e.g., "anthropic", "openai", "vertex_ai")
|
||||
|
||||
Returns:
|
||||
True if the batch response was successful, False otherwise
|
||||
Get the response from the batch job output file
|
||||
"""
|
||||
# Anthropic format: {"result": {"type": "succeeded", "message": {...}}}
|
||||
if custom_llm_provider == "anthropic":
|
||||
result = batch_job_output_file.get("result", {})
|
||||
return result.get("type") == "succeeded"
|
||||
|
||||
# Vertex AI format: {"status": "JOB_STATE_SUCCEEDED", "response": {...}}
|
||||
if custom_llm_provider == "vertex_ai":
|
||||
return batch_job_output_file.get("status") == "JOB_STATE_SUCCEEDED"
|
||||
|
||||
# OpenAI/Azure format: {"response": {"status_code": 200, "body": {...}}}
|
||||
# Default to OpenAI format for openai, azure, hosted_vllm, etc.
|
||||
_response: dict = batch_job_output_file.get("response", None) or {}
|
||||
return _response.get("status_code", None) == 200
|
||||
_response_body = _response.get("body", None) or {}
|
||||
return _response_body
|
||||
|
||||
|
||||
def _batch_response_was_successful(batch_job_output_file: dict) -> bool:
|
||||
"""
|
||||
Check if the batch job response status == 200
|
||||
"""
|
||||
_response: dict = batch_job_output_file.get("response", None) or {}
|
||||
return _response.get("status_code", None) == 200
|
||||
Loading…
Add table
Reference in a new issue