mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
feat(proxy): tool policies - auto-discover tools + policy enforcement guardrail (#22041)
* feat(proxy): tool policies - auto-discover tools, manage policies, guardrail enforcement
- New LiteLLM_ToolTable in schema.prisma to store discovered tools
- Auto-discovery: tools seen in LLM responses get upserted via ToolDiscoveryQueue
(hooks into DBSpendUpdateWriter, same pipeline as spend tracking)
- Management endpoints: GET /v1/tool/list, GET /v1/tool/{name}, POST /v1/tool/policy
- ToolPolicyGuardrail: blocks tool_calls in responses based on policy setting
- UI: Tool Policies page under Guardrails section with policy selector,
filters by policy/team/key, live tail, sortable table
- Unit tests for queue, writer, endpoints, guardrail
* feat(tool-policies): track call_count + discover tools from request body and /messages API
- Add call_count column to LiteLLM_ToolTable; incremented on every flush
- Extract tools from request body too (not just response tool_calls):
- OpenAI /chat/completions: tools[].function.name
- Anthropic /messages pass-through: request_body.tools[].name
- Show call_count column in UI table (sortable)
- UI: drop dual_llm option, keep only trusted/blocked
* fix: address greptile review feedback
- Remove redundant @@index([tool_name]) from schema.prisma (tool_name has @unique which already creates an index)
- Replace gen_random_uuid()::text with str(uuid.uuid4()) for portability
- Rewrite test_tool_registry_writer.py to mock execute_raw/query_raw (actual implementation) instead of Prisma model methods
- Fix test patches in test_tool_management_endpoints.py to target source modules since imports are inside function bodies
- Add "Tool Policies" page title to ToolPolicies.tsx
* fix: address greptile review round 2
- Replace NOW() with Python datetime parameter in tool_registry_writer (SQLite portability)
- Fix cache key collision in tool_policy_guardrail: use null-byte separator instead of colon
- Remove type==function filter from request-side tool extraction to match response-side behavior
- Clear seen_tool_names on flush so call_count increments per batch cycle not per pod lifetime
* fix: address greptile review round 3
- Fix test_seen_names_persist_across_flushes to match actual per-flush-cycle behavior
- Update module docstring in tool_discovery_queue.py to accurately describe flush behavior
- Add created_at/updated_at to raw SQL INSERT in batch_upsert_tools and update_tool_policy
* fix: cache tool policies per tool name not per combination
Previously the cache key was built from the full set of tool names in a
request, so each unique combination of tools got its own cold cache entry
and triggered a separate DB query. With N distinct tools across requests
this was effectively a DB hit on every request.
Now each tool name is cached individually. Cache hits are checked per
tool, only missing tools are fetched from DB in a single batch query,
and each result is cached separately. Once a tool's policy is warm,
any subsequent request using that tool benefits from the cache regardless
of what other tools are in the request.
* Update ui/litellm-dashboard/src/components/ToolPolicies.tsx
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
---------
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
This commit is contained in:
Ishaan Jaff
Sameer Kankute
parent
f13a8b6003
commit
fef13a161b
19 changed files with 1939 additions and 188 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
179
litellm/proxy/db/tool_registry_writer.py
Normal file
179
litellm/proxy/db/tool_registry_writer.py
Normal file
|
|
@ -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 {}
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
149
litellm/proxy/management_endpoints/tool_management_endpoints.py
Normal file
149
litellm/proxy/management_endpoints/tool_management_endpoints.py
Normal file
|
|
@ -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))
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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())
|
||||
|
|
|
|||
42
litellm/types/tool_management.py
Normal file
42
litellm/types/tool_management.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
197
tests/test_litellm/proxy/db/test_tool_registry_writer.py
Normal file
197
tests/test_litellm/proxy/db/test_tool_registry_writer.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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"
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
@ -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() {
|
|||
<AccessGroupsPage />
|
||||
) : page == "vector-stores" ? (
|
||||
<VectorStoreManagement accessToken={accessToken} userRole={userRole} userID={userID} />
|
||||
) : page == "tool-policies" ? (
|
||||
<ToolPolicies accessToken={accessToken} userRole={userRole} />
|
||||
) : page == "guardrails-monitor" ? (
|
||||
<GuardrailsMonitorView accessToken={accessToken} />
|
||||
) : page == "new_usage" ? (
|
||||
|
|
|
|||
415
ui/litellm-dashboard/src/components/ToolPolicies.tsx
Normal file
415
ui/litellm-dashboard/src/components/ToolPolicies.tsx
Normal file
|
|
@ -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 (
|
||||
<Select
|
||||
size="small"
|
||||
value={value}
|
||||
disabled={saving}
|
||||
loading={saving}
|
||||
onChange={(v) => onChange(toolName, v)}
|
||||
onClick={(e) => e.stopPropagation()}
|
||||
style={{
|
||||
minWidth: 110,
|
||||
fontWeight: 500,
|
||||
}}
|
||||
styles={{
|
||||
selector: {
|
||||
backgroundColor: style.bg,
|
||||
borderColor: style.border,
|
||||
color: style.color,
|
||||
borderRadius: 999,
|
||||
fontSize: 11,
|
||||
fontWeight: 600,
|
||||
paddingLeft: 8,
|
||||
paddingRight: 4,
|
||||
},
|
||||
}}
|
||||
popupMatchSelectWidth={false}
|
||||
options={POLICY_OPTIONS.map((o) => ({
|
||||
value: o.value,
|
||||
label: (
|
||||
<span
|
||||
style={{
|
||||
display: "inline-flex",
|
||||
alignItems: "center",
|
||||
gap: 6,
|
||||
fontSize: 12,
|
||||
fontWeight: 500,
|
||||
color: o.color,
|
||||
}}
|
||||
>
|
||||
<span
|
||||
style={{
|
||||
width: 8,
|
||||
height: 8,
|
||||
borderRadius: "50%",
|
||||
backgroundColor: o.color,
|
||||
display: "inline-block",
|
||||
flexShrink: 0,
|
||||
}}
|
||||
/>
|
||||
{o.label}
|
||||
</span>
|
||||
),
|
||||
}))}
|
||||
/>
|
||||
);
|
||||
};
|
||||
|
||||
export const ToolPolicies: React.FC<ToolPoliciesProps> = ({ accessToken }) => {
|
||||
const [tools, setTools] = useState<ToolRow[]>([]);
|
||||
const [loading, setLoading] = useState(true);
|
||||
const [isFetching, setIsFetching] = useState(false);
|
||||
const [error, setError] = useState<string | null>(null);
|
||||
const [saving, setSaving] = useState<string | null>(null);
|
||||
|
||||
const [searchTerm, setSearchTerm] = useState("");
|
||||
const [sortField, setSortField] = useState<SortField>("created_at");
|
||||
const [sortOrder, setSortOrder] = useState<"asc" | "desc">("desc");
|
||||
const [currentPage, setCurrentPage] = useState(1);
|
||||
const [isLiveTail, setIsLiveTail] = useState(true);
|
||||
const [activeFilters, setActiveFilters] = useState<FilterValues>({});
|
||||
const pageSize = 50;
|
||||
|
||||
const isFetchingDeferred = useDeferredValue(isFetching);
|
||||
const isButtonLoading = isFetching || isFetchingDeferred;
|
||||
|
||||
const load = useCallback(async () => {
|
||||
if (!accessToken) return;
|
||||
setIsFetching(true);
|
||||
setError(null);
|
||||
try {
|
||||
const rows = await fetchToolsList(accessToken);
|
||||
setTools(rows);
|
||||
} catch (e: any) {
|
||||
setError(e.message ?? "Failed to load tools");
|
||||
} finally {
|
||||
setIsFetching(false);
|
||||
setLoading(false);
|
||||
}
|
||||
}, [accessToken]);
|
||||
|
||||
useEffect(() => { load(); }, [load]);
|
||||
|
||||
useEffect(() => {
|
||||
if (!isLiveTail) return;
|
||||
const id = setInterval(load, 15000);
|
||||
return () => clearInterval(id);
|
||||
}, [isLiveTail, load]);
|
||||
|
||||
const handlePolicyChange = async (toolName: string, newPolicy: string) => {
|
||||
if (!accessToken) return;
|
||||
setSaving(toolName);
|
||||
try {
|
||||
await updateToolPolicy(accessToken, toolName, newPolicy);
|
||||
setTools((prev) =>
|
||||
prev.map((t) => (t.tool_name === toolName ? { ...t, call_policy: newPolicy } : t))
|
||||
);
|
||||
} catch (e: any) {
|
||||
alert(`Failed to update policy: ${e.message}`);
|
||||
} finally {
|
||||
setSaving(null);
|
||||
}
|
||||
};
|
||||
|
||||
const handleSortChange = (field: SortField, newState: SortState) => {
|
||||
if (newState === false) {
|
||||
setSortField("created_at");
|
||||
setSortOrder("desc");
|
||||
} else {
|
||||
setSortField(field);
|
||||
setSortOrder(newState);
|
||||
}
|
||||
setCurrentPage(1);
|
||||
};
|
||||
|
||||
const handleApplyFilters = (filters: FilterValues) => {
|
||||
setActiveFilters(filters);
|
||||
setCurrentPage(1);
|
||||
};
|
||||
|
||||
const handleResetFilters = () => {
|
||||
setActiveFilters({});
|
||||
setCurrentPage(1);
|
||||
};
|
||||
|
||||
// Build unique team/key options from loaded data
|
||||
const teamOptions = Array.from(new Set(tools.map((t) => t.team_id).filter(Boolean))).map(
|
||||
(v) => ({ label: v as string, value: v as string })
|
||||
);
|
||||
const keyAliasOptions = Array.from(new Set(tools.map((t) => t.key_alias).filter(Boolean))).map(
|
||||
(v) => ({ label: v as string, value: v as string })
|
||||
);
|
||||
|
||||
const filterOptions: FilterOption[] = [
|
||||
{
|
||||
name: "Policy",
|
||||
label: "Policy",
|
||||
options: POLICY_OPTIONS.map((o) => ({ label: o.label, value: o.value })),
|
||||
},
|
||||
{
|
||||
name: "Team Name",
|
||||
label: "Team Name",
|
||||
options: teamOptions,
|
||||
},
|
||||
{
|
||||
name: "Key Name",
|
||||
label: "Key Name",
|
||||
options: keyAliasOptions,
|
||||
},
|
||||
];
|
||||
|
||||
const SortHeader = ({ label, field }: { label: string; field: SortField }) => (
|
||||
<div className="flex items-center gap-1">
|
||||
<span>{label}</span>
|
||||
<TableHeaderSortDropdown
|
||||
sortState={sortField === field ? sortOrder : false}
|
||||
onSortChange={(s) => handleSortChange(field, s)}
|
||||
/>
|
||||
</div>
|
||||
);
|
||||
|
||||
const filtered = tools.filter((t) => {
|
||||
if (searchTerm) {
|
||||
const q = searchTerm.toLowerCase();
|
||||
const matchesSearch =
|
||||
t.tool_name.toLowerCase().includes(q) ||
|
||||
(t.team_id ?? "").toLowerCase().includes(q) ||
|
||||
(t.key_alias ?? "").toLowerCase().includes(q) ||
|
||||
(t.key_hash ?? "").toLowerCase().includes(q) ||
|
||||
t.call_policy.toLowerCase().includes(q);
|
||||
if (!matchesSearch) return false;
|
||||
}
|
||||
if (activeFilters["Policy"] && t.call_policy !== activeFilters["Policy"]) return false;
|
||||
if (activeFilters["Team Name"] && t.team_id !== activeFilters["Team Name"]) return false;
|
||||
if (activeFilters["Key Name"] && t.key_alias !== activeFilters["Key Name"]) return false;
|
||||
return true;
|
||||
});
|
||||
|
||||
const sorted = [...filtered].sort((a, b) => {
|
||||
const av = (a as any)[sortField] ?? "";
|
||||
const bv = (b as any)[sortField] ?? "";
|
||||
if (av < bv) return sortOrder === "desc" ? 1 : -1;
|
||||
if (av > bv) return sortOrder === "desc" ? -1 : 1;
|
||||
return 0;
|
||||
});
|
||||
|
||||
const totalPages = Math.max(1, Math.ceil(sorted.length / pageSize));
|
||||
const paginated = sorted.slice((currentPage - 1) * pageSize, currentPage * pageSize);
|
||||
|
||||
return (
|
||||
<div className="p-6 w-full">
|
||||
<h1 className="text-2xl font-semibold text-gray-900 mb-6">Tool Policies</h1>
|
||||
<div className="bg-white rounded-lg shadow w-full max-w-full box-border">
|
||||
|
||||
{/* Toolbar */}
|
||||
<div className="border-b px-6 py-4 w-full max-w-full box-border">
|
||||
<div className="flex flex-col md:flex-row items-start md:items-center justify-between space-y-4 md:space-y-0 w-full max-w-full box-border">
|
||||
<div className="flex flex-wrap items-center gap-3">
|
||||
<div className="relative w-64">
|
||||
<input
|
||||
type="text"
|
||||
placeholder="Search by Tool Name"
|
||||
className="w-full px-3 py-2 pl-8 border rounded-md text-sm focus:outline-none focus:ring-2 focus:ring-blue-500 focus:border-blue-500"
|
||||
value={searchTerm}
|
||||
onChange={(e) => { setSearchTerm(e.target.value); setCurrentPage(1); }}
|
||||
/>
|
||||
<svg className="absolute left-2.5 top-2.5 h-4 w-4 text-gray-500" fill="none" stroke="currentColor" viewBox="0 0 24 24">
|
||||
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M21 21l-6-6m2-5a7 7 0 11-14 0 7 7 0 0114 0z" />
|
||||
</svg>
|
||||
</div>
|
||||
|
||||
<div className="flex items-center gap-2">
|
||||
<span className="text-sm font-medium text-gray-900">Live Tail</span>
|
||||
<Switch color="green" checked={isLiveTail} onChange={setIsLiveTail} />
|
||||
</div>
|
||||
|
||||
<button
|
||||
onClick={load}
|
||||
disabled={isButtonLoading}
|
||||
className="flex items-center gap-1.5 px-3 py-2 text-sm border rounded-md hover:bg-gray-50 disabled:opacity-60"
|
||||
>
|
||||
<svg className={`w-4 h-4 ${isButtonLoading ? "animate-spin" : ""}`} fill="none" stroke="currentColor" viewBox="0 0 24 24">
|
||||
<path strokeLinecap="round" strokeLinejoin="round" strokeWidth={2} d="M4 4v5h.582m15.356 2A8.001 8.001 0 004.582 9m0 0H9m11 11v-5h-.581m0 0a8.003 8.003 0 01-15.357-2m15.357 2H15" />
|
||||
</svg>
|
||||
{isButtonLoading ? "Fetching" : "Fetch"}
|
||||
</button>
|
||||
</div>
|
||||
|
||||
<div className="flex items-center gap-4 text-sm text-gray-600 whitespace-nowrap">
|
||||
<span>
|
||||
Showing {filtered.length === 0 ? 0 : (currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, filtered.length)} of {filtered.length} results
|
||||
</span>
|
||||
<span>Page {currentPage} of {totalPages}</span>
|
||||
<div className="flex gap-1">
|
||||
<button onClick={() => setCurrentPage((p) => Math.max(1, p - 1))} disabled={currentPage === 1}
|
||||
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40">Previous</button>
|
||||
<button onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))} disabled={currentPage === totalPages}
|
||||
className="px-3 py-1.5 border rounded-md text-sm hover:bg-gray-50 disabled:opacity-40">Next</button>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Filter row */}
|
||||
<div className="mt-3">
|
||||
<FilterComponent
|
||||
options={filterOptions}
|
||||
onApplyFilters={handleApplyFilters}
|
||||
onResetFilters={handleResetFilters}
|
||||
buttonLabel="Filters"
|
||||
/>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
{/* Auto-refresh banner */}
|
||||
{isLiveTail && (
|
||||
<div className="bg-green-50 border-b border-green-100 px-6 py-2 flex items-center justify-between">
|
||||
<span className="text-sm text-green-700">Auto-refreshing every 15 seconds</span>
|
||||
<button onClick={() => setIsLiveTail(false)} className="text-xs text-green-600 underline">Stop</button>
|
||||
</div>
|
||||
)}
|
||||
|
||||
{error && (
|
||||
<div className="mx-6 mt-4 p-3 bg-red-50 border border-red-200 rounded text-sm text-red-700">{error}</div>
|
||||
)}
|
||||
|
||||
{/* Table */}
|
||||
<Table className="[&_td]:py-0.5 [&_th]:py-1 w-full">
|
||||
<TableHead>
|
||||
<TableRow>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="Discovered" field="created_at" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="Tool Name" field="tool_name" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="Policy" field="call_policy" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="# Calls" field="call_count" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="Team Name" field="team_id" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8">Key Hash</TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8"><SortHeader label="Key Name" field="key_alias" /></TableHeaderCell>
|
||||
<TableHeaderCell className="py-1 h-8">Origin</TableHeaderCell>
|
||||
</TableRow>
|
||||
</TableHead>
|
||||
<TableBody>
|
||||
{loading ? (
|
||||
<TableRow>
|
||||
<TableCell colSpan={8} className="h-8 text-center text-gray-500">Loading tools…</TableCell>
|
||||
</TableRow>
|
||||
) : paginated.length === 0 ? (
|
||||
<TableRow>
|
||||
<TableCell colSpan={8} className="h-8 text-center text-gray-500">
|
||||
No tools discovered yet. Make a chat completion that returns tool_calls to start auto-discovery.
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
) : (
|
||||
paginated.map((tool) => (
|
||||
<TableRow key={tool.tool_id} className="h-8 hover:bg-gray-50">
|
||||
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
|
||||
<TimeCell utcTime={tool.created_at ?? ""} />
|
||||
</TableCell>
|
||||
<TableCell className="py-0.5 max-h-8 overflow-hidden">
|
||||
<Tooltip title={tool.tool_name}>
|
||||
<span className="font-mono text-xs max-w-[20ch] truncate block font-medium">
|
||||
{tool.tool_name}
|
||||
</span>
|
||||
</Tooltip>
|
||||
</TableCell>
|
||||
<TableCell className="py-0.5 max-h-8">
|
||||
<PolicySelect
|
||||
value={tool.call_policy}
|
||||
toolName={tool.tool_name}
|
||||
saving={saving === tool.tool_name}
|
||||
onChange={handlePolicyChange}
|
||||
/>
|
||||
</TableCell>
|
||||
<TableCell className="py-0.5 max-h-8 text-right tabular-nums text-sm font-mono text-gray-700">
|
||||
{(tool.call_count ?? 0).toLocaleString()}
|
||||
</TableCell>
|
||||
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
|
||||
<Tooltip title={tool.team_id ?? "-"}>
|
||||
<span className="max-w-[15ch] truncate block">{tool.team_id ?? "-"}</span>
|
||||
</Tooltip>
|
||||
</TableCell>
|
||||
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
|
||||
<Tooltip title={tool.key_hash ?? "-"}>
|
||||
<span className="font-mono max-w-[15ch] truncate block text-blue-600">
|
||||
{tool.key_hash ?? "-"}
|
||||
</span>
|
||||
</Tooltip>
|
||||
</TableCell>
|
||||
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
|
||||
<Tooltip title={tool.key_alias ?? "-"}>
|
||||
<span className="max-w-[15ch] truncate block">{tool.key_alias ?? "-"}</span>
|
||||
</Tooltip>
|
||||
</TableCell>
|
||||
<TableCell className="py-0.5 max-h-8 overflow-hidden whitespace-nowrap">
|
||||
<Tooltip title={tool.origin ?? "-"}>
|
||||
<span className="max-w-[15ch] truncate block">{tool.origin ?? "-"}</span>
|
||||
</Tooltip>
|
||||
</TableCell>
|
||||
</TableRow>
|
||||
))
|
||||
)}
|
||||
</TableBody>
|
||||
</Table>
|
||||
|
||||
{/* Bottom pagination (only when > 1 page) */}
|
||||
{totalPages > 1 && (
|
||||
<div className="border-t px-6 py-3 flex items-center justify-between text-sm text-gray-600">
|
||||
<span>Showing {(currentPage - 1) * pageSize + 1} - {Math.min(currentPage * pageSize, sorted.length)} of {sorted.length}</span>
|
||||
<div className="flex gap-1">
|
||||
<button onClick={() => setCurrentPage((p) => Math.max(1, p - 1))} disabled={currentPage === 1}
|
||||
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40">Previous</button>
|
||||
<button onClick={() => setCurrentPage((p) => Math.min(totalPages, p + 1))} disabled={currentPage === totalPages}
|
||||
className="px-3 py-1.5 border rounded-md hover:bg-gray-50 disabled:opacity-40">Next</button>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
);
|
||||
};
|
||||
|
||||
export default ToolPolicies;
|
||||
|
|
@ -134,6 +134,12 @@ const menuGroups: MenuGroup[] = [
|
|||
label: "Vector Stores",
|
||||
icon: <DatabaseOutlined />,
|
||||
},
|
||||
{
|
||||
key: "tool-policies",
|
||||
page: "tool-policies",
|
||||
label: "Tool Policies",
|
||||
icon: <SafetyOutlined />,
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
|
|
|
|||
|
|
@ -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<string, any>;
|
||||
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<ToolRow[]> => {
|
||||
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<ToolRow> => {
|
||||
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();
|
||||
};
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue