Litellm Unified File ID output file id support (#10713)

* fix(router.py): write file to all deployments

allows unified file id to work across multiple deployments

* fix(view_logs/index.tsx): show call type in request logs

* fix(router.py): pass a deep copy of kwargs to avoid conflict across multiple runs

* fix(batch_utils.py): broaden check

* fix(router_utils.py): handle null type for function name

* fix(proxy_track_cost_callback.py): fix ruff check error

* fix(router.py): handle healthy_deployments as a dict

* feat(managed_files.py): support encoding / decoding unified batch id … (#10711)

* feat(managed_files.py): support encoding / decoding unified batch id when using managed files

allows routing retrieve batch to the right model id

* fix: fix linting error

* feat(managed_files.py): support unified output file id

enables batch output file id to be used to retrieve the actual file

* fix(managed_files.py): attempt to fix ci/cd linting error

* fix: fix ruff check

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
         )

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

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

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,

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)
             )

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

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:{}"

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"
+