diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 52a7d02fc2b..d7f075bc9ae 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -2266,15 +2266,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): """ ## check if file object if hasattr(response, "data") and isinstance(response.data, list): - if all(isinstance(file_object, FileObject) for file_object in response.data): - ## Get all file id's - ## Check which file id's were created by the user - ## Filter the response to only include the files created by the user - ## Return the filtered response - file_ids = [ - file_object.id - for file_object in cast(list[FileObject], response.data) # type: ignore - ] + page_items: Final = TypeAdapter(tuple[object, ...]).validate_python(response.data) + file_objects: Final = tuple(item for item in page_items if isinstance(item, FileObject)) + if len(file_objects) == len(page_items): + file_ids = [file_object.id for file_object in file_objects] user_created_file_ids = await self.get_user_created_file_ids(user_api_key_dict, file_ids) ## Filter the response to only include the files created by the user response.data = user_created_file_ids # type: ignore diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index 48a3791e93c..6a9e6a984a7 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -799,7 +799,7 @@ class DeploymentAffinityCheck(CustomLogger): verbose_router_logger.debug( "DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s", deployment_model_name, - model_id, + pin_value["model_id"], session_affinity_ttl, session_id, ) @@ -808,7 +808,7 @@ class DeploymentAffinityCheck(CustomLogger): "DeploymentAffinityCheck: session pin already claimed model_map_key=%s existing=%s ours=%s session_id=%s", deployment_model_name, claimed_session_pin, - model_id, + pin_value["model_id"], session_id, ) except Exception as e: diff --git a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py index c721e808ed5..ea881d5cac2 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_success_handler.py +++ b/tests/unit/proxy/pass_through_endpoints/test_success_handler.py @@ -659,6 +659,8 @@ async def test_completed_passthrough_cancellation_preserves_accounting( from unittest.mock import create_autospec import anyio + import fakeredis + from redis.asyncio.client import Pipeline from litellm.caching.dual_cache import DualCache from litellm.proxy import proxy_server @@ -667,20 +669,29 @@ async def test_completed_passthrough_cancellation_preserves_accounting( from litellm.proxy.spend_tracking.budget_reservation import release_unbound_budget_reservation from litellm.proxy.utils import ProxyLogging from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage + from tests._support.fake_redis import fake_redis_cache - class AllowanceCounter: - def __init__(self) -> None: - self.started = asyncio.Event() - self.release = asyncio.Event() - self.calls = 0 - self.units = 0.0 + allowance_started: Final = asyncio.Event() + allowance_release: Final = asyncio.Event() + executions: Final[list[int]] = [] # mutable-ok: counts pipeline flushes observed by the gated fake client - async def async_increment(self, key: str, value: float, ttl: int | None = None) -> float: - self.started.set() - await self.release.wait() - self.calls += 1 - self.units += value - return self.units + class GatedClient(fakeredis.FakeAsyncRedis): + def pipeline(self, transaction: bool = True, shard_hint: str | None = None) -> "Pipeline[bytes]": + pipe: Final = super().pipeline(transaction=transaction, shard_hint=shard_hint) # pyright: ignore[reportUnknownVariableType] # fakeredis pipeline is untyped + native_execute: Final = pipe.execute # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # redis-py pipeline.execute reply is untyped + + async def gated_execute(raise_on_error: bool = True) -> list[object]: + allowance_started.set() + await allowance_release.wait() + executions.append(1) + return await native_execute(raise_on_error=raise_on_error) # pyright: ignore[reportUnknownVariableType] # redis-py erases the reply type + + monkeypatch.setattr(pipe, "execute", gated_execute) # pyright: ignore[reportUnknownArgumentType] # fakeredis pipeline is untyped + return pipe # pyright: ignore[reportUnknownVariableType] # fakeredis pipeline is untyped + + server: Final = fakeredis.FakeServer() + allowance: Final = fake_redis_cache(server, client=GatedClient) # pyright: ignore[reportUnknownArgumentType] # fakeredis client class is partially untyped + verifier: Final = fakeredis.FakeAsyncRedis(server=server) now: Final = datetime(2026, 10, 20, tzinfo=timezone.utc) model: Final = "passthrough-cancellation-test" @@ -714,7 +725,6 @@ async def test_completed_passthrough_cancellation_preserves_accounting( writer: Final = create_autospec(DBSpendUpdateWriter, instance=True) proxy_logging: Final = ProxyLogging(user_api_key_cache=DualCache()) proxy_logging.db_spend_update_writer = writer - allowance: Final = AllowanceCounter() monkeypatch.setattr(proxy_logging.internal_usage_cache.dual_cache, "redis_cache", allowance) monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging) logging_obj: Final = LiteLLMLoggingObj( @@ -771,7 +781,7 @@ async def test_completed_passthrough_cancellation_preserves_accounting( pending: Final = asyncio.create_task(dispatch()) await entered.wait() if cancel_phase == "allocation": - await allowance.started.wait() + await allowance_started.wait() if cancel_mode == "scope": scope.cancel() elif cancel_mode != "none": @@ -781,7 +791,7 @@ async def test_completed_passthrough_cancellation_preserves_accounting( pending.cancel() if cancel_phase == "lock": logging_obj.cost_finalization_lock.release() - allowance.release.set() + allowance_release.set() if cancel_mode in ("once", "twice"): with pytest.raises(asyncio.CancelledError): await pending @@ -791,8 +801,9 @@ async def test_completed_passthrough_cancellation_preserves_accounting( assert reservation["callback_bound"] is True await release_unbound_budget_reservation(reservation) await GLOBAL_LOGGING_WORKER.flush() - assert allowance.calls == 1 - assert allowance.units == 3 + assert executions == [1] + (allocation_key,) = await verifier.keys("*") # pyright: ignore[reportUnknownVariableType] # fakeredis replies are untyped + assert await verifier.get(allocation_key) == b"3" # pyright: ignore[reportUnknownArgumentType] # fakeredis replies are untyped assert writer.update_database.await_count == 1 recorded: Final = writer.update_database.call_args.kwargs assert recorded["response_cost"] == pytest.approx(0.48)