tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-24 10:27:13 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(base_routing_strategy.py): fix function names

Browse Source
This commit is contained in:
Krrish Dholakia
2025-03-18 22:41:02 -07:00
parent 3033c40739
commit 084e8c425c
2 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/router_strategy/base_routing_strategy.py
+6 -6
View File
@@ -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)}"
tests/litellm/router_strategy/test_base_routing_strategy.py
+6 -6
View File
@@ -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