refactor(redis): move the tool discovery buffer behind a typed store

litellm/proxy/db/db_transaction_queue/redis_update_buffer.py

@@ -20,10 +20,8 @@ from litellm.caching import RedisCache
 from litellm.constants import (
     MAX_REDIS_BUFFER_DEQUEUE_COUNT,
     REDIS_SPEND_LOGS_BUFFER_MAX_ROWS,
-    REDIS_TOOL_DISCOVERY_BUFFER_KEY,
     REDIS_UPDATE_BUFFER_KEY,
 )
-from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.proxy._types import (
     BaseDailySpendTransaction,
     DailyAgentSpendTransaction,
@@ -59,12 +57,9 @@ from litellm.redis.stores.spend_update_buffer import (
     TransportSpendUpdateBufferStore,
     spend_buffer_key,
 )
+from litellm.redis.stores.tool_discovery_buffer import ToolDiscoveryBufferStore, TransportToolDiscoveryBufferStore
 from litellm.redis.transport.contract import RedisFailure, raise_redis_failure
 from litellm.secret_managers.main import str_to_bool
-from litellm.types.caching import (
-    RedisPipelineLpopOperation,
-    RedisPipelineRpushOperation,
-)
 from litellm.types.services import ServiceTypes
 
 SPEND_QUEUE_TARGET: Final = "spend_queue"
@@ -125,6 +120,10 @@ def _parked_spend_logs(redis_cache: RedisCache) -> SpendLogsParkingStore:
     return TransportSpendLogsParkingStore(redis_cache.transport())
 
 
+def _tool_discoveries(redis_cache: RedisCache) -> ToolDiscoveryBufferStore:
+    return TransportToolDiscoveryBufferStore(redis_cache.transport())
+
+
 def _zero_for_non_finite_constant(constant: str) -> float:
     verbose_proxy_logger.warning(
         "Redis spend buffer - reading a parked %s spend as 0, the spend it held is lost", constant
@@ -160,9 +159,6 @@ def _merged_entity_transactions(
     )
 
 
-_TOOL_DISCOVERY_ITEMS: Final = TypeAdapter(list[ToolDiscoveryQueueItem])
-
-
 class RedisUpdateBuffer:
     """
     Handles buffering database `UPDATE` transactions in Redis before committing them to the database
@@ -575,46 +571,38 @@ class RedisUpdateBuffer:
         items: Final = tool_discovery_queue.flush()
         if not items:
             return
-        try:
-            _ = await self.redis_cache.async_rpush_pipeline(
-                rpush_list=[
-                    RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(items)])
-                ]
-            )
-        except (RedisError, OSError, asyncio.TimeoutError) as e:
-            verbose_proxy_logger.error(
-                "Tool discovery - failed to push %d tool rows to Redis, restoring them for the next tick. Error: %s",
-                len(items),
-                str(e),
-            )
-            for item in items:
-                tool_discovery_queue.add_update(item)
+        pushed: Final = await _tool_discoveries(self.redis_cache).push(items)
+        if not isinstance(pushed, RedisFailure):
+            return
+        if pushed.kind not in ("transport", "server"):
+            raise_redis_failure(pushed)
+        verbose_proxy_logger.error(
+            "Tool discovery - failed to push %d tool rows to Redis, restoring them for the next tick. Error: %s",
+            len(items),
+            pushed.message,
+        )
+        for item in items:
+            tool_discovery_queue.add_update(item)
 
     @with_service_target(SPEND_QUEUE_TARGET)
     async def restore_tool_discoveries_to_redis(self, items: Sequence[ToolDiscoveryQueueItem]) -> None:
         """Put drained tool rows back for the next drain after a write that provably did not apply."""
         if self.redis_cache is None or not items:
             return
-        _ = await self.redis_cache.async_rpush_pipeline(
-            rpush_list=[
-                RedisPipelineRpushOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, values=[safe_dumps(list(items))])
-            ]
-        )
+        restored: Final = await _tool_discoveries(self.redis_cache).push(items)
+        if isinstance(restored, RedisFailure):
+            raise_redis_failure(restored)
 
     @with_service_target(SPEND_QUEUE_TARGET)
     async def get_all_tool_discoveries_from_redis_buffer(self) -> tuple[ToolDiscoveryQueueItem, ...]:
         """Drain every worker's tool rows and fold them to one row per tool_name."""
         if self.redis_cache is None:
             return ()
-        popped: Final = await self.redis_cache.async_lpop_pipeline(
-            lpop_list=[
-                RedisPipelineLpopOperation(key=REDIS_TOOL_DISCOVERY_BUFFER_KEY, count=MAX_REDIS_BUFFER_DEQUEUE_COUNT)
-            ]
-        )
-        entries: Final = popped[0] if popped else None
-        if not entries:
+        batches: Final = await _tool_discoveries(self.redis_cache).drain(count=MAX_REDIS_BUFFER_DEQUEUE_COUNT)
+        if isinstance(batches, RedisFailure):
+            raise_redis_failure(batches)
+        if not batches:
             return ()
-        batches: Final = tuple(_TOOL_DISCOVERY_ITEMS.validate_json(entry) for entry in entries)
         return aggregate_tool_discoveries(chain.from_iterable(batches))
 
     @with_service_target(SPEND_QUEUE_TARGET)

litellm/redis/stores/tool_discovery_buffer.py

@@ -0,0 +1,36 @@
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Final, Protocol
+
+from pydantic import TypeAdapter, ValidationError
+
+from litellm.constants import REDIS_TOOL_DISCOVERY_BUFFER_KEY
+from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+from litellm.proxy._types import ToolDiscoveryQueueItem
+from litellm.redis.transport.contract import RedisFailure, RedisTransport
+
+_ITEMS: Final = TypeAdapter(list[ToolDiscoveryQueueItem])
+
+
+class ToolDiscoveryBufferStore(Protocol):
+    async def push(self, items: Sequence[ToolDiscoveryQueueItem]) -> int | RedisFailure: ...
+
+    async def drain(self, *, count: int) -> tuple[tuple[ToolDiscoveryQueueItem, ...], ...] | RedisFailure: ...
+
+
+@dataclass(frozen=True, slots=True)
+class TransportToolDiscoveryBufferStore:
+    transport: RedisTransport
+
+    async def push(self, items: Sequence[ToolDiscoveryQueueItem]) -> int | RedisFailure:
+        return await self.transport.rpush(REDIS_TOOL_DISCOVERY_BUFFER_KEY, safe_dumps(list(items)))
+
+    async def drain(self, *, count: int) -> tuple[tuple[ToolDiscoveryQueueItem, ...], ...] | RedisFailure:
+        popped: Final = await self.transport.lpop(REDIS_TOOL_DISCOVERY_BUFFER_KEY, count)
+        if isinstance(popped, RedisFailure):
+            return popped
+        entries: Final = () if popped is None else (popped,) if isinstance(popped, bytes) else popped
+        try:
+            return tuple(tuple(_ITEMS.validate_json(entry)) for entry in entries)
+        except ValidationError as error:
+            return RedisFailure(kind="decode", message="Invalid tool discovery buffer entry", exception=error)

tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py

@@ -242,11 +242,10 @@ async def test_store_in_memory_spend_updates_restores_on_rpush_failure(
 
 
 @pytest.mark.asyncio
-async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_update_buffer, mock_redis_cache):
+async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_update_buffer, mock_redis_cache, raw_redis):
     """
     When all queues are empty, pipeline should never be called.
     """
-    mock_redis_cache.async_rpush_pipeline = AsyncMock()
 
     # All queues return empty
     empty_queue = AsyncMock()
@@ -263,7 +262,7 @@ async def test_store_in_memory_spend_updates_all_empty_returns_early(redis_updat
         daily_agent_spend_update_queue=empty_daily_queue,
     )
 
-    mock_redis_cache.async_rpush_pipeline.assert_not_called()
+    assert await raw_redis.keys("*") == []
 
 
 @pytest.mark.asyncio
@@ -540,11 +539,10 @@ async def test_restored_window_spend_transactions_drain_back_unchanged(
 
 
 @pytest.mark.asyncio
-async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache):
+async def test_restore_transactions_to_redis_noop_when_empty(redis_update_buffer, mock_redis_cache, raw_redis):
     """Nothing to restore -> no Redis call."""
-    mock_redis_cache.async_rpush_pipeline = AsyncMock()
     await redis_update_buffer.restore_transactions_to_redis()
-    mock_redis_cache.async_rpush_pipeline.assert_not_called()
+    assert await raw_redis.keys("*") == []
 
 
 @pytest.mark.asyncio
@@ -865,8 +863,9 @@ async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_
 async def test_tool_discoveries_go_back_on_the_queue_when_the_redis_push_fails():
     from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue
 
-    redis_cache = MagicMock()
-    redis_cache.async_rpush_pipeline = AsyncMock(side_effect=ConnectionError("redis down"))
+    server: Final = fakeredis.FakeServer()
+    redis_cache: Final = fake_redis_cache(server)
+    server.connected = False
     queue = ToolDiscoveryQueue()
     queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})
     queue.add_update({"tool_name": "get_weather", "origin": "user_defined"})

