diff --git a/litellm/constants.py b/litellm/constants.py index ee79f2fa56f..b1a0021bcc6 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -244,6 +244,7 @@ REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY = "litellm_daily_tag_spend_update_buffer MAX_REDIS_BUFFER_DEQUEUE_COUNT = int(os.getenv("MAX_REDIS_BUFFER_DEQUEUE_COUNT", 100)) # Bounds asyncio.Queue() instances (log queues, spend update queues, etc.) to prevent unbounded memory growth LITELLM_ASYNCIO_QUEUE_MAXSIZE = int(os.getenv("LITELLM_ASYNCIO_QUEUE_MAXSIZE", 1000)) +TOOL_POLICY_CACHE_TTL_SECONDS = int(os.getenv("TOOL_POLICY_CACHE_TTL_SECONDS", 60)) # Aggregation threshold: default to 80% of the asyncio queue maxsize so the check can always trigger. # Must be < LITELLM_ASYNCIO_QUEUE_MAXSIZE; if set higher the aggregation logic will never fire. MAX_SIZE_IN_MEMORY_QUEUE = int( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f354e28acd7..75b9f91acd9 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -183,6 +183,7 @@ class LitellmTableNames(str, enum.Enum): KEY_TABLE_NAME = "LiteLLM_VerificationToken" PROXY_MODEL_TABLE_NAME = "LiteLLM_ProxyModelTable" MANAGED_FILE_TABLE_NAME = "LiteLLM_ManagedFileTable" + TOOL_TABLE_NAME = "LiteLLM_ToolTable" class Litellm_EntityType(enum.Enum): @@ -4123,6 +4124,15 @@ class SpendUpdateQueueItem(TypedDict, total=False): response_cost: Optional[float] +class ToolDiscoveryQueueItem(TypedDict, total=False): + tool_name: str + origin: Optional[str] # MCP server name or "user_defined" + created_by: Optional[str] + key_hash: Optional[str] # hash of virtual key that triggered discovery + team_id: Optional[str] # team that triggered discovery + key_alias: Optional[str] # human-readable key alias + + class LiteLLM_ManagedFileTable(LiteLLMPydanticObjectBase): unified_file_id: str file_object: Optional[OpenAIFileObject] = None diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4f308952395..edf0cf0d397 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -13,7 +13,17 @@ import random import time import traceback from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + Dict, + List, + Literal, + Optional, + Union, + cast, + overload, +) import litellm from litellm._logging import verbose_proxy_logger @@ -23,18 +33,19 @@ from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.proxy._types import ( DB_CONNECTION_ERROR_TYPES, BaseDailySpendTransaction, - DailyTagSpendTransaction, - DailyOrganizationSpendTransaction, - DailyTeamSpendTransaction, - DailyEndUserSpendTransaction, - DailyUserSpendTransaction, DailyAgentSpendTransaction, + DailyEndUserSpendTransaction, + DailyOrganizationSpendTransaction, + DailyTagSpendTransaction, + DailyTeamSpendTransaction, + DailyUserSpendTransaction, DBSpendUpdateTransactions, Litellm_EntityType, LiteLLM_UserTable, SpendLogsMetadata, SpendLogsPayload, SpendUpdateQueueItem, + ToolDiscoveryQueueItem, ) from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( DailySpendUpdateQueue, @@ -42,6 +53,9 @@ from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( 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.spend_update_queue import SpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, +) from litellm.proxy.route_llm_request import ROUTE_ENDPOINT_MAPPING if TYPE_CHECKING: @@ -67,6 +81,7 @@ class DBSpendUpdateWriter: self.redis_update_buffer = RedisUpdateBuffer(redis_cache=self.redis_cache) self.pod_lock_manager = PodLockManager() self.spend_update_queue = SpendUpdateQueue() + self.tool_discovery_queue = ToolDiscoveryQueue() self.daily_spend_update_queue = DailySpendUpdateQueue() self.daily_team_spend_update_queue = DailySpendUpdateQueue() self.daily_end_user_spend_update_queue = DailySpendUpdateQueue() @@ -124,20 +139,53 @@ class DBSpendUpdateWriter: payload["startTime"] = payload["startTime"].isoformat() if isinstance(payload["endTime"], datetime): payload["endTime"] = payload["endTime"].isoformat() - + if org_id is not None and org_id != "": payload["organization_id"] = org_id if team_id is not None and team_id != "": payload["team_id"] = team_id - # One deepcopy shared by all 6 daily spend helpers (was 5, fixes agent bug) - payload_copy = copy.deepcopy(payload) + asyncio.create_task( + self._update_user_db( + response_cost=response_cost, + user_id=user_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + litellm_proxy_budget_name=litellm_proxy_budget_name, + end_user_id=end_user_id, + ) + ) + asyncio.create_task( + self._update_key_db( + response_cost=response_cost, + hashed_token=hashed_token, + prisma_client=prisma_client, + ) + ) + asyncio.create_task( + self._update_team_db( + response_cost=response_cost, + team_id=team_id, + user_id=user_id, + prisma_client=prisma_client, + ) + ) + asyncio.create_task( + self._update_org_db( + response_cost=response_cost, + org_id=org_id, + prisma_client=prisma_client, + ) + ) + asyncio.create_task( + self._update_tag_db( + response_cost=response_cost, + request_tags=copy.deepcopy(payload.get("request_tags")), + prisma_client=prisma_client, + ) + ) - # Deepcopy request_tags for _update_tag_db - request_tags = copy.deepcopy(payload.get("request_tags")) - - # Keep _insert_spend_log_to_db awaited inline (not a task, preserve current behavior) if disable_spend_logs is False: await self._insert_spend_log_to_db( payload=copy.deepcopy(payload), @@ -148,23 +196,54 @@ class DBSpendUpdateWriter: "disable_spend_logs=True. Skipping writing spend logs to db. Other spend updates - Key/User/Team table will still occur." ) - # Single task replaces 11 create_task() calls asyncio.create_task( - self._batch_database_updates( - response_cost=response_cost, - user_id=user_id, - hashed_token=hashed_token, - team_id=team_id, - org_id=org_id, - end_user_id=end_user_id, + self.add_spend_log_transaction_to_daily_user_transaction( + payload=copy.deepcopy(payload), prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - litellm_proxy_budget_name=litellm_proxy_budget_name, - payload_copy=payload_copy, - request_tags=request_tags, ) ) + asyncio.create_task( + self.add_spend_log_transaction_to_daily_end_user_transaction( + payload=copy.deepcopy(payload), + prisma_client=prisma_client, + ) + ) + + asyncio.create_task( + self.add_spend_log_transaction_to_daily_agent_transaction( + payload=payload, + prisma_client=prisma_client, + ) + ) + + asyncio.create_task( + self.add_spend_log_transaction_to_daily_team_transaction( + payload=copy.deepcopy(payload), + prisma_client=prisma_client, + ) + ) + asyncio.create_task( + self.add_spend_log_transaction_to_daily_org_transaction( + payload=copy.deepcopy(payload), + org_id=org_id, + prisma_client=prisma_client, + ) + ) + asyncio.create_task( + self.add_spend_log_transaction_to_daily_tag_transaction( + payload=copy.deepcopy(payload), + prisma_client=prisma_client, + ) + ) + + self._enqueue_tool_registry_upsert( + kwargs=kwargs, + completion_response=completion_response, + hashed_token=hashed_token, + team_id=team_id, + ) + verbose_proxy_logger.debug("Runs spend update on all tables") except Exception: verbose_proxy_logger.error( @@ -180,155 +259,102 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) - async def _batch_database_updates( + def _enqueue_tool_registry_upsert( self, - *, - response_cost: Optional[float], - user_id: Optional[str], - hashed_token: Optional[str], - team_id: Optional[str], - org_id: Optional[str], - end_user_id: Optional[str], - prisma_client: Optional[PrismaClient], - user_api_key_cache: DualCache, - litellm_proxy_budget_name: Optional[str], - payload_copy: dict, - request_tags: Optional[Any], - ): + kwargs: Optional[dict], + completion_response: Optional[Any], + hashed_token: Optional[str] = None, + team_id: Optional[str] = None, + ) -> None: """ - Runs all 11 spend-update helpers sequentially inside a single asyncio task. + Extract tool names from the LLM request and response and enqueue them + for upsert into LiteLLM_ToolTable via ToolDiscoveryQueue. - Each helper is wrapped in try/except so one failure doesn't prevent the others. + Handles four sources: + - MCP tools: standard_logging_object.mcp_tool_call_metadata.namespaced_tool_name + - Response tool_calls (OpenAI / Anthropic pass-through converted to OpenAI format): + completion_response.choices[].message.tool_calls[].function.name + - Request tools array (OpenAI format): kwargs["tools"][].function.name + - Request tools array (Anthropic /messages format): kwargs["passthrough_logging_payload"] + ["request_body"]["tools"][].name """ try: - await self._update_user_db( - response_cost=response_cost, - user_id=user_id, - prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, - litellm_proxy_budget_name=litellm_proxy_budget_name, - end_user_id=end_user_id, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_user_db failed: %s", - traceback.format_exc(), - ) + if kwargs is None: + return - try: - await self._update_key_db( - response_cost=response_cost, - hashed_token=hashed_token, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_key_db failed: %s", - traceback.format_exc(), - ) + # Extract key_alias from kwargs metadata if available + key_alias: Optional[str] = None + _litellm_params = kwargs.get("litellm_params") or {} + _metadata = _litellm_params.get("metadata") or {} + key_alias = _metadata.get("user_api_key_alias") or None - try: - await self._update_team_db( - response_cost=response_cost, - team_id=team_id, - user_id=user_id, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_team_db failed: %s", - traceback.format_exc(), - ) + def _enqueue(tool_name: str, origin: str = "user_defined") -> None: + self.tool_discovery_queue.add_update( + ToolDiscoveryQueueItem( + tool_name=tool_name, + origin=origin, + key_hash=hashed_token, + team_id=team_id, + key_alias=key_alias, + ) + ) - try: - await self._update_org_db( - response_cost=response_cost, - org_id=org_id, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_org_db failed: %s", - traceback.format_exc(), - ) + # --- MCP tool calls --- + sl_object = kwargs.get("standard_logging_object") + if sl_object is not None: + mcp_metadata = ( + sl_object.get("metadata", {}) or {} + ).get("mcp_tool_call_metadata") + if mcp_metadata and isinstance(mcp_metadata, dict): + tool_name = mcp_metadata.get("namespaced_tool_name") or mcp_metadata.get("name") + mcp_server_name = mcp_metadata.get("mcp_server_name") + if tool_name: + _enqueue(tool_name, origin=mcp_server_name or "user_defined") - try: - await self._update_tag_db( - response_cost=response_cost, - request_tags=request_tags, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: _update_tag_db failed: %s", - traceback.format_exc(), - ) + # --- Tools from request body (OpenAI format: tools[].function.name) --- + request_tools = kwargs.get("tools") or [] + for tool_def in request_tools: + if not isinstance(tool_def, dict): + continue + fn = tool_def.get("function") or {} + name = fn.get("name") if isinstance(fn, dict) else None + if name: + _enqueue(name) - try: - await self.add_spend_log_transaction_to_daily_user_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_user_transaction failed: %s", - traceback.format_exc(), - ) + # --- Tools from Anthropic /messages pass-through request body + # (Anthropic format: tools[].name, no "function" wrapper) --- + passthrough_payload = kwargs.get("passthrough_logging_payload") or {} + request_body = ( + passthrough_payload.get("request_body") + if isinstance(passthrough_payload, dict) + else None + ) or {} + for tool_def in request_body.get("tools") or []: + if not isinstance(tool_def, dict): + continue + name = tool_def.get("name") + if name: + _enqueue(name) - try: - await self.add_spend_log_transaction_to_daily_end_user_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: + # --- Response tool_calls (OpenAI format; Anthropic pass-through converts tool_use here) --- + if completion_response is not None and hasattr(completion_response, "choices"): + for choice in completion_response.choices or []: + message = getattr(choice, "message", None) + if message is None: + continue + tool_calls = getattr(message, "tool_calls", None) + if not tool_calls: + continue + for tc in tool_calls: + fn = getattr(tc, "function", None) + if fn is None: + continue + tool_name = getattr(fn, "name", None) + if tool_name: + _enqueue(tool_name) + except Exception as e: verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_end_user_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_agent_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_agent_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_team_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_team_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_org_transaction( - payload=payload_copy, - org_id=org_id, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_org_transaction failed: %s", - traceback.format_exc(), - ) - - try: - await self.add_spend_log_transaction_to_daily_tag_transaction( - payload=payload_copy, - prisma_client=prisma_client, - ) - except Exception: - verbose_proxy_logger.debug( - "_batch_database_updates: add_spend_log_transaction_to_daily_tag_transaction failed: %s", - traceback.format_exc(), + "_enqueue_tool_registry_upsert error (non-blocking): %s", e ) async def _update_key_db( @@ -846,6 +872,25 @@ class DBSpendUpdateWriter: daily_spend_transactions=daily_agent_spend_update_transactions, ) + ################## Tool Registry Upserts ################## + await self._flush_tool_discovery_queue(prisma_client=prisma_client) + + 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 = 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 + ) + async def _commit_spend_updates_to_db( # noqa: PLR0915 self, prisma_client: PrismaClient, @@ -1027,7 +1072,7 @@ class DBSpendUpdateWriter: team_id = key.split("::")[1] user_id = key.split("::")[3] team_memberships_to_invalidate.append((user_id, team_id)) - + for i in range(n_retry_times + 1): start_time = time.time() try: @@ -1064,13 +1109,11 @@ class DBSpendUpdateWriter: _raise_failed_update_spend_exception( e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj ) - + # Invalidate cache for updated team memberships # This ensures budget checks read fresh spend data from the database if team_memberships_to_invalidate and proxy_logging_obj is not None: - user_api_key_cache = proxy_logging_obj.call_details.get( - "user_api_key_cache" - ) + user_api_key_cache = proxy_logging_obj.call_details.get("user_api_key_cache") if user_api_key_cache is not None: for user_id, team_id in team_memberships_to_invalidate: cache_key = "team_membership:{}:{}".format(user_id, team_id) @@ -1382,9 +1425,7 @@ class DBSpendUpdateWriter: ), "endpoint": transaction.get("endpoint") or "", "prompt_tokens": transaction["prompt_tokens"], - "completion_tokens": transaction[ - "completion_tokens" - ], + "completion_tokens": transaction["completion_tokens"], "spend": transaction["spend"], "api_requests": transaction["api_requests"], "successful_requests": transaction[ @@ -1395,14 +1436,12 @@ class DBSpendUpdateWriter: # Add cache-related fields if they exist if "cache_read_input_tokens" in transaction: - common_data[ - "cache_read_input_tokens" - ] = transaction.get("cache_read_input_tokens", 0) + common_data["cache_read_input_tokens"] = ( + transaction.get("cache_read_input_tokens", 0) + ) if "cache_creation_input_tokens" in transaction: - common_data[ - "cache_creation_input_tokens" - ] = transaction.get( - "cache_creation_input_tokens", 0 + common_data["cache_creation_input_tokens"] = ( + transaction.get("cache_creation_input_tokens", 0) ) if entity_type == "tag" and "request_id" in transaction: @@ -1445,14 +1484,10 @@ class DBSpendUpdateWriter: } if entity_type == "tag" and "request_id" in transaction: - update_data["request_id"] = transaction.get( - "request_id" - ) + update_data["request_id"] = transaction.get("request_id") # Add endpoint to update_data so existing rows get their endpoint field updated - update_data["endpoint"] = ( - transaction.get("endpoint") or "" - ) + update_data["endpoint"] = transaction.get("endpoint") or "" table.upsert( where=where_clause, @@ -1636,9 +1671,7 @@ class DBSpendUpdateWriter: self, payload: Union[dict, SpendLogsPayload], prisma_client: PrismaClient, - type: Literal[ - "user", "team", "org", "request_tags", "end_user", "agent" - ] = "user", + type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user", ) -> Optional[BaseDailySpendTransaction]: common_expected_keys = ["startTime", "api_key"] if type == "user": @@ -1697,7 +1730,7 @@ class DBSpendUpdateWriter: endpoint = None if call_type: endpoint = ROUTE_ENDPOINT_MAPPING.get(call_type, None) - + daily_transaction = BaseDailySpendTransaction( date=date, api_key=payload["api_key"], @@ -1909,7 +1942,7 @@ class DBSpendUpdateWriter: endpoint_str = base_daily_transaction.get("endpoint") or "" daily_transaction_key = f"{payload['agent_id']}_{base_daily_transaction['date']}_{payload_with_agent_id['api_key']}_{payload_with_agent_id['model']}_{payload_with_agent_id['custom_llm_provider']}_{endpoint_str}" daily_transaction = DailyAgentSpendTransaction( - agent_id=payload["agent_id"], **base_daily_transaction + agent_id=payload['agent_id'], **base_daily_transaction ) await self.daily_agent_spend_update_queue.add_update( update={daily_transaction_key: daily_transaction} diff --git a/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py new file mode 100644 index 00000000000..16a3ada40f2 --- /dev/null +++ b/litellm/proxy/db/db_transaction_queue/tool_discovery_queue.py @@ -0,0 +1,54 @@ +""" +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. +""" + +from typing import List, Set + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ToolDiscoveryQueueItem + + +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. + """ + + def __init__(self) -> None: + self._seen_tool_names: Set[str] = set() + self._pending: List[ToolDiscoveryQueueItem] = [] + + def add_update(self, item: ToolDiscoveryQueueItem) -> None: + """Enqueue a tool discovery item if tool_name has not been seen before.""" + tool_name = 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) + verbose_proxy_logger.debug( + "ToolDiscoveryQueue: queued new tool %s (origin=%s)", + tool_name, + 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 items diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py new file mode 100644 index 00000000000..4e0a8095a08 --- /dev/null +++ b/litellm/proxy/db/tool_registry_writer.py @@ -0,0 +1,179 @@ +""" +DB helpers for LiteLLM_ToolTable — the global tool registry. + +Tools are auto-discovered from LLM responses and upserted here. +Admins use the management endpoints to read and update call_policy. + +NOTE: Uses raw SQL (query_raw / execute_raw) instead of Prisma model methods +because the generated Prisma Python client may not have LiteLLM_ToolTable +when running against an older generated schema. +""" + +import uuid +from datetime import datetime, timezone +from typing import TYPE_CHECKING, Dict, List, Optional + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import ToolDiscoveryQueueItem +from litellm.types.tool_management import LiteLLM_ToolTableRow, ToolCallPolicy + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + + +def _row_to_model(row: dict) -> LiteLLM_ToolTableRow: + return LiteLLM_ToolTableRow( + tool_id=row.get("tool_id", ""), + tool_name=row.get("tool_name", ""), + origin=row.get("origin"), + call_policy=row.get("call_policy", "untrusted"), + call_count=int(row.get("call_count") or 0), + assignments=row.get("assignments"), + key_hash=row.get("key_hash"), + team_id=row.get("team_id"), + key_alias=row.get("key_alias"), + created_at=row.get("created_at"), + updated_at=row.get("updated_at"), + created_by=row.get("created_by"), + updated_by=row.get("updated_by"), + ) + + +async def batch_upsert_tools( + prisma_client: "PrismaClient", + items: List[ToolDiscoveryQueueItem], +) -> None: + """ + Batch-upsert tool registry rows via raw SQL. + + On first insert: sets call_policy = "untrusted" (schema default), call_count = 1. + On conflict: increments call_count; preserves existing call_policy. + """ + if not items: + return + try: + data = [item for item in items if item.get("tool_name")] + if not data: + return + 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") + now = datetime.now(timezone.utc).isoformat() + await prisma_client.db.execute_raw( + 'INSERT INTO "LiteLLM_ToolTable" ' + "(tool_id, tool_name, origin, call_policy, call_count, created_by, updated_by, key_hash, team_id, key_alias, created_at, updated_at) " + "VALUES ($7, $1, $2, 'untrusted', 1, $3, $3, $4, $5, $6, $8, $8) " + "ON CONFLICT (tool_name) DO UPDATE SET " + "call_count = \"LiteLLM_ToolTable\".call_count + 1, " + "updated_at = $8", + tool_name, + origin, + created_by, + key_hash, + team_id, + key_alias, + str(uuid.uuid4()), + now, + ) + 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) + + +async def list_tools( + prisma_client: "PrismaClient", + call_policy: Optional[ToolCallPolicy] = None, +) -> List[LiteLLM_ToolTableRow]: + """Return all tools, optionally filtered by call_policy.""" + try: + if call_policy is not None: + rows = await prisma_client.db.query_raw( + 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' + 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' + 'FROM "LiteLLM_ToolTable" WHERE call_policy = $1 ORDER BY created_at DESC', + call_policy, + ) + else: + rows = await prisma_client.db.query_raw( + 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' + 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' + 'FROM "LiteLLM_ToolTable" ORDER BY created_at DESC', + ) + return [_row_to_model(row) for row in rows] + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer list_tools error: %s", e) + return [] + + +async def get_tool( + prisma_client: "PrismaClient", + tool_name: str, +) -> Optional[LiteLLM_ToolTableRow]: + """Return a single tool row by tool_name.""" + try: + rows = await prisma_client.db.query_raw( + 'SELECT tool_id, tool_name, origin, call_policy, call_count, assignments, ' + 'key_hash, team_id, key_alias, created_at, updated_at, created_by, updated_by ' + 'FROM "LiteLLM_ToolTable" WHERE tool_name = $1', + tool_name, + ) + if not rows: + return None + return _row_to_model(rows[0]) + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer get_tool error: %s", e) + return None + + +async def update_tool_policy( + prisma_client: "PrismaClient", + tool_name: str, + call_policy: ToolCallPolicy, + updated_by: Optional[str], +) -> Optional[LiteLLM_ToolTableRow]: + """Update the call_policy for a tool. Upserts the row if it does not exist yet.""" + try: + _updated_by = updated_by or "system" + now = datetime.now(timezone.utc).isoformat() + await prisma_client.db.execute_raw( + 'INSERT INTO "LiteLLM_ToolTable" (tool_id, tool_name, call_policy, created_by, updated_by, created_at, updated_at) ' + "VALUES ($4, $1, $2, $3, $3, $5, $5) " + "ON CONFLICT (tool_name) DO UPDATE SET call_policy = $2, updated_by = $3, updated_at = $5", + tool_name, + call_policy, + _updated_by, + str(uuid.uuid4()), + now, + ) + return await get_tool(prisma_client, tool_name) + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer update_tool_policy error: %s", e) + return None + + +async def get_tools_by_names( + prisma_client: "PrismaClient", + tool_names: List[str], +) -> Dict[str, str]: + """ + Return a {tool_name: call_policy} map for the given tool names. + Used by the policy enforcement guardrail — single batch query, never N+1. + """ + if not tool_names: + return {} + try: + placeholders = ", ".join(f"${i+1}" for i in range(len(tool_names))) + rows = await prisma_client.db.query_raw( + f'SELECT tool_name, call_policy FROM "LiteLLM_ToolTable" WHERE tool_name IN ({placeholders})', + *tool_names, + ) + return {row["tool_name"]: row["call_policy"] for row in rows} + except Exception as e: + verbose_proxy_logger.error("tool_registry_writer get_tools_by_names error: %s", e) + return {} diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py new file mode 100644 index 00000000000..5a43006e23c --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/__init__.py @@ -0,0 +1,16 @@ +import litellm +from litellm.types.guardrails import Guardrail, LitellmParams + + +def initialize_guardrail(litellm_params: LitellmParams, guardrail: Guardrail): + from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import ( + ToolPolicyGuardrail, + ) + + _callback = ToolPolicyGuardrail( + guardrail_name=guardrail.get("guardrail_name", "tool_policy"), + event_hook=litellm_params.mode, + default_on=litellm_params.default_on, + ) + litellm.logging_callback_manager.add_litellm_callback(_callback) + return _callback diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py new file mode 100644 index 00000000000..87558566c42 --- /dev/null +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -0,0 +1,163 @@ +""" +Tool Policy Guardrail + +Reads call_policy from LiteLLM_ToolTable and enforces it on LLM requests/responses. + +Policy values: + "trusted" - allow through (no action) + "untrusted" - allow through (no action; default for newly discovered tools) + "blocked" - raise HTTPException, preventing the tool call + "dual_llm" - (Phase 3) send to second LLM for verification; currently treated as allowed + +Configuration in proxy config YAML: + guardrails: + - guardrail_name: "tool_policy" + litellm_params: + guardrail: tool_policy + mode: post_call + +or both pre and post call: + - guardrail_name: "tool_policy" + litellm_params: + guardrail: tool_policy + mode: during_call # runs before LLM and on response +""" + +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional + +from fastapi import HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.constants import TOOL_POLICY_CACHE_TTL_SECONDS +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + log_guardrail_information, +) +from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import GenericGuardrailAPIInputs + +if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + +GUARDRAIL_NAME = "tool_policy" + + +class ToolPolicyGuardrail(CustomGuardrail): + """ + Guardrail that enforces per-tool call policies stored in LiteLLM_ToolTable. + + Tools with call_policy="blocked" are rejected before/after the LLM call. + Tools with call_policy="trusted" or "untrusted" pass through unchanged. + """ + + def __init__(self, **kwargs: Any) -> None: + if "supported_event_hooks" not in kwargs: + kwargs["supported_event_hooks"] = [ + GuardrailEventHooks.pre_call, + GuardrailEventHooks.post_call, + GuardrailEventHooks.during_call, + ] + super().__init__(**kwargs) + self._policy_cache: DualCache = DualCache() + + @log_guardrail_information + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional["LiteLLMLoggingObj"] = None, + ) -> GenericGuardrailAPIInputs: + """ + Enforce tool policies on both request tools and response tool_calls. + + - input_type="request": check inputs["tools"] (tool definitions in the LLM request) + - input_type="response": check inputs["tool_calls"] (tool_calls in the LLM response) + + Raises HTTPException (400) if any tool is "blocked". + """ + if input_type == "request": + tools = inputs.get("tools") or [] + tool_names = [ + t["function"]["name"] + for t in tools + if isinstance(t, dict) + and isinstance(t.get("function"), dict) + and t["function"].get("name") + ] + else: # response + tool_calls = inputs.get("tool_calls") or [] + tool_names = [] + for tc in tool_calls: + fn = None + if isinstance(tc, dict): + fn = (tc.get("function") or {}).get("name") + elif hasattr(tc, "function"): + fn = getattr(tc.function, "name", None) + if fn: + tool_names.append(fn) + + if not tool_names: + return inputs + + policy_map = await self._get_policies_cached(tool_names) + + blocked = [name for name in tool_names if policy_map.get(name) == "blocked"] + if blocked: + verbose_proxy_logger.warning( + "ToolPolicyGuardrail: blocking tool(s) %s (policy=blocked)", blocked + ) + raise HTTPException( + status_code=400, + detail={ + "error": "Violated tool policy", + "blocked_tools": blocked, + "message": f"Tool(s) {blocked} are blocked by policy.", + }, + ) + + return inputs + + async def _get_policies_cached(self, tool_names: List[str]) -> Dict[str, str]: + """ + Batch-fetch call_policy for the given tool names. + + Caches per individual tool name (not per combination) so that adding + a new tool to a request doesn't invalidate the cached policies for all + the other tools already in the cache. + """ + from litellm.proxy.db.tool_registry_writer import get_tools_by_names + from litellm.proxy.proxy_server import prisma_client + + if not tool_names or prisma_client is None: + return {} + + result: Dict[str, str] = {} + cache_misses: List[str] = [] + + for name in tool_names: + cached = await self._policy_cache.async_get_cache(f"tool_policy:{name}") + if cached is not None and isinstance(cached, str): + result[name] = cached + else: + cache_misses.append(name) + + if cache_misses: + fetched = await get_tools_by_names( + prisma_client=prisma_client, tool_names=cache_misses + ) + for name, policy in fetched.items(): + result[name] = policy + await self._policy_cache.async_set_cache( + key=f"tool_policy:{name}", + value=policy, + ttl=TOOL_POLICY_CACHE_TTL_SECONDS, + ) + verbose_proxy_logger.debug( + "ToolPolicyGuardrail: fetched %d policies from DB (cache hits: %d)", + len(cache_misses), + len(tool_names) - len(cache_misses), + ) + + return result diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py new file mode 100644 index 00000000000..89880c9a4ec --- /dev/null +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -0,0 +1,149 @@ +""" +TOOL POLICY MANAGEMENT + +All /tool management endpoints + +GET /v1/tool/list - List all discovered tools and their policies +GET /v1/tool/{tool_name} - Get a single tool's details +POST /v1/tool/policy - Update the call_policy for a tool +""" + +from typing import Optional + +from fastapi import APIRouter, Depends, HTTPException + +from litellm._logging import verbose_proxy_logger +from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.types.tool_management import ( + LiteLLM_ToolTableRow, + ToolCallPolicy, + ToolListResponse, + ToolPolicyUpdateRequest, + ToolPolicyUpdateResponse, +) + +router = APIRouter() + + +@router.get( + "/v1/tool/list", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolListResponse, +) +async def list_tools( + call_policy: Optional[ToolCallPolicy] = None, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + List all auto-discovered tools and their call policies. + + Parameters: + - call_policy: Optional filter — one of "trusted", "untrusted", "dual_llm", "blocked" + """ + from litellm.proxy.db.tool_registry_writer import list_tools as db_list_tools + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + tools = await db_list_tools(prisma_client=prisma_client, call_policy=call_policy) + return ToolListResponse(tools=tools, total=len(tools)) + except Exception as e: + verbose_proxy_logger.exception("Error listing tools: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.get( + "/v1/tool/{tool_name:path}", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=LiteLLM_ToolTableRow, +) +async def get_tool( + tool_name: str, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Get details for a single tool. + + Parameters: + - tool_name: The tool name (supports namespaced names with slashes) + """ + from litellm.proxy.db.tool_registry_writer import get_tool as db_get_tool + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + tool = await db_get_tool(prisma_client=prisma_client, tool_name=tool_name) + if tool is None: + raise HTTPException( + status_code=404, detail=f"Tool '{tool_name}' not found" + ) + return tool + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error getting tool: %s", e) + raise HTTPException(status_code=500, detail=str(e)) + + +@router.post( + "/v1/tool/policy", + tags=["tool management"], + dependencies=[Depends(user_api_key_auth)], + response_model=ToolPolicyUpdateResponse, +) +async def update_tool_policy( + data: ToolPolicyUpdateRequest, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Set the call policy for a tool. + + Parameters: + - tool_name: str - The tool to update + - call_policy: "trusted" | "untrusted" | "dual_llm" | "blocked" + + Setting a tool to "blocked" will cause the ToolPolicyGuardrail to remove + that tool_call from LLM responses before returning them to the client. + """ + from litellm.proxy.db.tool_registry_writer import ( + update_tool_policy as db_update_tool_policy, + ) + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise HTTPException( + status_code=500, detail=CommonProxyErrors.db_not_connected_error.value + ) + + try: + updated = await db_update_tool_policy( + prisma_client=prisma_client, + tool_name=data.tool_name, + call_policy=data.call_policy, + updated_by=user_api_key_dict.user_id, + ) + if updated is None: + raise HTTPException( + status_code=500, detail=f"Failed to update policy for tool '{data.tool_name}'" + ) + return ToolPolicyUpdateResponse( + tool_name=updated.tool_name, + call_policy=updated.call_policy, + updated=True, + ) + except HTTPException: + raise + except Exception as e: + verbose_proxy_logger.exception("Error updating tool policy: %s", e) + raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 77cb12f608e..f3c309f2998 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -410,6 +410,9 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team, validate_membership, ) +from litellm.proxy.management_endpoints.tool_management_endpoints import ( + router as tool_management_router, +) from litellm.proxy.management_endpoints.ui_sso import ( get_disabled_non_admin_personal_key_creation, ) @@ -12882,6 +12885,7 @@ app.include_router(budget_management_router) app.include_router(model_management_router) app.include_router(model_access_group_management_router) app.include_router(tag_management_router) +app.include_router(tool_management_router) app.include_router(cost_tracking_settings_router) app.include_router(router_settings_router) app.include_router(fallback_management_router) diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 50c0a55a875..23917cf7c7f 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1051,6 +1051,26 @@ model LiteLLM_PolicyAttachmentTable { updated_by String? } +// Global tool registry - auto-discovered from LLM responses; admins set call_policy here +model LiteLLM_ToolTable { + tool_id String @id @default(uuid()) + tool_name String @unique // e.g. "huggingface_remote-mcp__dynamic_space" + origin String? // MCP server name or "user_defined" + call_policy String @default("untrusted") // "trusted" | "untrusted" | "dual_llm" | "blocked" + call_count Int @default(0) // cumulative number of times this tool was seen + assignments Json? @default("{}") + key_hash String? // hash of the virtual key that first called this tool + team_id String? // team that first called this tool + key_alias String? // human-readable alias of the virtual key + created_at DateTime @default(now()) + created_by String? + updated_at DateTime @default(now()) @updatedAt + updated_by String? + + @@index([call_policy]) + @@index([team_id]) +} + //Unified Access Groups table for storing unified access groups model LiteLLM_AccessGroupTable { access_group_id String @id @default(uuid()) diff --git a/litellm/types/tool_management.py b/litellm/types/tool_management.py new file mode 100644 index 00000000000..8704ff27759 --- /dev/null +++ b/litellm/types/tool_management.py @@ -0,0 +1,42 @@ +""" +Pydantic models for Tool Policy management endpoints. +""" + +from datetime import datetime +from typing import Dict, List, Literal, Optional + +from pydantic import BaseModel + +ToolCallPolicy = Literal["trusted", "untrusted", "dual_llm", "blocked"] + + +class LiteLLM_ToolTableRow(BaseModel): + tool_id: str + tool_name: str + origin: Optional[str] = None + call_policy: ToolCallPolicy = "untrusted" + call_count: int = 0 + assignments: Optional[Dict] = None + key_hash: Optional[str] = None + team_id: Optional[str] = None + key_alias: Optional[str] = None + created_at: Optional[datetime] = None + updated_at: Optional[datetime] = None + created_by: Optional[str] = None + updated_by: Optional[str] = None + + +class ToolListResponse(BaseModel): + tools: List[LiteLLM_ToolTableRow] + total: int + + +class ToolPolicyUpdateRequest(BaseModel): + tool_name: str + call_policy: ToolCallPolicy + + +class ToolPolicyUpdateResponse(BaseModel): + tool_name: str + call_policy: ToolCallPolicy + updated: bool diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py new file mode 100644 index 00000000000..defdb3834d8 --- /dev/null +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_tool_discovery_queue.py @@ -0,0 +1,75 @@ +""" +Unit tests for ToolDiscoveryQueue. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.db_transaction_queue.tool_discovery_queue import ( + ToolDiscoveryQueue, +) + + +@pytest.fixture +def queue(): + return ToolDiscoveryQueue() + + +def test_add_single_tool(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" + + +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_deduplication_different_names(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" + + +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 + + +def test_flush_returns_list(queue): + result = queue.flush() + assert isinstance(result, list) diff --git a/tests/test_litellm/proxy/db/test_tool_registry_writer.py b/tests/test_litellm/proxy/db/test_tool_registry_writer.py new file mode 100644 index 00000000000..44f9e32058a --- /dev/null +++ b/tests/test_litellm/proxy/db/test_tool_registry_writer.py @@ -0,0 +1,197 @@ +""" +Unit tests for tool_registry_writer.py — uses a mock prisma client +that exposes execute_raw / query_raw (matching the actual raw-SQL implementation). +""" + +import os +import sys +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.db.tool_registry_writer import ( + batch_upsert_tools, + get_tool, + get_tools_by_names, + list_tools, + update_tool_policy, +) + + +def _make_prisma(query_rows=None): + """Return a minimal mock prisma_client with execute_raw / query_raw.""" + default_row = { + "tool_id": "uuid-1", + "tool_name": "my_tool", + "origin": "user_defined", + "call_policy": "untrusted", + "call_count": 1, + "assignments": {}, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": None, + } + rows = query_rows if query_rows is not None else [default_row] + + prisma = MagicMock() + prisma.db.execute_raw = AsyncMock(return_value=None) + prisma.db.query_raw = AsyncMock(return_value=rows) + return prisma + + +@pytest.mark.asyncio +async def test_batch_upsert_tools_calls_execute_raw(): + prisma = _make_prisma() + items = [{"tool_name": "tool_a", "origin": "mcp_server", "created_by": None}] + await batch_upsert_tools(prisma, items) + prisma.db.execute_raw.assert_awaited_once() + call_args = prisma.db.execute_raw.call_args + sql = call_args.args[0] + assert "LiteLLM_ToolTable" in sql + assert "ON CONFLICT" in sql + + +@pytest.mark.asyncio +async def test_batch_upsert_tools_empty_list(): + prisma = _make_prisma() + await batch_upsert_tools(prisma, []) + prisma.db.execute_raw.assert_not_awaited() + + +@pytest.mark.asyncio +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.execute_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_batch_upsert_multiple_tools_calls_execute_raw_per_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"}, + ] + await batch_upsert_tools(prisma, items) + assert prisma.db.execute_raw.await_count == 2 + + +@pytest.mark.asyncio +async def test_list_tools_no_filter(): + row = { + "tool_id": "id1", + "tool_name": "tool_a", + "origin": "mcp", + "call_policy": "untrusted", + "call_count": 5, + "assignments": {}, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": None, + } + prisma = _make_prisma(query_rows=[row]) + result = await list_tools(prisma) + assert len(result) == 1 + assert result[0].tool_name == "tool_a" + assert result[0].call_count == 5 + prisma.db.query_raw.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_list_tools_with_policy_filter(): + row = { + "tool_id": "id1", + "tool_name": "blocked_tool", + "origin": None, + "call_policy": "blocked", + "call_count": 2, + "assignments": None, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": None, + } + prisma = _make_prisma(query_rows=[row]) + result = await list_tools(prisma, call_policy="blocked") + assert result[0].call_policy == "blocked" + call_args = prisma.db.query_raw.call_args + sql = call_args.args[0] + assert "WHERE call_policy" in sql + + +@pytest.mark.asyncio +async def test_get_tool_found(): + prisma = _make_prisma() + result = await get_tool(prisma, "my_tool") + assert result is not None + assert result.tool_name == "my_tool" + prisma.db.query_raw.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_get_tool_not_found(): + prisma = _make_prisma(query_rows=[]) + result = await get_tool(prisma, "nonexistent") + assert result is None + + +@pytest.mark.asyncio +async def test_update_tool_policy_calls_execute_raw(): + row = { + "tool_id": "uuid-1", + "tool_name": "my_tool", + "origin": "user_defined", + "call_policy": "blocked", + "call_count": 1, + "assignments": {}, + "key_hash": None, + "team_id": None, + "key_alias": None, + "created_at": datetime.now(timezone.utc), + "updated_at": datetime.now(timezone.utc), + "created_by": None, + "updated_by": "admin", + } + prisma = _make_prisma(query_rows=[row]) + result = await update_tool_policy(prisma, "my_tool", "blocked", "admin") + assert result is not None + assert result.call_policy == "blocked" + prisma.db.execute_raw.assert_awaited_once() + call_args = prisma.db.execute_raw.call_args + sql = call_args.args[0] + assert "ON CONFLICT" in sql + assert "call_policy" in sql + + +@pytest.mark.asyncio +async def test_get_tools_by_names_returns_policy_map(): + rows = [ + {"tool_name": "tool_a", "call_policy": "trusted"}, + {"tool_name": "tool_b", "call_policy": "blocked"}, + ] + prisma = _make_prisma(query_rows=rows) + result = await get_tools_by_names(prisma, ["tool_a", "tool_b"]) + assert result == {"tool_a": "trusted", "tool_b": "blocked"} + + +@pytest.mark.asyncio +async def test_get_tools_by_names_empty_list(): + prisma = _make_prisma() + result = await get_tools_by_names(prisma, []) + assert result == {} + prisma.db.query_raw.assert_not_awaited() diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py new file mode 100644 index 00000000000..c6a81efbf0b --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py @@ -0,0 +1,181 @@ +""" +Unit tests for ToolPolicyGuardrail. +""" + +import os +import sys +from typing import Any +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../../../../../..")) + +from litellm.proxy.guardrails.guardrail_hooks.tool_policy.tool_policy_guardrail import ( + ToolPolicyGuardrail, +) +from litellm.types.guardrails import GuardrailEventHooks + + +@pytest.fixture +def guardrail(): + return ToolPolicyGuardrail() + + +# --- helpers --- + +def _tool_request_inputs(tool_names: list) -> dict: + return { + "tools": [ + {"type": "function", "function": {"name": name, "description": ""}} + for name in tool_names + ] + } + + +def _tool_response_inputs(tool_names: list) -> dict: + return { + "tool_calls": [ + {"type": "function", "function": {"name": name}} + for name in tool_names + ] + } + + +# --- tests --- + + +def test_guardrail_supports_pre_and_post_call(guardrail): + hooks = guardrail.supported_event_hooks + assert GuardrailEventHooks.pre_call in hooks + assert GuardrailEventHooks.post_call in hooks + + +@pytest.mark.asyncio +async def test_no_tools_in_request_passes_through(guardrail): + inputs: Any = {"tools": []} + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_no_tool_calls_in_response_passes_through(guardrail): + inputs: Any = {"tool_calls": []} + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_untrusted_tools_pass_through(guardrail): + policy_map = {"search": "untrusted", "read_file": "trusted"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_request_inputs(["search", "read_file"]) + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_blocked_tool_in_request_raises_http_exception(guardrail): + policy_map = {"dangerous_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_request_inputs(["dangerous_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert exc_info.value.status_code == 400 + assert "dangerous_tool" in exc_info.value.detail["blocked_tools"] + + +@pytest.mark.asyncio +async def test_blocked_tool_in_response_raises_http_exception(guardrail): + policy_map = {"exfil_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_response_inputs(["exfil_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) + assert exc_info.value.status_code == 400 + assert "exfil_tool" in exc_info.value.detail["blocked_tools"] + + +@pytest.mark.asyncio +async def test_mixed_blocked_and_allowed_raises_for_blocked(guardrail): + policy_map = {"safe_tool": "trusted", "bad_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + inputs: Any = _tool_request_inputs(["safe_tool", "bad_tool"]) + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + blocked = exc_info.value.detail["blocked_tools"] + assert "bad_tool" in blocked + assert "safe_tool" not in blocked + + +@pytest.mark.asyncio +async def test_tool_not_in_db_passes_through(guardrail): + """Tools not found in the DB (no entry) should not be blocked.""" + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value={})): + inputs: Any = _tool_request_inputs(["unknown_tool"]) + result = await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="request" + ) + assert result is inputs + + +@pytest.mark.asyncio +async def test_get_policies_cached_uses_cache(guardrail): + """Second call with same tool names should return the cached result.""" + policy_map = {"tool_a": "trusted"} + with patch( + "litellm.proxy.db.tool_registry_writer.get_tools_by_names", + new=AsyncMock(return_value=policy_map), + ) as mock_db, patch( + "litellm.proxy.proxy_server.prisma_client", + new=MagicMock(), + ): + # first call — should hit DB + result1 = await guardrail._get_policies_cached(["tool_a"]) + assert result1 == policy_map + + # second call — should hit cache, not DB again + result2 = await guardrail._get_policies_cached(["tool_a"]) + assert result2 == policy_map + + assert mock_db.call_count == 1 + + +@pytest.mark.asyncio +async def test_get_policies_cached_no_prisma(guardrail): + """Without a prisma client, returns empty dict.""" + with patch( + "litellm.proxy.proxy_server.prisma_client", + None, + ): + result = await guardrail._get_policies_cached(["tool_a"]) + assert result == {} + + +@pytest.mark.asyncio +async def test_response_tool_calls_as_objects(guardrail): + """tool_calls that are objects (not dicts) with .function.name should work.""" + policy_map = {"obj_tool": "blocked"} + with patch.object(guardrail, "_get_policies_cached", new=AsyncMock(return_value=policy_map)): + fn = MagicMock() + fn.name = "obj_tool" + tc = MagicMock() + tc.function = fn + inputs: Any = {"tool_calls": [tc]} + with pytest.raises(HTTPException): + await guardrail.apply_guardrail( + inputs=inputs, request_data={}, input_type="response" + ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py new file mode 100644 index 00000000000..6f1d373fdee --- /dev/null +++ b/tests/test_litellm/proxy/management_endpoints/test_tool_management_endpoints.py @@ -0,0 +1,149 @@ +""" +Unit tests for tool management endpoints (/v1/tool/*). +Uses FastAPI TestClient with mocked DB functions. + +Patches target the source modules (litellm.proxy.db.tool_registry_writer.* +and litellm.proxy.proxy_server.prisma_client) because the endpoint code +imports these inside function bodies to avoid circular imports. +""" + +import os +import sys +from datetime import datetime, timezone +from typing import Optional +from unittest.mock import AsyncMock, MagicMock, patch + +from fastapi import FastAPI +from fastapi.testclient import TestClient + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.proxy.management_endpoints.tool_management_endpoints import router +from litellm.types.tool_management import LiteLLM_ToolTableRow + +# --- helpers --- + + +def _make_tool_row( + tool_name: str = "my_tool", + call_policy: str = "untrusted", + origin: Optional[str] = None, +) -> LiteLLM_ToolTableRow: + now = datetime.now(timezone.utc) + return LiteLLM_ToolTableRow( + tool_id="uuid-1", + tool_name=tool_name, + origin=origin, + call_policy=call_policy, # type: ignore[arg-type] + assignments={}, + created_at=now, + updated_at=now, + ) + + +def _make_app() -> FastAPI: + """Build a minimal FastAPI app with the tool management router.""" + app = FastAPI() + app.include_router(router) + return app + + +# Stub the auth dependency so we don't need a real proxy running. +def _override_auth(): + from litellm.proxy._types import UserAPIKeyAuth + + return UserAPIKeyAuth(api_key="sk-test", user_id="admin") + + +# A real (non-None) prisma stub for truthiness checks. +_MOCK_PRISMA = MagicMock() + + +# --- test class --- + + +class TestToolManagementEndpoints: + def setup_method(self): + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + + app = _make_app() + app.dependency_overrides[user_api_key_auth] = _override_auth + self.client = TestClient(app, raise_server_exceptions=True) + + @patch( + "litellm.proxy.db.tool_registry_writer.list_tools", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_list_tools_returns_200(self, mock_db_list): + mock_db_list.return_value = [_make_tool_row()] + + resp = self.client.get("/v1/tool/list") + assert resp.status_code == 200 + body = resp.json() + assert body["total"] == 1 + assert body["tools"][0]["tool_name"] == "my_tool" + + @patch( + "litellm.proxy.db.tool_registry_writer.list_tools", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_list_tools_with_policy_filter(self, mock_db_list): + mock_db_list.return_value = [_make_tool_row(call_policy="blocked")] + + resp = self.client.get("/v1/tool/list?call_policy=blocked") + assert resp.status_code == 200 + assert resp.json()["tools"][0]["call_policy"] == "blocked" + + @patch( + "litellm.proxy.db.tool_registry_writer.get_tool", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_get_tool_found(self, mock_db_get): + mock_db_get.return_value = _make_tool_row(tool_name="tool_a") + + resp = self.client.get("/v1/tool/tool_a") + assert resp.status_code == 200 + assert resp.json()["tool_name"] == "tool_a" + + @patch( + "litellm.proxy.db.tool_registry_writer.get_tool", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_get_tool_not_found_returns_404(self, mock_db_get): + mock_db_get.return_value = None + + resp = self.client.get("/v1/tool/nonexistent", follow_redirects=True) + assert resp.status_code == 404 + + @patch( + "litellm.proxy.db.tool_registry_writer.update_tool_policy", + new_callable=AsyncMock, + ) + @patch("litellm.proxy.proxy_server.prisma_client", _MOCK_PRISMA) + def test_update_tool_policy_blocked(self, mock_db_update): + mock_db_update.return_value = _make_tool_row(call_policy="blocked") + + resp = self.client.post( + "/v1/tool/policy", + json={"tool_name": "my_tool", "call_policy": "blocked"}, + ) + assert resp.status_code == 200 + body = resp.json() + assert body["call_policy"] == "blocked" + assert body["updated"] is True + + @patch("litellm.proxy.proxy_server.prisma_client", None) + def test_list_tools_no_db_returns_500(self): + resp = self.client.get("/v1/tool/list") + assert resp.status_code == 500 + + def test_update_tool_policy_invalid_policy_returns_422(self): + resp = self.client.post( + "/v1/tool/policy", + json={"tool_name": "my_tool", "call_policy": "invalid_value"}, + ) + assert resp.status_code == 422 diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index fb749d7afb0..258c2ccb0e0 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -38,6 +38,7 @@ import Usage from "@/components/usage"; import UserDashboard from "@/components/user_dashboard"; import { AccessGroupsPage } from "@/components/AccessGroups/AccessGroupsPage"; import VectorStoreManagement from "@/components/vector_store_management"; +import ToolPolicies from "@/components/ToolPolicies"; import SpendLogsTable from "@/components/view_logs"; import ViewUserDashboard from "@/components/view_users"; import { ThemeProvider } from "@/contexts/ThemeContext"; @@ -548,6 +549,8 @@ function CreateKeyPageContent() { ) : page == "vector-stores" ? ( + ) : page == "tool-policies" ? ( + ) : page == "guardrails-monitor" ? ( ) : page == "new_usage" ? ( diff --git a/ui/litellm-dashboard/src/components/ToolPolicies.tsx b/ui/litellm-dashboard/src/components/ToolPolicies.tsx new file mode 100644 index 00000000000..0e3f5434e7f --- /dev/null +++ b/ui/litellm-dashboard/src/components/ToolPolicies.tsx @@ -0,0 +1,415 @@ +"use client"; + +import React, { useCallback, useDeferredValue, useEffect, useState } from "react"; +import { Select, Switch, Tooltip } from "antd"; +import { Select, Tooltip } from "antd"; +import { + Table, + TableHead, + TableHeaderCell, + TableBody, + TableRow, + TableCell, +} from "@tremor/react"; +import { TimeCell } from "./view_logs/time_cell"; +import { TableHeaderSortDropdown } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; +import type { SortState } from "./common_components/TableHeaderSortDropdown/TableHeaderSortDropdown"; +import FilterComponent, { FilterOption } from "./molecules/filter"; +import { fetchToolsList, updateToolPolicy, ToolRow } from "./networking"; + +const POLICY_OPTIONS = [ + { value: "trusted", label: "trusted", color: "#065f46", bg: "#d1fae5", border: "#6ee7b7" }, + { value: "blocked", label: "blocked", color: "#991b1b", bg: "#fee2e2", border: "#fca5a5" }, +] as const; + +type PolicyValue = "trusted" | "blocked"; + +const policyStyle = (p: string) => + POLICY_OPTIONS.find((o) => o.value === p) ?? POLICY_OPTIONS[1]; + +type SortField = "tool_name" | "call_policy" | "team_id" | "key_alias" | "created_at" | "call_count"; + +interface FilterValues { + [key: string]: string; +} + +interface ToolPoliciesProps { + accessToken: string | null; + userRole?: string; +} + +const PolicySelect: React.FC<{ + value: string; + toolName: string; + saving: boolean; + onChange: (toolName: string, policy: string) => void; +}> = ({ value, toolName, saving, onChange }) => { + const style = policyStyle(value); + return ( + { setSearchTerm(e.target.value); setCurrentPage(1); }} + /> + + + + + +
+ Live Tail + +
+ + + + +
+ + Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results + + Page {currentPage} of {totalPages} +
+ + +
+
+ + + {/* Filter row */} +
+ +
+ + + {/* Auto-refresh banner */} + {isLiveTail && ( +
+ Auto-refreshing every 15 seconds + +
+ )} + + {error && ( +
{error}
+ )} + + {/* Table */} + + + + + + + + + Key Hash + + Origin + + + + {loading ? ( + + Loading tools… + + ) : paginated.length === 0 ? ( + + + No tools discovered yet. Make a chat completion that returns tool_calls to start auto-discovery. + + + ) : ( + paginated.map((tool) => ( + + + + + + + + {tool.tool_name} + + + + + + + + {(tool.call_count ?? 0).toLocaleString()} + + + + {tool.team_id ?? "-"} + + + + + + {tool.key_hash ?? "-"} + + + + + + {tool.key_alias ?? "-"} + + + + + {tool.origin ?? "-"} + + + + )) + )} + +
+ + {/* Bottom pagination (only when > 1 page) */} + {totalPages > 1 && ( +
+ Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of {sorted.length} +
+ + +
+
+ )} + + + ); +}; + +export default ToolPolicies; diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index da3ca2a8bae..2cbeb22ec81 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -134,6 +134,12 @@ const menuGroups: MenuGroup[] = [ label: "Vector Stores", icon: , }, + { + key: "tool-policies", + page: "tool-policies", + label: "Tool Policies", + icon: , + }, ], }, ], diff --git a/ui/litellm-dashboard/src/components/networking.tsx b/ui/litellm-dashboard/src/components/networking.tsx index 6ffd744cce9..8536a584ee1 100644 --- a/ui/litellm-dashboard/src/components/networking.tsx +++ b/ui/litellm-dashboard/src/components/networking.tsx @@ -9854,3 +9854,57 @@ export const checkGdprCompliance = async ( } return response.json(); }; + +export interface ToolRow { + tool_id: string; + tool_name: string; + origin?: string; + call_policy: string; + call_count?: number; + assignments?: Record; + key_hash?: string; + team_id?: string; + key_alias?: string; + created_at?: string; + updated_at?: string; + created_by?: string; + updated_by?: string; +} + +export const fetchToolsList = async (accessToken: string): Promise => { + const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/tool/list` : `/v1/tool/list`; + const response = await fetch(url, { + method: "GET", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + }); + if (!response.ok) { + const errorData = await response.text(); + throw new Error(errorData); + } + const data = await response.json(); + return data.tools ?? []; +}; + +export const updateToolPolicy = async ( + accessToken: string, + toolName: string, + callPolicy: string +): Promise => { + const url = proxyBaseUrl ? `${proxyBaseUrl}/v1/tool/policy` : `/v1/tool/policy`; + const response = await fetch(url, { + method: "POST", + headers: { + [globalLitellmHeaderName]: `Bearer ${accessToken}`, + "Content-Type": "application/json", + }, + body: JSON.stringify({ tool_name: toolName, call_policy: callPolicy }), + }); + if (!response.ok) { + const errorData = await response.text(); + throw new Error(errorData); + } + return response.json(); +};