diff --git a/enterprise/enterprise_hooks/managed_files.py b/enterprise/enterprise_hooks/managed_files.py index 3819a58756..ebd066e842 100644 --- a/enterprise/enterprise_hooks/managed_files.py +++ b/enterprise/enterprise_hooks/managed_files.py @@ -137,6 +137,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "rerank", "acreate_batch", "aretrieve_batch", + "afile_content", ], ) -> Union[Exception, str, Dict, None]: """ @@ -154,6 +155,20 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) data["model_file_id_mapping"] = model_file_id_mapping + elif call_type == CallTypes.afile_content.value: + retrieve_file_id = cast(Optional[str], data.get("file_id")) + potential_file_id = ( + _is_base64_encoded_unified_file_id(retrieve_file_id) + if retrieve_file_id + else False + ) + if potential_file_id: + model_id = self.get_model_id_from_unified_file_id(potential_file_id) + if model_id: + data["model"] = model_id + data["file_id"] = self.get_output_file_id_from_unified_file_id( + potential_file_id + ) elif call_type == CallTypes.acreate_batch.value: input_file_id = cast(Optional[str], data.get("input_file_id")) if input_file_id: @@ -171,8 +186,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) 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 + potential_model_id = self.get_model_id_from_unified_batch_id( + potential_batch_id + ) + if potential_model_id is None: + raise Exception( + f"LiteLLM Managed Batch ID with id={retrieve_batch_id} is invalid - does not contain encoded model_id." + ) + data["model"] = potential_model_id data["batch_id"] = self.get_batch_id_from_unified_batch_id( potential_batch_id ) @@ -333,8 +354,15 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): file_type = file_data["content_type"] + output_file_id = file_objects[0].id + model_id = file_objects[0]._hidden_params.get("model_id") + unified_file_id = SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( - file_type, str(uuid.uuid4()), ",".join(target_model_names_list) + file_type, + str(uuid.uuid4()), + ",".join(target_model_names_list), + output_file_id, + model_id, ) # Convert to URL-safe base64 and strip padding @@ -362,9 +390,41 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) return base64.urlsafe_b64encode(unified_batch_id.encode()).decode().rstrip("=") - def get_model_id_from_unified_batch_id(self, file_id: str) -> str: + def get_unified_output_file_id( + self, output_file_id: str, model_id: str, model_name: str + ) -> str: + unified_output_file_id = ( + SpecialEnums.LITELLM_MANAGED_FILE_COMPLETE_STR.value.format( + "application/json", + str(uuid.uuid4()), + model_name, + output_file_id, + model_id, + ) + ) + return ( + base64.urlsafe_b64encode(unified_output_file_id.encode()) + .decode() + .rstrip("=") + ) + + def get_model_id_from_unified_file_id(self, file_id: str) -> str: + return file_id.split("llm_output_file_model_id,")[1].split(";")[0] + + def get_output_file_id_from_unified_file_id(self, file_id: str) -> str: + return file_id.split("llm_output_file_id,")[1].split(";")[0] + + def get_model_id_from_unified_batch_id(self, file_id: str) -> Optional[str]: + """ + Get the model_id from the file_id + + Expected format: litellm_proxy;model_id:{};llm_batch_id:{};llm_output_file_id:{} + """ ## use regex to get the model_id from the file_id - return file_id.split("model_id:")[1].split(";")[0] + try: + return file_id.split("model_id:")[1].split(";")[0] + except Exception: + return None def get_batch_id_from_unified_batch_id(self, file_id: str) -> str: ## use regex to get the batch_id from the file_id @@ -375,15 +435,28 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) -> Any: if isinstance(response, LiteLLMBatch): ## Check if unified_file_id is in the response - unified_batch_id = response._hidden_params.get( + unified_file_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: + unified_batch_id = response._hidden_params.get( + "unified_batch_id" + ) # managed batch id + model_id = cast(Optional[str], response._hidden_params.get("model_id")) + model_name = cast(Optional[str], response._hidden_params.get("model_name")) + if (unified_batch_id or unified_file_id) and model_id: response.id = self.get_unified_batch_id( batch_id=response.id, model_id=model_id ) + if ( + response.output_file_id and model_name and model_id + ): # return a file id with the model_id and output_file_id + response.output_file_id = self.get_unified_output_file_id( + output_file_id=response.output_file_id, + model_id=model_id, + model_name=model_name, + ) + return await super().async_post_call_success_hook( data, user_api_key_dict, response ) diff --git a/litellm/files/main.py b/litellm/files/main.py index ded74cc653..5d0dc05771 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -15,6 +15,7 @@ import httpx import litellm from litellm import get_secret_str +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.files.handler import AzureOpenAIFilesAPI from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler @@ -743,11 +744,13 @@ async def afile_content( try: loop = asyncio.get_event_loop() kwargs["afile_content"] = True + model = kwargs.pop("model", None) # Use a partial function to pass your keyword arguments func = partial( file_content, file_id, + model, custom_llm_provider, extra_headers, extra_body, @@ -770,7 +773,10 @@ async def afile_content( def file_content( file_id: str, - custom_llm_provider: Literal["openai", "azure"] = "openai", + model: Optional[str] = None, + custom_llm_provider: Optional[ + Union[Literal["openai", "azure", "vertex_ai"], str] + ] = None, extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -788,10 +794,18 @@ def file_content( client = kwargs.get("client") # set timeout for 10 minutes by default + try: + if model is not None: + _, custom_llm_provider, _, _ = get_llm_provider( + model, custom_llm_provider + ) + except Exception: + pass + if ( timeout is not None and isinstance(timeout, httpx.Timeout) - and supports_httpx_timeout(custom_llm_provider) is False + and supports_httpx_timeout(cast(str, custom_llm_provider)) is False ): read_timeout = timeout.read or 600 timeout = read_timeout # default 10 min timeout diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index bf5173f893..02f6e5d275 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -253,6 +253,7 @@ async def retrieve_batch( ) data = cast(dict, _retrieve_batch_request) + unified_batch_id = _is_base64_encoded_unified_file_id(batch_id) base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) ( @@ -268,10 +269,7 @@ async def retrieve_batch( route_type="aretrieve_batch", ) - if ( - litellm.enable_loadbalancing_on_batch_endpoints is True - or data.get("model") is not None - ): + if litellm.enable_loadbalancing_on_batch_endpoints is True or unified_batch_id: if llm_router is None: raise HTTPException( status_code=500, @@ -281,6 +279,7 @@ async def retrieve_batch( ) response = await llm_router.aretrieve_batch(**data) # type: ignore + response._hidden_params["unified_batch_id"] = unified_batch_id else: custom_llm_provider = ( provider diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 8fd1b4ddc0..2fd56af0b9 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -116,6 +116,7 @@ class ProxyBaseLLMRequestProcessing: "adelete_responses", "acreate_batch", "aretrieve_batch", + "afile_content", ], version: Optional[str] = None, user_model: Optional[str] = None, diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 568dff7cc1..3c2c3d80dc 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -413,7 +413,6 @@ async def get_file_content( ``` """ from litellm.proxy.proxy_server import ( - add_litellm_data_to_request, general_settings, llm_router, proxy_config, @@ -421,16 +420,21 @@ async def get_file_content( version, ) - data: Dict = {} + data: Dict = {"file_id": file_id} try: # 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="afile_content", ) custom_llm_provider = ( @@ -464,15 +468,32 @@ async def get_file_content( param="None", code=500, ) - response = await managed_files_obj.afile_content( - file_id=file_id, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, - llm_router=llm_router, - **data, - ) + model = cast(Optional[str], data.get("model")) + if model: + response = await llm_router.afile_content( + **{ + "model": model, + "file_id": file_id, + **data, + } + ) # type: ignore + + else: + response = await managed_files_obj.afile_content( + **{ + "file_id": file_id, + "litellm_parent_otel_span": user_api_key_dict.parent_otel_span, + "llm_router": llm_router, + **data, + } + ) else: response = await litellm.afile_content( - custom_llm_provider=custom_llm_provider, file_id=file_id, **data # type: ignore + **{ + "custom_llm_provider": custom_llm_provider, + "file_id": file_id, + **data, + } # type: ignore ) ### ALERTING ### @@ -514,7 +535,7 @@ async def get_file_content( await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, request_data=data ) - verbose_proxy_logger.error( + verbose_proxy_logger.exception( "litellm.proxy.proxy_server.retrieve_file_content(): Exception occured - {}".format( str(e) ) diff --git a/litellm/router.py b/litellm/router.py index 19943ba504..e0068cacbc 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -2980,6 +2980,11 @@ class Router: model=model ) new_kwargs = copy.deepcopy(kwargs) + self._update_kwargs_with_deployment( + deployment=cast(dict, model_name), + kwargs=new_kwargs, + function_name="aretrieve_batch", + ) new_kwargs.pop("custom_llm_provider", None) return await litellm.aretrieve_batch( custom_llm_provider=custom_llm_provider, **new_kwargs # type: ignore diff --git a/litellm/types/utils.py b/litellm/types/utils.py index b6ac371850..d4e4ab4041 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2308,9 +2308,7 @@ class ExtractedFileData(TypedDict): class SpecialEnums(Enum): LITELM_MANAGED_FILE_ID_PREFIX = "litellm_proxy" - LITELLM_MANAGED_FILE_COMPLETE_STR = ( - "litellm_proxy:{};unified_id,{};target_model_names,{}" - ) + LITELLM_MANAGED_FILE_COMPLETE_STR = "litellm_proxy:{};unified_id,{};target_model_names,{};llm_output_file_id,{};llm_output_file_model_id,{}" LITELLM_MANAGED_RESPONSE_COMPLETE_STR = ( "litellm:custom_llm_provider:{};model_id:{};response_id:{}" diff --git a/tests/litellm/test_router.py b/tests/litellm/test_router.py index 0de5bbbb42..e645e7f217 100644 --- a/tests/litellm/test_router.py +++ b/tests/litellm/test_router.py @@ -148,7 +148,6 @@ async def test_router_acreate_file(): # 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(): """ @@ -175,3 +174,4 @@ async def test_router_async_get_healthy_deployments(): assert len(result) == 1 assert result[0]["model_name"] == "gpt-3.5-turbo" assert result[0]["litellm_params"]["model"] == "gpt-3.5-turbo" +