From 084e8c425c2b8f5a336d946ff6f41cfea128bdb9 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Tue, 18 Mar 2025 22:41:02 -0700 Subject: [PATCH] refactor(base_routing_strategy.py): fix function names --- litellm/router_strategy/base_routing_strategy.py | 12 ++++++------ .../router_strategy/test_base_routing_strategy.py | 12 ++++++------ 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 41d3ebc8f1..10969a809b 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -61,7 +61,7 @@ class BaseRoutingStrategy(ABC): ttl=ttl, ) self.redis_increment_operation_queue.append(increment_op) - self.add_to_cache_keys(key=key) + self.add_to_in_memory_keys_to_update(key=key) return result async def periodic_sync_in_memory_spend_with_redis( @@ -116,13 +116,13 @@ class BaseRoutingStrategy(ABC): ) self.redis_increment_operation_queue = [] - def add_to_cache_keys(self, key: str): + def add_to_in_memory_keys_to_update(self, key: str): self.in_memory_keys_to_update.add(key) - def get_cache_keys(self) -> Set[str]: + def get_in_memory_keys_to_update(self) -> Set[str]: return self.in_memory_keys_to_update - def reset_cache_keys(self): + def reset_in_memory_keys_to_update(self): self.in_memory_keys_to_update = set() async def _sync_in_memory_spend_with_redis(self): @@ -147,7 +147,7 @@ class BaseRoutingStrategy(ABC): await self._push_in_memory_increments_to_redis() # 2. Fetch all current provider spend from Redis to update in-memory cache - cache_keys = self.get_cache_keys() + cache_keys = self.get_in_memory_keys_to_update() cache_keys_list = list(cache_keys) @@ -167,7 +167,7 @@ class BaseRoutingStrategy(ABC): f"Updated in-memory cache for {key}: {value}" ) - self.reset_cache_keys() + self.reset_in_memory_keys_to_update() except Exception as e: verbose_router_logger.exception( f"Error syncing in-memory cache with Redis: {str(e)}" diff --git a/tests/litellm/router_strategy/test_base_routing_strategy.py b/tests/litellm/router_strategy/test_base_routing_strategy.py index de3796352c..b47a2f1c90 100644 --- a/tests/litellm/router_strategy/test_base_routing_strategy.py +++ b/tests/litellm/router_strategy/test_base_routing_strategy.py @@ -120,15 +120,15 @@ async def test_sync_in_memory_spend_with_redis(base_strategy, mock_dual_cache): def test_cache_keys_management(base_strategy): # Test adding and getting cache keys - base_strategy.add_to_cache_keys("key1") - base_strategy.add_to_cache_keys("key2") - base_strategy.add_to_cache_keys("key1") # Duplicate should be ignored + base_strategy.add_to_in_memory_keys_to_update("key1") + base_strategy.add_to_in_memory_keys_to_update("key2") + base_strategy.add_to_in_memory_keys_to_update("key1") # Duplicate should be ignored - cache_keys = base_strategy.get_cache_keys() + cache_keys = base_strategy.get_in_memory_keys_to_update() assert len(cache_keys) == 2 assert "key1" in cache_keys assert "key2" in cache_keys # Test resetting cache keys - base_strategy.reset_cache_keys() - assert len(base_strategy.get_cache_keys()) == 0 + base_strategy.reset_in_memory_keys_to_update() + assert len(base_strategy.get_in_memory_keys_to_update()) == 0