diff --git a/enterprise/enterprise_hooks/managed_files.py b/enterprise/enterprise_hooks/managed_files.py index 3d5f6f6096..3819a58756 100644 --- a/enterprise/enterprise_hooks/managed_files.py +++ b/enterprise/enterprise_hooks/managed_files.py @@ -23,7 +23,7 @@ from litellm.types.llms.openai import ( OpenAIFileObject, OpenAIFilesPurpose, ) -from litellm.types.utils import SpecialEnums +from litellm.types.utils import LiteLLMBatch, LLMResponseTypes, SpecialEnums if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -136,6 +136,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "pass_through_endpoint", "rerank", "acreate_batch", + "aretrieve_batch", ], ) -> Union[Exception, str, Dict, None]: """ @@ -161,6 +162,21 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) data["model_file_id_mapping"] = model_file_id_mapping + elif call_type == CallTypes.aretrieve_batch.value: + retrieve_batch_id = cast(Optional[str], data.get("batch_id")) + potential_batch_id = ( + _is_base64_encoded_unified_file_id(retrieve_batch_id) + if retrieve_batch_id + else False + ) + if potential_batch_id: + ## for managed batch id - get the model id + model_id = self.get_model_id_from_unified_batch_id(potential_batch_id) + data["model"] = model_id + data["batch_id"] = self.get_batch_id_from_unified_batch_id( + potential_batch_id + ) + return data async def async_pre_call_deployment_hook( @@ -340,6 +356,38 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return response + def get_unified_batch_id(self, batch_id: str, model_id: str) -> str: + unified_batch_id = SpecialEnums.LITELLM_MANAGED_BATCH_COMPLETE_STR.value.format( + model_id, batch_id + ) + return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=") + + def get_model_id_from_unified_batch_id(self, file_id: str) -> str: + ## use regex to get the model_id from the file_id + return file_id.split("model_id:")[1].split(";")[0] + + def get_batch_id_from_unified_batch_id(self, file_id: str) -> str: + ## use regex to get the batch_id from the file_id + return file_id.split("llm_batch_id:")[1].split(",")[0] + + async def async_post_call_success_hook( + self, data: Dict, user_api_key_dict: UserAPIKeyAuth, response: LLMResponseTypes + ) -> Any: + if isinstance(response, LiteLLMBatch): + ## Check if unified_file_id is in the response + unified_batch_id = response._hidden_params.get( + "unified_file_id" + ) # managed file id + model_id = response._hidden_params.get("model_id") + if unified_batch_id and model_id: + response.id = self.get_unified_batch_id( + batch_id=response.id, model_id=model_id + ) + + return await super().async_post_call_success_hook( + data, user_api_key_dict, response + ) + async def afile_retrieve( self, file_id: str, litellm_parent_otel_span: Optional[Span] ) -> OpenAIFileObject: diff --git a/litellm/batches/main.py b/litellm/batches/main.py index 0be9667790..9852755622 100644 --- a/litellm/batches/main.py +++ b/litellm/batches/main.py @@ -51,7 +51,7 @@ async def acreate_batch( extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, -) -> Batch: +) -> LiteLLMBatch: """ Async: Creates and executes a batch from an uploaded file of request diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 354cde5739..0789d7669d 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -3674,6 +3674,7 @@ def get_standard_logging_object_payload( or litellm_params.get("metadata", None) or {} ) + completion_start_time = kwargs.get("completion_start_time", end_time) call_type = kwargs.get("call_type") cache_hit = kwargs.get("cache_hit", False) diff --git a/litellm/proxy/_experimental/out/onboarding.html b/litellm/proxy/_experimental/out/onboarding.html deleted file mode 100644 index 4b81c88e56..0000000000 --- a/litellm/proxy/_experimental/out/onboarding.html +++ /dev/null @@ -1 +0,0 @@ -LiteLLM Dashboard \ No newline at end of file diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 732265548d..eacfe88d6d 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -56,4 +56,17 @@ model_list: model: gpt-image-1 api_key: os.environ/OPENAI_API_KEY # drop_params: true - + - model_name: "gpt-4o-batch" + litellm_params: + model: azure/gpt-4o-mini + api_base: os.environ/AZURE_API_BASE + api_key: os.environ/AZURE_API_KEY + model_info: + id: my-general-azure-deployment + - model_name: "gpt-4o-batch" + litellm_params: + model: azure/gpt-4o-mini + api_base: https://krris-m2f9a9i7-eastus2.openai.azure.com + api_key: 04d22fb7e9ad4d9c8afe7c6abf97a6fc + model_info: + id: my-unique-azure-deployment diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index d8f3c58849..bf5173f893 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -126,7 +126,7 @@ async def create_batch( response = await llm_router.acreate_batch(**_create_batch_data) # type: ignore elif ( - unified_file_id + unified_file_id and input_file_id ): # litellm_proxy:application/octet-stream;unified_id,c4843482-b176-4901-8292-7523fd0f2c6e;target_model_names,gpt-4o-mini target_model_names = get_models_from_unified_file_id(unified_file_id) ## EXPECTS 1 MODEL @@ -150,11 +150,18 @@ async def create_batch( ) response = await llm_router.acreate_batch(**_create_batch_data) + response.input_file_id = input_file_id + response._hidden_params["unified_file_id"] = unified_file_id else: response = await litellm.acreate_batch( custom_llm_provider=custom_llm_provider, **_create_batch_data # type: ignore ) + ### CALL HOOKS ### - modify outgoing data + response = await proxy_logging_obj.post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + ### ALERTING ### asyncio.create_task( proxy_logging_obj.update_request_status( @@ -231,7 +238,6 @@ async def retrieve_batch( ``` """ from litellm.proxy.proxy_server import ( - add_litellm_data_to_request, general_settings, llm_router, proxy_config, @@ -248,22 +254,24 @@ async def retrieve_batch( data = cast(dict, _retrieve_batch_request) - # setup logging - data["litellm_call_id"] = request.headers.get( - "x-litellm-call-id", str(uuid.uuid4()) - ) - - # Include original request and headers in the data - data = await add_litellm_data_to_request( - data=data, + base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) + ( + data, + litellm_logging_obj, + ) = await base_llm_response_processor.common_processing_pre_call_logic( request=request, general_settings=general_settings, user_api_key_dict=user_api_key_dict, version=version, + proxy_logging_obj=proxy_logging_obj, proxy_config=proxy_config, + route_type="aretrieve_batch", ) - if litellm.enable_loadbalancing_on_batch_endpoints is True: + if ( + litellm.enable_loadbalancing_on_batch_endpoints is True + or data.get("model") is not None + ): if llm_router is None: raise HTTPException( status_code=500, @@ -283,6 +291,11 @@ async def retrieve_batch( custom_llm_provider=custom_llm_provider, **data # type: ignore ) + ### CALL HOOKS ### - modify outgoing data + response = await proxy_logging_obj.post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + ### ALERTING ### asyncio.create_task( proxy_logging_obj.update_request_status( diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3c9a3f4a8f..8fd1b4ddc0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -115,6 +115,7 @@ class ProxyBaseLLMRequestProcessing: "aget_responses", "adelete_responses", "acreate_batch", + "aretrieve_batch", ], version: Optional[str] = None, user_model: Optional[str] = None, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 4b8447fb03..22e023d9ed 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -96,11 +96,7 @@ class _ProxyDBLogger(CustomLogger): start_time=None, end_time=None, # start/end time for completion ): - from litellm.proxy.proxy_server import ( - prisma_client, - proxy_logging_obj, - update_cache, - ) + from litellm.proxy.proxy_server import proxy_logging_obj, update_cache verbose_proxy_logger.debug("INSIDE _PROXY_track_cost_callback") try: @@ -134,7 +130,7 @@ class _ProxyDBLogger(CustomLogger): ) verbose_proxy_logger.debug( - f"user_api_key {user_api_key}, prisma_client: {prisma_client}" + f"user_api_key {user_api_key}, user_id {user_id}, team_id {team_id}, end_user_id {end_user_id}" ) if _should_track_cost_callback( user_api_key=user_api_key, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 097f798de2..e81501eebe 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -67,7 +67,9 @@ def _get_metadata_variable_name(request: Request) -> str: "batches", "/v1/messages", "responses", + "files", ] + if any( [ litellm_metadata_route in request.url.path diff --git a/litellm/router.py b/litellm/router.py index 3e9087a100..19943ba504 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -106,12 +106,7 @@ from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_successes_for_current_minute, ) from litellm.scheduler import FlowItem, Scheduler -from litellm.types.llms.openai import ( - AllMessageValues, - Batch, - FileTypes, - OpenAIFileObject, -) +from litellm.types.llms.openai import AllMessageValues, FileTypes, OpenAIFileObject from litellm.types.router import ( CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, VALID_LITELLM_ENVIRONMENTS, @@ -135,7 +130,7 @@ from litellm.types.router import ( RoutingStrategy, ) from litellm.types.services import ServiceTypes -from litellm.types.utils import GenericBudgetConfigType +from litellm.types.utils import GenericBudgetConfigType, LiteLLMBatch from litellm.types.utils import ModelInfo from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.types.utils import StandardLoggingPayload @@ -2743,68 +2738,87 @@ class Router: f"Inside _atext_completion()- model: {model}; kwargs: {kwargs}" ) parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) - deployment = await self.async_get_available_deployment( + healthy_deployments = await self.async_get_healthy_deployments( model=model, messages=[{"role": "user", "content": "files-api-fake-text"}], specific_deployment=kwargs.pop("specific_deployment", None), request_kwargs=kwargs, - ) - self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs) - - data = deployment["litellm_params"].copy() - model_name = data["model"] - - model_client = self._get_async_openai_model_client( - deployment=deployment, - kwargs=kwargs, - ) - self.total_calls[model_name] += 1 - - ## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ## - stripped_model, custom_llm_provider, _, _ = get_llm_provider( - model=data["model"] + parent_otel_span=parent_otel_span, ) - response = litellm.acreate_file( - **{ - **data, - "custom_llm_provider": custom_llm_provider, - "caching": self.cache_responses, - "client": model_client, - **kwargs, - } - ) + async def create_file_for_deployment(deployment: dict) -> OpenAIFileObject: + kwargs_copy = copy.deepcopy(kwargs) + self._update_kwargs_with_deployment( + deployment=deployment, + kwargs=kwargs_copy, + function_name="acreate_file", + ) + data = deployment["litellm_params"].copy() + model_name = data["model"] - rpm_semaphore = self._get_client( - deployment=deployment, - kwargs=kwargs, - client_type="max_parallel_requests", - ) + model_client = self._get_async_openai_model_client( + deployment=deployment, + kwargs=kwargs_copy, + ) + self.total_calls[model_name] += 1 - if rpm_semaphore is not None and isinstance( - rpm_semaphore, asyncio.Semaphore - ): - async with rpm_semaphore: - """ - - Check rpm limits before making the call - - If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe) - """ + ## REPLACE MODEL IN FILE WITH SELECTED DEPLOYMENT ## + stripped_model, custom_llm_provider, _, _ = get_llm_provider( + model=data["model"] + ) + response = litellm.acreate_file( + **{ + **data, + "custom_llm_provider": custom_llm_provider, + "caching": self.cache_responses, + "client": model_client, + **kwargs_copy, + } + ) + + rpm_semaphore = self._get_client( + deployment=deployment, + kwargs=kwargs_copy, + client_type="max_parallel_requests", + ) + + if rpm_semaphore is not None and isinstance( + rpm_semaphore, asyncio.Semaphore + ): + async with rpm_semaphore: + """ + - Check rpm limits before making the call + - If allowed, increment the rpm limit (allows global value to be updated, concurrency-safe) + """ + await self.async_routing_strategy_pre_call_checks( + deployment=deployment, parent_otel_span=parent_otel_span + ) + response = await response # type: ignore + else: await self.async_routing_strategy_pre_call_checks( deployment=deployment, parent_otel_span=parent_otel_span ) response = await response # type: ignore - else: - await self.async_routing_strategy_pre_call_checks( - deployment=deployment, parent_otel_span=parent_otel_span + + self.success_calls[model_name] += 1 + verbose_router_logger.info( + f"litellm.acreate_file(model={model_name})\033[32m 200 OK\033[0m" ) - response = await response # type: ignore - self.success_calls[model_name] += 1 - verbose_router_logger.info( - f"litellm.acreate_file(model={model_name})\033[32m 200 OK\033[0m" - ) + return response - return response # type: ignore + tasks = [] + if isinstance(healthy_deployments, dict): + tasks.append(create_file_for_deployment(healthy_deployments)) + else: + for deployment in healthy_deployments: + tasks.append(create_file_for_deployment(deployment)) + + responses = await asyncio.gather(*tasks) + + if len(responses) == 0: + raise Exception("No healthy deployments found.") + return responses[0] except Exception as e: verbose_router_logger.exception( f"litellm.acreate_file(model={model}, {kwargs})\033[31m Exception {str(e)}\033[0m" @@ -2817,7 +2831,7 @@ class Router: self, model: str, **kwargs, - ) -> Batch: + ) -> LiteLLMBatch: try: kwargs["model"] = model kwargs["original_function"] = self._acreate_batch @@ -2848,7 +2862,7 @@ class Router: self, model: str, **kwargs, - ) -> Batch: + ) -> LiteLLMBatch: try: verbose_router_logger.debug( f"Inside _acreate_batch()- model: {model}; kwargs: {kwargs}" @@ -2861,7 +2875,6 @@ class Router: request_kwargs=kwargs, ) - kwargs["model_info"] = deployment.get("model_info", {}) data = deployment["litellm_params"].copy() model_name = data["model"] self._update_kwargs_with_deployment( @@ -2913,8 +2926,9 @@ class Router: self.success_calls[model_name] += 1 verbose_router_logger.info( - f"litellm.acreate_file(model={model_name})\033[32m 200 OK\033[0m" + f"litellm.acreate_batch(model={model_name})\033[32m 200 OK\033[0m" ) + return response # type: ignore except Exception as e: verbose_router_logger.exception( @@ -2926,26 +2940,44 @@ class Router: async def aretrieve_batch( self, + model: Optional[str] = None, **kwargs, - ) -> Batch: + ) -> LiteLLMBatch: """ Iterate through all models in a model group to check for batch Future Improvement - cache the result. """ try: - filtered_model_list = self.get_model_list() + parent_otel_span = _get_parent_otel_span_from_kwargs(kwargs) + if model is not None: + filtered_model_list: Optional[ + Union[List[DeploymentTypedDict], List[Dict], Dict] + ] = await self.async_get_healthy_deployments( + model=model, + messages=[{"role": "user", "content": "retrieve-api-fake-text"}], + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + parent_otel_span=parent_otel_span, + ) + else: + filtered_model_list = self.get_model_list() if filtered_model_list is None: raise Exception("Router not yet initialized.") receieved_exceptions = [] - async def try_retrieve_batch(model_name): + async def try_retrieve_batch(model_name: DeploymentTypedDict): try: + model = model_name["litellm_params"].get("model") + if model is None: + raise Exception( + f"Model not found in litellm_params for deployment: {model_name}" + ) # Update kwargs with the current model name or any other model-specific adjustments ## SET CUSTOM PROVIDER TO SELECTED DEPLOYMENT ## _, custom_llm_provider, _, _ = get_llm_provider( # type: ignore - model=model_name["litellm_params"]["model"] + model=model ) new_kwargs = copy.deepcopy(kwargs) new_kwargs.pop("custom_llm_provider", None) @@ -2957,15 +2989,35 @@ class Router: return None # Check all models in parallel - results = await asyncio.gather( - *[try_retrieve_batch(model) for model in filtered_model_list], - return_exceptions=True, - ) + if ( + filtered_model_list is not None + and isinstance(filtered_model_list, list) + and len(filtered_model_list) > 0 + ): + results = await asyncio.gather( + *[ + try_retrieve_batch(cast(DeploymentTypedDict, model)) + for model in filtered_model_list + ], + return_exceptions=True, + ) + elif filtered_model_list is not None and isinstance( + filtered_model_list, dict + ): + results = await try_retrieve_batch( + cast(DeploymentTypedDict, filtered_model_list) + ) + else: + raise Exception("No healthy deployments found.") # Check for successful responses and handle exceptions - for result in results: - if isinstance(result, Batch): - return result + if results is not None: + if isinstance(results, LiteLLMBatch): + return results + elif isinstance(results, list): + for result in results: + if isinstance(result, LiteLLMBatch): + return result # If no valid Batch response was found, raise the first encountered exception if receieved_exceptions: @@ -5936,6 +5988,80 @@ class Router: return model, healthy_deployments + async def async_get_healthy_deployments( + self, + model: str, + request_kwargs: Dict, + messages: Optional[List[Dict[str, str]]] = None, + input: Optional[Union[str, List]] = None, + specific_deployment: Optional[bool] = False, + parent_otel_span: Optional[Span] = None, + ) -> Union[List[Dict], Dict]: + """ + Get the healthy deployments for a model. + + Returns: + - List[Dict], if multiple models chosen + *OR* + - Dict, if specific model chosen + """ + model, healthy_deployments = self._common_checks_available_deployment( + model=model, + messages=messages, + input=input, + specific_deployment=specific_deployment, + ) # type: ignore + + if isinstance(healthy_deployments, dict): + return healthy_deployments + + cooldown_deployments = await _async_get_cooldown_deployments( + litellm_router_instance=self, parent_otel_span=parent_otel_span + ) + verbose_router_logger.debug( + f"async cooldown deployments: {cooldown_deployments}" + ) + verbose_router_logger.debug(f"cooldown_deployments: {cooldown_deployments}") + healthy_deployments = self._filter_cooldown_deployments( + healthy_deployments=healthy_deployments, + cooldown_deployments=cooldown_deployments, + ) + + healthy_deployments = await self.async_callback_filter_deployments( + model=model, + healthy_deployments=healthy_deployments, + messages=( + cast(List[AllMessageValues], messages) if messages is not None else None + ), + request_kwargs=request_kwargs, + parent_otel_span=parent_otel_span, + ) + + if self.enable_pre_call_checks and messages is not None: + healthy_deployments = self._pre_call_checks( + model=model, + healthy_deployments=cast(List[Dict], healthy_deployments), + messages=messages, + request_kwargs=request_kwargs, + ) + # check if user wants to do tag based routing + healthy_deployments = await get_deployments_for_tag( # type: ignore + llm_router_instance=self, + model=model, + request_kwargs=request_kwargs, + healthy_deployments=healthy_deployments, + ) + + if len(healthy_deployments) == 0: + exception = await async_raise_no_deployment_exception( + litellm_router_instance=self, + model=model, + parent_otel_span=parent_otel_span, + ) + raise exception + + return healthy_deployments + async def async_get_available_deployment( self, model: str, @@ -5965,61 +6091,17 @@ class Router: ) try: parent_otel_span = _get_parent_otel_span_from_kwargs(request_kwargs) - model, healthy_deployments = self._common_checks_available_deployment( + healthy_deployments = await self.async_get_healthy_deployments( model=model, + request_kwargs=request_kwargs, messages=messages, input=input, specific_deployment=specific_deployment, - ) # type: ignore + parent_otel_span=parent_otel_span, + ) if isinstance(healthy_deployments, dict): return healthy_deployments - cooldown_deployments = await _async_get_cooldown_deployments( - litellm_router_instance=self, parent_otel_span=parent_otel_span - ) - verbose_router_logger.debug( - f"async cooldown deployments: {cooldown_deployments}" - ) - verbose_router_logger.debug(f"cooldown_deployments: {cooldown_deployments}") - healthy_deployments = self._filter_cooldown_deployments( - healthy_deployments=healthy_deployments, - cooldown_deployments=cooldown_deployments, - ) - - healthy_deployments = await self.async_callback_filter_deployments( - model=model, - healthy_deployments=healthy_deployments, - messages=( - cast(List[AllMessageValues], messages) - if messages is not None - else None - ), - request_kwargs=request_kwargs, - parent_otel_span=parent_otel_span, - ) - - if self.enable_pre_call_checks and messages is not None: - healthy_deployments = self._pre_call_checks( - model=model, - healthy_deployments=cast(List[Dict], healthy_deployments), - messages=messages, - request_kwargs=request_kwargs, - ) - # check if user wants to do tag based routing - healthy_deployments = await get_deployments_for_tag( # type: ignore - llm_router_instance=self, - model=model, - request_kwargs=request_kwargs, - healthy_deployments=healthy_deployments, - ) - - if len(healthy_deployments) == 0: - exception = await async_raise_no_deployment_exception( - litellm_router_instance=self, - model=model, - parent_otel_span=parent_otel_span, - ) - raise exception start_time = time.time() if ( self.routing_strategy == "usage-based-routing-v2" diff --git a/litellm/router_strategy/lowest_tpm_rpm.py b/litellm/router_strategy/lowest_tpm_rpm.py index 8658793973..121df00d30 100644 --- a/litellm/router_strategy/lowest_tpm_rpm.py +++ b/litellm/router_strategy/lowest_tpm_rpm.py @@ -106,6 +106,8 @@ class LowestTPMLoggingHandler(CustomLogger): elif isinstance(id, int): id = str(id) + if "usage" not in response_obj: + return total_tokens = response_obj["usage"]["total_tokens"] # ------------ @@ -144,7 +146,7 @@ class LowestTPMLoggingHandler(CustomLogger): if self.test_flag: self.logged_success += 1 except Exception as e: - verbose_router_logger.error( + verbose_router_logger.exception( "litellm.router_strategy.lowest_tpm_rpm.py::async_log_success_event(): Exception occured - {}".format( str(e) ) diff --git a/litellm/router_utils/batch_utils.py b/litellm/router_utils/batch_utils.py index 50b24ec363..a55be22913 100644 --- a/litellm/router_utils/batch_utils.py +++ b/litellm/router_utils/batch_utils.py @@ -48,7 +48,7 @@ def replace_model_in_jsonl( return None -def _get_router_metadata_variable_name(function_name) -> str: +def _get_router_metadata_variable_name(function_name: Optional[str]) -> str: """ Helper to return what the "metadata" field should be called in the request data @@ -57,9 +57,11 @@ def _get_router_metadata_variable_name(function_name) -> str: For ALL other endpoints we call this "metadata """ ROUTER_METHODS_USING_LITELLM_METADATA = set( - ["batch", "generic_api_call", "_acreate_batch"] + ["batch", "generic_api_call", "_acreate_batch", "file"] ) - if function_name in ROUTER_METHODS_USING_LITELLM_METADATA: + if function_name and any( + method in function_name for method in ROUTER_METHODS_USING_LITELLM_METADATA + ): return "litellm_metadata" else: return "metadata" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index a3eb6f7fda..b6ac371850 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2316,9 +2316,11 @@ class SpecialEnums(Enum): "litellm:custom_llm_provider:{};model_id:{};response_id:{}" ) + LITELLM_MANAGED_BATCH_COMPLETE_STR = "litellm_proxy;model_id:{};llm_batch_id:{}" + LLMResponseTypes = Union[ - ModelResponse, EmbeddingResponse, ImageResponse, OpenAIFileObject + ModelResponse, EmbeddingResponse, ImageResponse, OpenAIFileObject, LiteLLMBatch ] diff --git a/tests/enterprise/enterprise_hooks/test_managed_files.py b/tests/enterprise/enterprise_hooks/test_managed_files.py index 89b332a506..e84034ed09 100644 --- a/tests/enterprise/enterprise_hooks/test_managed_files.py +++ b/tests/enterprise/enterprise_hooks/test_managed_files.py @@ -40,6 +40,24 @@ def test_get_file_ids_from_messages(): ] +@pytest.mark.asyncio +async def test_async_pre_call_hook_batch_retrieve(): + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=MagicMock() + ) + data = { + "user_api_key_dict": {"parent_otel_span": MagicMock()}, + "data": { + "batch_id": "bGl0ZWxsbV9wcm94eTttb2RlbF9pZDpteS1nZW5lcmFsLWF6dXJlLWRlcGxveW1lbnQ7bGxtX2JhdGNoX2lkOmJhdGNoX2EzMjJiNmJhLWFjN2UtNDg4OC05MjljLTFhZDM0NDJmMDZlZA", + }, + "call_type": "aretrieve_batch", + "cache": MagicMock(), + } + response = await proxy_managed_files.async_pre_call_hook(**data) + assert response["batch_id"] == "batch_a322b6ba-ac7e-4888-929c-1ad3442f06ed" + assert response["model"] == "my-general-azure-deployment" + + # def test_list_managed_files(): # proxy_managed_files = _PROXY_LiteLLMManagedFiles(DualCache()) diff --git a/tests/litellm/test_router.py b/tests/litellm/test_router.py index 95bdfccebb..0de5bbbb42 100644 --- a/tests/litellm/test_router.py +++ b/tests/litellm/test_router.py @@ -118,3 +118,60 @@ async def test_router_with_tags_and_fallbacks(): mock_testing_fallbacks=True, metadata={"tags": ["test"]}, ) + + +@pytest.mark.asyncio +async def test_router_acreate_file(): + """ + Write to all deployments of a model + """ + from unittest.mock import MagicMock, call, patch + + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + }, + {"model_name": "gpt-3.5-turbo", "litellm_params": {"model": "gpt-4o-mini"}}, + ], + ) + + with patch("litellm.acreate_file", return_value=MagicMock()) as mock_acreate_file: + mock_acreate_file.return_value = MagicMock() + response = await router.acreate_file( + model="gpt-3.5-turbo", + purpose="test", + file=MagicMock(), + ) + + # assert that the mock_acreate_file was called twice + assert mock_acreate_file.call_count == 2 + + +@pytest.mark.asyncio +async def test_router_async_get_healthy_deployments(): + """ + Test that afile_content returns the correct file content + """ + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": {"model": "gpt-3.5-turbo"}, + }, + ], + ) + + result = await router.async_get_healthy_deployments( + model="gpt-3.5-turbo", + request_kwargs={}, + messages=None, + input=None, + specific_deployment=False, + parent_otel_span=None, + ) + + assert len(result) == 1 + assert result[0]["model_name"] == "gpt-3.5-turbo" + assert result[0]["litellm_params"]["model"] == "gpt-3.5-turbo" diff --git a/tests/router_unit_tests/test_router_batch_utils.py b/tests/router_unit_tests/test_router_batch_utils.py index 3d1bc92101..40a491a84a 100644 --- a/tests/router_unit_tests/test_router_batch_utils.py +++ b/tests/router_unit_tests/test_router_batch_utils.py @@ -82,3 +82,9 @@ def test_router_metadata_variable_name(): assert ( _get_router_metadata_variable_name(function_name="batch") == "litellm_metadata" ) + assert ( + _get_router_metadata_variable_name(function_name="acreate_file") == "litellm_metadata" + ) + assert ( + _get_router_metadata_variable_name(function_name="aget_file") == "litellm_metadata" + ) diff --git a/ui/litellm-dashboard/src/components/view_logs/index.tsx b/ui/litellm-dashboard/src/components/view_logs/index.tsx index 4c913bffc4..a8c47945a2 100644 --- a/ui/litellm-dashboard/src/components/view_logs/index.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/index.tsx @@ -769,6 +769,10 @@ export function RequestViewer({ row }: { row: Row }) { Model ID: {row.original.model_id} +
+ Call Type: + {row.original.call_type} +
Provider: {row.original.custom_llm_provider || "-"} @@ -779,10 +783,12 @@ export function RequestViewer({ row }: { row: Row }) { {row.original.api_base || "-"}
-
- Start Time: - {row.original.startTime} -
+ {row?.original?.requester_ip_address && ( +
+ IP Address: + {row?.original?.requester_ip_address} +
+ )}
@@ -798,12 +804,7 @@ export function RequestViewer({ row }: { row: Row }) { Cache Hit: {row.original.cache_hit}
- {row?.original?.requester_ip_address && ( -
- IP Address: - {row?.original?.requester_ip_address} -
- )} +
Status: }) {
+
+ Start Time: + {row.original.startTime} +
End Time: {row.original.endTime}