tests/unit/redis/stores/test_tool_discovery_buffer.py

@@ -0,0 +1,50 @@
+from typing import Final
+
+import fakeredis
+import pytest
+
+from litellm.constants import REDIS_TOOL_DISCOVERY_BUFFER_KEY
+from litellm.redis.stores.tool_discovery_buffer import TransportToolDiscoveryBufferStore
+from litellm.redis.transport.contract import RedisFailure
+from tests._support.fake_redis import fake_redis_cache
+
+KEY: Final = f"ns:{REDIS_TOOL_DISCOVERY_BUFFER_KEY}"
+
+
+def _store(server: fakeredis.FakeServer) -> TransportToolDiscoveryBufferStore:
+    return TransportToolDiscoveryBufferStore(fake_redis_cache(server, namespace="ns").transport())
+
+
+@pytest.mark.asyncio
+async def test_each_push_is_one_json_list_entry_and_a_drain_returns_every_batch(server: fakeredis.FakeServer) -> None:
+    store: Final = _store(server)
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+
+    assert await store.push(({"tool_name": "a", "call_count": 2},)) == 1
+    assert await store.push(({"tool_name": "b", "origin": None, "call_count": 1},)) == 2
+
+    assert await raw.lrange(KEY, 0, -1) == [
+        b'[{"tool_name": "a", "call_count": 2}]',
+        b'[{"tool_name": "b", "origin": null, "call_count": 1}]',
+    ]
+    assert await store.drain(count=10) == (
+        ({"tool_name": "a", "call_count": 2},),
+        ({"tool_name": "b", "origin": None, "call_count": 1},),
+    )
+    assert await store.drain(count=10) == ()
+
+
+@pytest.mark.asyncio
+async def test_bad_entries_rejections_and_outages_are_failures(server: fakeredis.FakeServer) -> None:
+    raw: Final = fakeredis.FakeAsyncRedis(server=server)
+    store: Final = _store(server)
+    await raw.rpush(KEY, b'{"not": "a list"}')
+    malformed: Final = await store.drain(count=10)
+    await raw.set(KEY, "not a list")
+    rejected: Final = await store.push(({"tool_name": "a", "call_count": 1},))
+    server.connected = False
+    outage: Final = await store.drain(count=10)
+
+    assert isinstance(malformed, RedisFailure) and malformed.kind == "decode"
+    assert isinstance(rejected, RedisFailure) and rejected.kind == "server"
+    assert isinstance(outage, RedisFailure) and outage.kind == "transport"