diff --git a/litellm/constants.py b/litellm/constants.py index ce002cc4ee8..57d5452385d 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -446,6 +446,7 @@ REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_end_user_sp REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_agent_spend_update_buffer" REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_daily_tag_spend_update_buffer" REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY: Final = "litellm_window_spend_update_buffer" +REDIS_TOOL_DISCOVERY_BUFFER_KEY: Final = "litellm_tool_discovery_buffer" MAX_REDIS_BUFFER_DEQUEUE_COUNT: Final = get_env_int("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100) REDIS_SPEND_LOGS_BUFFER_KEY: Final = "litellm_spend_logs_buffer" REDIS_SPEND_LOGS_BUFFER_MAX_ROWS: Final = 100000 @@ -1810,6 +1811,7 @@ LITELLM_UI_SESSION_DURATION: Final = os.getenv("LITELLM_UI_SESSION_DURATION", "2 ########################### DB CRON JOB NAMES ########################### DB_SPEND_UPDATE_JOB_NAME: Final = "db_spend_update_job" +TOOL_DISCOVERY_COMMIT_JOB_NAME: Final = "tool_discovery_commit_job" DB_DAILY_TAG_SPEND_UPDATE_JOB_NAME: Final = "db_daily_tag_spend_update_job" PROMETHEUS_EMIT_BUDGET_METRICS_JOB_NAME: Final = "prometheus_emit_budget_metrics" CLOUDZERO_EXPORT_USAGE_DATA_JOB_NAME: Final = "cloudzero_export_usage_data" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 4272bfde3f5..5e812eec477 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -5198,6 +5198,7 @@ class ToolDiscoveryQueueItem(TypedDict, total=False): team_id: str | None # team that triggered discovery key_alias: str | None # human-readable key alias user_agent: str | None # HTTP User-Agent of the caller + call_count: ReadOnly[int] # invocations folded into this row since the previous flush from litellm.models.managed_files import ( # noqa: E402 diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 3d8678d70fb..127376b98c4 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -37,7 +37,6 @@ from litellm.constants import ( from litellm.litellm_core_utils.litellm_logging import coerce_model_access_groups from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( - DB_CONNECTION_ERROR_TYPES, DB_RETRY_SAFE_ERROR_TYPES, BaseDailySpendTransaction, DailyAgentSpendTransaction, @@ -74,7 +73,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, ) -from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler, db_write_failure_is_requeue_safe from litellm.proxy.db.model_usage_rollup import build_model_usage_transaction from litellm.proxy.db.request_error_tracking import request_error_accumulator from litellm.proxy.db.rollup_lock_timeout import apply_rollup_lock_timeout @@ -231,16 +230,6 @@ class _DailySpendCommit(Protocol[_DailySpendTransactionT]): ) -> None: ... -_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) - - -def _spend_commit_failure_is_requeue_safe(e: Exception) -> bool: - if isinstance(e, DB_CONNECTION_ERROR_TYPES): - return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) - sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) - return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES - - _SpendTableName = Literal[ "user_list_transactions", "end_user_list_transactions", @@ -278,7 +267,7 @@ def _spend_tables_left_to_send( _SPEND_TABLE_COMMIT_ORDER[len(committed)] if len(committed) < len(_SPEND_TABLE_COMMIT_ORDER) else None ) dropped: Final[frozenset[_SpendTableName]] = ( - frozenset() if in_flight is None or _spend_commit_failure_is_requeue_safe(failure) else frozenset({in_flight}) + frozenset() if in_flight is None or db_write_failure_is_requeue_safe(failure) else frozenset({in_flight}) ) if dropped and in_flight is not None: spend_log_error( @@ -1757,7 +1746,7 @@ class DBSpendUpdateWriter: window_spend_transactions=window_spend_update_transactions, ) except Exception as e: - if not _spend_commit_failure_is_requeue_safe(e): + if not db_write_failure_is_requeue_safe(e): uncommitted.pop("window_spend_update_transactions", None) spend_log_error( "Spend tracking - dropped %d budget window increments: the failed statement may have " @@ -1923,7 +1912,7 @@ class DBSpendUpdateWriter: window_spend_transactions=window_spend_update_transactions, ) except Exception as e: # noqa: BLE001 # the increments go back on the queue; the rest of the flush must run - if _spend_commit_failure_is_requeue_safe(e): + if db_write_failure_is_requeue_safe(e): spend_log_error( "Spend tracking - failed to commit budget window spend updates. " "Re-queued %d window increments for retry on next tick. Error: %s", @@ -1941,9 +1930,6 @@ class DBSpendUpdateWriter: exc=e, ) - ################## Tool Registry Upserts ################## - await self._flush_tool_discovery_queue(prisma_client=prisma_client) - async def commit_daily_tag_spend_to_db( self, prisma_client: PrismaClient, @@ -2075,20 +2061,6 @@ class DBSpendUpdateWriter: ) raise - async def _flush_tool_discovery_queue( - self, - prisma_client: PrismaClient, - ) -> None: - """Flush ToolDiscoveryQueue and batch-upsert new tools into LiteLLM_ToolTable.""" - from litellm.proxy.db.tool_registry_writer import batch_upsert_tools - - try: - items: Final = self.tool_discovery_queue.flush() - if items: - await batch_upsert_tools(prisma_client=prisma_client, items=items) - except Exception as e: - verbose_proxy_logger.debug("_flush_tool_discovery_queue error (non-blocking): %s", e) - @staticmethod async def handle_spend_update_failure( e: Exception, @@ -2100,7 +2072,6 @@ class DBSpendUpdateWriter: """Retry a failed spend-update transaction on connection errors, deadlocks or a rollup ``lock_timeout`` (55P03), else re-raise. All three roll the transaction back before any increment applied, so re-sending the same batch cannot double-count.""" - from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.utils import raise_failed_update_spend_exception is_retryable = ( @@ -2609,7 +2580,7 @@ class DBSpendUpdateWriter: _mark_daily_spend_commit_started() _mark_daily_spend_commit_finished() except Exception as batch_error: - if _spend_commit_failure_is_requeue_safe(batch_error): + if db_write_failure_is_requeue_safe(batch_error): spend_log_error( "Daily %s spend batch upsert failed. Table: %s, Rows: %d, Error: %s", entity_type, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 898aee493df..22f55ea9153 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -9,8 +9,11 @@ import json from collections.abc import Mapping, Sequence from datetime import datetime from functools import reduce +from itertools import chain from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast +from pydantic import TypeAdapter + from redis.exceptions import RedisError from litellm._internal_context import with_service_target @@ -26,6 +29,7 @@ from litellm.constants import ( REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_KEY, REDIS_SPEND_LOGS_BUFFER_MAX_ROWS, + REDIS_TOOL_DISCOVERY_BUFFER_KEY, REDIS_UPDATE_BUFFER_KEY, REDIS_WINDOW_SPEND_UPDATE_BUFFER_KEY, ) @@ -41,12 +45,17 @@ from litellm.proxy._types import ( DBSpendUpdateTransactions, Litellm_EntityType, SpendUpdateQueueItem, + ToolDiscoveryQueueItem, ) from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, ) from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, + aggregate_tool_discoveries, +) from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import ( WindowSpendTransaction, WindowSpendUpdateQueue, @@ -133,6 +142,9 @@ 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 @@ -551,6 +563,56 @@ class RedisUpdateBuffer: str(e), ) + @with_service_target(SPEND_QUEUE_TARGET) + async def store_tool_discoveries_in_redis(self, tool_discovery_queue: ToolDiscoveryQueue) -> None: + """Move this worker's pending tool rows to Redis so only the leader writes LiteLLM_ToolTable.""" + if self.redis_cache is None: + return + 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) + + @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))]) + ] + ) + + @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: + 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) async def store_spend_logs_in_redis( self, diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py index d5f6a447131..e9ec040023c 100644 --- a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -1,52 +1,57 @@ """ -In-memory buffer for tool registry upserts. - -Unlike SpendUpdateQueue (which aggregates increments), ToolDiscoveryQueue -uses set-deduplication: each unique tool_name is only queued once per flush -cycle (~30s). The seen-set is cleared on every flush so that call_count -increments in subsequent cycles rather than stopping after the first flush. +In-memory buffer for tool registry upserts, aggregated by tool_name so one flush carries one row per +tool with the number of times it was invoked since the previous flush. """ +from collections.abc import Iterable from typing import Final +from pydantic import TypeAdapter + from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem +_ITEM: Final = TypeAdapter(ToolDiscoveryQueueItem) + + +def _merged(existing: ToolDiscoveryQueueItem | None, item: ToolDiscoveryQueueItem) -> ToolDiscoveryQueueItem: + """The first sighting keeps its origin and caller fields, later sightings only add to call_count.""" + increment: Final = item.get("call_count", 1) + if existing is None: + return _ITEM.validate_python({**item, "call_count": increment}) + return _ITEM.validate_python({**existing, "call_count": existing.get("call_count", 1) + increment}) + + +def aggregate_tool_discoveries(items: Iterable[ToolDiscoveryQueueItem]) -> tuple[ToolDiscoveryQueueItem, ...]: + """Collapse items to one per tool_name, summing call_count, in first-seen order.""" + by_name: Final[dict[str, ToolDiscoveryQueueItem]] = {} # mutable-ok: single pass fold keyed by tool_name + for item in items: + if not (tool_name := item.get("tool_name", "")): + continue + by_name[tool_name] = _merged(by_name.get(tool_name), item) + return tuple(by_name.values()) + class ToolDiscoveryQueue: - """ - In-memory buffer for tool registry upserts. - - Deduplicates by tool_name within each flush cycle: a tool is only queued - once per ~30s batch, so call_count increments once per flush cycle the - tool appears in (not once per invocation, but not once per pod lifetime - either). The seen-set is cleared on flush so subsequent batches can - re-count the same tool. - """ + """In-memory buffer for tool registry upserts, one pending row per tool_name with a running call_count.""" def __init__(self) -> None: - self._seen_tool_names: set[str] = set() - self._pending: list[ToolDiscoveryQueueItem] = [] + self._pending: dict[str, ToolDiscoveryQueueItem] = {} # mutable-ok: folded per add_update, cleared on flush def add_update(self, item: ToolDiscoveryQueueItem) -> None: - """Enqueue a tool discovery item if tool_name has not been seen before.""" + """Count one more invocation of item's tool_name, keeping the first sighting's origin and caller fields.""" tool_name: Final = item.get("tool_name", "") if not tool_name: return - if tool_name in self._seen_tool_names: - verbose_proxy_logger.debug("ToolDiscoveryQueue: skipping already-seen tool %s", tool_name) - return - self._seen_tool_names.add(tool_name) - self._pending.append(item) + self._pending[tool_name] = _merged(self._pending.get(tool_name), item) verbose_proxy_logger.debug( - "ToolDiscoveryQueue: queued new tool %s (origin=%s)", + "ToolDiscoveryQueue: tool %s call_count=%s (origin=%s)", tool_name, + self._pending[tool_name].get("call_count"), item.get("origin"), ) def flush(self) -> list[ToolDiscoveryQueueItem]: - """Return and clear all pending items. Resets seen-set so the next - flush cycle can re-count the same tools.""" - items, self._pending = self._pending, [] - self._seen_tool_names.clear() + """Return and clear all pending rows.""" + items, self._pending = list(self._pending.values()), {} return items diff --git a/litellm/proxy/db/exception_handler.py b/litellm/proxy/db/exception_handler.py index 379e0012ea8..5dfcaa4b2f8 100644 --- a/litellm/proxy/db/exception_handler.py +++ b/litellm/proxy/db/exception_handler.py @@ -8,6 +8,7 @@ from pydantic import TypeAdapter, ValidationError from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, + DB_RETRY_SAFE_ERROR_TYPES, ProxyErrorTypes, ProxyException, ) @@ -637,3 +638,14 @@ async def call_with_db_reconnect_retry( # propagate — repeated reconnect-loops are the watchdog's job, not # this helper's. return await coro_factory() + + +_DATA_REJECTED_SQLSTATE_CLASSES: Final = frozenset({"22", "23"}) + + +def db_write_failure_is_requeue_safe(e: Exception) -> bool: + """Whether a failed write provably never applied, so its rows can be sent again without double counting.""" + if isinstance(e, DB_CONNECTION_ERROR_TYPES): + return isinstance(e, DB_RETRY_SAFE_ERROR_TYPES) + sqlstate: Final = PrismaDBExceptionHandler.postgres_sqlstate(e) + return sqlstate is None or sqlstate[:2] not in _DATA_REJECTED_SQLSTATE_CLASSES diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index 2e3530d7613..f7ec88945a1 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -14,8 +14,17 @@ from typing import TYPE_CHECKING, Final, Protocol from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger +from litellm.constants import TOOL_DISCOVERY_COMMIT_JOB_NAME from litellm.proxy._types import ToolDiscoveryQueueItem -from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry +from litellm.proxy.db.db_span import db_span +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, + aggregate_tool_discoveries, +) +from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry, db_write_failure_is_requeue_safe +from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ToolRepository @@ -105,58 +114,75 @@ def _row_to_model(row: object) -> LiteLLM_ToolTableRow: async def batch_upsert_tools( prisma_client: "PrismaClient", - items: list[ToolDiscoveryQueueItem], + items: Sequence[ToolDiscoveryQueueItem], ) -> None: - """ - Batch-upsert tool registry rows via Prisma. + """Write a batch of discovered tools to LiteLLM_ToolTable in a single statement. Raises on failure so the + caller decides whether the rows can be retried.""" + rows: Final = aggregate_tool_discoveries(items) + if not rows: + return + sql, params = build_tool_bulk_upsert(rows) + async with db_span("commit_tool_discoveries", "LiteLLM_ToolTable"): + await prisma_client.db.execute_raw(sql, *params) # pyright: ignore[reportAny] # untyped prisma client + verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(rows)) - On first insert: sets input_policy/output_policy = "untrusted" (default), call_count = 1. - On conflict: increments call_count; preserves existing policies. - """ + +async def flush_tool_discovery_queue(prisma_client: "PrismaClient", tool_discovery_queue: ToolDiscoveryQueue) -> None: + """One worker's own rows straight to LiteLLM_ToolTable, the path without the Redis buffer.""" + items: Final = tool_discovery_queue.flush() if not items: return try: - data: Final = [item for item in items if item.get("tool_name")] - if not data: - return - now: Final = datetime.now(timezone.utc) - table: Final = _tool_table_actions(prisma_client) - for item in data: - tool_name = item.get("tool_name", "") - origin = item.get("origin") or "user_defined" - created_by = item.get("created_by") or "system" - key_hash = item.get("key_hash") - team_id = item.get("team_id") - key_alias = item.get("key_alias") - user_agent = item.get("user_agent") - await table.upsert( - where={"tool_name": tool_name}, - data={ - "create": { - "tool_id": str(uuid.uuid4()), - "tool_name": tool_name, - "origin": origin, - "input_policy": "untrusted", - "output_policy": "untrusted", - "call_count": 1, - "created_by": created_by, - "updated_by": created_by, - "key_hash": key_hash, - "team_id": team_id, - "key_alias": key_alias, - "user_agent": user_agent, - "last_used_at": now, - }, - "update": { - "call_count": {"increment": 1}, - "updated_at": now, - "last_used_at": now, - }, - }, + await batch_upsert_tools(prisma_client=prisma_client, items=items) + except Exception as e: # noqa: BLE001 # a lost tool row only delays its discovery, the scheduler job must not fail + verbose_proxy_logger.debug("flush_tool_discovery_queue error (non-blocking): %s", e) + + +async def commit_tool_discoveries( + prisma_client: "PrismaClient", + tool_discovery_queue: ToolDiscoveryQueue, + redis_update_buffer: RedisUpdateBuffer, + pod_lock_manager: PodLockManager, + use_redis_buffer: bool, +) -> None: + """Scheduled alongside update_spend on its own lock, so tool table writes never ride the spend commit. + + With the Redis buffer every worker hands its rows to Redis and the holder of the tool lock writes them all in + one statement. Without it each worker writes its own rows.""" + if not use_redis_buffer: + await flush_tool_discovery_queue(prisma_client=prisma_client, tool_discovery_queue=tool_discovery_queue) + return + await redis_update_buffer.store_tool_discoveries_in_redis(tool_discovery_queue=tool_discovery_queue) + if not await pod_lock_manager.acquire_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME): + return + try: + await _commit_tool_discoveries_from_redis(prisma_client=prisma_client, redis_update_buffer=redis_update_buffer) + finally: + await pod_lock_manager.release_lock(cronjob_id=TOOL_DISCOVERY_COMMIT_JOB_NAME) + + +async def _commit_tool_discoveries_from_redis( + prisma_client: "PrismaClient", redis_update_buffer: RedisUpdateBuffer +) -> None: + """Tool lock holder only: drain every worker's tool rows from Redis and write them in one statement.""" + items: Final = await redis_update_buffer.get_all_tool_discoveries_from_redis_buffer() + if not items: + return + try: + await batch_upsert_tools(prisma_client=prisma_client, items=items) + except Exception as e: # noqa: BLE001 # the rows go back to redis or are dropped, the scheduler job must not fail + if db_write_failure_is_requeue_safe(e): + await redis_update_buffer.restore_tool_discoveries_to_redis(items) + verbose_proxy_logger.error( + "Tool discovery - %d tool rows went back to Redis for the next tick. Error: %s", len(items), str(e) ) - verbose_proxy_logger.debug("tool_registry_writer: upserted %d tool(s)", len(data)) - except Exception as e: - verbose_proxy_logger.error("tool_registry_writer batch_upsert_tools error: %s", e) + return + verbose_proxy_logger.error( + "Tool discovery - dropped %d tool rows: the failed statement may have applied or the database " + "refused the data, so re-sending it is not safe. Error: %s", + len(items), + str(e), + ) _NO_OWNERS: Final[Mapping[str, ToolDiscoveryUser]] = MappingProxyType({}) diff --git a/litellm/proxy/db/tool_table_bulk_upsert.py b/litellm/proxy/db/tool_table_bulk_upsert.py new file mode 100644 index 00000000000..e586042ae7c --- /dev/null +++ b/litellm/proxy/db/tool_table_bulk_upsert.py @@ -0,0 +1,71 @@ +"""One multi-row ``INSERT ... ON CONFLICT ("tool_name") DO UPDATE`` per batch of discovered tools. + +LiteLLM_ToolTable holds one row per tool name the proxy has seen. A batch arrives already folded to one +entry per tool name with a summed call_count, and this statement inserts the names the table has not seen +with the untrusted policies the columns default to, and for the names it has adds the batch's call_count +while leaving the policies and the discovering caller's fields alone. Same shape as daily_spend_bulk_upsert. +""" + +import uuid +from collections.abc import Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final, TypeAlias + +from litellm.proxy._types import ToolDiscoveryQueueItem + +SqlValue: TypeAlias = str | int | None + +_TOOL_ROW_COLUMNS: Final[tuple[str, ...]] = ( + "tool_id", + "tool_name", + "origin", + "call_count", + "created_by", + "updated_by", + "key_hash", + "team_id", + "key_alias", + "user_agent", +) +_CASTS: Final[Mapping[str, str]] = MappingProxyType({"call_count": "integer"}) + + +def _row_params(item: ToolDiscoveryQueueItem) -> tuple[SqlValue, ...]: + created_by: Final = item.get("created_by") or "system" + return ( + str(uuid.uuid4()), + item.get("tool_name", ""), + item.get("origin") or "user_defined", + item.get("call_count", 1), + created_by, + created_by, + item.get("key_hash"), + item.get("team_id"), + item.get("key_alias"), + item.get("user_agent"), + ) + + +def build_tool_bulk_upsert(items: Sequence[ToolDiscoveryQueueItem]) -> tuple[str, tuple[SqlValue, ...]]: + """The single statement writing one batch of distinct tool_names, plus its positional arguments.""" + width: Final = len(_TOOL_ROW_COLUMNS) + rows: Final = ", ".join( + "(" + + ", ".join( + f"${row_index * width + offset + 1}::{_CASTS.get(column, 'text')}" + for offset, column in enumerate(_TOOL_ROW_COLUMNS) + ) + + ", (NOW() AT TIME ZONE 'UTC'), (NOW() AT TIME ZONE 'UTC'))" + for row_index in range(len(items)) + ) + columns: Final = ", ".join(f'"{column}"' for column in _TOOL_ROW_COLUMNS) + sql: Final = ( + f'INSERT INTO "LiteLLM_ToolTable" ({columns}, "last_used_at", "updated_at")\n' + f"VALUES {rows}\n" + 'ON CONFLICT ("tool_name") DO UPDATE SET\n' + ' "call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count",\n' + ' "last_used_at" = EXCLUDED."last_used_at",\n' + ' "updated_at" = EXCLUDED."updated_at"' + ) + return sql, tuple(chain.from_iterable(map(_row_params, items))) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6dcde1aacae..80970cba90b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -10888,6 +10888,19 @@ class ProxyStartupEvent: misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) + ### TOOL DISCOVERY (own lock, same interval, never rides the spend commit) ### + from litellm.proxy.utils import commit_tool_discoveries + + scheduler.add_job( # pyright: ignore[reportUnknownMemberType] # apscheduler is untyped + commit_tool_discoveries, + "interval", + seconds=batch_writing_interval, + args=[prisma_client, proxy_logging_obj], + id="commit_tool_discoveries_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + ### UPDATE DAILY TAG SPEND (separate scheduler job with longer interval) ### ## Reduces QPS as there are more tags for a single request tag_spend_update_interval: Final = int(batch_writing_interval * DAILY_TAG_SPEND_BATCH_MULTIPLIER) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index c06bf54fbe7..5351c37b0e6 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -180,6 +180,7 @@ from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header from litellm.proxy.common_utils.config_sync_pubsub import publish_config_param_change from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.db import tool_registry_writer from litellm.proxy.db.create_views import ( create_missing_views, create_view_tolerating_race, @@ -187,6 +188,7 @@ from litellm.proxy.db.create_views import ( ) from litellm.proxy.db.db_span import db_span from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.db.db_url_settings import ( DatabaseURLSettings, add_missing_query_params, @@ -8289,6 +8291,18 @@ class ProxyUpdateSpend: return False +async def commit_tool_discoveries(prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None: + """Scheduler job writing discovered tools to LiteLLM_ToolTable, on its own lock, next to update_spend.""" + writer: Final = proxy_logging_obj.db_spend_update_writer + await tool_registry_writer.commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=writer.tool_discovery_queue, + redis_update_buffer=writer.redis_update_buffer, + pod_lock_manager=writer.pod_lock_manager, + use_redis_buffer=RedisUpdateBuffer.should_commit_spend_updates_to_redis(), + ) + + async def update_spend( prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, diff --git a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py index 5ecc8784384..76a2dbf73f1 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_redis_update_buffer.py @@ -5,7 +5,6 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest - from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer from litellm.proxy.proxy_server import ProxyStartupEvent @@ -701,3 +700,18 @@ async def test_store_spend_logs_in_redis_is_off_unless_transaction_buffering_is_ assert await buffer.store_spend_logs_in_redis([{"request_id": "a"}]) is False assert redis.rows == [] + + +@pytest.mark.asyncio +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")) + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + + await RedisUpdateBuffer(redis_cache=redis_cache).store_tool_discoveries_in_redis(queue) + + assert queue.flush() == [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 2}] diff --git a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py index e400ad16e84..0a015255f2e 100644 --- a/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py +++ b/tests/unit/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -5,9 +5,9 @@ Unit tests for ToolDiscoveryQueue. import pytest - from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( ToolDiscoveryQueue, + aggregate_tool_discoveries, ) @@ -16,57 +16,50 @@ def queue(): return ToolDiscoveryQueue() -def test_add_single_tool(queue): +def test_a_tool_seen_once_flushes_with_call_count_one(queue): queue.add_update({"tool_name": "my_tool", "origin": "user_defined"}) - items = queue.flush() - assert len(items) == 1 - assert items[0]["tool_name"] == "my_tool" - assert items[0]["origin"] == "user_defined" + assert queue.flush() == [{"tool_name": "my_tool", "origin": "user_defined", "call_count": 1}] -def test_deduplication_same_name(queue): - """Adding the same tool_name twice should only keep the first.""" - queue.add_update({"tool_name": "tool_a", "origin": "mcp_server"}) - queue.add_update({"tool_name": "tool_a", "origin": "user_defined"}) - items = queue.flush() - assert len(items) == 1 - assert items[0]["origin"] == "mcp_server" # first wins +def test_repeated_sightings_fold_into_one_row_counting_every_call(queue): + queue.add_update({"tool_name": "tool_a", "origin": "mcp_server", "created_by": "alice"}) + queue.add_update({"tool_name": "tool_a", "origin": "user_defined", "created_by": "bob"}) + queue.add_update({"tool_name": "tool_a"}) + assert queue.flush() == [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": "alice", "call_count": 3}] -def test_deduplication_different_names(queue): +def test_distinct_tools_each_get_their_own_row(queue): queue.add_update({"tool_name": "tool_a"}) queue.add_update({"tool_name": "tool_b"}) - items = queue.flush() - assert len(items) == 2 - names = {i["tool_name"] for i in items} - assert names == {"tool_a", "tool_b"} - - -def test_flush_clears_pending(queue): - queue.add_update({"tool_name": "tool_x"}) - items1 = queue.flush() - assert len(items1) == 1 - items2 = queue.flush() - assert len(items2) == 0 - - -def test_seen_names_reset_after_flush(queue): - """Seen-set is cleared on flush so the same tool can re-enter the next cycle.""" queue.add_update({"tool_name": "tool_a"}) - queue.flush() - queue.add_update({"tool_name": "tool_a"}) # same tool, new cycle - items = queue.flush() - assert len(items) == 1 - assert items[0]["tool_name"] == "tool_a" + assert queue.flush() == [ + {"tool_name": "tool_a", "call_count": 2}, + {"tool_name": "tool_b", "call_count": 1}, + ] + + +def test_flush_clears_pending_and_the_next_cycle_counts_from_zero(queue): + queue.add_update({"tool_name": "tool_x"}) + queue.add_update({"tool_name": "tool_x"}) + assert queue.flush() == [{"tool_name": "tool_x", "call_count": 2}] + assert queue.flush() == [] + queue.add_update({"tool_name": "tool_x"}) + assert queue.flush() == [{"tool_name": "tool_x", "call_count": 1}] def test_empty_tool_name_ignored(queue): queue.add_update({"tool_name": ""}) queue.add_update({"tool_name": None}) # type: ignore[arg-type] - items = queue.flush() - assert len(items) == 0 + assert queue.flush() == [] -def test_flush_returns_list(queue): - result = queue.flush() - assert isinstance(result, list) +def test_aggregate_sums_call_counts_across_workers_flushes(): + worker_a = [{"tool_name": "get_weather", "origin": "user_defined", "call_count": 3}] + worker_b = [ + {"tool_name": "get_weather", "origin": "mcp_server", "call_count": 2}, + {"tool_name": "search", "call_count": 1}, + ] + assert aggregate_tool_discoveries([*worker_a, *worker_b]) == ( + {"tool_name": "get_weather", "origin": "user_defined", "call_count": 5}, + {"tool_name": "search", "call_count": 1}, + ) diff --git a/tests/unit/proxy/db/test_db_spend_update_writer.py b/tests/unit/proxy/db/test_db_spend_update_writer.py index d440b99a632..b2f006aa52f 100644 --- a/tests/unit/proxy/db/test_db_spend_update_writer.py +++ b/tests/unit/proxy/db/test_db_spend_update_writer.py @@ -12,6 +12,8 @@ from unittest.mock import AsyncMock, MagicMock, call, patch import httpx import pytest + +from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation from prisma.errors import DataError as PrismaDataError from prisma.errors import RawQueryError from redis.exceptions import DataError @@ -2365,6 +2367,33 @@ async def test_commit_spend_updates_uses_pipeline(): mock_redis_update_buffer.get_all_daily_tag_spend_update_transactions_from_redis_buffer.assert_not_called() +class _InMemoryRedis: + def __init__(self) -> None: + self.store: dict[str, str] = {} + self.lists: dict[str, list[object]] = {} + + async def async_set_cache(self, key: str, value: str, nx: bool = False, **kwargs: object) -> bool | None: + if nx and key in self.store: + return None + self.store[key] = json.dumps(value) + return True + + async def async_get_cache(self, key: str, **kwargs: object) -> str | None: + raw = self.store.get(key) + return json.loads(raw) if raw is not None else None + + async def async_delete_cache(self, key: str) -> int: + return 1 if self.store.pop(key, None) is not None else 0 + + async def async_rpush_pipeline(self, rpush_list: list[RedisPipelineRpushOperation]) -> list[int]: + for op in rpush_list: + self.lists.setdefault(op["key"], []).extend(op["values"]) + return [len(self.lists[op["key"]]) for op in rpush_list] + + async def async_lpop_pipeline(self, lpop_list: list[RedisPipelineLpopOperation]) -> list[list[object] | None]: + return [self.lists.pop(op["key"], None) for op in lpop_list] + + @pytest.mark.asyncio async def test_commit_with_redis_requeues_all_on_db_failure(): """ @@ -3213,7 +3242,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFakeDB() db.query_raw = AsyncMock(side_effect=Exception("connection reset")) - db_writer._flush_tool_discovery_queue = AsyncMock() await db_writer._commit_spend_updates_to_db_without_redis_buffer( prisma_client=_WindowSpendFakePrisma(db), @@ -3221,7 +3249,6 @@ async def test_failed_window_spend_commit_requeues_the_increments_and_continues_ proxy_logging_obj=MagicMock(), ) - db_writer._flush_tool_discovery_queue.assert_called_once() requeued = await db_writer.window_spend_update_queue.flush_and_get_aggregated_window_spend_transactions() assert requeued == (transaction,) @@ -3297,7 +3324,6 @@ async def test_failed_daily_spend_commit_is_requeued_only_when_the_rows_are_prov db_writer = DBSpendUpdateWriter() await db_writer.daily_spend_update_queue.add_update({"user-key": _daily_txn(user_id="user-1")}) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=failure) - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3323,7 +3349,6 @@ async def test_failed_daily_spend_commit_drops_only_the_batch_that_was_sent(): {f"user-{i:03d}": _daily_txn(user_id=f"user-{i:03d}") for i in range(150)} ) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend", failure=httpx.ReadTimeout("no reply")) - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3351,7 +3376,6 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other team_txn = {key: value for key, value in _daily_txn().items() if key != "user_id"} | {"team_id": "team-1"} await db_writer.daily_team_spend_update_queue.add_update({"team-key": team_txn}) db = _DailySpendFakeDB(failing_table="LiteLLM_DailyUserSpend") - db_writer._flush_tool_discovery_queue = AsyncMock() proxy_logging_obj = MagicMock() proxy_logging_obj.failure_handler = AsyncMock() @@ -3362,7 +3386,6 @@ async def test_failed_daily_spend_commit_requeues_the_rows_and_flushes_the_other assert _daily_upserts(db, "LiteLLM_DailyUserSpend") == [] (team_upsert,) = _daily_upserts(db, "LiteLLM_DailyTeamSpend") assert _row_values(team_upsert, "team_id") == ["team-1"] - db_writer._flush_tool_discovery_queue.assert_called_once() db.failing_table = None await db_writer._commit_spend_updates_to_db_without_redis_buffer( @@ -3589,7 +3612,6 @@ async def test_failed_window_spend_commit_is_requeued_only_when_the_rows_are_pro ) await db_writer.window_spend_update_queue.add_update(transaction) db = _WindowSpendFailureDB(failure) - db_writer._flush_tool_discovery_queue = AsyncMock() await db_writer._commit_spend_updates_to_db_without_redis_buffer( prisma_client=_WindowSpendFakePrisma(db), diff --git a/tests/unit/proxy/db/test_tool_registry_writer.py b/tests/unit/proxy/db/test_tool_registry_writer.py index c9df665741d..76898b0eeab 100644 --- a/tests/unit/proxy/db/test_tool_registry_writer.py +++ b/tests/unit/proxy/db/test_tool_registry_writer.py @@ -3,22 +3,28 @@ Unit tests for tool_registry_writer.py — uses a mock prisma client that exposes litellm_tooltable.upsert / find_many / find_unique. """ +import json from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock import pytest -from prisma.errors import PrismaError - +from prisma.errors import PrismaError, RawQueryError +from litellm.constants import DB_SPEND_UPDATE_JOB_NAME, REDIS_TOOL_DISCOVERY_BUFFER_KEY, TOOL_DISCOVERY_COMMIT_JOB_NAME +from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager +from litellm.proxy.db.db_transaction_queue.redis_update_buffer import RedisUpdateBuffer +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ToolDiscoveryQueue from litellm.proxy.db.tool_registry_writer import ( ToolPolicyRegistry, batch_upsert_tools, + commit_tool_discoveries, get_tool, get_tool_policy_registry, get_tools_by_names, list_tools, update_tool_policy, ) +from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation def _mock_row(**kwargs): @@ -62,6 +68,7 @@ def _make_prisma( prisma = MagicMock() prisma.db.litellm_tooltable = MagicMock() prisma.db.litellm_tooltable.upsert = AsyncMock(return_value=upsert_return) + prisma.db.execute_raw = AsyncMock(return_value=1) prisma.db.litellm_tooltable.find_many = AsyncMock( return_value=find_many_rows if find_many_rows is not None else [] ) @@ -73,28 +80,32 @@ def _make_prisma( return prisma +def _statement(prisma) -> tuple[str, tuple[object, ...]]: + prisma.db.execute_raw.assert_awaited_once() + sql, *params = prisma.db.execute_raw.call_args.args + return sql, tuple(params) + + @pytest.mark.asyncio -async def test_batch_upsert_tools_calls_upsert(): +async def test_batch_upsert_tools_inserts_a_new_tool_with_the_batch_call_count_and_adds_it_on_conflict(): prisma = _make_prisma() - items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}] - await batch_upsert_tools(prisma, items) - prisma.db.litellm_tooltable.upsert.assert_awaited_once() - call_kw = prisma.db.litellm_tooltable.upsert.call_args.kwargs - assert call_kw["where"] == {"tool_name": "tool_a"} - assert call_kw["data"]["create"]["tool_name"] == "tool_a" - assert call_kw["data"]["create"]["origin"] == "mcp_server" - assert call_kw["data"]["create"]["input_policy"] == "untrusted" - assert call_kw["data"]["create"]["output_policy"] == "untrusted" - assert call_kw["data"]["create"]["call_count"] == 1 - assert call_kw["data"]["update"]["call_count"] == {"increment": 1} - assert "updated_at" in call_kw["data"]["update"] + await batch_upsert_tools( + prisma, [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 4}] + ) + sql, params = _statement(prisma) + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') + assert 'ON CONFLICT ("tool_name") DO UPDATE SET' in sql + assert '"call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count"' in sql + assert '"input_policy"' not in sql, "policies come from the table defaults and are never overwritten" + assert params[1:] == ("tool_a", "mcp_server", 4, "system", "system", None, None, None, None) + prisma.db.litellm_tooltable.upsert.assert_not_awaited() @pytest.mark.asyncio async def test_batch_upsert_tools_empty_list(): prisma = _make_prisma() await batch_upsert_tools(prisma, []) - prisma.db.litellm_tooltable.upsert.assert_not_awaited() + prisma.db.execute_raw.assert_not_awaited() @pytest.mark.asyncio @@ -102,21 +113,22 @@ async def test_batch_upsert_tools_skips_empty_names(): prisma = _make_prisma() items = [{"tool_name": "", "origin": None}, {"tool_name": None}] # type: ignore[list-item] await batch_upsert_tools(prisma, items) - prisma.db.litellm_tooltable.upsert.assert_not_awaited() + prisma.db.execute_raw.assert_not_awaited() @pytest.mark.asyncio -async def test_batch_upsert_multiple_tools_calls_upsert_per_tool(): +async def test_batch_upsert_many_tools_is_one_statement_with_one_row_per_distinct_tool(): prisma = _make_prisma() items = [ - {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}, - {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice"}, + {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 2}, + {"tool_name": "tool_b", "origin": "user_defined", "created_by": "alice", "call_count": 1}, + {"tool_name": "tool_a", "origin": "user_defined", "created_by": "bob", "call_count": 5}, ] await batch_upsert_tools(prisma, items) - assert prisma.db.litellm_tooltable.upsert.await_count == 2 - calls = prisma.db.litellm_tooltable.upsert.call_args_list - assert calls[0].kwargs["where"]["tool_name"] == "tool_a" - assert calls[1].kwargs["where"]["tool_name"] == "tool_b" + sql, params = _statement(prisma) + assert sql.count("$1::text") == 1 and "$11::text" in sql and "$21::text" not in sql + assert params[1:10] == ("tool_a", "mcp_server", 7, "system", "system", None, None, None, None) + assert params[11:20] == ("tool_b", "user_defined", 1, "alice", "alice", None, None, None, None) @pytest.mark.asyncio @@ -264,20 +276,14 @@ async def test_update_tool_policy_calls_upsert_then_get_tool(): assert call_kw["where"] == {"tool_name": "my_tool"} assert call_kw["data"]["update"]["input_policy"] == "blocked" assert call_kw["data"]["update"]["updated_by"] == "admin" - prisma.db.litellm_tooltable.find_unique.assert_awaited_with( - where={"tool_name": "my_tool"} - ) + prisma.db.litellm_tooltable.find_unique.assert_awaited_with(where={"tool_name": "my_tool"}) @pytest.mark.asyncio async def test_get_tools_by_names_returns_policy_map(): rows = [ - _mock_row( - tool_name="tool_a", input_policy="trusted", output_policy="untrusted" - ), - _mock_row( - tool_name="tool_b", input_policy="blocked", output_policy="untrusted" - ), + _mock_row(tool_name="tool_a", input_policy="trusted", output_policy="untrusted"), + _mock_row(tool_name="tool_b", input_policy="blocked", output_policy="untrusted"), ] prisma = _make_prisma(find_many_rows=rows) result = await get_tools_by_names(prisma, ["tool_a", "tool_b"]) @@ -285,9 +291,7 @@ async def test_get_tools_by_names_returns_policy_map(): "tool_a": ("trusted", "untrusted"), "tool_b": ("blocked", "untrusted"), } - prisma.db.litellm_tooltable.find_many.assert_awaited_once_with( - where={"tool_name": {"in": ["tool_a", "tool_b"]}} - ) + prisma.db.litellm_tooltable.find_many.assert_awaited_once_with(where={"tool_name": {"in": ["tool_a", "tool_b"]}}) @pytest.mark.asyncio @@ -382,12 +386,8 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read(): return [] mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock( - side_effect=_flaky_find_many - ) - mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( - return_value=[] - ) + mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(side_effect=_flaky_find_many) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(return_value=[]) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 @@ -397,10 +397,7 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_first_read(): assert len(invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "sync_tool_policy_from_db_tools_lookup_failure" - ) + assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_tools_lookup_failure" assert registry.is_initialized() @@ -421,9 +418,7 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read() mock_prisma_client = MagicMock() mock_prisma_client.db.litellm_tooltable.find_many = AsyncMock(return_value=[]) - mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock( - side_effect=_flaky_perms_find_many - ) + mock_prisma_client.db.litellm_objectpermissiontable.find_many = AsyncMock(side_effect=_flaky_perms_find_many) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 @@ -433,7 +428,142 @@ async def test_sync_tool_policy_from_db_retries_on_transport_error_second_read() assert len(perms_invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "sync_tool_policy_from_db_perms_lookup_failure" + assert reconnect_kwargs["reason"] == "sync_tool_policy_from_db_perms_lookup_failure" + + +class _InMemoryRedis: + def __init__(self) -> None: + self.store: dict[str, str] = {} + self.lists: dict[str, list[object]] = {} + + async def async_set_cache(self, key: str, value: str, nx: bool = False, **kwargs: object) -> bool | None: + if nx and key in self.store: + return None + self.store[key] = json.dumps(value) + return True + + async def async_get_cache(self, key: str, **kwargs: object) -> str | None: + raw = self.store.get(key) + return json.loads(raw) if raw is not None else None + + async def async_delete_cache(self, key: str) -> int: + return 1 if self.store.pop(key, None) is not None else 0 + + async def async_rpush_pipeline(self, rpush_list: list[RedisPipelineRpushOperation]) -> list[int]: + for op in rpush_list: + self.lists.setdefault(op["key"], []).extend(op["values"]) + return [len(self.lists[op["key"]]) for op in rpush_list] + + async def async_lpop_pipeline(self, lpop_list: list[RedisPipelineLpopOperation]) -> list[list[object] | None]: + return [self.lists.pop(op["key"], None) for op in lpop_list] + + +def _postgres_rejection(sqlstate: str) -> RawQueryError: + return RawQueryError( + data={"user_facing_error": {"error_code": "P2010", "meta": {"code": sqlstate, "message": "db error"}}} ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "lock_held_by_another_pod", + [None, DB_SPEND_UPDATE_JOB_NAME, TOOL_DISCOVERY_COMMIT_JOB_NAME], + ids=["free", "spend_lock_held", "tool_lock_held"], +) +async def test_commit_with_redis_writes_every_workers_tools_in_one_statement_from_the_tool_lock_holder_only( + lock_held_by_another_pod: str | None, +): + redis = _InMemoryRedis() + if lock_held_by_another_pod is not None: + assert await PodLockManager(redis_cache=redis).acquire_lock(cronjob_id=lock_held_by_another_pod) + other_worker = ToolDiscoveryQueue() + other_worker.add_update({"tool_name": "get_weather", "origin": "mcp_server"}) + other_worker.add_update({"tool_name": "search", "origin": "user_defined"}) + await RedisUpdateBuffer(redis_cache=redis).store_tool_discoveries_in_redis(other_worker) + + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(return_value=2) + prisma_client.db.litellm_tooltable.upsert = AsyncMock() + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=True, + ) + + assert queue.flush() == [], "every worker hands its tools to redis each tick" + prisma_client.db.litellm_tooltable.upsert.assert_not_awaited() + if lock_held_by_another_pod == TOOL_DISCOVERY_COMMIT_JOB_NAME: + prisma_client.db.execute_raw.assert_not_awaited() + assert len(redis.lists[REDIS_TOOL_DISCOVERY_BUFFER_KEY]) == 2, "non-holders leave their tools for the holder" + return + assert REDIS_TOOL_DISCOVERY_BUFFER_KEY not in redis.lists, "the spend lock must not gate the tool drain" + prisma_client.db.execute_raw.assert_awaited_once() + sql, *params = prisma_client.db.execute_raw.call_args.args + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') and 'ON CONFLICT ("tool_name")' in sql + rows = {params[i + 1]: params[i + 3] for i in range(0, len(params), 10)} + assert rows == {"get_weather": 3, "search": 1}, f"one row per tool summed across workers, got {params}" + + +@pytest.mark.asyncio +async def test_commit_without_redis_writes_only_this_workers_tools_in_one_statement_and_never_touches_redis(): + redis = _InMemoryRedis() + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + queue.add_update({"tool_name": "search", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(return_value=2) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=False, + ) + + assert redis.lists == {}, "without the buffer flag nothing goes through redis" + assert queue.flush() == [] + prisma_client.db.execute_raw.assert_awaited_once() + sql, *params = prisma_client.db.execute_raw.call_args.args + assert 'ON CONFLICT ("tool_name")' in sql + rows = {params[i + 1]: params[i + 3] for i in range(0, len(params), 10)} + assert rows == {"get_weather": 2, "search": 1}, f"got {params}" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "requeued"), + [ + (Exception("connection reset"), True), + (_postgres_rejection("23502"), False), + ], + ids=["connection_error_goes_back_to_redis", "rejected_row_is_dropped"], +) +async def test_failed_tool_table_write_requeues_the_rows_only_when_they_provably_did_not_apply( + failure: Exception, requeued: bool +): + redis = _InMemoryRedis() + queue = ToolDiscoveryQueue() + queue.add_update({"tool_name": "get_weather", "origin": "user_defined"}) + prisma_client = MagicMock() + prisma_client.db.execute_raw = AsyncMock(side_effect=failure) + + await commit_tool_discoveries( + prisma_client=prisma_client, + tool_discovery_queue=queue, + redis_update_buffer=RedisUpdateBuffer(redis_cache=redis), + pod_lock_manager=PodLockManager(redis_cache=redis), + use_redis_buffer=True, + ) + + prisma_client.db.execute_raw.assert_awaited_once() + left_in_redis = [json.loads(entry) for entry in redis.lists.get(REDIS_TOOL_DISCOVERY_BUFFER_KEY, [])] + expected = [[{"tool_name": "get_weather", "origin": "user_defined", "call_count": 1}]] if requeued else [] + assert left_in_redis == expected, f"requeued={requeued} but redis holds {left_in_redis}" diff --git a/tests/unit/proxy/db/test_tool_table_bulk_upsert.py b/tests/unit/proxy/db/test_tool_table_bulk_upsert.py new file mode 100644 index 00000000000..d1e79030f9d --- /dev/null +++ b/tests/unit/proxy/db/test_tool_table_bulk_upsert.py @@ -0,0 +1,38 @@ +"""Tests for the single-statement LiteLLM_ToolTable upsert (LIT-9997).""" + +from litellm.proxy.db.tool_table_bulk_upsert import build_tool_bulk_upsert + + +def test_one_row_per_item_and_the_conflict_update_only_moves_the_counters(): + sql, params = build_tool_bulk_upsert( + [ + {"tool_name": "tool_a", "origin": "mcp_server", "created_by": None, "call_count": 4}, + { + "tool_name": "tool_b", + "origin": "user_defined", + "created_by": "alice", + "call_count": 1, + "key_hash": "hash", + "team_id": "team", + "key_alias": "alias", + "user_agent": "agent", + }, + ] + ) + assert sql.startswith('INSERT INTO "LiteLLM_ToolTable"') + assert sql.count("$1::text") == 1 and "$11::text" in sql and "$21::text" not in sql + assert "$4::integer" in sql and "$14::integer" in sql + assert 'ON CONFLICT ("tool_name") DO UPDATE SET' in sql + assert '"call_count" = "LiteLLM_ToolTable"."call_count" + EXCLUDED."call_count"' in sql + assert '"last_used_at" = EXCLUDED."last_used_at"' in sql + assert '"input_policy"' not in sql and '"output_policy"' not in sql, "policies stay on the table defaults" + assert '"origin" =' not in sql and '"key_hash" =' not in sql, "an existing row keeps who discovered it" + assert len(params) == 20 + assert params[1:10] == ("tool_a", "mcp_server", 4, "system", "system", None, None, None, None) + assert params[11:20] == ("tool_b", "user_defined", 1, "alice", "alice", "hash", "team", "alias", "agent") + assert params[0] != params[10], "each row gets its own tool_id" + + +def test_defaults_fill_origin_created_by_and_call_count(): + _, params = build_tool_bulk_upsert([{"tool_name": "tool_a"}]) + assert params[1:6] == ("tool_a", "user_defined", 1, "system", "system")