From 854183e3b9895b361da83df2357e4c93e773429b Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 10 Dec 2025 17:01:21 +0530 Subject: [PATCH] Revert batch utils with original logic --- litellm/batches/batch_utils.py | 414 ++++++++++++--------------------- 1 file changed, 154 insertions(+), 260 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index b7d147d8e19..8a078eeaca1 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -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 \ No newline at end of file