mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(proxy): buffer tool discoveries in Redis and write LiteLLM_ToolTable in one statement per drain, on its own lock (#48089)
Backport of #48089 to lts/1.108.x.
Cherry-picked from e6a32bbeeb (main).
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu
kerry
parent
4bb0c86b2d
commit
6becdf304e
15 changed files with 583 additions and 209 deletions
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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({})
|
||||
|
|
|
|||
71
litellm/proxy/db/tool_table_bulk_upsert.py
Normal file
71
litellm/proxy/db/tool_table_bulk_upsert.py
Normal file
|
|
@ -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)))
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}]
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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}"
|
||||
|
|
|
|||
38
tests/unit/proxy/db/test_tool_table_bulk_upsert.py
Normal file
38
tests/unit/proxy/db/test_tool_table_bulk_upsert.py
Normal file
|
|
@ -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")
|
||||
Loading…
Add table
Reference in a new issue