fix(managed_files.py): support for DELETE endpoint for files

This commit is contained in:
Krrish Dholakia 2025-04-11 15:26:51 -07:00
parent cbcf028da5
commit 522ffd6e7c
3 changed files with 90 additions and 8 deletions

View file

@ -8,7 +8,7 @@ from abc import ABC, abstractmethod
from datetime import datetime
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast
from litellm import verbose_logger
from litellm import Router, verbose_logger
from litellm.caching.caching import DualCache
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
@ -59,7 +59,7 @@ class BaseFileEndpoints(ABC):
pass
class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
class _PROXY_LiteLLMManagedFiles(CustomLogger):
# Class variables or attributes
def __init__(
self, internal_usage_cache: InternalUsageCache, prisma_client: PrismaClient
@ -92,6 +92,25 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
litellm_parent_otel_span=litellm_parent_otel_span,
)
async def delete_unified_file_id(
self, file_id: str, litellm_parent_otel_span: Optional[Span] = None
) -> OpenAIFileObject:
key = f"litellm_proxy/{file_id}"
## get old value
old_value = await self.internal_usage_cache.async_get_cache(
key=key,
litellm_parent_otel_span=litellm_parent_otel_span,
)
if old_value is None or not isinstance(old_value, OpenAIFileObject):
raise Exception(f"LiteLLM Managed File object with id={file_id} not found")
## delete old value
await self.internal_usage_cache.async_set_cache(
key=key,
value=None,
litellm_parent_otel_span=litellm_parent_otel_span,
)
return old_value
async def async_pre_call_hook(
self,
user_api_key_dict: UserAPIKeyAuth,
@ -322,6 +341,24 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
return []
async def afile_delete(
self, custom_llm_provider: str, file_id: str, **data: Dict
self,
file_id: str,
litellm_parent_otel_span: Optional[Span],
llm_router: Router,
**data: Dict,
) -> OpenAIFileObject:
raise NotImplementedError("afile_delete not implemented")
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
stored_file_object = await self.delete_unified_file_id(
file_id, litellm_parent_otel_span
)
if stored_file_object:
return stored_file_object
else:
raise Exception(f"LiteLLM Managed File object with id={file_id} not found")

View file

@ -565,7 +565,8 @@ async def get_file(
code=500,
)
response = await managed_files_obj.afile_retrieve(
file_id=file_id, litellm_parent_otel_span=user_api_key_dict.parent_otel_span, **data # type: ignore
file_id=file_id,
litellm_parent_otel_span=user_api_key_dict.parent_otel_span,
)
else:
response = await litellm.afile_retrieve(
@ -663,6 +664,7 @@ async def delete_file(
from litellm.proxy.proxy_server import (
add_litellm_data_to_request,
general_settings,
llm_router,
proxy_config,
proxy_logging_obj,
version,
@ -685,10 +687,41 @@ async def delete_file(
proxy_config=proxy_config,
)
response = await litellm.afile_delete(
custom_llm_provider=custom_llm_provider, file_id=file_id, **data # type: ignore
## check if file_id is a litellm managed file
is_base64_unified_file_id = (
_PROXY_LiteLLMManagedFiles._is_base64_encoded_unified_file_id(file_id)
)
if is_base64_unified_file_id:
managed_files_obj = cast(
Optional[_PROXY_LiteLLMManagedFiles],
proxy_logging_obj.get_proxy_hook("managed_files"),
)
if managed_files_obj is None:
raise ProxyException(
message="Managed files hook not found",
type="None",
param="None",
code=500,
)
if llm_router is None:
raise ProxyException(
message="LLM Router not found",
type="None",
param="None",
code=500,
)
response = await managed_files_obj.afile_delete(
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_delete(
custom_llm_provider=custom_llm_provider, file_id=file_id, **data # type: ignore
)
### ALERTING ###
asyncio.create_task(
proxy_logging_obj.update_request_status(

View file

@ -729,6 +729,9 @@ class Router:
self.aresponses = self.factory_function(
litellm.aresponses, call_type="aresponses"
)
self.afile_delete = self.factory_function(
litellm.afile_delete, call_type="afile_delete"
)
self.responses = self.factory_function(litellm.responses, call_type="responses")
def validate_fallbacks(self, fallback_param: Optional[List]):
@ -2435,6 +2438,8 @@ class Router:
model_name = data["model"]
self.total_calls[model_name] += 1
### get custom
response = original_function(
**{
**data,
@ -2514,9 +2519,15 @@ class Router:
# Perform pre-call checks for routing strategy
self.routing_strategy_pre_call_checks(deployment=deployment)
try:
_, custom_llm_provider, _, _ = get_llm_provider(model=data["model"])
except Exception:
custom_llm_provider = None
response = original_function(
**{
**data,
"custom_llm_provider": custom_llm_provider,
"caching": self.cache_responses,
**kwargs,
}
@ -3058,6 +3069,7 @@ class Router:
"anthropic_messages",
"aresponses",
"responses",
"afile_delete",
] = "assistants",
):
"""
@ -3102,7 +3114,7 @@ class Router:
return await self._pass_through_moderation_endpoint_factory(
original_function=original_function, **kwargs
)
elif call_type in ("anthropic_messages", "aresponses"):
elif call_type in ("anthropic_messages", "aresponses", "afile_delete"):
return await self._ageneric_api_call_with_fallbacks(
original_function=original_function,
**kwargs,