feat(managed_files.py): support /delete for files + feat(managed_batches): support /cancel for batches

enterprise/litellm_enterprise/proxy/hooks/managed_files.py

@@ -498,7 +498,6 @@ async def get_model_file_id_mapping(
         for file_id in file_ids:
             ## CHECK IF FILE ID IS MANAGED BY LITELM
             is_base64_unified_file_id = _is_base64_encoded_unified_file_id(file_id)
-
             if is_base64_unified_file_id:
                 litellm_managed_file_ids.append(file_id)
 
@@ -509,6 +508,7 @@ async def get_model_file_id_mapping(
                 unified_file_object = await self.get_unified_file_id(
                     file_id, litellm_parent_otel_span
                 )
+
                 if unified_file_object:
                     file_id_mapping[file_id] = unified_file_object.model_mappings
 
@@ -784,18 +784,21 @@ async def afile_delete(
         llm_router: Router,
         **data: Dict,
     ) -> OpenAIFileObject:
-        file_id = convert_b64_uid_to_unified_uid(file_id)
+
+        # file_id = convert_b64_uid_to_unified_uid(file_id)
         model_file_id_mapping = await self.get_model_file_id_mapping(
             [file_id], litellm_parent_otel_span
         )
+
         specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
         if specific_model_file_id_mapping:
-            for model_id, file_id in specific_model_file_id_mapping.items():
-                await llm_router.afile_delete(model=model_id, file_id=file_id, **data)  # type: ignore
+            for model_id, model_file_id in specific_model_file_id_mapping.items():
+                await llm_router.afile_delete(model=model_id, file_id=model_file_id, **data)  # type: ignore
 
         stored_file_object = await self.delete_unified_file_id(
             file_id, litellm_parent_otel_span
         )
+
         if stored_file_object:
             return stored_file_object
         else:
@@ -816,6 +819,7 @@ async def afile_content(
             model_file_id_mapping
             or await self.get_model_file_id_mapping([file_id], litellm_parent_otel_span)
         )
+
         specific_model_file_id_mapping = model_file_id_mapping.get(file_id)
 
         if specific_model_file_id_mapping:

litellm/batches/main.py

@@ -17,6 +17,8 @@
 from typing import Any, Coroutine, Dict, Literal, Optional, Union, cast
 
 import httpx
+from openai.types.batch import BatchRequestCounts
+from openai.types.batch import Metadata as BatchMetadata
 
 import litellm
 from litellm._logging import verbose_logger
@@ -223,10 +225,12 @@ def create_batch(
                 api_key=optional_params.api_key,
                 logging_obj=litellm_logging_obj,
                 _is_async=_is_async,
-                client=client
-                if client is not None
-                and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
-                else None,
+                client=(
+                    client
+                    if client is not None
+                    and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
+                    else None
+                ),
                 timeout=timeout,
                 model=model,
             )
@@ -609,10 +613,12 @@ def retrieve_batch(
                     function_id="batch_retrieve",
                 ),
                 _is_async=_is_async,
-                client=client
-                if client is not None
-                and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
-                else None,
+                client=(
+                    client
+                    if client is not None
+                    and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
+                    else None
+                ),
                 timeout=timeout,
                 model=model,
             )
@@ -799,6 +805,7 @@ def list_batches(
 
 async def acancel_batch(
     batch_id: str,
+    model: Optional[str] = None,
     custom_llm_provider: Literal["openai", "azure"] = "openai",
     metadata: Optional[Dict[str, str]] = None,
     extra_headers: Optional[Dict[str, str]] = None,
@@ -813,11 +820,13 @@ async def acancel_batch(
     try:
         loop = asyncio.get_event_loop()
         kwargs["acancel_batch"] = True
+        model = kwargs.pop("model", None)
 
         # Use a partial function to pass your keyword arguments
         func = partial(
             cancel_batch,
             batch_id,
+            model,
             custom_llm_provider,
             metadata,
             extra_headers,
@@ -840,7 +849,8 @@ async def acancel_batch(
 
 def cancel_batch(
     batch_id: str,
-    custom_llm_provider: Literal["openai", "azure"] = "openai",
+    model: Optional[str] = None,
+    custom_llm_provider: Union[Literal["openai", "azure"], str] = "openai",
     metadata: Optional[Dict[str, str]] = None,
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
@@ -852,6 +862,17 @@ def cancel_batch(
     LiteLLM Equivalent of POST https://api.openai.com/v1/batches/{batch_id}/cancel
     """
     try:
+
+        try:
+            if model is not None:
+                _, custom_llm_provider, _, _ = get_llm_provider(
+                    model=model,
+                    custom_llm_provider=custom_llm_provider,
+                )
+        except Exception as e:
+            verbose_logger.exception(
+                f"litellm.batches.main.py::cancel_batch() - Error inferring custom_llm_provider - {str(e)}"
+            )
         optional_params = GenericLiteLLMParams(**kwargs)
         litellm_params = get_litellm_params(
             custom_llm_provider=custom_llm_provider,
@@ -1005,21 +1026,28 @@ async def _async_get_status():
             created_at=status_response["submitTime"],
             in_progress_at=status_response["lastModifiedTime"],
             completed_at=status_response.get("endTime"),
-            failed_at=status_response.get("endTime")
-            if status_response["status"] == "failed"
-            else None,
-            request_counts={
-                "total": 1,
-                "completed": 1 if status_response["status"] == "completed" else 0,
-                "failed": 1 if status_response["status"] == "failed" else 0,
-            },
-            metadata={
-                "output_file_id": status_response["outputDataConfig"][
-                    "s3OutputDataConfig"
-                ]["s3Uri"],
-                "failure_message": status_response.get("failureMessage"),
-                "model_arn": status_response["modelArn"],
-            },
+            failed_at=(
+                status_response.get("endTime")
+                if status_response["status"] == "failed"
+                else None
+            ),
+            request_counts=BatchRequestCounts(
+                total=1,
+                completed=1 if status_response["status"] == "completed" else 0,
+                failed=1 if status_response["status"] == "failed" else 0,
+            ),
+            metadata=dict(
+                **{
+                    "output_file_id": status_response["outputDataConfig"][
+                        "s3OutputDataConfig"
+                    ]["s3Uri"],
+                    "failure_message": status_response.get("failureMessage") or "",
+                    "model_arn": status_response["modelArn"],
+                }
+            ),
+            completion_window="24h",
+            endpoint="/v1/embeddings",
+            input_file_id="",
         )
 
         return result

litellm/files/main.py

@@ -95,7 +95,9 @@ async def acreate_file(
 def create_file(
     file: FileTypes,
     purpose: Literal["assistants", "batch", "fine-tune"],
-    custom_llm_provider: Optional[Literal["openai", "azure", "vertex_ai", "bedrock"]] = None,
+    custom_llm_provider: Optional[
+        Literal["openai", "azure", "vertex_ai", "bedrock"]
+    ] = None,
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
     **kwargs,
@@ -155,10 +157,12 @@ def create_file(
                 api_key=optional_params.api_key,
                 logging_obj=logging_obj,
                 _is_async=_is_async,
-                client=client
-                if client is not None
-                and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
-                else None,
+                client=(
+                    client
+                    if client is not None
+                    and isinstance(client, (HTTPHandler, AsyncHTTPHandler))
+                    else None
+                ),
                 timeout=timeout,
             )
         elif custom_llm_provider == "openai":
@@ -441,12 +445,14 @@ async def afile_delete(
     """
     try:
         loop = asyncio.get_event_loop()
+        model = kwargs.pop("model", None)
         kwargs["is_async"] = True
 
         # Use a partial function to pass your keyword arguments
         func = partial(
             file_delete,
             file_id,
+            model,
             custom_llm_provider,
             extra_headers,
             extra_body,
@@ -470,7 +476,8 @@ async def afile_delete(
 @client
 def file_delete(
     file_id: str,
-    custom_llm_provider: Literal["openai", "azure"] = "openai",
+    model: Optional[str] = None,
+    custom_llm_provider: Union[Literal["openai", "azure"], str] = "openai",
     extra_headers: Optional[Dict[str, str]] = None,
     extra_body: Optional[Dict[str, str]] = None,
     **kwargs,
@@ -481,6 +488,13 @@ def file_delete(
     LiteLLM Equivalent of DELETE https://api.openai.com/v1/files
     """
     try:
+        try:
+            if model is not None:
+                _, custom_llm_provider, _, _ = get_llm_provider(
+                    model, custom_llm_provider
+                )
+        except Exception:
+            pass
         optional_params = GenericLiteLLMParams(**kwargs)
         litellm_params_dict = get_litellm_params(**kwargs)
         ### TIMEOUT LOGIC ###
@@ -566,7 +580,7 @@ def file_delete(
             )
         else:
             raise litellm.exceptions.BadRequestError(
-                message="LiteLLM doesn't support {} for 'create_batch'. Only 'openai' is supported.".format(
+                message="LiteLLM doesn't support {} for 'delete_batch'. Only 'openai' is supported.".format(
                     custom_llm_provider
                 ),
                 model="n/a",