Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py: 48%

1""" 

2In-memory buffer for tool registry upserts. 

3 

4Unlike SpendUpdateQueue (which aggregates increments), ToolDiscoveryQueue 

5uses set-deduplication: each unique tool_name is only queued once per flush 

6cycle (~30s). The seen-set is cleared on every flush so that call_count 

7increments in subsequent cycles rather than stopping after the first flush. 

8""" 

9 

10from typing import Final 

11 

12from litellm._logging import verbose_proxy_logger 

13from litellm.proxy._types import ToolDiscoveryQueueItem 

14 

15 

16class ToolDiscoveryQueue: 

17 """ 

18 In-memory buffer for tool registry upserts. 

19 

20 Deduplicates by tool_name within each flush cycle: a tool is only queued 

21 once per ~30s batch, so call_count increments once per flush cycle the 

22 tool appears in (not once per invocation, but not once per pod lifetime 

23 either). The seen-set is cleared on flush so subsequent batches can 

24 re-count the same tool. 

25 """ 

26 

27 def __init__(self) -> None: 

28 self._seen_tool_names: set[str] = set() 

29 self._pending: list[ToolDiscoveryQueueItem] = [] 

30 

31 def add_update(self, item: ToolDiscoveryQueueItem) -> None: 

32 """Enqueue a tool discovery item if tool_name has not been seen before.""" 

33 tool_name: Final = item.get("tool_name", "") 

34 if not tool_name: 

35 return 

36 if tool_name in self._seen_tool_names: 

37 verbose_proxy_logger.debug("ToolDiscoveryQueue: skipping already-seen tool %s", tool_name) 

38 return 

39 self._seen_tool_names.add(tool_name) 

40 self._pending.append(item) 

41 verbose_proxy_logger.debug( 

42 "ToolDiscoveryQueue: queued new tool %s (origin=%s)", 

43 tool_name, 

44 item.get("origin"), 

45 ) 

46 

47 def flush(self) -> list[ToolDiscoveryQueueItem]: 

48 """Return and clear all pending items. Resets seen-set so the next 

49 flush cycle can re-count the same tools.""" 

50 items, self._pending = self._pending, [] 

51 self._seen_tool_names.clear() 

52 return items