From a5dd01d8aca49a94433d4a430958fe369ac18a2f Mon Sep 17 00:00:00 2001 From: kothamah <104782493+kothamah@users.noreply.github.com> Date: Thu, 19 Mar 2026 13:52:23 -0400 Subject: [PATCH 001/226] added bedrock guardrail API exception --- .../guardrail_hooks/bedrock_guardrails.py | 37 ++++++++++++------- 1 file changed, 23 insertions(+), 14 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index 8ef188bb23..02d87a4962 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -82,7 +82,12 @@ def _redact_pii_matches(response_json: dict) -> dict: redacted_response = copy.deepcopy(response_json) # Get assessments from the response - assessments = redacted_response.get("assessments", []) + # NOTE: We use `.get("key") or []` instead of `.get("key", [])` because + # the Bedrock API can return explicit `null` for list fields (e.g. "regexes": null). + # In Python, dict.get("key", []) returns None (not []) when the key exists + # with a None/null value. The `or []` ensures we always get an iterable, + # preventing "TypeError: 'NoneType' object is not iterable". + assessments = redacted_response.get("assessments") or [] if not assessments: return redacted_response @@ -90,13 +95,13 @@ def _redact_pii_matches(response_json: dict) -> dict: # Redact PII entities in sensitive information policy sensitive_info_policy = assessment.get("sensitiveInformationPolicy") if sensitive_info_policy: - pii_entities = sensitive_info_policy.get("piiEntities", []) + pii_entities = sensitive_info_policy.get("piiEntities") or [] for pii_entity in pii_entities: if "match" in pii_entity: pii_entity["match"] = "[REDACTED]" # Redact regex matches - regexes = sensitive_info_policy.get("regexes", []) + regexes = sensitive_info_policy.get("regexes") or [] for regex_match in regexes: if "match" in regex_match: regex_match["match"] = "[REDACTED]" @@ -104,12 +109,12 @@ def _redact_pii_matches(response_json: dict) -> dict: # Redact custom word matches in word policy word_policy = assessment.get("wordPolicy") if word_policy: - custom_words = word_policy.get("customWords", []) + custom_words = word_policy.get("customWords") or [] for custom_word in custom_words: if "match" in custom_word: custom_word["match"] = "[REDACTED]" - managed_words = word_policy.get("managedWordLists", []) + managed_words = word_policy.get("managedWordLists") or [] for managed_word in managed_words: if "match" in managed_word: managed_word["match"] = "[REDACTED]" @@ -682,7 +687,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): return False # Check assessments to determine if any actions were BLOCKED (vs ANONYMIZED) - assessments = response.get("assessments", []) + # NOTE: Use `or []` instead of default param to handle explicit null from Bedrock API. + # See _redact_pii_matches() for detailed explanation of the null safety pattern. + assessments = response.get("assessments") or [] if not assessments: return False @@ -690,7 +697,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check topic policy topic_policy = assessment.get("topicPolicy") if topic_policy: - topics = topic_policy.get("topics", []) + topics = topic_policy.get("topics") or [] for topic in topics: if topic.get("action") == "BLOCKED": return True @@ -698,7 +705,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check content policy content_policy = assessment.get("contentPolicy") if content_policy: - filters = content_policy.get("filters", []) + filters = content_policy.get("filters") or [] for filter_item in filters: if filter_item.get("action") == "BLOCKED": return True @@ -706,11 +713,11 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check word policy word_policy = assessment.get("wordPolicy") if word_policy: - custom_words = word_policy.get("customWords", []) + custom_words = word_policy.get("customWords") or [] for custom_word in custom_words: if custom_word.get("action") == "BLOCKED": return True - managed_words = word_policy.get("managedWordLists", []) + managed_words = word_policy.get("managedWordLists") or [] for managed_word in managed_words: if managed_word.get("action") == "BLOCKED": return True @@ -718,12 +725,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check sensitive information policy sensitive_info_policy = assessment.get("sensitiveInformationPolicy") if sensitive_info_policy: - pii_entities = sensitive_info_policy.get("piiEntities", []) + pii_entities = sensitive_info_policy.get("piiEntities") or [] if pii_entities: for pii_entity in pii_entities: if pii_entity.get("action") == "BLOCKED": return True - regexes = sensitive_info_policy.get("regexes", []) + regexes = sensitive_info_policy.get("regexes") or [] if regexes: for regex in regexes: if regex.get("action") == "BLOCKED": @@ -732,7 +739,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): # Check contextual grounding policy contextual_grounding_policy = assessment.get("contextualGroundingPolicy") if contextual_grounding_policy: - grounding_filters = contextual_grounding_policy.get("filters", []) + grounding_filters = contextual_grounding_policy.get("filters") or [] for grounding_filter in grounding_filters: if grounding_filter.get("action") == "BLOCKED": return True @@ -1391,7 +1398,9 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Raises: Exception: If content is blocked by Bedrock guardrail """ - texts = inputs.get("texts", []) + # NOTE: Use `or []` to handle case where inputs["texts"] is explicitly None. + # dict.get("texts", []) would return None if the key exists with a None value. + texts = inputs.get("texts") or [] try: verbose_proxy_logger.debug( f"Bedrock Guardrail: Applying guardrail to {len(texts)} text(s)" From ead822b6985443827d3accb7c0d6cfb7fce40df7 Mon Sep 17 00:00:00 2001 From: kothamah <104782493+kothamah@users.noreply.github.com> Date: Thu, 19 Mar 2026 13:55:36 -0400 Subject: [PATCH 002/226] Added test cases for the null type handling --- .../test_bedrock_guardrails.py | 2972 ++++++++--------- 1 file changed, 1472 insertions(+), 1500 deletions(-) diff --git a/tests/guardrails_tests/test_bedrock_guardrails.py b/tests/guardrails_tests/test_bedrock_guardrails.py index cb594c221c..03fe63b307 100644 --- a/tests/guardrails_tests/test_bedrock_guardrails.py +++ b/tests/guardrails_tests/test_bedrock_guardrails.py @@ -1,1090 +1,213 @@ -import sys +""" +Unit tests for Bedrock Guardrails +""" +import json import os -import io, asyncio +import sys +from unittest.mock import AsyncMock, MagicMock, patch + import pytest +from fastapi import HTTPException + +sys.path.insert(0, os.path.abspath("../../../../../..")) -sys.path.insert(0, os.path.abspath("../..")) -import litellm -from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import BedrockGuardrail from litellm.proxy._types import UserAPIKeyAuth -from litellm.caching import DualCache -from unittest.mock import MagicMock, AsyncMock, patch +from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( + BedrockGuardrail, + _redact_pii_matches, +) @pytest.mark.asyncio -async def test_bedrock_guardrails_pii_masking(): - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() +async def test__redact_pii_matches_function(): + """Test the _redact_pii_matches function directly""" - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, - {"role": "assistant", "content": "Hello, how can I help you today?"}, - {"role": "user", "content": "I need to cancel my order"}, - { - "role": "user", - "content": "ok, my credit card number is 1234-5678-9012-3456", - }, - ], - } - - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - print("response after moderation hook", response) - - if response: # Only assert if response is not None - assert response["messages"][0]["content"] == "Hello, my phone number is {PHONE}" - assert response["messages"][1]["content"] == "Hello, how can I help you today?" - assert response["messages"][2]["content"] == "I need to cancel my order" - assert ( - response["messages"][3]["content"] - == "ok, my credit card number is {CREDIT_DEBIT_CARD_NUMBER}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_pii_masking_content_list(): - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "gpt-4o", - "messages": [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Hello, my phone number is +1 412 555 1212", - }, - {"type": "text", "text": "what time is it?"}, - ], - }, - {"role": "assistant", "content": "Hello, how can I help you today?"}, - {"role": "user", "content": "who is the president of the united states?"}, - ], - } - - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - print(response) - - if response: # Only assert if response is not None - # Verify that the list content is properly masked - assert isinstance(response["messages"][0]["content"], list) - assert ( - response["messages"][0]["content"][0]["text"] - == "Hello, my phone number is {PHONE}" - ) - assert response["messages"][0]["content"][1]["text"] == "what time is it?" - assert response["messages"][1]["content"] == "Hello, how can I help you today?" - assert ( - response["messages"][2]["content"] - == "who is the president of the united states?" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_block_messages_api(): - """ - Test that guardrails block messages API requests containing 'coffee' and raise the expected exception. - """ - from fastapi import HTTPException - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "claude-3-5-sonnet-20240620", - "messages": [ - { - "role": "user", - "content": [ - { - "type": "text", - "text": "Hello, my phone number is +1 412 555 1212", - }, - {"type": "text", "text": "what time is it?"}, - ], - }, - {"role": "user", "content": "tell me about coffee"}, - ], - } - - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_pre_call_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="anthropic_messages", - cache=MagicMock(spec=DualCache), - ) - - exception = exc_info.value - assert exception.status_code == 400 - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - assert ( - detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question. coffee guardrail applied " - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_block_responses_api(): - """ - Test that guardrails block responses API requests containing 'coffee' and raise the expected exception. - """ - from fastapi import HTTPException - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - ) - - request_data = { - "model": "gpt-4.1", - "input": "Tell me a three sentence bedtime story about a unicorn drinking coffee", - "stream": False, - } - - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_pre_call_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="responses", - cache=MagicMock(spec=DualCache), - ) - - exception = exc_info.value - assert exception.status_code == 400 - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - assert ( - detail["bedrock_guardrail_response"] - == "Sorry, the model cannot answer this question. coffee guardrail applied " - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_with_streaming(): - from litellm.proxy.utils import ProxyLogging - from litellm.types.guardrails import GuardrailEventHooks - - # Create proper mock objects - mock_user_api_key_cache = MagicMock(spec=DualCache) - mock_user_api_key_dict = UserAPIKeyAuth() - - with pytest.raises(Exception): # Assert that this raises an exception - proxy_logging_obj = ProxyLogging( - user_api_key_cache=mock_user_api_key_cache, - premium_user=True, - ) - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - litellm.callbacks.append(guardrail) - - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "Hi I like coffee"}], - "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - - response = await litellm.acompletion( - **request_data, - ) - - response = proxy_logging_obj.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=response, - request_data=request_data, - ) - - async for chunk in response: - print(chunk) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_with_streaming_no_violation(): - from litellm.proxy.utils import ProxyLogging - from litellm.types.guardrails import GuardrailEventHooks - - # Create proper mock objects - mock_user_api_key_cache = MagicMock(spec=DualCache) - mock_user_api_key_dict = UserAPIKeyAuth() - - proxy_logging_obj = ProxyLogging( - user_api_key_cache=mock_user_api_key_cache, - premium_user=True, - ) - - guardrail = BedrockGuardrail( - guardrailIdentifier="ff6ujrregl1q", - guardrailVersion="DRAFT", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - litellm.callbacks.append(guardrail) - - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "hi"}], - "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - - response = await litellm.acompletion( - **request_data, - ) - - response = proxy_logging_obj.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=response, - request_data=request_data, - ) - - async for chunk in response: - print(chunk) - - -@pytest.mark.asyncio -async def test_bedrock_guardrails_streaming_request_body_mock(): - """Test that the exact request body sent to Bedrock matches expected format when using streaming""" - import json - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.caching import DualCache - from litellm.types.guardrails import GuardrailEventHooks - - # Create mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - mock_cache = MagicMock(spec=DualCache) - - # Create the guardrail - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - # Mock the assembled response from streaming - mock_response = litellm.ModelResponse( - id="test-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", content="The capital of Spain is Madrid." - ), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion", - ) - - # Mock Bedrock API response - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = {"action": "NONE", "outputs": []} - - # Patch the async_handler.post method to capture the request body - with patch.object(guardrail, "async_handler") as mock_async_handler: - mock_async_handler.post = AsyncMock(return_value=mock_bedrock_response) - - # Test data - simulating request data and assembled response - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "what's the capital of spain?"}], - "stream": True, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - - # Call the method that should make the Bedrock API request - await guardrail.make_bedrock_api_request( - source="OUTPUT", response=mock_response, request_data=request_data - ) - - # Verify the API call was made - mock_async_handler.post.assert_called_once() - - # Get the request data that was passed - call_args = mock_async_handler.post.call_args - - # The data should be in the 'data' parameter of the prepared request - # We need to parse the JSON from the prepared request body - prepared_request_body = call_args.kwargs.get("data") - - # Parse the JSON body - if isinstance(prepared_request_body, bytes): - actual_body = json.loads(prepared_request_body.decode("utf-8")) - else: - actual_body = json.loads(prepared_request_body) - - # Expected body based on the convert_to_bedrock_format method behavior - expected_body = { - "source": "OUTPUT", - "content": [{"text": {"text": "The capital of Spain is Madrid."}}], - } - - print("Actual Bedrock request body:", json.dumps(actual_body, indent=2)) - print("Expected Bedrock request body:", json.dumps(expected_body, indent=2)) - - # Assert the request body matches exactly - assert ( - actual_body == expected_body - ), f"Request body mismatch. Expected: {expected_body}, Got: {actual_body}" - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_aws_param_persistence(): - """Test that AWS auth params set on init are used for every request and not popped out.""" - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.guardrails import GuardrailEventHooks - - guardrail = BedrockGuardrail( - guardrailIdentifier="wf0hkdb5x07f", - guardrailVersion="DRAFT", - aws_access_key_id="test-access-key", - aws_secret_access_key="test-secret-key", - aws_region_name="us-east-1", - supported_event_hooks=[GuardrailEventHooks.post_call], - guardrail_name="bedrock-post-guard", - ) - - with patch.object( - guardrail, "get_credentials", wraps=guardrail.get_credentials - ) as mock_get_creds: - for i in range(3): - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": f"request {i}"}], - "stream": False, - "metadata": {"guardrails": ["bedrock-post-guard"]}, - } - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - # Configure the mock response properly - mock_response = AsyncMock() - mock_response.status_code = 200 - mock_response.json = MagicMock( - return_value={"action": "NONE", "outputs": []} - ) - mock_post.return_value = mock_response - await guardrail.make_bedrock_api_request( - source="INPUT", - messages=request_data.get("messages"), - request_data=request_data, - ) - - assert mock_get_creds.call_count == 3 - for call in mock_get_creds.call_args_list: - kwargs = call.kwargs - print("used the following kwargs to get credentials=", kwargs) - assert kwargs["aws_access_key_id"] == "test-access-key" - assert kwargs["aws_secret_access_key"] == "test-secret-key" - assert kwargs["aws_region_name"] == "us-east-1" - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_vs_anonymized_actions(): - """Test that BLOCKED actions raise exceptions but ANONYMIZED actions do not""" - from unittest.mock import MagicMock - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrail, - ) - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrailResponse, - ) - - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Test 1: ANONYMIZED action should NOT raise exception - anonymized_response: BedrockGuardrailResponse = { + # Test case 1: Response with PII entities + response_with_pii = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Hello, my phone number is {PHONE}"}], "assessments": [ { "sensitiveInformationPolicy": { "piiEntities": [ + {"type": "NAME", "match": "John Smith", "action": "BLOCKED"}, { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - } - ] - } - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception( - anonymized_response - ) - assert should_raise is False, "ANONYMIZED actions should not raise exceptions" - - # Test 2: BLOCKED action should raise exception - blocked_response: BedrockGuardrailResponse = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(blocked_response) - assert should_raise is True, "BLOCKED actions should raise exceptions" - - # Test 3: Mixed actions - should raise if ANY action is BLOCKED - mixed_response: BedrockGuardrailResponse = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - } - ] - }, - "topicPolicy": { - "topics": [ - {"name": "Blocked Topic", "type": "DENY", "action": "BLOCKED"} - ] - }, - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(mixed_response) - assert ( - should_raise is True - ), "Mixed actions with any BLOCKED should raise exceptions" - - # Test 4: NONE action should not raise exception - none_response: BedrockGuardrailResponse = { - "action": "NONE", - "outputs": [], - "assessments": [], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception(none_response) - assert should_raise is False, "NONE actions should not raise exceptions" - - # Test 5: Test other policy types with BLOCKED actions - content_blocked_response: BedrockGuardrailResponse = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "contentPolicy": { - "filters": [ - {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} - ] - } - } - ], - } - - should_raise = guardrail._should_raise_guardrail_blocked_exception( - content_blocked_response - ) - assert ( - should_raise is True - ), "Content policy BLOCKED actions should raise exceptions" - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_masking_with_anonymized_response(): - """Test that masking works correctly when guardrail returns ANONYMIZED actions""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.caching import DualCache - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - mask_request_content=True, - ) - - # Mock the Bedrock API response with ANONYMIZED action - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Hello, my phone number is {PHONE}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - } - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # This should NOT raise an exception since action is ANONYMIZED - try: - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - # Should succeed and return data with masked content - assert response is not None - assert ( - response["messages"][0]["content"] - == "Hello, my phone number is {PHONE}" - ) - except Exception as e: - pytest.fail( - f"Should not raise exception for ANONYMIZED actions, but got: {e}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_uses_masked_output_without_masking_flags(): - """Test that masked output from guardrails is used even when masking flags are not enabled""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail WITHOUT masking flags enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - # Note: No mask_request_content=True or mask_response_content=True - ) - - # Mock the Bedrock API response with ANONYMIZED action and masked output - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Hello, my phone number is {PHONE} and email is {EMAIL}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", + "type": "US_SOCIAL_SECURITY_NUMBER", + "match": "324-12-3212", + "action": "BLOCKED", }, + {"type": "PHONE", "match": "607-456-7890", "action": "BLOCKED"}, + ] + } + } + ], + "outputs": [{"text": "Input blocked by PII policy"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_pii) + + # Verify that PII matches are redacted + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + + assert pii_entities[0]["match"] == "[REDACTED]", "Name should be redacted" + assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" + assert pii_entities[2]["match"] == "[REDACTED]", "Phone should be redacted" + + # Verify other fields remain unchanged + assert pii_entities[0]["type"] == "NAME" + assert pii_entities[1]["type"] == "US_SOCIAL_SECURITY_NUMBER" + assert pii_entities[2]["type"] == "PHONE" + assert redacted_response["action"] == "GUARDRAIL_INTERVENED" + assert redacted_response["outputs"][0]["text"] == "Input blocked by PII policy" + + print("PII redaction function test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_no_pii(): + """Test _redact_pii_matches with response that has no PII""" + + response_no_pii = {"action": "NONE", "assessments": [], "outputs": []} + + # Call the redaction function + redacted_response = _redact_pii_matches(response_no_pii) + + # Should return the same response unchanged + assert redacted_response == response_no_pii + print("No PII redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_empty_assessments(): + """Test _redact_pii_matches with empty assessments""" + + response_empty_assessments = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [{"sensitiveInformationPolicy": {"piiEntities": []}}], + "outputs": [{"text": "Some output"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_empty_assessments) + + # Should return the same response unchanged + assert redacted_response == response_empty_assessments + print("Empty assessments redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_malformed_response(): + """Test _redact_pii_matches with malformed response (should not crash)""" + + # Test with completely malformed response + malformed_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": "not_a_list", # This should cause an exception + } + + # Should not crash and return original response + redacted_response = _redact_pii_matches(malformed_response) + assert redacted_response == malformed_response + + # Test with missing keys + missing_keys_response = { + "action": "GUARDRAIL_INTERVENED" + # Missing assessments key + } + + redacted_response = _redact_pii_matches(missing_keys_response) + assert redacted_response == missing_keys_response + + print("Malformed response redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_multiple_assessments(): + """Test _redact_pii_matches with multiple assessments containing PII""" + + response_multiple_assessments = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ { "type": "EMAIL", - "match": "user@example.com", + "match": "john@example.com", "action": "ANONYMIZED", - }, + } ] } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - { - "role": "user", - "content": "Hello, my phone number is +1 412 555 1212 and email is user@example.com", }, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # This should use the masked output even without masking flags - response = await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Should use the masked content from guardrail output - assert response is not None - assert ( - response["messages"][0]["content"] - == "Hello, my phone number is {PHONE} and email is {EMAIL}" - ) - print("✅ Masked output was applied even without masking flags enabled") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_response_pii_masking_non_streaming(): - """Test that PII masking is applied to response content in non-streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail with response masking enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - ) - - # Mock the Bedrock API response with ANONYMIZED PII - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [ - { - "text": "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" - } - ], - "assessments": [ { "sensitiveInformationPolicy": { "piiEntities": [ { "type": "CREDIT_DEBIT_CARD_NUMBER", "match": "1234-5678-9012-3456", + "action": "BLOCKED", + }, + { + "type": "ADDRESS", + "match": "123 Main St, Anytown USA", "action": "ANONYMIZED", }, + ] + } + }, + ], + "outputs": [{"text": "Multiple PII detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_multiple_assessments) + + # Verify all PII in all assessments are redacted + assessment1_pii = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + assessment2_pii = redacted_response["assessments"][1]["sensitiveInformationPolicy"][ + "piiEntities" + ] + + assert assessment1_pii[0]["match"] == "[REDACTED]", "Email should be redacted" + assert assessment2_pii[0]["match"] == "[REDACTED]", "Credit card should be redacted" + assert assessment2_pii[1]["match"] == "[REDACTED]", "Address should be redacted" + + # Verify types remain unchanged + assert assessment1_pii[0]["type"] == "EMAIL" + assert assessment2_pii[0]["type"] == "CREDIT_DEBIT_CARD_NUMBER" + assert assessment2_pii[1]["type"] == "ADDRESS" + + print("Multiple assessments redaction test passed") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_logging_uses_redacted_response(): + """Test that the Bedrock guardrail uses redacted response for logging""" + + # Create proper mock objects + mock_user_api_key_dict = UserAPIKeyAuth() + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock the Bedrock API response with PII + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ { "type": "PHONE", - "match": "+1 412 555 1212", - "action": "ANONYMIZED", - }, - ] - } - } - ], - } - - # Create a mock response that contains PII - mock_response = litellm.ModelResponse( - id="test-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", - content="My credit card number is 1234-5678-9012-3456 and my phone is +1 412 555 1212", - ), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion", - ) - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "What's your credit card and phone number?"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Call the post-call success hook - await guardrail.async_post_call_success_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - response=mock_response, - ) - - # Verify that the response content was masked - assert ( - mock_response.choices[0].message.content - == "My credit card number is {CREDIT_DEBIT_CARD_NUMBER} and my phone is {PHONE}" - ) - print("✓ Non-streaming response PII masking test passed") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_response_pii_masking_streaming(): - """Test that PII masking is applied to response content in streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail with response masking enabled - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - ) - - # Mock the Bedrock API response with ANONYMIZED PII - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "Sure! My email is {EMAIL} and SSN is {US_SSN}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "EMAIL", - "match": "john@example.com", - "action": "ANONYMIZED", - }, - { - "type": "US_SSN", - "match": "123-45-6789", - "action": "ANONYMIZED", - }, - ] - } - } - ], - } - - # Create mock streaming chunks - async def mock_streaming_response(): - chunks = [ - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="Sure! My email is "), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta( - content="john@example.com and SSN is " - ), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="123-45-6789"), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ] - for chunk in chunks: - yield chunk - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "What's your email and SSN?"}, - ], - "stream": True, - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Call the streaming hook - masked_stream = guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data, - ) - - # Collect all chunks from the masked stream - masked_chunks = [] - async for chunk in masked_stream: - masked_chunks.append(chunk) - - # Verify that we got chunks back - assert len(masked_chunks) > 0 - - # Reconstruct the full response from chunks to verify masking - full_content = "" - for chunk in masked_chunks: - if hasattr(chunk, "choices") and chunk.choices: - if hasattr(chunk.choices[0], "delta") and chunk.choices[0].delta: - if ( - hasattr(chunk.choices[0].delta, "content") - and chunk.choices[0].delta.content - ): - full_content += chunk.choices[0].delta.content - - # Verify that the reconstructed content contains the masked PII - assert "Sure! My email is {EMAIL} and SSN is {US_SSN}" == full_content - print("✓ Streaming response PII masking test passed") - - -@pytest.mark.asyncio -async def test_convert_to_bedrock_format_input_source(): - """Test convert_to_bedrock_format with INPUT source and mock messages""" - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrail, - ) - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockRequest, - ) - from unittest.mock import patch - - # Create the guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock messages - mock_messages = [ - {"role": "user", "content": "Hello, how are you?"}, - {"role": "assistant", "content": "I'm doing well, thank you!"}, - { - "role": "user", - "content": [ - {"type": "text", "text": "What's the weather like?"}, - {"type": "text", "text": "Is it sunny today?"}, - ], - }, - ] - - # Call the method - result = guardrail.convert_to_bedrock_format(source="INPUT", messages=mock_messages) - - # Verify the result structure - assert isinstance(result, dict) - assert result.get("source") == "INPUT" - assert "content" in result - assert isinstance(result.get("content"), list) - - # Verify content items - expected_content_items = [ - {"text": {"text": "Hello, how are you?"}}, - {"text": {"text": "I'm doing well, thank you!"}}, - {"text": {"text": "What's the weather like?"}}, - {"text": {"text": "Is it sunny today?"}}, - ] - - assert result.get("content") == expected_content_items - print("✅ INPUT source test passed - result:", result) - - -@pytest.mark.asyncio -async def test_convert_to_bedrock_format_output_source(): - """Test convert_to_bedrock_format with OUTPUT source and mock ModelResponse""" - from litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrail, - ) - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockRequest, - ) - import litellm - from unittest.mock import patch - - # Create the guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock ModelResponse - mock_response = litellm.ModelResponse( - id="test-response-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", content="This is a test response from the model." - ), - finish_reason="stop", - ), - litellm.Choices( - index=1, - message=litellm.Message( - role="assistant", content="This is a second choice response." - ), - finish_reason="stop", - ), - ], - created=1234567890, - model="gpt-4o", - object="chat.completion", - ) - - # Call the method - result = guardrail.convert_to_bedrock_format( - source="OUTPUT", response=mock_response - ) - - # Verify the result structure - assert isinstance(result, dict) - assert result.get("source") == "OUTPUT" - assert "content" in result - assert isinstance(result.get("content"), list) - - # Verify content items - should contain both choice contents - expected_content_items = [ - {"text": {"text": "This is a test response from the model."}}, - {"text": {"text": "This is a second choice response."}}, - ] - - assert result.get("content") == expected_content_items - print("✅ OUTPUT source test passed - result:", result) - - -@pytest.mark.asyncio -async def test_convert_to_bedrock_format_post_call_streaming_hook(): - """Test async_post_call_streaming_iterator_hook makes OUTPUT bedrock request and applies masking""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - import litellm - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock streaming chunks that contain PII - async def mock_streaming_response(): - chunks = [ - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="My email is "), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="john@example.com"), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ] - for chunk in chunks: - yield chunk - - # Mock Bedrock API response with PII masking - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "My email is {EMAIL}"}], - "assessments": [ - { - "sensitiveInformationPolicy": { - "piiEntities": [ - { - "type": "EMAIL", - "match": "john@example.com", + "match": "+1 412 555 1212", # This should be redacted in logs "action": "ANONYMIZED", } ] @@ -1095,99 +218,82 @@ async def test_convert_to_bedrock_format_post_call_streaming_hook(): request_data = { "model": "gpt-4o", - "messages": [{"role": "user", "content": "What's your email?"}], - "stream": True, + "messages": [ + {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, + ], } - # Track which bedrock API calls were made - bedrock_calls = [] + # Mock AWS credentials to avoid credential loading issues in CI + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None - # Mock the make_bedrock_api_request method to track calls - async def mock_make_bedrock_api_request( - source, messages=None, response=None, request_data=None - ): - bedrock_calls.append( - { - "source": source, - "messages": messages, - "response": response, - "request_data": request_data, - } - ) - # Return the mock bedrock response - from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( - BedrockGuardrailResponse, - ) - - return BedrockGuardrailResponse(**mock_bedrock_response.json()) - - # Patch the bedrock API request method + # Mock AWS-related methods to ensure test runs without external dependencies with patch.object( - guardrail, "make_bedrock_api_request", side_effect=mock_make_bedrock_api_request - ): + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch( + "litellm.proxy.guardrails.guardrail_hooks.bedrock_guardrails.verbose_proxy_logger.debug" + ) as mock_debug, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ) as mock_load_creds, patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ) as mock_prepare_request: - # Call the streaming hook - result_generator = guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), + mock_post.return_value = mock_bedrock_response + + # Call the method that should log the redacted response + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), request_data=request_data, ) - # Collect all chunks from the result - result_chunks = [] - async for chunk in result_generator: - result_chunks.append(chunk) + # Verify that debug logging was called + mock_debug.assert_called() + + # Get the logged response (second argument to debug call) + logged_calls = mock_debug.call_args_list + bedrock_response_log_call = None + + for call in logged_calls: + args, kwargs = call + if len(args) >= 2 and "Bedrock AI response" in str(args[0]): + bedrock_response_log_call = call + break - # Verify bedrock API calls were made - # Note: When event_hook is None (default), the guardrail is considered enabled for all hooks. - # In post_call, INPUT validation is skipped if pre_call/during_call is already enabled - # to avoid redundant validation. Since event_hook=None means all hooks are enabled, - # only OUTPUT validation should be performed in post_call. assert ( - len(bedrock_calls) == 1 - ), f"Expected 1 bedrock call (OUTPUT only), got {len(bedrock_calls)}" + bedrock_response_log_call is not None + ), "Should have logged Bedrock AI response" - # Verify the OUTPUT call - output_call = bedrock_calls[0] - assert output_call["source"] == "OUTPUT" - assert output_call["response"] is not None - assert output_call["messages"] is None # OUTPUT calls don't need messages + # Extract the logged response data + logged_response = bedrock_response_log_call[0][ + 1 + ] # Second argument to debug call - # Verify that the response content was masked - # The streaming chunks should now contain the masked content - full_content = "" - for chunk in result_chunks: - if hasattr(chunk, "choices") and chunk.choices: - if ( - hasattr(chunk.choices[0], "delta") - and chunk.choices[0].delta.content - ): - full_content += chunk.choices[0].delta.content - - # The content should be masked (contains {EMAIL} instead of john@example.com) + # Verify that the logged response has redacted PII assert ( - "{EMAIL}" in full_content - ), f"Expected masked content with {{EMAIL}}, got: {full_content}" - assert ( - "john@example.com" not in full_content - ), f"Original email should be masked, got: {full_content}" - - print( - "✅ Post-call streaming hook test passed - OUTPUT source used for masking" + logged_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] + == "[REDACTED]" ) - print( - f"✅ Bedrock calls made: {[call['source'] for call in bedrock_calls]} " - "(INPUT validation skipped due to event_hook=None implying pre_call/during_call enabled)" + + # Verify other fields are preserved + assert logged_response["action"] == "GUARDRAIL_INTERVENED" + assert ( + logged_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["type"] + == "PHONE" ) - print(f"✅ Final masked content: {full_content}") + + print("Bedrock guardrail logging redaction test passed") @pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_action_shows_output_text(): - """Test that BLOCKED actions raise HTTPException with the output text in the detail""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from fastapi import HTTPException +async def test_bedrock_guardrail_original_response_not_modified(): + """Test that the original response is not modified by redaction, only the logged version""" # Create proper mock objects mock_user_api_key_dict = UserAPIKeyAuth() @@ -1196,408 +302,1274 @@ async def test_bedrock_guardrail_blocked_action_shows_output_text(): guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" ) - # Mock the Bedrock API response with BLOCKED action and output text - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { + # Mock the Bedrock API response with PII + original_response_data = { "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "this violates litellm corporate guardrail policy"}], + "outputs": [{"text": "Hello, my phone number is {PHONE}"}], "assessments": [ { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "PHONE", + "match": "+1 412 555 1212", # This should NOT be modified in original + "action": "ANONYMIZED", + } ] } } ], } + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = original_response_data + request_data = { "model": "gpt-4o", "messages": [ - {"role": "user", "content": "Tell me how to make explosives"}, + {"role": "user", "content": "Hello, my phone number is +1 412 555 1212"}, ], } - # Patch the async_handler.post method + # Mock AWS credentials to avoid credential loading issues in CI + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Mock AWS-related methods to ensure test runs without external dependencies with patch.object( guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ) as mock_load_creds, patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ) as mock_prepare_request: + mock_post.return_value = mock_bedrock_response - # This should raise HTTPException due to BLOCKED action - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Verify the exception details - exception = exc_info.value - assert exception.status_code == 400 - assert "detail" in exception.__dict__ - - # Check that the detail contains the expected structure - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - - # Verify that the output text from both outputs is included - expected_output_text = "this violates litellm corporate guardrail policy" - assert detail["bedrock_guardrail_response"] == expected_output_text - - print( - "✅ BLOCKED action HTTPException test passed - output text properly included" + # Call the method + result = await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), + request_data=request_data, ) + # Verify that the original response data was not modified + # (The json() method should return the original data) + original_data = mock_bedrock_response.json() + assert ( + original_data["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ][0]["match"] + == "+1 412 555 1212" + ) + + # Verify that the returned BedrockGuardrailResponse contains original data + assert ( + result["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][ + "match" + ] + == "+1 412 555 1212" + ) + + print("Original response not modified test passed") + @pytest.mark.asyncio -async def test_bedrock_guardrail_blocked_action_empty_outputs(): - """Test that BLOCKED actions with empty outputs still raise HTTPException""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from fastapi import HTTPException +async def test__redact_pii_matches_preserves_non_pii_entities(): + """Test that _redact_pii_matches only affects PII-related entities and preserves other assessment data""" - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Mock the Bedrock API response with BLOCKED action but no outputs - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { + response_with_mixed_data = { "action": "GUARDRAIL_INTERVENED", - "outputs": [], # Empty outputs "assessments": [ { - "contentPolicy": { - "filters": [ - {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Violent content here"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # This should raise HTTPException due to BLOCKED action - with pytest.raises(HTTPException) as exc_info: - await guardrail.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Verify the exception details - exception = exc_info.value - assert exception.status_code == 400 - - # Check that the detail contains the expected structure with empty output text - detail = exception.detail - assert isinstance(detail, dict) - assert detail["error"] == "Violated guardrail policy" - assert detail["bedrock_guardrail_response"] == "" # Empty string for no outputs - - print("✅ BLOCKED action with empty outputs test passed") - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_disable_exception_on_block_non_streaming(): - """Test that disable_exception_on_block=True prevents exceptions in non-streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from fastapi import HTTPException - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Test 1: disable_exception_on_block=False (default) - should raise exception - guardrail_default = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=False, - ) - - # Mock the Bedrock API response with BLOCKED action - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "topicPolicy": { - "topics": [ - {"name": "Sensitive Topic", "type": "DENY", "action": "BLOCKED"} - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [ - {"role": "user", "content": "Tell me how to make explosives"}, - ], - } - - # Patch the async_handler.post method - with patch.object( - guardrail_default.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should raise HTTPException when disable_exception_on_block=False - with pytest.raises(HTTPException) as exc_info: - await guardrail_default.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - - # Verify the exception details - exception = exc_info.value - assert exception.status_code == 400 - assert "Violated guardrail policy" in str(exception.detail) - - # Test 2: disable_exception_on_block=True - should NOT raise exception - guardrail_disabled = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=True, - ) - - with patch.object( - guardrail_disabled.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should NOT raise exception when disable_exception_on_block=True - try: - response = await guardrail_disabled.async_moderation_hook( - data=request_data, - user_api_key_dict=mock_user_api_key_dict, - call_type="completion", - ) - # Should succeed and return data (even though content was blocked) - assert response is not None - print("✅ No exception raised when disable_exception_on_block=True") - except Exception as e: - pytest.fail( - f"Should not raise exception when disable_exception_on_block=True, but got: {e}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_disable_exception_on_block_streaming(): - """Test that disable_exception_on_block=True prevents exceptions in streaming scenarios""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - from fastapi import HTTPException - import litellm - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Mock streaming chunks that would normally trigger a block - async def mock_streaming_response(): - chunks = [ - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta( - content="Here's how to make explosives: " - ), - finish_reason=None, - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ModelResponseStream( - id="test-id", - choices=[ - litellm.utils.StreamingChoices( - index=0, - delta=litellm.utils.Delta(content="step 1, step 2..."), - finish_reason="stop", - ) - ], - created=1234567890, - model="gpt-4o", - object="chat.completion.chunk", - ), - ] - for chunk in chunks: - yield chunk - - # Mock Bedrock API response with BLOCKED action - mock_bedrock_response = MagicMock() - mock_bedrock_response.status_code = 200 - mock_bedrock_response.json.return_value = { - "action": "GUARDRAIL_INTERVENED", - "outputs": [{"text": "I can't provide that information."}], - "assessments": [ - { - "contentPolicy": { - "filters": [ - {"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"} - ] - } - } - ], - } - - request_data = { - "model": "gpt-4o", - "messages": [{"role": "user", "content": "Tell me how to make explosives"}], - "stream": True, - } - - # Test 1: disable_exception_on_block=False (default) - should raise exception - guardrail_default = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=False, - ) - - with patch.object( - guardrail_default.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should raise exception during streaming processing - with pytest.raises(HTTPException): - result_generator = ( - guardrail_default.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data, - ) - ) - - # Try to consume the generator - should raise exception - async for chunk in result_generator: - pass - - # Test 2: disable_exception_on_block=True - should NOT raise exception - guardrail_disabled = BedrockGuardrail( - guardrailIdentifier="test-guardrail", - guardrailVersion="DRAFT", - disable_exception_on_block=True, - ) - - with patch.object( - guardrail_disabled.async_handler, "post", new_callable=AsyncMock - ) as mock_post: - mock_post.return_value = mock_bedrock_response - - # Should NOT raise exception when disable_exception_on_block=True - try: - result_generator = ( - guardrail_disabled.async_post_call_streaming_iterator_hook( - user_api_key_dict=mock_user_api_key_dict, - response=mock_streaming_response(), - request_data=request_data, - ) - ) - - # Consume the generator - should succeed without exceptions - result_chunks = [] - async for chunk in result_generator: - result_chunks.append(chunk) - - # Should have received chunks back even though content was blocked - assert len(result_chunks) > 0 - print( - "✅ Streaming completed without exception when disable_exception_on_block=True" - ) - - except Exception as e: - pytest.fail( - f"Should not raise exception when disable_exception_on_block=True in streaming, but got: {e}" - ) - - -@pytest.mark.asyncio -async def test_bedrock_guardrail_post_call_success_hook_no_output_text(): - """Test that async_post_call_success_hook skips when there's no output text""" - from unittest.mock import AsyncMock, MagicMock, patch - from litellm.proxy._types import UserAPIKeyAuth - from litellm.types.utils import ModelResponseStream - import litellm - - # Create proper mock objects - mock_user_api_key_dict = UserAPIKeyAuth() - - # Create guardrail instance - guardrail = BedrockGuardrail( - guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" - ) - - # Create a ModelResponse with tool calls (no text content) - # This simulates a response where the LLM is making a tool call - mock_response = litellm.ModelResponse( - id="test-id", - choices=[ - litellm.Choices( - index=0, - message=litellm.Message( - role="assistant", - content=None, # No text content - tool_calls=[ - litellm.utils.ChatCompletionMessageToolCall( - id="tooluse_kZJMlvQmRJ6eAyJE5GIl7Q", - function=litellm.utils.Function( - name="top_song", arguments='{"sign": "WZPZ"}' - ), - type="function", - ) + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "EMAIL", + "match": "user@example.com", + "action": "ANONYMIZED", + "confidence": "HIGH", + } ], - ), - finish_reason="tool_calls", - ) + "regexes": [ + { + "name": "custom_pattern", + "match": "some_pattern_match", + "action": "BLOCKED", + } + ], + }, + "contentPolicy": { + "filters": [ + { + "type": "VIOLENCE", + "confidence": "MEDIUM", + "action": "BLOCKED", + } + ] + }, + "topicPolicy": { + "topics": [ + { + "name": "Restricted Topic", + "type": "DENY", + "action": "BLOCKED", + } + ] + }, + } ], - created=1234567890, - model="gpt-4o", - object="chat.completion", + "outputs": [{"text": "Content blocked"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_mixed_data) + + # Verify that PII entity matches are redacted + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + assert pii_entities[0]["match"] == "[REDACTED]", "PII match should be redacted" + assert pii_entities[0]["type"] == "EMAIL", "PII type should be preserved" + assert pii_entities[0]["action"] == "ANONYMIZED", "PII action should be preserved" + assert pii_entities[0]["confidence"] == "HIGH", "PII confidence should be preserved" + + # Verify that regex matches are also redacted (updated behavior) + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "regexes" + ] + assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" + assert regexes[0]["name"] == "custom_pattern", "Regex name should be preserved" + assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" + + # Verify that other policies are completely unchanged + content_policy = redacted_response["assessments"][0]["contentPolicy"] + assert content_policy["filters"][0]["type"] == "VIOLENCE" + assert content_policy["filters"][0]["confidence"] == "MEDIUM" + + topic_policy = redacted_response["assessments"][0]["topicPolicy"] + assert topic_policy["topics"][0]["name"] == "Restricted Topic" + + # Verify top-level fields are unchanged + assert redacted_response["action"] == "GUARDRAIL_INTERVENED" + assert redacted_response["outputs"][0]["text"] == "Content blocked" + + print("Preserves non-PII entities test passed") + + +@pytest.mark.asyncio +async def test_pii_redaction_matches_debug_output_format(): + """Test that demonstrates the exact behavior shown in your debug output""" + + # This matches the structure from your debug output + original_response = { + "action": "GUARDRAIL_INTERVENED", + "actionReason": "Guardrail blocked.", + "assessments": [ + { + "invocationMetrics": { + "guardrailCoverage": { + "textCharacters": {"guarded": 84, "total": 84} + }, + "guardrailProcessingLatency": 322, + "usage": { + "contentPolicyImageUnits": 0, + "contentPolicyUnits": 0, + "contextualGroundingPolicyUnits": 0, + "sensitiveInformationPolicyFreeUnits": 0, + "sensitiveInformationPolicyUnits": 1, + "topicPolicyUnits": 0, + "wordPolicyUnits": 0, + }, + }, + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "action": "BLOCKED", + "detected": True, + "match": "John Smith", + "type": "NAME", + }, + { + "action": "BLOCKED", + "detected": True, + "match": "324-12-3212", + "type": "US_SOCIAL_SECURITY_NUMBER", + }, + { + "action": "BLOCKED", + "detected": True, + "match": "607-456-7890", + "type": "PHONE", + }, + ] + }, + } + ], + "blockedResponse": "Input blocked by PII policy", + "guardrailCoverage": {"textCharacters": {"guarded": 84, "total": 84}}, + "output": [{"text": "Input blocked by PII policy"}], + "outputs": [{"text": "Input blocked by PII policy"}], + "usage": { + "contentPolicyImageUnits": 0, + "contentPolicyUnits": 0, + "contextualGroundingPolicyUnits": 0, + "sensitiveInformationPolicyFreeUnits": 0, + "sensitiveInformationPolicyUnits": 1, + "topicPolicyUnits": 0, + "wordPolicyUnits": 0, + }, + } + + # Apply redaction + redacted_response = _redact_pii_matches(original_response) + + # Verify the redacted response matches your expected debug output + pii_entities = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "piiEntities" + ] + + # All PII matches should be redacted + assert pii_entities[0]["match"] == "[REDACTED]", "NAME should be redacted" + assert pii_entities[1]["match"] == "[REDACTED]", "SSN should be redacted" + assert pii_entities[2]["match"] == "[REDACTED]", "PHONE should be redacted" + + # But all other fields should be preserved + assert pii_entities[0]["type"] == "NAME" + assert pii_entities[1]["type"] == "US_SOCIAL_SECURITY_NUMBER" + assert pii_entities[2]["type"] == "PHONE" + assert pii_entities[0]["action"] == "BLOCKED" + assert pii_entities[0]["detected"] == True + + # Verify that the original response is unchanged + original_pii_entities = original_response["assessments"][0][ + "sensitiveInformationPolicy" + ]["piiEntities"] + assert ( + original_pii_entities[0]["match"] == "John Smith" + ), "Original should be unchanged" + assert ( + original_pii_entities[1]["match"] == "324-12-3212" + ), "Original should be unchanged" + assert ( + original_pii_entities[2]["match"] == "607-456-7890" + ), "Original should be unchanged" + + # Verify all other metadata is preserved in redacted response + assert redacted_response["action"] == "GUARDRAIL_INTERVENED" + assert redacted_response["actionReason"] == "Guardrail blocked." + assert redacted_response["blockedResponse"] == "Input blocked by PII policy" + assert ( + redacted_response["assessments"][0]["invocationMetrics"][ + "guardrailProcessingLatency" + ] + == 322 ) - data = { + print("PII redaction matches debug output format test passed") + print( + f"Original PII values preserved: {[e['match'] for e in original_pii_entities]}" + ) + print(f"Redacted PII values: {[e['match'] for e in pii_entities]}") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_with_regex_matches(): + """Test redaction of regex matches in sensitive information policy""" + + response_with_regex = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "regexes": [ + { + "name": "SSN_PATTERN", + "match": "123-45-6789", + "action": "BLOCKED", + }, + { + "name": "CREDIT_CARD_PATTERN", + "match": "4111-1111-1111-1111", + "action": "ANONYMIZED", + }, + ] + } + } + ], + "outputs": [{"text": "Regex patterns detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_regex) + + # Verify that regex matches are redacted + regexes = redacted_response["assessments"][0]["sensitiveInformationPolicy"][ + "regexes" + ] + + assert regexes[0]["match"] == "[REDACTED]", "SSN regex match should be redacted" + assert ( + regexes[1]["match"] == "[REDACTED]" + ), "Credit card regex match should be redacted" + + # Verify other fields are preserved + assert regexes[0]["name"] == "SSN_PATTERN", "Regex name should be preserved" + assert regexes[0]["action"] == "BLOCKED", "Regex action should be preserved" + assert regexes[1]["name"] == "CREDIT_CARD_PATTERN", "Regex name should be preserved" + assert regexes[1]["action"] == "ANONYMIZED", "Regex action should be preserved" + + # Verify original response is unchanged + original_regexes = response_with_regex["assessments"][0][ + "sensitiveInformationPolicy" + ]["regexes"] + assert original_regexes[0]["match"] == "123-45-6789", "Original should be unchanged" + assert ( + original_regexes[1]["match"] == "4111-1111-1111-1111" + ), "Original should be unchanged" + + print("Regex matches redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_with_custom_words(): + """Test redaction of custom word matches in word policy""" + + response_with_custom_words = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "wordPolicy": { + "customWords": [ + { + "match": "confidential_data", + "action": "BLOCKED", + }, + { + "match": "secret_information", + "action": "ANONYMIZED", + }, + ] + } + } + ], + "outputs": [{"text": "Custom words detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_custom_words) + + # Verify that custom word matches are redacted + custom_words = redacted_response["assessments"][0]["wordPolicy"]["customWords"] + + assert ( + custom_words[0]["match"] == "[REDACTED]" + ), "First custom word match should be redacted" + assert ( + custom_words[1]["match"] == "[REDACTED]" + ), "Second custom word match should be redacted" + + # Verify other fields are preserved + assert ( + custom_words[0]["action"] == "BLOCKED" + ), "Custom word action should be preserved" + assert ( + custom_words[1]["action"] == "ANONYMIZED" + ), "Custom word action should be preserved" + + # Verify original response is unchanged + original_custom_words = response_with_custom_words["assessments"][0]["wordPolicy"][ + "customWords" + ] + assert ( + original_custom_words[0]["match"] == "confidential_data" + ), "Original should be unchanged" + assert ( + original_custom_words[1]["match"] == "secret_information" + ), "Original should be unchanged" + + print("Custom words redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_with_managed_words(): + """Test redaction of managed word matches in word policy""" + + response_with_managed_words = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "wordPolicy": { + "managedWordLists": [ + { + "match": "inappropriate_word", + "action": "BLOCKED", + "type": "PROFANITY", + }, + { + "match": "offensive_term", + "action": "ANONYMIZED", + "type": "HATE_SPEECH", + }, + ] + } + } + ], + "outputs": [{"text": "Managed words detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(response_with_managed_words) + + # Verify that managed word matches are redacted + managed_words = redacted_response["assessments"][0]["wordPolicy"][ + "managedWordLists" + ] + + assert ( + managed_words[0]["match"] == "[REDACTED]" + ), "First managed word match should be redacted" + assert ( + managed_words[1]["match"] == "[REDACTED]" + ), "Second managed word match should be redacted" + + # Verify other fields are preserved + assert ( + managed_words[0]["action"] == "BLOCKED" + ), "Managed word action should be preserved" + assert ( + managed_words[0]["type"] == "PROFANITY" + ), "Managed word type should be preserved" + assert ( + managed_words[1]["action"] == "ANONYMIZED" + ), "Managed word action should be preserved" + assert ( + managed_words[1]["type"] == "HATE_SPEECH" + ), "Managed word type should be preserved" + + # Verify original response is unchanged + original_managed_words = response_with_managed_words["assessments"][0][ + "wordPolicy" + ]["managedWordLists"] + assert ( + original_managed_words[0]["match"] == "inappropriate_word" + ), "Original should be unchanged" + assert ( + original_managed_words[1]["match"] == "offensive_term" + ), "Original should be unchanged" + + print("Managed words redaction test passed") + + +@pytest.mark.asyncio +async def test__redact_pii_matches_comprehensive_coverage(): + """Test redaction across all supported policy types in a single response""" + + comprehensive_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "EMAIL", + "match": "user@example.com", + "action": "ANONYMIZED", + } + ], + "regexes": [ + { + "name": "PHONE_PATTERN", + "match": "555-123-4567", + "action": "BLOCKED", + } + ], + }, + "wordPolicy": { + "customWords": [ + { + "match": "confidential", + "action": "BLOCKED", + } + ], + "managedWordLists": [ + { + "match": "inappropriate", + "action": "ANONYMIZED", + "type": "PROFANITY", + } + ], + }, + } + ], + "outputs": [{"text": "Multiple policy violations detected"}], + } + + # Call the redaction function + redacted_response = _redact_pii_matches(comprehensive_response) + + # Verify all match fields are redacted + assessment = redacted_response["assessments"][0] + + # PII entities + pii_entities = assessment["sensitiveInformationPolicy"]["piiEntities"] + assert ( + pii_entities[0]["match"] == "[REDACTED]" + ), "PII entity match should be redacted" + + # Regex matches + regexes = assessment["sensitiveInformationPolicy"]["regexes"] + assert regexes[0]["match"] == "[REDACTED]", "Regex match should be redacted" + + # Custom words + custom_words = assessment["wordPolicy"]["customWords"] + assert ( + custom_words[0]["match"] == "[REDACTED]" + ), "Custom word match should be redacted" + + # Managed words + managed_words = assessment["wordPolicy"]["managedWordLists"] + assert ( + managed_words[0]["match"] == "[REDACTED]" + ), "Managed word match should be redacted" + + # Verify all other fields are preserved + assert pii_entities[0]["type"] == "EMAIL" + assert regexes[0]["name"] == "PHONE_PATTERN" + assert managed_words[0]["type"] == "PROFANITY" + + # Verify original response is unchanged + original_assessment = comprehensive_response["assessments"][0] + assert ( + original_assessment["sensitiveInformationPolicy"]["piiEntities"][0]["match"] + == "user@example.com" + ) + assert ( + original_assessment["sensitiveInformationPolicy"]["regexes"][0]["match"] + == "555-123-4567" + ) + assert ( + original_assessment["wordPolicy"]["customWords"][0]["match"] == "confidential" + ) + assert ( + original_assessment["wordPolicy"]["managedWordLists"][0]["match"] + == "inappropriate" + ) + + print("Comprehensive coverage redaction test passed") + +@pytest.mark.asyncio +async def test_bedrock_guardrail_respects_custom_runtime_endpoint(monkeypatch): + """Test that BedrockGuardrail respects aws_bedrock_runtime_endpoint when set""" + + # Clear any existing environment variable to ensure clean test + monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) + + # Create guardrail with custom runtime endpoint + custom_endpoint = "https://custom-bedrock.example.com" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + aws_bedrock_runtime_endpoint=custom_endpoint, + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-east-1" + + # Mock the _load_credentials method to avoid actual AWS credential loading + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the custom endpoint is used in the URL + expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected URL to contain custom endpoint. Got: {prepped_request.url}" + + print(f"Custom runtime endpoint test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_respects_env_runtime_endpoint(monkeypatch): + """Test that BedrockGuardrail respects AWS_BEDROCK_RUNTIME_ENDPOINT environment variable""" + + custom_endpoint = "https://env-bedrock.example.com" + + # Set the environment variable + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", custom_endpoint) + + # Create guardrail without explicit aws_bedrock_runtime_endpoint + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-east-1" + + # Mock the _load_credentials method + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the custom endpoint from environment is used in the URL + expected_url = f"{custom_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected URL to contain env endpoint. Got: {prepped_request.url}" + + print(f"Environment runtime endpoint test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_uses_default_endpoint_when_no_custom_set(monkeypatch): + """Test that BedrockGuardrail uses default endpoint when no custom endpoint is set""" + + # Ensure no environment variable is set + monkeypatch.delenv("AWS_BEDROCK_RUNTIME_ENDPOINT", raising=False) + + # Create guardrail without any custom endpoint + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-west-2" + + # Mock the _load_credentials method + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the default endpoint is used + expected_url = f"https://bedrock-runtime.{aws_region_name}.amazonaws.com/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected default URL. Got: {prepped_request.url}" + + print(f"Default endpoint test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_parameter_takes_precedence_over_env(monkeypatch): + """Test that aws_bedrock_runtime_endpoint parameter takes precedence over environment variable + + This test verifies the corrected behavior where the parameter should take precedence + over the environment variable, consistent with the endpoint_url logic. + """ + + param_endpoint = "https://param-bedrock.example.com" + env_endpoint = "https://env-bedrock.example.com" + + # Set environment variable + monkeypatch.setenv("AWS_BEDROCK_RUNTIME_ENDPOINT", env_endpoint) + + # Create guardrail with explicit aws_bedrock_runtime_endpoint + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + aws_bedrock_runtime_endpoint=param_endpoint, + ) + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + # Test data + data = {"source": "INPUT", "content": [{"text": {"text": "test content"}}]} + optional_params = {} + aws_region_name = "us-east-1" + + # Mock the _load_credentials method + with patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, aws_region_name) + ): + # Call _prepare_request which internally calls get_runtime_endpoint + prepped_request = guardrail._prepare_request( + credentials=mock_credentials, + data=data, + optional_params=optional_params, + aws_region_name=aws_region_name, + ) + + # Verify that the parameter takes precedence over environment variable + expected_url = f"{param_endpoint}/guardrail/{guardrail.guardrailIdentifier}/version/{guardrail.guardrailVersion}/apply" + assert ( + prepped_request.url == expected_url + ), f"Expected parameter endpoint to take precedence. Got: {prepped_request.url}" + + print(f"Parameter precedence test passed. URL: {prepped_request.url}") + + +@pytest.mark.asyncio +async def test_bedrock_apply_guardrail_with_only_tool_calls_response(): + """Test that apply_guardrail handles response with tool_calls (no text content) without calling Bedrock API""" + # Create a BedrockGuardrail instance + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + # Mock the make_bedrock_api_request method + with patch.object( + guardrail, "make_bedrock_api_request", new_callable=AsyncMock + ) as mock_api_request: + # Test the apply_guardrail method with tool_calls in response + inputs = { + "texts": [], + "tool_calls": [ + { + "id": "call_eFSCWFsyL7MclHYnzKrcQnMK", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"location":"São Paulo"}', + }, + } + ], + } + + guardrailed_inputs = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="response", + logging_obj=None, + ) + + # Verify the result - should succeed without errors + assert guardrailed_inputs is not None + assert "tool_calls" in guardrailed_inputs + assert len(guardrailed_inputs["tool_calls"]) == 1 + assert ( + guardrailed_inputs["tool_calls"][0]["id"] + == "call_eFSCWFsyL7MclHYnzKrcQnMK" + ) + assert guardrailed_inputs["tool_calls"][0]["function"]["name"] == "get_weather" + assert ( + guardrailed_inputs["tool_calls"][0]["function"]["arguments"] + == '{"location":"São Paulo"}' + ) + # Verify that the Bedrock API was NOT called since there's no text to process + mock_api_request.assert_not_called() + print("✅ apply_guardrail with tool_calls test passed - no API call made") + + +@pytest.mark.asyncio +async def test_bedrock_guardrail_blocked_content_with_masking_enabled(): + """Test that BLOCKED content raises exception even when masking is enabled + + This test verifies the bug fix where previously mask_request_content=True or + mask_response_content=True would bypass all BLOCKED content checks. Now it + properly distinguishes between BLOCKED (raise exception) and ANONYMIZED (apply masking). + """ + + # Create guardrail with masking enabled + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + mask_request_content=True, # Masking enabled + mask_response_content=True, # Masking enabled + ) + + # Mock Bedrock response with BLOCKED content (hate speech) + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "contentPolicy": { + "filters": [ + { + "type": "HATE", + "confidence": "HIGH", + "action": "BLOCKED", # Should raise exception + } + ] + }, + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "type": "NAME", + "match": "John Doe", + "action": "ANONYMIZED", # Should be masked + } + ] + }, + } + ], + "outputs": [{"text": "Content blocked due to policy violation"}], + } + + mock_bedrock_response = MagicMock() + mock_bedrock_response.status_code = 200 + mock_bedrock_response.json.return_value = blocked_response + + # Mock credentials + mock_credentials = MagicMock() + mock_credentials.access_key = "test-access-key" + mock_credentials.secret_key = "test-secret-key" + mock_credentials.token = None + + request_data = { "model": "gpt-4o", "messages": [ - {"role": "user", "content": "Hello"}, + {"role": "user", "content": "Test message with PII and hate speech"}, ], } - mock_user_api_key_dict = UserAPIKeyAuth() + + # Mock AWS-related methods + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ): + mock_post.return_value = mock_bedrock_response + + # Should raise HTTPException for BLOCKED content + with pytest.raises(HTTPException) as exc_info: + await guardrail.make_bedrock_api_request( + source="INPUT", + messages=request_data.get("messages"), + request_data=request_data, + ) + + # Verify exception details + assert exc_info.value.status_code == 400 + assert "Violated guardrail policy" in str(exc_info.value.detail) + + print("✅ BLOCKED content with masking enabled raises exception correctly") - result = await guardrail.async_post_call_success_hook( - data=data, - response=mock_response, - user_api_key_dict=mock_user_api_key_dict, - ) - # If no error is raised and result is None, then the test passes - assert result is None - print("✅ No output text in response test passed") + +# ────────────────────────────────────────────────────────────────────────────── +# Null-safety tests for Bedrock guardrail responses +# +# The Bedrock ApplyGuardrail API can return explicit null/None for list fields +# such as "regexes", "piiEntities", "topics", "filters", "customWords", and +# "managedWordLists" when a particular policy category is present in the +# assessment but has no matches. +# +# Python's dict.get("key", []) returns None (NOT []) when the key exists with +# a None value. The `or []` fallback ensures we always iterate over a list. +# +# Without the fix, iterating over None raises: +# TypeError: 'NoneType' object is not iterable +# which surfaces to callers as: +# openai.InternalServerError: Error code: 500 +# {'error': {'message': "Bedrock guardrail failed: 'NoneType' object is not iterable", ...}} +# ────────────────────────────────────────────────────────────────────────────── + + +class TestRedactPiiMatchesNullSafety: + """Tests for _redact_pii_matches handling of null/None list fields from Bedrock API.""" + + @pytest.mark.asyncio + async def test_should_handle_null_regexes_in_sensitive_info_policy(self): + """Bedrock can return regexes: null while piiEntities has data. + + Real-world scenario: guardrail detects PII (e.g. EMAIL) but has no + custom regex patterns configured, so the API returns regexes: null. + """ + response = { + "action": "NONE", + "actionReason": "No action.", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": [ + { + "action": "NONE", + "detected": True, + "match": "joebloggs@gmail.com", + "type": "EMAIL", + } + ], + "regexes": None, # Explicit null from Bedrock API + }, + } + ], + } + + # Should not raise TypeError: 'NoneType' object is not iterable + redacted = _redact_pii_matches(response) + + # PII match should be redacted + pii = redacted["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"] + assert pii[0]["match"] == "[REDACTED]" + assert pii[0]["type"] == "EMAIL" + + @pytest.mark.asyncio + async def test_should_handle_null_pii_entities_in_sensitive_info_policy(self): + """Bedrock can return piiEntities: null while regexes has data.""" + response = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, # null from Bedrock API + "regexes": [ + { + "name": "CUSTOM_PATTERN", + "match": "secret-abc-123", + "action": "BLOCKED", + } + ], + }, + } + ], + } + + redacted = _redact_pii_matches(response) + + regexes = redacted["assessments"][0]["sensitiveInformationPolicy"]["regexes"] + assert regexes[0]["match"] == "[REDACTED]" + + @pytest.mark.asyncio + async def test_should_handle_null_custom_words_and_managed_words(self): + """Bedrock can return null for customWords and managedWordLists in wordPolicy.""" + response = { + "action": "NONE", + "assessments": [ + { + "wordPolicy": { + "customWords": None, # null from Bedrock API + "managedWordLists": None, # null from Bedrock API + }, + } + ], + } + + # Should not raise TypeError + redacted = _redact_pii_matches(response) + + # Values should remain None (no crash) + assert redacted["assessments"][0]["wordPolicy"]["customWords"] is None + assert redacted["assessments"][0]["wordPolicy"]["managedWordLists"] is None + + @pytest.mark.asyncio + async def test_should_handle_null_assessments_list(self): + """Bedrock can return assessments: null.""" + response = { + "action": "NONE", + "assessments": None, # null from Bedrock API + } + + # Should not raise TypeError + redacted = _redact_pii_matches(response) + assert redacted["assessments"] is None + + @pytest.mark.asyncio + async def test_should_handle_all_null_policy_sub_lists_together(self): + """All sub-list fields are null at the same time — worst-case scenario.""" + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + }, + "wordPolicy": { + "customWords": None, + "managedWordLists": None, + }, + "topicPolicy": None, + "contentPolicy": None, + "contextualGroundingPolicy": None, + } + ], + } + + # Should not raise any exception + redacted = _redact_pii_matches(response) + assert redacted is not None + + +class TestShouldRaiseGuardrailBlockedExceptionNullSafety: + """Tests for _should_raise_guardrail_blocked_exception handling of null list fields.""" + + def _create_guardrail(self) -> BedrockGuardrail: + return BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + @pytest.mark.asyncio + async def test_should_handle_all_null_policy_sub_lists(self): + """All policy sub-lists are null — should not crash, should return False.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, # null from Bedrock API + }, + "contentPolicy": { + "filters": None, # null + }, + "wordPolicy": { + "customWords": None, # null + "managedWordLists": None, # null + }, + "sensitiveInformationPolicy": { + "piiEntities": None, # null + "regexes": None, # null + }, + "contextualGroundingPolicy": { + "filters": None, # null + }, + } + ], + } + + # No BLOCKED actions found (all lists null) → should return False + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_detect_blocked_despite_other_null_lists(self): + """A mix of null lists and a real BLOCKED action — should still detect it.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, # null — should not crash + }, + "contentPolicy": { + "filters": [ + { + "type": "HATE", + "confidence": "HIGH", + "action": "BLOCKED", + } + ], + }, + "wordPolicy": { + "customWords": None, # null + "managedWordLists": None, # null + }, + "sensitiveInformationPolicy": { + "piiEntities": None, # null + "regexes": None, # null + }, + "contextualGroundingPolicy": None, # entire policy is null + } + ], + } + + # Should return True because contentPolicy has a BLOCKED filter + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_assessments_list(self): + """assessments itself is null — should return False.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": None, # null from Bedrock API + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_handle_null_topics_with_blocked_word_policy(self): + """topics is null but wordPolicy has a BLOCKED customWord.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "topicPolicy": { + "topics": None, + }, + "wordPolicy": { + "customWords": [ + {"match": "badword", "action": "BLOCKED"} + ], + "managedWordLists": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_pii_with_blocked_regex(self): + """piiEntities is null but regexes has a BLOCKED match.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": [ + {"name": "SSN", "match": "123-45-6789", "action": "BLOCKED"} + ], + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is True + + @pytest.mark.asyncio + async def test_should_handle_null_grounding_filters(self): + """contextualGroundingPolicy.filters is null — should not crash.""" + guardrail = self._create_guardrail() + + response = { + "action": "GUARDRAIL_INTERVENED", + "assessments": [ + { + "contextualGroundingPolicy": { + "filters": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + @pytest.mark.asyncio + async def test_should_not_crash_when_action_is_not_intervened(self): + """If action != GUARDRAIL_INTERVENED, null lists should never be reached.""" + guardrail = self._create_guardrail() + + response = { + "action": "NONE", + "assessments": [ + { + "sensitiveInformationPolicy": { + "piiEntities": None, + "regexes": None, + }, + } + ], + } + + result = guardrail._should_raise_guardrail_blocked_exception(response) + assert result is False + + +class TestApplyGuardrailNullSafety: + """Tests for apply_guardrail handling of null/None texts input.""" + + @pytest.mark.asyncio + async def test_should_handle_none_texts_in_inputs(self): + """inputs[\"texts\"] is explicitly None — should not crash.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + inputs = {"texts": None} # Explicit None + + mock_credentials = MagicMock() + + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ): + # With empty texts (from None → []), no Bedrock API call should be made + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + # Should return empty texts without crashing + assert result.get("texts") == [] + # No Bedrock API call should be made for empty input + mock_post.assert_not_called() + + @pytest.mark.asyncio + async def test_should_handle_missing_texts_key(self): + """inputs has no \"texts\" key at all — should not crash.""" + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", guardrailVersion="DRAFT" + ) + + inputs = {} # No "texts" key + + mock_credentials = MagicMock() + + with patch.object( + guardrail.async_handler, "post", new_callable=AsyncMock + ) as mock_post, patch.object( + guardrail, "_load_credentials", return_value=(mock_credentials, "us-east-1") + ), patch.object( + guardrail, "_prepare_request", return_value=MagicMock() + ): + result = await guardrail.apply_guardrail( + inputs=inputs, + request_data={}, + input_type="request", + ) + + assert result.get("texts") == [] + mock_post.assert_not_called() From c2f486acfe09051f69a7d56ddaabb5ce1d330829 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Mon, 6 Apr 2026 13:26:49 -0700 Subject: [PATCH 003/226] fix: batch-limit stale managed object cleanup to prevent 300K row UPDATE (#25227) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add STALE_OBJECT_CLEANUP_BATCH_SIZE constant Configurable batch limit (default 1000) for stale managed object cleanup, preventing unbounded UPDATE queries from hitting 300K+ rows at once. * Batch-limit stale managed object cleanup with single bounded SQL query Two fixes to _cleanup_stale_managed_objects: 1. Replace unbounded update_many with a single execute_raw using a subquery LIMIT, capping each poll cycle to STALE_OBJECT_CLEANUP_BATCH_SIZE rows. Zero rows loaded into Python memory — everything stays in Postgres. Uses the same PostgreSQL raw-SQL pattern as spend_log_cleanup.py (the proxy requires PostgreSQL per schema.prisma). 2. Extract _expire_stale_rows as a separate method for testability. Keeps the file_purpose='response' filter to avoid incorrectly expiring long-running batch or fine-tune jobs that legitimately exceed the staleness cutoff. --- .../common_utils/check_responses_cost.py | 45 +++++++++++++++---- litellm/constants.py | 3 ++ 2 files changed, 40 insertions(+), 8 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py index 54fbc7abcc..dc0168683c 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_responses_cost.py @@ -11,6 +11,7 @@ from litellm._logging import verbose_proxy_logger from litellm.constants import ( MANAGED_OBJECT_STALENESS_CUTOFF_DAYS, MAX_OBJECTS_PER_POLL_CYCLE, + STALE_OBJECT_CLEANUP_BATCH_SIZE, ) if TYPE_CHECKING: @@ -32,21 +33,49 @@ class CheckResponsesCost: self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router + async def _expire_stale_rows( + self, cutoff: datetime, batch_size: int + ) -> int: + """Execute the bounded UPDATE that marks stale rows as 'stale_expired'. + + Isolated so it can be swapped / mocked in tests without touching the + orchestration logic in ``_cleanup_stale_managed_objects``. + + Uses PostgreSQL syntax (``$1::timestamptz``, ``LIMIT``, double-quoted + identifiers) which is the only dialect the proxy supports — every + ``schema.prisma`` in the repo sets ``provider = "postgresql"``. + Same pattern as ``spend_log_cleanup.py``. + """ + return await self.prisma_client.db.execute_raw( + """ + UPDATE "LiteLLM_ManagedObjectTable" + SET "status" = 'stale_expired' + WHERE "id" IN ( + SELECT "id" FROM "LiteLLM_ManagedObjectTable" + WHERE "file_purpose" = 'response' + AND "status" NOT IN ('completed', 'complete', 'failed', 'expired', 'cancelled', 'stale_expired') + AND "created_at" < $1::timestamptz + ORDER BY "created_at" ASC + LIMIT $2 + ) + """, + cutoff, + batch_size, + ) + async def _cleanup_stale_managed_objects(self) -> None: """ Mark managed objects older than MANAGED_OBJECT_STALENESS_CUTOFF_DAYS days in non-terminal states as 'stale_expired'. These will never complete and should not be polled. + + Runs as a single DB query with a subquery LIMIT so no rows are loaded + into Python memory. Processes at most STALE_OBJECT_CLEANUP_BATCH_SIZE + rows per invocation to avoid overwhelming the DB when there is a large + backlog. """ cutoff = datetime.now(timezone.utc) - timedelta(days=MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) - result = await self.prisma_client.db.litellm_managedobjecttable.update_many( - where={ - "file_purpose": "response", - "status": {"not_in": ["completed", "complete", "failed", "expired", "cancelled", "stale_expired"]}, - "created_at": {"lt": cutoff}, - }, - data={"status": "stale_expired"}, - ) + result = await self._expire_stale_rows(cutoff, STALE_OBJECT_CLEANUP_BATCH_SIZE) if result > 0: verbose_proxy_logger.warning( f"CheckResponsesCost: marked {result} stale managed objects " diff --git a/litellm/constants.py b/litellm/constants.py index 1af53b2dae..28c6c0cc0e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1367,6 +1367,9 @@ MAX_OBJECTS_PER_POLL_CYCLE = max(1, int(os.getenv("MAX_OBJECTS_PER_POLL_CYCLE", MANAGED_OBJECT_STALENESS_CUTOFF_DAYS = max( 1, int(os.getenv("MANAGED_OBJECT_STALENESS_CUTOFF_DAYS", 7)) ) +STALE_OBJECT_CLEANUP_BATCH_SIZE = max( + 1, int(os.getenv("STALE_OBJECT_CLEANUP_BATCH_SIZE", 1000)) +) # Set PROXY_BATCH_POLLING_ENABLED=false to disable the CheckBatchCost and # CheckResponsesCost background polling jobs entirely (e.g. to avoid DB load on # installations with large numbers of stale managed objects). From 8b0fadf99c95d3f3dd51a957b3716af9aff7585a Mon Sep 17 00:00:00 2001 From: Otavio Brito <69211663+otaviofbrito@users.noreply.github.com> Date: Mon, 6 Apr 2026 17:42:24 -0300 Subject: [PATCH 004/226] [bug-fix]return actual status code - /v1/messages/count_tokens endpoint (#21352) * return actual status code - /count_tokens endpoint * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> * fix greptile suggestion * rollback file * add test case --------- Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> --- litellm/proxy/proxy_server.py | 26 +++++--- .../test_proxy_token_counter.py | 64 +++++++++++++++++++ 2 files changed, 81 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9738ae4f1a..715fc44114 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9117,15 +9117,23 @@ async def token_counter(request: TokenCountRequest, call_endpoint: bool = False) ) is True ): - result = await provider_counter.count_tokens( - model_to_use=model_to_use or "", - messages=messages, # type: ignore - contents=contents, - deployment=deployment, - request_model=request.model, - tools=tools, - system=system, - ) + try: + result = await provider_counter.count_tokens( + model_to_use=model_to_use or "", + messages=messages, # type: ignore + contents=contents, + deployment=deployment, + request_model=request.model, + ) + except httpx.HTTPStatusError as e: + error_message = getattr(e, "message", None) or str(e) + status_code = getattr(e, "status_code", None) or e.response.status_code + raise ProxyException( + message=error_message, + type="token_counting_error", + param="model", + code=status_code + ) ######################################################### # Transfrom the Response to the well known format ######################################################### diff --git a/tests/proxy_unit_tests/test_proxy_token_counter.py b/tests/proxy_unit_tests/test_proxy_token_counter.py index cca1e84360..ea25369d8f 100644 --- a/tests/proxy_unit_tests/test_proxy_token_counter.py +++ b/tests/proxy_unit_tests/test_proxy_token_counter.py @@ -994,6 +994,70 @@ async def test_bedrock_handler_httpx_error_status_code_propagation(): assert exc_info.value.message == "Forbidden - Invalid credentials" +@pytest.mark.asyncio +async def test_token_counter_httpx_status_error_raises_proxy_exception(): + """ + When provider_counter.count_tokens() raises httpx.HTTPStatusError, + the token_counter endpoint should catch it and raise a ProxyException + with the upstream status code and error message. + """ + + upstream_status = 429 + upstream_message = "Rate limit exceeded" + response = httpx.Response( + status_code=upstream_status, + request=httpx.Request("POST", "https://provider.example.com/count"), + ) + http_error = httpx.HTTPStatusError( + message=upstream_message, + request=response.request, + response=response, + ) + + mock_counter = MagicMock() + mock_counter.should_use_token_counting_api.return_value = True + mock_counter.count_tokens = AsyncMock(side_effect=http_error) + + # Save originals + original_get_provider_token_counter = litellm.proxy.proxy_server._get_provider_token_counter + original_router = litellm.proxy.proxy_server.llm_router + + try: + def mock_get_provider_token_counter(deployment, model_to_use): + return (mock_counter, "claude-4-6-sonnet", "vertex_ai") + + litellm.proxy.proxy_server._get_provider_token_counter = mock_get_provider_token_counter + + mock_router = MagicMock() + mock_router.async_get_available_deployment = AsyncMock( + return_value={ + "litellm_params": { + "model": "vertex_ai/claude-4-6-sonnet", + "api_key": "fake-key", + }, + "model_info": {}, + } + ) + litellm.proxy.proxy_server.llm_router = mock_router + + with pytest.raises(ProxyException) as exc_info: + await token_counter( + request=TokenCountRequest( + model="claude-4-6-sonnet", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) + + assert exc_info.value.code == str(upstream_status) + assert upstream_message in exc_info.value.message + assert exc_info.value.type == "token_counting_error" + assert exc_info.value.param == "model" + finally: + litellm.proxy.proxy_server._get_provider_token_counter = original_get_provider_token_counter + litellm.proxy.proxy_server.llm_router = original_router + + @pytest.mark.asyncio async def test_proxy_token_counter_error_raises_exception_when_disabled(): """ From 414d3966bfe212d5d2376c3b6ae1c3adbb0f67ef Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Mon, 6 Apr 2026 13:48:43 -0700 Subject: [PATCH 005/226] feat(teams): per-member model scope + team default_team_member_models (#24950) * fix(bedrock): strip [1m]/[200k] context window suffixes before cost lookup * test(bedrock): add test for [1m] context window suffix stripping in cost lookup * schema: add allowed_models to BudgetTable, default_team_member_models to TeamTable * migration: add allowed_models and default_team_member_models columns * types: add allowed_models to TeamMemberAddRequest, TeamMemberUpdateRequest, UpdateTeamRequest * utils: add allowed_models param to add_new_member, persist to budget table * common_utils: add allowed_models to _upsert_budget_and_membership * team endpoints: seed allowed_models on member_add, persist on member_update and team/update * auth: enforce per-member allowed_models at request time * networking: add allowed_models to Member type and teamMemberUpdateCall * TeamMemberTab: add Model Scope column showing per-member allowed_models * EditMembership: add Allowed Models multi-select field * TeamInfo: add default_team_member_models field in Settings tab * chore: sync schema.prisma copies from root * fix(team_member_update): update existing budget in-place instead of creating new one When a member already has a budget_id, patch only the fields the caller provided rather than always creating a fresh budget record. The old code ignored existing_budget_id entirely, so updating only allowed_models silently dropped the stored max_budget / tpm_limit / rpm_limit values. * fix(auth): pass llm_router to _check_team_member_model_access Without the router, _can_object_call_model cannot resolve wildcard model names (e.g. openai/*) or access-group names in allowed_models, causing legitimate requests to be denied. Thread the existing llm_router from _run_common_checks through to the new member-scope check. * feat(ui): add Team Member Settings accordion to Create Team modal Groups default_team_member_models, member budget/key duration, and tpm/rpm defaults into a single collapsible section. The model picker is filtered to only show the models selected for the team, and the copy distinguishes it from the team-level Models field. * feat(ui): consolidate Team Member Settings into accordion in edit team form Moves default_team_member_models + per-member budget/key/tpm/rpm fields into a collapsible "Team Member Settings" panel. Keeps the top-level form focused on team-wide settings (team models, team budget, tpm/rpm). * fix(ui): use tremor Accordion for Team Member Settings in edit team form * fix(ui): move Team Member Settings accordion above budget fields in Create Team * chore: fixes --------- Co-authored-by: Ishaan Jaffer Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: Yuneng Jiang --- .../migration.sql | 9 ++ .../litellm_proxy_extras/schema.prisma | 4 +- litellm/llms/bedrock/common_utils.py | 16 +- litellm/proxy/_types.py | 18 +++ litellm/proxy/auth/auth_checks.py | 74 ++++++++- .../management_endpoints/common_utils.py | 38 ++++- .../management_endpoints/team_endpoints.py | 67 ++++---- litellm/proxy/management_helpers/utils.py | 24 +-- litellm/proxy/schema.prisma | 4 +- schema.prisma | 4 +- .../llms/bedrock/test_bedrock_common_utils.py | 43 +++++- .../components/modals/CreateTeamModal.tsx | 105 +++++++++---- .../src/components/networking.tsx | 4 + .../src/components/team/EditMembership.tsx | 15 +- .../src/components/team/TeamInfo.tsx | 146 +++++++++++++----- .../src/components/team/TeamMemberTab.tsx | 39 +++++ 16 files changed, 480 insertions(+), 130 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_team_member_model_scope/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_team_member_model_scope/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_team_member_model_scope/migration.sql new file mode 100644 index 0000000000..d34482b339 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_team_member_model_scope/migration.sql @@ -0,0 +1,9 @@ +-- Add per-member model scope to LiteLLM_BudgetTable +-- allowed_models: empty array = inherit team models; non-empty = enforce member-level restriction +ALTER TABLE "LiteLLM_BudgetTable" + ADD COLUMN IF NOT EXISTS "allowed_models" TEXT[] DEFAULT ARRAY[]::TEXT[]; + +-- Add default_team_member_models to LiteLLM_TeamTable +-- Seeds allowed_models for newly added team members; empty = no per-member restriction +ALTER TABLE "LiteLLM_TeamTable" + ADD COLUMN IF NOT EXISTS "default_team_member_models" TEXT[] DEFAULT ARRAY[]::TEXT[]; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index fce95465b5..03760695f3 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -17,8 +17,9 @@ model LiteLLM_BudgetTable { tpm_limit BigInt? rpm_limit BigInt? model_max_budget Json? - budget_duration String? + budget_duration String? budget_reset_at DateTime? + allowed_models String[] @default([]) // per-member model scope; empty = inherit team models created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@ -140,6 +141,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + default_team_member_models String[] @default([]) // default allowed_models for newly added team members; empty = no per-member restriction model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 9666aa68c9..a7cb6ed9f1 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -432,12 +432,16 @@ def strip_bedrock_routing_prefix(model: str) -> str: def strip_bedrock_throughput_suffix(model: str) -> str: - """Strip throughput tier suffixes from Bedrock model names.""" + """Strip throughput tier suffixes and context window suffixes from Bedrock model names.""" import re # Pattern matches model:version:throughput where throughput is like 51k, 18k, etc. # Keep the model:version part, strip the :throughput suffix - return re.sub(r"(:\d+):\d+k$", r"\1", model) + model = re.sub(r"(:\d+):\d+k$", r"\1", model) + # Strip context window suffixes like [1m], [200k], etc. + # e.g. "us.anthropic.claude-opus-4-6-v1[1m]" -> "us.anthropic.claude-opus-4-6-v1" + model = re.sub(r"\[\w+\]$", "", model) + return model def get_bedrock_base_model(model: str) -> str: @@ -1062,9 +1066,11 @@ class CommonBatchFilesUtils: return ( dict(prepped.headers), - request_data.encode("utf-8") - if isinstance(request_data, str) - else request_data, + ( + request_data.encode("utf-8") + if isinstance(request_data, str) + else request_data + ), ) def generate_unique_job_name(self, model: str, prefix: str = "litellm") -> str: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 441b3b836a..452ba7f799 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1649,6 +1649,9 @@ class TeamBase(LiteLLMPydanticObjectBase): blocked: bool = False router_settings: Optional[dict] = None access_group_ids: Optional[List[str]] = None + default_team_member_models: Optional[List[str]] = ( + None # default allowed_models seeded onto new team members + ) class NewTeamRequest(TeamBase): @@ -1742,6 +1745,9 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): enforced_file_expires_after: Optional[dict] = None router_settings: Optional[dict] = None access_group_ids: Optional[List[str]] = None + default_team_member_models: Optional[List[str]] = ( + None # default allowed_models seeded onto new team members + ) class ResetTeamBudgetRequest(LiteLLMPydanticObjectBase): @@ -1943,6 +1949,9 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None model_max_budget: Optional[dict] = None budget_duration: Optional[str] = None + allowed_models: Optional[List[str]] = ( + None # per-member model scope; empty = inherit team models + ) model_config = ConfigDict(protected_namespaces=()) @@ -3735,6 +3744,10 @@ class TeamMemberAddRequest(MemberAddRequest): default=None, description="Maximum budget allocated to this user within the team. If not set, user has unlimited budget within team limits", ) + allowed_models: Optional[List[str]] = Field( + default=None, + description="List of models this team member can access. If not set, inherits the team's default_team_member_models or all team models.", + ) class TeamMemberDeleteRequest(MemberDeleteRequest): @@ -3750,6 +3763,10 @@ class TeamMemberUpdateRequest(TeamMemberDeleteRequest): rpm_limit: Optional[int] = Field( default=None, description="Requests per minute limit for this team member" ) + allowed_models: Optional[List[str]] = Field( + default=None, + description="List of models this team member can access. Pass an empty list to remove per-member model restrictions.", + ) class TeamMemberUpdateResponse(MemberUpdateResponse): @@ -3757,6 +3774,7 @@ class TeamMemberUpdateResponse(MemberUpdateResponse): max_budget_in_team: Optional[float] = None tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None + allowed_models: Optional[List[str]] = None class TeamModelAddRequest(BaseModel): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 68bde8434a..bf0de3d3ca 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -8,6 +8,7 @@ Run checks for: 2. If user is in budget 3. If end_user ('user' passed to /chat/completions, /embeddings endpoint) is in budget """ + import asyncio import re import time @@ -456,9 +457,9 @@ async def common_checks( # noqa: PLR0915 model=_model, team_object=team_object, llm_router=llm_router, - team_model_aliases=valid_token.team_model_aliases - if valid_token - else None, + team_model_aliases=( + valid_token.team_model_aliases if valid_token else None + ), ): raise ProxyException( message=f"Team not allowed to access model. Team={team_object.team_id}, Model={_model}. Allowed team models = {team_object.models}", @@ -467,6 +468,21 @@ async def common_checks( # noqa: PLR0915 code=status.HTTP_401_UNAUTHORIZED, ) + # 2.2. If team member has per-member model scope, enforce it + if _model and team_object and valid_token and valid_token.user_id: + with tracer.trace( + "litellm.proxy.auth.common_checks.check_team_member_model_access" + ): + await _check_team_member_model_access( + model=_model, + team_object=team_object, + valid_token=valid_token, + llm_router=llm_router, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent if valid_token is not None and valid_token.agent_id: from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry @@ -3108,6 +3124,58 @@ async def _check_team_member_budget( ) +async def _check_team_member_model_access( + model: Union[str, List[str]], + team_object: LiteLLM_TeamTable, + valid_token: UserAPIKeyAuth, + llm_router: Optional[Router], + prisma_client: Optional["PrismaClient"], + user_api_key_cache: DualCache, + proxy_logging_obj: ProxyLogging, +) -> None: + """ + Check if a team member's per-member model scope allows access to the requested model. + + Only enforced when the member's budget table has a non-empty allowed_models list. + If allowed_models is empty or absent, the team-level models list applies (no extra restriction). + """ + if valid_token.user_id is None or team_object.team_id is None: + return + + team_membership = await get_team_membership( + user_id=valid_token.user_id, + team_id=team_object.team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + proxy_logging_obj=proxy_logging_obj, + ) + + if ( + team_membership is None + or team_membership.litellm_budget_table is None + or not team_membership.litellm_budget_table.allowed_models + ): + return # no per-member restriction — inherit team-level check + + member_allowed_models: List[str] = ( + team_membership.litellm_budget_table.allowed_models + ) + try: + _can_object_call_model( + model=model, + llm_router=llm_router, + models=member_allowed_models, + object_type="team", + ) + except ProxyException: + raise ProxyException( + message=f"Team member not allowed to access model. User={valid_token.user_id}, Team={team_object.team_id}, Model={model}. Allowed member models = {member_allowed_models}", + type=ProxyErrorTypes.team_model_access_denied, + param="model", + code=status.HTTP_401_UNAUTHORIZED, + ) + + async def _team_max_budget_check( team_object: Optional[LiteLLM_TeamTable], valid_token: Optional[UserAPIKeyAuth], diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index efc42d3355..07286b4fa8 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any, Dict, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union from litellm._logging import verbose_proxy_logger from litellm.caching import DualCache @@ -354,6 +354,7 @@ async def _upsert_budget_and_membership( user_api_key_dict: UserAPIKeyAuth, tpm_limit: Optional[int] = None, rpm_limit: Optional[int] = None, + allowed_models: Optional[List[str]] = None, ): """ Helper function to Create/Update or Delete the budget within the team membership @@ -366,11 +367,17 @@ async def _upsert_budget_and_membership( user_api_key_dict: User API Key dictionary containing user information tpm_limit: Tokens per minute limit for the team member rpm_limit: Requests per minute limit for the team member + allowed_models: Per-member model scope. None = don't change. [] = remove restrictions. Non-empty list = enforce. - If max_budget, tpm_limit, and rpm_limit are all None, the user's budget is removed from the team membership. + If max_budget, tpm_limit, rpm_limit, and allowed_models are all None, the user's budget is removed from the team membership. If any of these values exist, a budget is updated or created and linked to the team membership. """ - if max_budget is None and tpm_limit is None and rpm_limit is None: + if ( + max_budget is None + and tpm_limit is None + and rpm_limit is None + and allowed_models is None + ): # disconnect the budget since all limits are None await tx.litellm_teammembership.update( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, @@ -378,7 +385,27 @@ async def _upsert_budget_and_membership( ) return - # create a new budget + if existing_budget_id is not None: + # Update the existing budget in-place to preserve fields not being changed. + # Only write fields that the caller explicitly provided (non-None). + update_data: Dict[str, Any] = { + "updated_by": user_api_key_dict.user_id or "", + } + if max_budget is not None: + update_data["max_budget"] = max_budget + if tpm_limit is not None: + update_data["tpm_limit"] = tpm_limit + if rpm_limit is not None: + update_data["rpm_limit"] = rpm_limit + if allowed_models is not None: + update_data["allowed_models"] = allowed_models + await tx.litellm_budgettable.update( + where={"budget_id": existing_budget_id}, + data=update_data, + ) + return + + # No existing budget — create a new one and link it to the membership. create_data: Dict[str, Any] = { "created_by": user_api_key_dict.user_id or "", "updated_by": user_api_key_dict.user_id or "", @@ -389,12 +416,13 @@ async def _upsert_budget_and_membership( create_data["tpm_limit"] = tpm_limit if rpm_limit is not None: create_data["rpm_limit"] = rpm_limit + if allowed_models is not None: + create_data["allowed_models"] = allowed_models new_budget = await tx.litellm_budgettable.create( data=create_data, include={"team_membership": True}, ) - # upsert the team membership with the new/updated budget await tx.litellm_teammembership.upsert( where={ "user_id_team_id": { diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index e153457378..c980fcb82f 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1545,12 +1545,12 @@ async def update_team( # noqa: PLR0915 updated_kv["router_settings"] = safe_dumps(updated_kv["router_settings"]) updated_kv = prisma_client.jsonify_team_object(db_data=updated_kv) - team_row: Optional[ - LiteLLM_TeamTable - ] = await prisma_client.db.litellm_teamtable.update( - where={"team_id": data.team_id}, - data=updated_kv, - include={"litellm_model_table": True}, # type: ignore + team_row: Optional[LiteLLM_TeamTable] = ( + await prisma_client.db.litellm_teamtable.update( + where={"team_id": data.team_id}, + data=updated_kv, + include={"litellm_model_table": True}, # type: ignore + ) ) if team_row is None or team_row.team_id is None: @@ -1764,6 +1764,11 @@ async def _process_team_members( else None ) + # Resolve allowed_models: explicit request value, or fall back to team's default_team_member_models + member_allowed_models = data.allowed_models + if member_allowed_models is None and complete_team_data.default_team_member_models: + member_allowed_models = complete_team_data.default_team_member_models + if isinstance(data.member, Member): try: updated_user, updated_tm = await add_new_member( @@ -1774,6 +1779,7 @@ async def _process_team_members( litellm_proxy_admin_name=litellm_proxy_admin_name, team_id=data.team_id, default_team_budget_id=default_team_budget_id, + allowed_models=member_allowed_models, ) except Exception as e: raise HTTPException( @@ -1798,6 +1804,7 @@ async def _process_team_members( litellm_proxy_admin_name=litellm_proxy_admin_name, team_id=data.team_id, default_team_budget_id=default_team_budget_id, + allowed_models=member_allowed_models, ) except Exception as e: raise HTTPException( @@ -2297,13 +2304,13 @@ async def team_member_delete( ) # Fetch keys before deletion to persist them - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={ - "user_id": {"in": list(user_ids_to_delete)}, - "team_id": data.team_id, - } + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={ + "user_id": {"in": list(user_ids_to_delete)}, + "team_id": data.team_id, + } + ) ) if keys_to_delete: @@ -2435,6 +2442,7 @@ async def team_member_update( user_api_key_dict=user_api_key_dict, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + allowed_models=data.allowed_models, ) ### update team member role @@ -2467,6 +2475,7 @@ async def team_member_update( max_budget_in_team=data.max_budget_in_team, tpm_limit=data.tpm_limit, rpm_limit=data.rpm_limit, + allowed_models=data.allowed_models, ) @@ -2687,10 +2696,10 @@ async def delete_team( team_rows: List[LiteLLM_TeamTable] = [] for team_id in data.team_ids: try: - team_row_base: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id} + team_row_base: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id} + ) ) if team_row_base is None: raise Exception @@ -2749,10 +2758,10 @@ async def delete_team( _persist_deleted_verification_tokens, ) - keys_to_delete: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={"team_id": {"in": data.team_ids}} + keys_to_delete: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"team_id": {"in": data.team_ids}} + ) ) if keys_to_delete: @@ -2991,11 +3000,11 @@ async def team_info( ) try: - team_info: Optional[ - BaseModel - ] = await prisma_client.db.litellm_teamtable.find_unique( - where={"team_id": team_id}, - include={"object_permission": True}, + team_info: Optional[BaseModel] = ( + await prisma_client.db.litellm_teamtable.find_unique( + where={"team_id": team_id}, + include={"object_permission": True}, + ) ) if team_info is None: raise Exception @@ -3316,7 +3325,7 @@ async def _build_team_list_where_conditions( user_object_correct_type = await get_user_object( user_id=user_id, prisma_client=prisma_client, - user_api_key_cache=user_api_key_cache, + user_api_key_cache=user_api_key_cache, # type: ignore[arg-type] user_id_upsert=False, proxy_logging_obj=proxy_logging_obj, ) @@ -3835,9 +3844,7 @@ async def list_team( except Exception as e: team_exception = """Invalid team object for team_id: {}. team_object={}. Error: {} - """.format( - team.team_id, team.model_dump(), str(e) - ) + """.format(team.team_id, team.model_dump(), str(e)) verbose_proxy_logger.exception(team_exception) continue # Sort the responses by team_alias diff --git a/litellm/proxy/management_helpers/utils.py b/litellm/proxy/management_helpers/utils.py index 7d485fdebd..3e42d39207 100644 --- a/litellm/proxy/management_helpers/utils.py +++ b/litellm/proxy/management_helpers/utils.py @@ -2,7 +2,7 @@ ## Helper utils for the management endpoints (keys/users/teams) from datetime import datetime from functools import wraps -from typing import Optional, Tuple +from typing import List, Optional, Tuple from fastapi import HTTPException, Request @@ -148,6 +148,7 @@ async def add_new_member( user_api_key_dict: UserAPIKeyAuth, litellm_proxy_admin_name: str, default_team_budget_id: Optional[str] = None, + allowed_models: Optional[List[str]] = None, ) -> Tuple[LiteLLM_UserTable, Optional[LiteLLM_TeamMembership]]: """ Add a new member to a team @@ -206,17 +207,18 @@ async def add_new_member( }, ) - # Check if trying to set a budget for team member - - if max_budget_in_team is not None: + # Check if trying to set a budget or model scope for team member + if max_budget_in_team is not None or allowed_models is not None: # create a new budget item for this member - response = await prisma_client.db.litellm_budgettable.create( - data={ - "max_budget": max_budget_in_team, - "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - } - ) + budget_data: dict = { + "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, + } + if max_budget_in_team is not None: + budget_data["max_budget"] = max_budget_in_team + if allowed_models is not None: + budget_data["allowed_models"] = allowed_models + response = await prisma_client.db.litellm_budgettable.create(data=budget_data) _budget_id = response.budget_id else: diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index fce95465b5..03760695f3 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -17,8 +17,9 @@ model LiteLLM_BudgetTable { tpm_limit BigInt? rpm_limit BigInt? model_max_budget Json? - budget_duration String? + budget_duration String? budget_reset_at DateTime? + allowed_models String[] @default([]) // per-member model scope; empty = inherit team models created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@ -140,6 +141,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + default_team_member_models String[] @default([]) // default allowed_models for newly added team members; empty = no per-member restriction model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) diff --git a/schema.prisma b/schema.prisma index fce95465b5..03760695f3 100644 --- a/schema.prisma +++ b/schema.prisma @@ -17,8 +17,9 @@ model LiteLLM_BudgetTable { tpm_limit BigInt? rpm_limit BigInt? model_max_budget Json? - budget_duration String? + budget_duration String? budget_reset_at DateTime? + allowed_models String[] @default([]) // per-member model scope; empty = inherit team models created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at") @@ -140,6 +141,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + default_team_member_models String[] @default([]) // default allowed_models for newly added team members; empty = no per-member restriction model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index 5ce291aa16..a4e621b0bb 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -29,25 +29,30 @@ def test_govcloud_cross_region_inference_prefix(): Test that GovCloud models with cross-region inference prefix (us-gov.) are parsed correctly """ bedrock_model_info = BedrockModelInfo - + # Test us-gov prefix is stripped correctly for Claude models base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.anthropic.claude-haiku-4-5-20251001-v1:0" ) +<<<<<<< worktree-rustling-wishing-kite + assert base_model == "anthropic.claude-3-5-sonnet-20240620-v1:0" + +======= assert base_model == "anthropic.claude-haiku-4-5-20251001-v1:0" +>>>>>>> main # Test us-gov prefix is stripped correctly for different Claude versions base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0" ) assert base_model == "anthropic.claude-sonnet-4-5-20250929-v1:0" - + # Test us-gov prefix is stripped correctly for Haiku models base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.anthropic.claude-3-haiku-20240307-v1:0" ) assert base_model == "anthropic.claude-3-haiku-20240307-v1:0" - + # Test us-gov prefix is stripped correctly for Meta models base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.meta.llama3-8b-instruct-v1:0" @@ -55,3 +60,35 @@ def test_govcloud_cross_region_inference_prefix(): assert base_model == "meta.llama3-8b-instruct-v1:0" +def test_context_window_suffix_stripped_for_cost_lookup(): + """ + Test that [1m], [200k] etc. context window suffixes are stripped from + Bedrock model names before cost lookup. + + Models configured like `bedrock/us.anthropic.claude-opus-4-6-v1[1m]` + should resolve to the base model name so pricing can be found. + """ + from litellm.llms.bedrock.common_utils import get_bedrock_base_model + + assert ( + get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1[1m]") + == "anthropic.claude-opus-4-6-v1" + ) + assert ( + get_bedrock_base_model("us.anthropic.claude-sonnet-4-6[1m]") + == "anthropic.claude-sonnet-4-6" + ) + assert ( + get_bedrock_base_model("global.anthropic.claude-opus-4-5-20251101-v1:0[1m]") + == "anthropic.claude-opus-4-5-20251101-v1:0" + ) + # Ensure models without suffix are unaffected + assert ( + get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1") + == "anthropic.claude-opus-4-6-v1" + ) + # Ensure :51k throughput suffix still works + assert ( + get_bedrock_base_model("anthropic.claude-3-5-sonnet-20241022-v2:0:51k") + == "anthropic.claude-3-5-sonnet-20241022-v2:0" + ) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx index 0aa42b69a0..b4f99f5ca8 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/teams/components/modals/CreateTeamModal.tsx @@ -391,6 +391,80 @@ const CreateTeamModal = ({ + + + Team Member Settings + + + + Optional defaults applied when members join this team. All fields can be overridden per member. + + prev.models !== cur.models} + > + {({ getFieldValue }) => { + const teamModels: string[] = getFieldValue("models") || []; + const opts = teamModels.length > 0 ? teamModels : modelsToPick; + return ( + + Default Model Access{" "} + + + + + } + name="default_team_member_models" + > + + {opts.map((m) => ( + + {getModelDisplayName(m)} + + ))} + + + ); + }} + + (value ? Number(value) : undefined)} + tooltip="Default spend budget for each member in this team." + > + + + + + + + + + + + + + + @@ -409,7 +483,7 @@ const CreateTeamModal = ({ { if (!mcpAccessGroupsLoaded) { fetchMcpAccessGroups(); @@ -432,35 +506,6 @@ const CreateTeamModal = ({ }} /> - (value ? Number(value) : undefined)} - tooltip="This is the individual budget for a user in the team." - > - - - - - - - - - - - { @@ -3898,6 +3899,9 @@ export const teamMemberUpdateCall = async ( if (formValues.rpm_limit !== undefined && formValues.rpm_limit !== null) { requestBody.rpm_limit = formValues.rpm_limit; } + if (formValues.allowed_models !== undefined) { + requestBody.allowed_models = formValues.allowed_models; + } console.log("Final request body:", requestBody); diff --git a/ui/litellm-dashboard/src/components/team/EditMembership.tsx b/ui/litellm-dashboard/src/components/team/EditMembership.tsx index e6d9ccad60..cc655ba52b 100644 --- a/ui/litellm-dashboard/src/components/team/EditMembership.tsx +++ b/ui/litellm-dashboard/src/components/team/EditMembership.tsx @@ -21,7 +21,7 @@ interface ModalConfig { additionalFields?: Array<{ name: string; label: string | React.ReactNode; - type: "input" | "select" | "numerical"; + type: "input" | "select" | "numerical" | "multi-select"; options?: Array<{ label: string; value: string }>; rules?: any[]; step?: number; @@ -65,6 +65,8 @@ const MemberModal = ({ max_budget_in_team: (initialData as any).max_budget_in_team || null, tpm_limit: (initialData as any).tpm_limit || null, rpm_limit: (initialData as any).rpm_limit || null, + // Keep array values for multi-select fields + allowed_models: (initialData as any).allowed_models || [], }; console.log("Setting form values:", formValues); form.setFieldsValue(formValues); @@ -115,7 +117,7 @@ const MemberModal = ({ const renderField = (field: { name: string; label: string | React.ReactNode; - type: "input" | "select" | "numerical"; + type: "input" | "select" | "numerical" | "multi-select"; options?: Array<{ label: string; value: string }>; rules?: any[]; step?: number; @@ -144,6 +146,15 @@ const MemberModal = ({ ))} ); + case "multi-select": + return ( + - - - - - - form.setFieldValue("team_member_budget_duration", value)} - value={form.getFieldValue("team_member_budget_duration")} - /> - - - - - - - - - - - - - + + + Team Member Settings + + + + Optional defaults applied when members join this team. All fields can be overridden per member. + + + Default Model Access{" "} + + + + + } + name="default_team_member_models" + > + prev.models !== cur.models}> + {({ getFieldValue }) => { + const teamModels = getFieldValue("models") || info.models || []; + return ( + @@ -1318,6 +1362,18 @@ const TeamInfoView: React.FC = ({ ))} + {info.default_team_member_models && info.default_team_member_models.length > 0 && ( +
+ Default Member Models +
+ {info.default_team_member_models.map((model, index) => ( + + {model} + + ))} +
+
+ )}
Rate Limits
TPM: {info.tpm_limit || "Unlimited"}
@@ -1483,6 +1539,20 @@ const TeamInfoView: React.FC = ({ min: 0, placeholder: "Requests per minute limit for this member in this team", }, + { + name: "allowed_models", + label: ( + + Allowed Models{" "} + + + + + ), + type: "multi-select" as const, + options: (info.models || []).map((m: string) => ({ label: m, value: m })), + placeholder: "Leave empty to inherit all team models", + }, ], }} /> diff --git a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx index 652d1dbcd9..e880aa49f6 100644 --- a/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamMemberTab.tsx @@ -82,7 +82,45 @@ export default function TeamMemberTab({ const isUserTeamAdmin = isUserTeamAdminForSingleTeam(teamData.team_info.members_with_roles, userId || ""); const isProxyAdmin = isProxyAdminRole(userRole || ""); + const getUserAllowedModels = (userId: string | null): string[] | null => { + if (!userId) return null; + const membership = teamData.team_memberships.find((tm) => tm.user_id === userId); + const models = membership?.litellm_budget_table?.allowed_models; + return models && models.length > 0 ? models : null; + }; + const extraColumns: ColumnsType = [ + { + title: ( + + Model Scope + + + + + ), + key: "model_scope", + render: (_: unknown, record: Member) => { + const models = getUserAllowedModels(record.user_id); + if (!models) { + return (all team models); + } + const displayed = models.slice(0, 2); + const remaining = models.length - displayed.length; + return ( + + {displayed.map((m) => ( + {m} + ))} + {remaining > 0 && ( + + +{remaining} more + + )} + + ); + }, + }, { title: ( @@ -138,6 +176,7 @@ export default function TeamMemberTab({ max_budget_in_team: membership?.litellm_budget_table?.max_budget || null, tpm_limit: membership?.litellm_budget_table?.tpm_limit || null, rpm_limit: membership?.litellm_budget_table?.rpm_limit || null, + allowed_models: membership?.litellm_budget_table?.allowed_models || [], }; setSelectedEditMember(enhancedMember); setIsEditMemberModalVisible(true); From 0afffe436686da1f89a5b83b07bb14e1bf0a5bb7 Mon Sep 17 00:00:00 2001 From: ishaan-berri <155045088+ishaan-berri@users.noreply.github.com> Date: Mon, 6 Apr 2026 14:02:04 -0700 Subject: [PATCH 006/226] feat: multiple concurrent budget windows per API key and team (#24883) (#25109) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: multiple concurrent budget windows per API key and team (#24883) * feat(proxy): add BudgetLimitEntry type and wire budget_limits into key/team models * feat(schema): add budget_limits Json column to VerificationToken and TeamTable * feat(migrations): add migration for budget_limits column on keys and teams * feat(keys): initialize budget_limits windows with reset_at on key create/update * feat(teams): initialize budget_limits windows with reset_at on team create/update * feat(auth): add _virtual_key_multi_budget_check and _team_multi_budget_check * feat(auth): call multi-budget checks from common_checks for keys and teams * feat(proxy): increment per-window Redis spend counters after each request * feat(budget): reset individual budget windows on schedule via reset_budget_job * feat(ui): add hourly option to BudgetDurationDropdown * feat(ui): add budget_limits field to KeyResponse type * feat(ui): add Budget Windows editor to key edit view * feat(ui): add Budget Windows editor to create key form * fix(proxy): strip budget_limits=None before Prisma upsert to fix login 500 Prisma rejects nullable JSON fields (Json? without @default) when passed as Python None — it needs the field omitted entirely so the DB stores NULL via the column's nullable constraint. This was breaking /v2/login because the UI session key creation path hit the upsert with budget_limits=None. * ui(key-edit): use antd InputNumber+Button for budget windows, add reset hints * ui(create-key): use antd InputNumber+Button for budget windows, add reset hints * docs(users): add multiple budget windows section with API + dashboard walkthrough * fix: BudgetExceededError returns HTTP 429 instead of 400 - Add status_code=429 to BudgetExceededError class - auth_exception_handler hardcoded code=400 → code=429 * fix: no-op else branch in multi-budget auth checks causes KeyError - BudgetLimitEntry objects must be coerced via model_dump() not left as-is - Move _virtual_key_multi_budget_check into common_checks (was asymmetric with _team_multi_budget_check which already lived there) * fix: len() on JSON string returns char count not window count Guard with isinstance check + json.loads() before iterating per-window Redis counters in increment_spend_counters * fix: silent except:pass hides Redis reset failures in reset_budget_windows Log Redis counter reset failures as warnings so they are observable * test: add unit tests for multi-budget window enforcement 5 tests covering: no budget_limits passes, under budget passes, over hourly window raises 429, over monthly window raises 429, BudgetLimitEntry objects coerced without KeyError * fix: key per-window counters stable across reorders (duration key, not index) * fix: team+key per-window spend increments use duration key, not index * fix: budget window reset uses duration key; log failures instead of swallowing * refactor: extract BudgetWindowsEditor to shared component * refactor: key_edit_view imports BudgetWindowsEditor from shared component * refactor: create_key_button imports BudgetWindowsEditor from shared component --------- Co-authored-by: Ishaan Jaffer * fix(reset_budget_job): extract _reset_expired_window helper to fix PLR0915 too many statements * feat(skills): Skills Registry & Hub — register skills, browse in AI Hub, public skill hub (#25118) * feat(skills): add domain and namespace fields to plugin types * feat(skills): store and return domain/namespace inside manifest_json * feat(skills): add /public/skill_hub endpoint for unauthenticated access * feat(skills): whitelist /public/skill_hub from auth requirements * feat(skills): add domain, namespace to Plugin and RegisterPluginRequest types * feat(skills): smart URL parser — paste github URL, auto-detect source type and name * feat(skills): replace enable toggle with Public badge, make rows clickable * feat(skills): add skill detail view with Overview and How to Use tabs * feat(skills): add MakeSkillPublicForm modal for publishing skills to the hub * feat(skills): rename panel to Skills, wire in skill detail view on row click * feat(skills): add skill hub table columns — name, description, domain, source, status * feat(skills): add SkillHubDashboard with stats row, domain dropdown filter, and table * feat(skills): add Skill Hub tab to AI Hub with Select Skills to Make Public button * feat(skills): move Skills to top-level nav item directly under MCP Servers * feat(skills): add skillHubPublicCall and NEXT_PUBLIC_BASE_URL support * feat(skills): add Skill Hub tab to public AI Hub page * feat(skills): add skills page routing in main app router * feat(skills): add /skills page route * chore: update package-lock after npm install * docs(skills): add Skills Gateway doc page with mermaid architecture diagram * docs(skills): add Skills Gateway to sidebar under Agent & MCP Gateway * docs(skills): add loom walkthrough video to Skills Gateway doc * chore: fixes --------- Co-authored-by: Ishaan Jaffer Co-authored-by: Yuneng Jiang --- docs/my-website/docs/proxy/users.md | 61 +++ docs/my-website/docs/skills_gateway.md | 111 ++++++ docs/my-website/sidebars.js | 7 + .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 2 + litellm/exceptions.py | 3 +- litellm/proxy/_types.py | 20 + .../claude_code_marketplace.py | 6 + litellm/proxy/auth/auth_checks.py | 81 ++++ litellm/proxy/auth/auth_exception_handler.py | 2 +- litellm/proxy/auth/user_api_key_auth.py | 12 +- .../proxy/common_utils/reset_budget_job.py | 109 +++++- .../key_management_endpoints.py | 65 +++- .../management_endpoints/team_endpoints.py | 25 ++ litellm/proxy/proxy_server.py | 40 ++ .../public_endpoints/public_endpoints.py | 43 ++ litellm/proxy/schema.prisma | 2 + litellm/proxy/utils.py | 22 +- litellm/types/proxy/claude_code_endpoints.py | 4 + .../proxy/auth/test_multi_budget_windows.py | 137 +++++++ ui/litellm-dashboard/package-lock.json | 82 +--- .../src/app/(dashboard)/skills/page.tsx | 17 + ui/litellm-dashboard/src/app/page.tsx | 2 +- .../src/components/AIHub/ModelHubTable.tsx | 63 ++- .../components/AIHub/SkillHubDashboard.tsx | 139 +++++++ .../src/components/claude_code_plugins.tsx | 115 +++--- .../MakeSkillPublicForm.tsx | 249 ++++++++++++ .../claude_code_plugins/add_plugin_form.tsx | 331 ++++++++-------- .../claude_code_plugins/plugin_table.tsx | 71 +--- .../claude_code_plugins/skill_detail.tsx | 366 ++++++++++++++++++ .../components/claude_code_plugins/types.ts | 14 +- .../budget_duration_dropdown.tsx | 2 + .../key_team_helpers/BudgetWindowsEditor.tsx | 84 ++++ .../components/key_team_helpers/key_list.tsx | 1 + .../src/components/leftnav.tsx | 15 +- .../src/components/networking.tsx | 26 +- .../organisms/create_key_button.tsx | 30 +- .../src/components/public_model_hub.tsx | 27 ++ .../components/skill_hub_table_columns.tsx | 114 ++++++ .../components/templates/key_edit_view.tsx | 27 +- 40 files changed, 2103 insertions(+), 429 deletions(-) create mode 100644 docs/my-website/docs/skills_gateway.md create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql create mode 100644 tests/test_litellm/proxy/auth/test_multi_budget_windows.py create mode 100644 ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx create mode 100644 ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx create mode 100644 ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx create mode 100644 ui/litellm-dashboard/src/components/claude_code_plugins/skill_detail.tsx create mode 100644 ui/litellm-dashboard/src/components/key_team_helpers/BudgetWindowsEditor.tsx create mode 100644 ui/litellm-dashboard/src/components/skill_hub_table_columns.tsx diff --git a/docs/my-website/docs/proxy/users.md b/docs/my-website/docs/proxy/users.md index 88a7a0f1e0..0e36e84c20 100644 --- a/docs/my-website/docs/proxy/users.md +++ b/docs/my-website/docs/proxy/users.md @@ -333,6 +333,67 @@ curl 'http://0.0.0.0:4000/key/generate' \ }' ``` +#### **Set multiple budget windows on a key** + +Apply multiple concurrent budget limits at different time scales on the same key — for example, cap a key at **$10/day** AND **$100/month**. + +**When is this useful?** + +A single `budget_duration` window can't prevent a bad day from burning your entire month. Multiple budget windows let you: + +- Block a runaway usage spike within the day while still allowing normal monthly spend. +- Give Claude Code rollouts a daily guardrail (`24h`) and a monthly ceiling (`30d`) so a single heavy session doesn't exhaust the whole month. +- Layer fine-grained hourly limits for bursty workloads on top of a weekly cap. + +:::info + +See [User Budget docs](https://docs.litellm.ai/docs/proxy/users) for more on how budgets work across keys, teams, and users. + +::: + +**Via API** + +Pass `budget_limits` as a list of `{budget_duration, max_budget}` objects: + +```bash +curl 'http://0.0.0.0:4000/key/generate' \ +--header 'Authorization: Bearer ' \ +--header 'Content-Type: application/json' \ +--data-raw '{ + "budget_limits": [ + {"budget_duration": "24h", "max_budget": 10}, + {"budget_duration": "30d", "max_budget": 100} + ] +}' +``` + +Each window is tracked independently and resets on its own schedule: + +| `budget_duration` | Resets | +|---|---| +| `1h` | Every hour | +| `24h` | Daily at midnight UTC | +| `7d` | Every Sunday at midnight UTC | +| `30d` | 1st of every month at midnight UTC | + +**Via Dashboard** + +Open **Virtual Keys → Create Key → Optional Settings → Budget Windows**. + +![Step 1 - open key settings](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/18930ba5-67c0-4031-afc0-57f37b4e59e4/ascreenshot_ef79d8a000bb41cdacf1bd9827732ee8_text_export.jpeg) + +Click **+ Add Budget Window** to add a row, choose the period from the dropdown, and enter the spend cap. + +![Step 2 - add a window](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/5ae8c0b3-2d03-41ad-a63c-47b20c350dfe/ascreenshot_1a7dc6c7d65544f38fd8a65604674f22_text_export.jpeg) + +Add a second row for a different time period (e.g. monthly $100 on top of a daily $10). + +![Step 3 - add second window](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/cbded3a7-1086-4e20-8f0f-de154b76146c/ascreenshot_c51c18752c3b4f8b976d28799b2638b6_text_export.jpeg) + +Each window shows the reset schedule below the input so it's always clear when spend resets. + +![Step 4 - reset hints](https://colony-recorder.s3.amazonaws.com/files/2026-04-01/8754f121-1640-4892-9dd0-fd4a870418bf/ascreenshot_8079eb0df2194e8f99e5258ba4b3c082_text_export.jpeg) + ### ✨ Virtual Key (Model Specific) diff --git a/docs/my-website/docs/skills_gateway.md b/docs/my-website/docs/skills_gateway.md new file mode 100644 index 0000000000..d0eb810757 --- /dev/null +++ b/docs/my-website/docs/skills_gateway.md @@ -0,0 +1,111 @@ +# Skills Gateway + + + +LiteLLM acts as a **Skills Registry** — a central place to register, manage, and discover Claude Code skills across your organization. Teams can publish skills once and have agents and developers find them through a single hub. + +## How it works + +```mermaid +graph TD + Dev["👨‍💻 Developer
registers a skill
(GitHub URL or subdir)"] -->|POST /claude-code/plugins| Proxy["LiteLLM Proxy
(Skills Registry)"] + + Admin["🔑 Admin
publishes skill
(marks as public)"] -->|enable via UI or API| Proxy + + Proxy -->|GET /public/skill_hub| SkillHub["🗂️ Skill Hub
(AI Hub → Skill Hub tab)"] + Proxy -->|GET /claude-code/marketplace.json| Marketplace["📦 Claude Code
Marketplace endpoint"] + + SkillHub --> Human["🧑 Human
browses & discovers skills
in AI Hub UI"] + Marketplace --> Agent["🤖 Agent / Claude Code
installs skill with
/plugin marketplace add <name>"] + + style Proxy fill:#1a73e8,color:#fff + style SkillHub fill:#e8f0fe,color:#1a73e8 + style Marketplace fill:#e8f0fe,color:#1a73e8 +``` + +## Quick start + +### 1. Register a skill + +Paste any GitHub URL into the Skills UI — LiteLLM auto-detects the source type and skill name. + +```bash +curl -X POST https://your-proxy/claude-code/plugins \ + -H "Authorization: Bearer $LITELLM_KEY" \ + -H "Content-Type: application/json" \ + -d '{ + "name": "grill-me", + "source": { + "source": "git-subdir", + "url": "https://github.com/mattpocock/skills", + "path": "grill-me" + }, + "description": "Interview skill for relentless questioning", + "domain": "Productivity", + "namespace": "interviews" + }' +``` + +Skills nested in subdirectories (e.g. `github.com/org/repo/tree/main/skill-name`) are supported — LiteLLM parses the URL automatically in the UI. + +### 2. Publish to hub + +In the Admin UI: **AI Hub → Skill Hub → Select Skills to Make Public**. + +Or via API: + +```bash +curl -X POST https://your-proxy/claude-code/plugins/grill-me/enable \ + -H "Authorization: Bearer $LITELLM_KEY" +``` + +### 3. Browse the hub + +Public skills appear at: +- **Admin UI**: AI Hub → Skill Hub tab +- **Public page**: `/ui/model_hub` → Skill Hub tab (no login required) +- **API**: `GET /public/skill_hub` + +### 4. Install in Claude Code + +Point Claude Code at your proxy marketplace once: + +```json title="~/.claude/settings.json" +{ + "extraKnownMarketplaces": { + "my-org": { + "source": "url", + "url": "https://your-proxy/claude-code/marketplace.json" + } + } +} +``` + +Then install any skill: + +``` +/plugin marketplace add grill-me +``` + +## Skill fields + +| Field | Description | +|-------|-------------| +| `name` | Unique skill identifier (used in `/plugin marketplace add`) | +| `source` | Git source — `github`, `url`, or `git-subdir` | +| `description` | Short description shown in the hub | +| `domain` | Category for grouping (e.g. `Engineering`, `Productivity`) | +| `namespace` | Subcategory within a domain (e.g. `quality`, `meetings`) | +| `keywords` | Tags for search and filtering | +| `version` | Semver string | + +## API reference + +| Endpoint | Auth | Description | +|----------|------|-------------| +| `POST /claude-code/plugins` | Required | Register a skill | +| `GET /claude-code/plugins` | Required | List all skills (admin) | +| `POST /claude-code/plugins/{name}/enable` | Required | Publish a skill | +| `POST /claude-code/plugins/{name}/disable` | Required | Unpublish a skill | +| `GET /public/skill_hub` | None | List public skills | +| `GET /claude-code/marketplace.json` | None | Claude Code marketplace manifest | diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index ab8f257c7d..bb863a0e77 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -333,6 +333,13 @@ const sidebars = { }, ], }, + { + type: "category", + label: "Skills Gateway", + items: [ + "skills_gateway", + ], + }, ], }, { diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql new file mode 100644 index 0000000000..fdd65543a3 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260401000000_add_budget_limits/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable: add budget_limits column to LiteLLM_VerificationToken +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "budget_limits" JSONB; + +-- AlterTable: add budget_limits column to LiteLLM_TeamTable +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "budget_limits" JSONB; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 03760695f3..d570b1dc35 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -141,6 +141,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] default_team_member_models String[] @default([]) // default allowed_models for newly added team members; empty = no per-member restriction model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team @@ -402,6 +403,7 @@ model LiteLLM_VerificationToken { rotation_interval String? // How often to rotate (e.g., "30d", "90d") last_rotation_at DateTime? // When this key was last rotated key_rotation_at DateTime? // When this key should next be rotated + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) diff --git a/litellm/exceptions.py b/litellm/exceptions.py index abdba09dd8..51810c5643 100644 --- a/litellm/exceptions.py +++ b/litellm/exceptions.py @@ -281,7 +281,7 @@ class Timeout(openai.APITimeoutError): # type: ignore return _message -class PermissionDeniedError(openai.PermissionDeniedError): # type:ignore +class PermissionDeniedError(openai.PermissionDeniedError): # type: ignore def __init__( self, message, @@ -847,6 +847,7 @@ class BudgetExceededError(Exception): ): self.current_cost = current_cost self.max_budget = max_budget + self.status_code = 429 message = ( message or f"Budget has been exceeded! Current cost: {current_cost}, Max budget: {max_budget}" diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 452ba7f799..12b0c54258 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -596,6 +596,7 @@ class LiteLLMRoutes(enum.Enum): "/public/model_hub", "/public/agent_hub", "/public/mcp_hub", + "/public/skill_hub", "/public/litellm_model_cost_map", ] ) @@ -868,6 +869,14 @@ class LiteLLM_ObjectPermissionBase(LiteLLMPydanticObjectBase): models: Optional[List[str]] = None +class BudgetLimitEntry(LiteLLMPydanticObjectBase): + """A single budget window with its own limit and independent reset schedule.""" + + budget_duration: str # e.g. "24h", "7d", "30d" + max_budget: float # max spend in USD for this window + reset_at: Optional[datetime] = None # populated at creation/reset time + + class GenerateRequestBase(LiteLLMPydanticObjectBase): """ Overlapping schema between key and user generate/update requests @@ -887,6 +896,9 @@ class GenerateRequestBase(LiteLLMPydanticObjectBase): rpm_limit: Optional[int] = None budget_duration: Optional[str] = None + budget_limits: Optional[List[BudgetLimitEntry]] = ( + None # multiple concurrent budget windows + ) allowed_cache_controls: Optional[list] = [] config: Optional[dict] = {} permissions: Optional[dict] = {} @@ -991,6 +1003,7 @@ class GenerateKeyResponse(KeyRequestBase): "permissions", "model_max_budget", "router_settings", + "budget_limits", ] for field in dict_fields: value = values.get(field) @@ -1644,6 +1657,9 @@ class TeamBase(LiteLLMPydanticObjectBase): max_budget: Optional[float] = None soft_budget: Optional[float] = None budget_duration: Optional[str] = None + budget_limits: Optional[List[BudgetLimitEntry]] = ( + None # multiple concurrent budget windows + ) models: list = [] blocked: bool = False @@ -1745,6 +1761,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): enforced_file_expires_after: Optional[dict] = None router_settings: Optional[dict] = None access_group_ids: Optional[List[str]] = None + budget_limits: Optional[List[BudgetLimitEntry]] = ( + None # multiple concurrent budget windows default_team_member_models: Optional[List[str]] = ( None # default allowed_models seeded onto new team members ) @@ -1893,6 +1911,7 @@ class LiteLLM_TeamTable(TeamBase): "model_max_budget", "model_aliases", "router_settings", + "budget_limits", ] if isinstance(values, BaseModel): @@ -2378,6 +2397,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): last_rotation_at: Optional[datetime] = None # When this key was last rotated key_rotation_at: Optional[datetime] = None # When this key should next be rotated router_settings: Optional[dict] = None + budget_limits: Optional[List[dict]] = None # multiple concurrent budget windows model_config = ConfigDict(protected_namespaces=()) diff --git a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py index abd3ce5661..78cce865a2 100644 --- a/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py +++ b/litellm/proxy/anthropic_endpoints/claude_code_endpoints/claude_code_marketplace.py @@ -260,6 +260,10 @@ async def register_plugin( manifest["keywords"] = request.keywords if request.category: manifest["category"] = request.category + if request.domain: + manifest["domain"] = request.domain + if request.namespace: + manifest["namespace"] = request.namespace # Check if plugin exists existing = await prisma_client.db.litellm_claudecodeplugintable.find_unique( @@ -362,6 +366,8 @@ async def list_plugins( homepage=manifest.get("homepage"), keywords=manifest.get("keywords"), category=manifest.get("category"), + domain=manifest.get("domain"), + namespace=manifest.get("namespace"), enabled=p.enabled, created_at=p.created_at.isoformat() if p.created_at else None, updated_at=p.updated_at.isoformat() if p.updated_at else None, diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index bf0de3d3ca..8e0cf62e56 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -534,6 +534,17 @@ async def common_checks( # noqa: PLR0915 valid_token=valid_token, ) + # 3.1. Multi-window budget check for team + with tracer.trace("litellm.proxy.auth.common_checks.team_multi_budget_check"): + await _team_multi_budget_check(team_object=team_object) + + # 3.2. Multi-window budget check for key + with tracer.trace( + "litellm.proxy.auth.common_checks.virtual_key_multi_budget_check" + ): + if valid_token is not None: + await _virtual_key_multi_budget_check(valid_token=valid_token) + # 3.0.5. If team is over soft budget (alert only, doesn't block) with tracer.trace("litellm.proxy.auth.common_checks.team_soft_budget_check"): await _team_soft_budget_check( @@ -2982,6 +2993,43 @@ async def _virtual_key_max_budget_check( ) +async def _virtual_key_multi_budget_check( + valid_token: UserAPIKeyAuth, +): + """ + Raises BudgetExceededError if any budget window in valid_token.budget_limits is exceeded. + + Each window has its own Redis counter keyed by spend:key:{token}:window:{budget_duration}. + Using budget_duration (not list index) keeps counters stable when windows are reordered + or removed during a key update. + + Note: counters are not seeded from DB on Redis cold-start. After a Redis flush, + per-window spend resets to zero within the current window period. This is an acceptable + trade-off: the DB stores reset_at timestamps but not per-window accumulated spend. + """ + if not valid_token.budget_limits: + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in valid_token.budget_limits: + w: dict = window if isinstance(window, dict) else window.model_dump() + counter_key = f"spend:key:{valid_token.token}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + ) + if window_spend >= w["max_budget"]: + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], + message=( + f"ExceededBudget: Key over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + ) + + async def _virtual_key_soft_budget_check( valid_token: UserAPIKeyAuth, proxy_logging_obj: ProxyLogging, @@ -3223,6 +3271,39 @@ async def _team_max_budget_check( ) +async def _team_multi_budget_check( + team_object: Optional[LiteLLM_TeamTable], +): + """ + Raises BudgetExceededError if any budget window in team_object.budget_limits is exceeded. + + Each window has its own Redis counter keyed by spend:team:{team_id}:window:{budget_duration}. + Using budget_duration (not list index) keeps counters stable when windows are reordered + or removed during a team update. + """ + if team_object is None or not team_object.budget_limits: + return + + from litellm.proxy.proxy_server import get_current_spend + + for window in team_object.budget_limits: + w: dict = window if isinstance(window, dict) else window.model_dump() + counter_key = f"spend:team:{team_object.team_id}:window:{w['budget_duration']}" + window_spend = await get_current_spend( + counter_key=counter_key, + fallback_spend=0.0, + ) + if window_spend >= w["max_budget"]: + raise litellm.BudgetExceededError( + current_cost=window_spend, + max_budget=w["max_budget"], + message=( + f"ExceededBudget: Team={team_object.team_id} over {w['budget_duration']} budget. " + f"Spend=${window_spend:.4f}, Limit=${w['max_budget']:.2f}" + ), + ) + + async def _team_soft_budget_check( team_object: Optional[LiteLLM_TeamTable], valid_token: Optional[UserAPIKeyAuth], diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index 9c306acd2c..b9b614994c 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -106,7 +106,7 @@ class UserAPIKeyAuthExceptionHandler: message=e.message, type=ProxyErrorTypes.budget_exceeded, param=None, - code=400, + code=429, ) if isinstance(e, HTTPException): raise ProxyException( diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 046c39a910..1d5b913f9b 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -984,9 +984,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 route=route, ) if _end_user_object is not None: - end_user_params[ - "allowed_model_region" - ] = _end_user_object.allowed_model_region + end_user_params["allowed_model_region"] = ( + _end_user_object.allowed_model_region + ) if _end_user_object.litellm_budget_table is not None: _apply_budget_limits_to_end_user_params( end_user_params=end_user_params, @@ -1540,9 +1540,9 @@ async def _user_api_key_auth_builder( # noqa: PLR0915 if _end_user_object is not None: valid_token_dict.update(end_user_params) - valid_token_dict[ - "end_user_object_permission" - ] = _end_user_object.object_permission + valid_token_dict["end_user_object_permission"] = ( + _end_user_object.object_permission + ) # check if token is from litellm-ui, litellm ui makes keys to allow users to login with sso. These keys can only be used for LiteLLM UI functions # sso/login, ui/login, /key functions and /user functions diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index bcfaed2439..695facd704 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -2,7 +2,7 @@ import asyncio import json import time from datetime import datetime, timedelta, timezone -from typing import List, Literal, Optional, Union +from typing import Any, List, Literal, Optional, Union from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -48,6 +48,9 @@ class ResetBudgetJob: ### RESET ENDUSER (Customer) BUDGET and corresponding Budget duration ### await self.reset_budget_for_litellm_budget_table() + ### RESET MULTI-WINDOW BUDGETS ### + await self.reset_budget_windows() + async def reset_budget_for_litellm_team_members( self, budgets_to_reset: List[LiteLLM_BudgetTableFull] ): @@ -549,6 +552,102 @@ class ResetBudgetJob: ) verbose_proxy_logger.exception("Failed to reset budget for teams: %s", e) + @staticmethod + async def _reset_expired_window( + window: dict, + counter_key: str, + spend_counter_cache: Any, + now: datetime, + ) -> bool: + """Reset a single budget window if expired. Returns True if the window was reset.""" + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + reset_at_str = window.get("reset_at") + if not reset_at_str: + return False + reset_at = datetime.fromisoformat( + reset_at_str.replace("Z", "+00:00") + ).replace(tzinfo=None) + if reset_at > now: + return False + spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.0) + if spend_counter_cache.redis_cache is not None: + try: + await spend_counter_cache.redis_cache.async_set_cache( + key=counter_key, value=0.0 + ) + except Exception as redis_err: + verbose_proxy_logger.warning( + "Failed to reset Redis counter %s: %s", counter_key, redis_err + ) + window["reset_at"] = get_budget_reset_time( + budget_duration=window["budget_duration"] + ).isoformat() + return True + + async def reset_budget_windows(self) -> None: + """ + For keys and teams with budget_limits, reset any individual windows where + reset_at <= now. Only the expired windows are reset; other windows are untouched. + """ + from litellm.proxy.proxy_server import spend_counter_cache + + now = datetime.utcnow() + + # --- Keys --- + try: + all_keys = await self.prisma_client.db.litellm_verificationtoken.find_many( + where={"budget_limits": {"not": None}} # type: ignore[arg-type] + ) + for key in all_keys: + raw = key.budget_limits # type: ignore[attr-defined] + if not raw: + continue + windows: list = raw if isinstance(raw, list) else json.loads(raw) + changed = False + for window in windows: + counter_key = f"spend:key:{key.token}:window:{window['budget_duration']}" + if await ResetBudgetJob._reset_expired_window( + window, counter_key, spend_counter_cache, now + ): + changed = True + if changed: + await self.prisma_client.db.litellm_verificationtoken.update( + where={"token": key.token}, + data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] + ) + except Exception as e: + verbose_proxy_logger.exception( + "Failed to reset budget windows for keys: %s", e + ) + + # --- Teams --- + try: + all_teams = await self.prisma_client.db.litellm_teamtable.find_many( + where={"budget_limits": {"not": None}} # type: ignore[arg-type] + ) + for team in all_teams: + raw = team.budget_limits # type: ignore[attr-defined] + if not raw: + continue + windows = raw if isinstance(raw, list) else json.loads(raw) + changed = False + for window in windows: + counter_key = f"spend:team:{team.team_id}:window:{window['budget_duration']}" + if await ResetBudgetJob._reset_expired_window( + window, counter_key, spend_counter_cache, now + ): + changed = True + if changed: + await self.prisma_client.db.litellm_teamtable.update( + where={"team_id": team.team_id}, + data={"budget_limits": json.dumps(windows)}, # type: ignore[arg-type] + ) + except Exception as e: + verbose_proxy_logger.exception( + "Failed to reset budget windows for teams: %s", e + ) + @staticmethod async def _reset_budget_common( item: Union[LiteLLM_TeamTable, LiteLLM_UserTable, LiteLLM_VerificationToken], @@ -570,14 +669,14 @@ class ResetBudgetJob: from litellm.proxy.proxy_server import spend_counter_cache counter_key = None - if item_type == "key" and hasattr(item, "token") and item.token is not None: - counter_key = f"spend:key:{item.token}" + if item_type == "key" and hasattr(item, "token") and item.token is not None: # type: ignore[union-attr] + counter_key = f"spend:key:{item.token}" # type: ignore[union-attr] elif ( item_type == "team" and hasattr(item, "team_id") - and item.team_id is not None + and item.team_id is not None # type: ignore[union-attr] ): - counter_key = f"spend:team:{item.team_id}" + counter_key = f"spend:team:{item.team_id}" # type: ignore[union-attr] if counter_key is not None: # Always reset in-memory (local fallback) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 00d8ce182e..89c7213176 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -744,9 +744,9 @@ async def _common_key_generation_helper( # noqa: PLR0915 request_type="key", **data_json, table_name="key" ) - response[ - "soft_budget" - ] = data.soft_budget # include the user-input soft budget in the response + response["soft_budget"] = ( + data.soft_budget + ) # include the user-input soft budget in the response response = GenerateKeyResponse(**response) @@ -1571,6 +1571,19 @@ async def prepare_key_update_data( non_default_values["budget_reset_at"] = key_reset_at non_default_values["budget_duration"] = budget_duration + if "budget_limits" in non_default_values and non_default_values["budget_limits"]: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + raw_windows = non_default_values["budget_limits"] + initialized_windows = [] + for window in raw_windows: + w = window if isinstance(window, dict) else window.model_dump() + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + non_default_values["budget_limits"] = json.dumps(initialized_windows) + if "object_permission" in non_default_values: non_default_values = await _handle_update_object_permission( data_json=non_default_values, @@ -2832,6 +2845,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 rotation_interval: Optional[str] = None, router_settings: Optional[dict] = None, access_group_ids: Optional[list] = None, + budget_limits: Optional[list] = None, # multiple concurrent budget windows ): from litellm.proxy.proxy_server import premium_user, prisma_client @@ -2863,6 +2877,18 @@ async def generate_key_helper_fn( # noqa: PLR0915 else: reset_at = get_budget_reset_time(budget_duration=budget_duration) + # Initialize reset_at for each budget window + budget_limits_json: Optional[str] = None + if budget_limits: + initialized_windows = [] + for window in budget_limits: + w = dict(window) if not isinstance(window, dict) else {**window} + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + budget_limits_json = json.dumps(initialized_windows) + aliases_json = json.dumps(aliases) config_json = json.dumps(config) permissions_json = json.dumps(permissions) @@ -2945,6 +2971,7 @@ async def generate_key_helper_fn( # noqa: PLR0915 "organization_id": organization_id, "budget_id": budget_id, "blocked": blocked, + "budget_limits": budget_limits_json, "created_by": created_by, "updated_by": updated_by, "allowed_routes": allowed_routes or [], @@ -3222,10 +3249,10 @@ async def delete_verification_tokens( try: if prisma_client: tokens = [_hash_token_if_needed(token=key) for key in tokens] - _keys_being_deleted: List[ - LiteLLM_VerificationToken - ] = await prisma_client.db.litellm_verificationtoken.find_many( - where={"token": {"in": tokens}} + _keys_being_deleted: List[LiteLLM_VerificationToken] = ( + await prisma_client.db.litellm_verificationtoken.find_many( + where={"token": {"in": tokens}} + ) ) if len(_keys_being_deleted) == 0: @@ -3425,9 +3452,9 @@ async def _rotate_master_key( # noqa: PLR0915 from litellm.proxy.proxy_server import proxy_config try: - models: Optional[ - List - ] = await prisma_client.db.litellm_proxymodeltable.find_many() + models: Optional[List] = ( + await prisma_client.db.litellm_proxymodeltable.find_many() + ) except Exception: models = None # 2. process model table @@ -4067,11 +4094,11 @@ async def validate_key_list_check( param="user_id", code=status.HTTP_403_FORBIDDEN, ) - complete_user_info_db_obj: Optional[ - BaseModel - ] = await prisma_client.db.litellm_usertable.find_unique( - where={"user_id": user_api_key_dict.user_id}, - include={"organization_memberships": True}, + complete_user_info_db_obj: Optional[BaseModel] = ( + await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id}, + include={"organization_memberships": True}, + ) ) if complete_user_info_db_obj is None: @@ -4154,10 +4181,10 @@ async def _fetch_user_team_objects( if complete_user_info is None or not complete_user_info.teams: return [] - teams: Optional[ - List[BaseModel] - ] = await prisma_client.db.litellm_teamtable.find_many( - where={"team_id": {"in": complete_user_info.teams}} + teams: Optional[List[BaseModel]] = ( + await prisma_client.db.litellm_teamtable.find_many( + where={"team_id": {"in": complete_user_info.teams}} + ) ) if teams is None: return [] diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index c980fcb82f..15b7f9c646 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -978,6 +978,19 @@ async def new_team( # noqa: PLR0915 budget_duration=complete_team_data.budget_duration, ) + # If budget_limits is set, initialize reset_at for each window + if complete_team_data.budget_limits: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + initialized_windows = [] + for window in complete_team_data.budget_limits: + w = window if isinstance(window, dict) else window.model_dump() + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + complete_team_data.budget_limits = initialized_windows + ## Add Team Member Budget Table members_with_roles: List[Member] = [] if complete_team_data.members_with_roles is not None: @@ -1593,6 +1606,18 @@ def _set_budget_reset_at(data: UpdateTeamRequest, updated_kv: dict) -> None: reset_at = get_budget_reset_time(budget_duration=data.budget_duration) updated_kv["budget_reset_at"] = reset_at + if data.budget_limits is not None and len(data.budget_limits) > 0: + from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time + + initialized_windows = [] + for window in data.budget_limits: + w = window if isinstance(window, dict) else window.model_dump() + w["reset_at"] = get_budget_reset_time( + budget_duration=w["budget_duration"] + ).isoformat() + initialized_windows.append(w) + updated_kv["budget_limits"] = json.dumps(initialized_windows) + async def handle_update_object_permission( data_json: dict, existing_team_row: LiteLLM_TeamTable diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 715fc44114..ef4807275e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1781,6 +1781,26 @@ async def increment_spend_counters( increment=response_cost, ) + # Increment per-window budget counters for multi-budget keys + key_obj = await user_api_key_cache.async_get_cache(key=hashed_token) + if key_obj is not None: + key_budget_limits = getattr(key_obj, "budget_limits", None) or ( + key_obj.get("budget_limits") if isinstance(key_obj, dict) else None + ) + if isinstance(key_budget_limits, str): + key_budget_limits = json.loads(key_budget_limits) + if isinstance(key_budget_limits, list): + for window in key_budget_limits: + duration = ( + window["budget_duration"] + if isinstance(window, dict) + else window.budget_duration + ) + await spend_counter_cache.async_increment_cache( + key=f"spend:key:{hashed_token}:window:{duration}", + value=response_cost, + ) + if team_id is not None: await _init_and_increment_spend_counter( counter_key=f"spend:team:{team_id}", @@ -1788,6 +1808,26 @@ async def increment_spend_counters( increment=response_cost, ) + # Increment per-window budget counters for multi-budget teams + team_obj = await user_api_key_cache.async_get_cache(key=f"team_id:{team_id}") + if team_obj is not None: + team_budget_limits = getattr(team_obj, "budget_limits", None) or ( + team_obj.get("budget_limits") if isinstance(team_obj, dict) else None + ) + if isinstance(team_budget_limits, str): + team_budget_limits = json.loads(team_budget_limits) + if isinstance(team_budget_limits, list): + for window in team_budget_limits: + duration = ( + window["budget_duration"] + if isinstance(window, dict) + else window.budget_duration + ) + await spend_counter_cache.async_increment_cache( + key=f"spend:team:{team_id}:window:{duration}", + value=response_cost, + ) + if user_id is not None and team_id is not None: await _init_and_increment_spend_counter( counter_key=f"spend:team_member:{user_id}:{team_id}", diff --git a/litellm/proxy/public_endpoints/public_endpoints.py b/litellm/proxy/public_endpoints/public_endpoints.py index eb9ed59055..10b0a6644e 100644 --- a/litellm/proxy/public_endpoints/public_endpoints.py +++ b/litellm/proxy/public_endpoints/public_endpoints.py @@ -247,6 +247,49 @@ async def get_mcp_servers(): ] +@router.get( + "/public/skill_hub", + tags=["public", "Claude Code Marketplace"], +) +async def public_skill_hub(): + """Return enabled (public) Claude Code skills — no auth required.""" + from litellm.proxy.anthropic_endpoints.claude_code_endpoints.claude_code_marketplace import ( + _get_prisma_client, + ) + from litellm.types.proxy.claude_code_endpoints import ListPluginsResponse, PluginListItem + + try: + prisma_client = await _get_prisma_client() + plugins = await prisma_client.db.litellm_claudecodeplugintable.find_many( + where={"enabled": True} + ) + items = [] + for plugin in plugins: + raw = plugin.manifest_json or {} + manifest = json.loads(raw) if isinstance(raw, str) else raw + items.append( + PluginListItem( + id=plugin.id, + name=plugin.name, + enabled=plugin.enabled, + created_at=str(plugin.created_at) if plugin.created_at else None, + updated_at=str(plugin.updated_at) if plugin.updated_at else None, + source=manifest.get("source", {}), + description=manifest.get("description"), + version=manifest.get("version"), + category=manifest.get("category"), + keywords=manifest.get("keywords"), + author=manifest.get("author"), + homepage=manifest.get("homepage"), + domain=manifest.get("domain"), + namespace=manifest.get("namespace"), + ) + ) + return ListPluginsResponse(plugins=items, count=len(items)) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + + @router.get( "/public/model_hub/info", tags=["public", "model management"], diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 03760695f3..d570b1dc35 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -141,6 +141,7 @@ model LiteLLM_TeamTable { team_member_permissions String[] @default([]) access_group_ids String[] @default([]) policies String[] @default([]) + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] default_team_member_models String[] @default([]) // default allowed_models for newly added team members; empty = no per-member restriction model_id Int? @unique // id for LiteLLM_ModelTable -> stores team-level model aliases allow_team_guardrail_config Boolean @default(false) // if true, team admin can configure guardrails for this team @@ -402,6 +403,7 @@ model LiteLLM_VerificationToken { rotation_interval String? // How often to rotate (e.g., "30d", "90d") last_rotation_at DateTime? // When this key was last rotated key_rotation_at DateTime? // When this key should next be rotated + budget_limits Json? // multiple concurrent budget windows [{budget_duration, max_budget, reset_at}] litellm_budget_table LiteLLM_BudgetTable? @relation(fields: [budget_id], references: [budget_id]) litellm_organization_table LiteLLM_OrganizationTable? @relation(fields: [organization_id], references: [organization_id]) litellm_project_table LiteLLM_ProjectTable? @relation(fields: [project_id], references: [project_id]) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ec98cfd4d1..8d1ac63191 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -1898,9 +1898,9 @@ class ProxyLogging: normalized_call_type = CallTypes.aembedding.value if normalized_call_type is not None: litellm_logging_obj.call_type = normalized_call_type - litellm_logging_obj.model_call_details[ - "call_type" - ] = normalized_call_type + litellm_logging_obj.model_call_details["call_type"] = ( + normalized_call_type + ) # Pass-through endpoints are logged via the callback loop's # async_post_call_failure_hook — skip pre_call and failure handlers. if litellm_logging_obj.call_type == CallTypes.pass_through.value: @@ -2524,8 +2524,7 @@ class PrismaClient: required_view = "LiteLLM_VerificationTokenView" expected_views_str = ", ".join(f"'{view}'" for view in expected_views) pg_schema = os.getenv("DATABASE_SCHEMA", "public") - ret = await self.db.query_raw( - f""" + ret = await self.db.query_raw(f""" WITH existing_views AS ( SELECT viewname FROM pg_views @@ -2537,8 +2536,7 @@ class PrismaClient: (SELECT COUNT(*) FROM existing_views) AS view_count, ARRAY_AGG(viewname) AS view_names FROM existing_views - """ - ) + """) expected_total_views = len(expected_views) if ret[0]["view_count"] == expected_total_views: verbose_proxy_logger.info("All necessary views exist!") @@ -2547,8 +2545,7 @@ class PrismaClient: ## check if required view exists ## if ret[0]["view_names"] and required_view not in ret[0]["view_names"]: await self.health_check() # make sure we can connect to db - await self.db.execute_raw( - """ + await self.db.execute_raw(""" CREATE VIEW "LiteLLM_VerificationTokenView" AS SELECT v.*, @@ -2558,8 +2555,7 @@ class PrismaClient: t.rpm_limit AS team_rpm_limit FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; - """ - ) + """) verbose_proxy_logger.info( "LiteLLM_VerificationTokenView Created in DB!" @@ -3142,6 +3138,10 @@ class PrismaClient: hashed_token = self.hash_token(token=token) db_data = self.jsonify_object(data=data) db_data["token"] = hashed_token + # Prisma rejects nullable JSON fields set to None (no default). + # Strip them so the DB stores NULL via the column's nullable constraint. + if db_data.get("budget_limits") is None: + db_data.pop("budget_limits", None) print_verbose( "PrismaClient: Before upsert into litellm_verificationtoken" ) diff --git a/litellm/types/proxy/claude_code_endpoints.py b/litellm/types/proxy/claude_code_endpoints.py index bdf4f122e0..e5d82a4503 100644 --- a/litellm/types/proxy/claude_code_endpoints.py +++ b/litellm/types/proxy/claude_code_endpoints.py @@ -49,6 +49,8 @@ class RegisterPluginRequest(BaseModel): homepage: Optional[str] = Field(None, description="Plugin homepage URL") keywords: Optional[List[str]] = Field(None, description="Search keywords") category: Optional[str] = Field(None, description="Plugin category") + domain: Optional[str] = Field(None, description="Skill domain (e.g., 'Productivity')") + namespace: Optional[str] = Field(None, description="Skill namespace within domain (e.g., 'workflows')") class PluginResponse(BaseModel): @@ -82,6 +84,8 @@ class PluginListItem(BaseModel): homepage: Optional[str] = None keywords: Optional[List[str]] = None category: Optional[str] = None + domain: Optional[str] = None + namespace: Optional[str] = None enabled: bool created_at: Optional[str] updated_at: Optional[str] diff --git a/tests/test_litellm/proxy/auth/test_multi_budget_windows.py b/tests/test_litellm/proxy/auth/test_multi_budget_windows.py new file mode 100644 index 0000000000..ed94fca837 --- /dev/null +++ b/tests/test_litellm/proxy/auth/test_multi_budget_windows.py @@ -0,0 +1,137 @@ +""" +Unit tests for multi-budget-window enforcement on API keys. +""" + +from unittest.mock import AsyncMock, patch + +import pytest + +import litellm +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.auth.auth_checks import _virtual_key_multi_budget_check + + +def _make_valid_token(**kwargs) -> UserAPIKeyAuth: + defaults = dict( + token="sk-test-token", + key_name="test", + spend=0.0, + max_budget=None, + budget_limits=[], + ) + defaults.update(kwargs) + return UserAPIKeyAuth(**defaults) + + +@pytest.mark.asyncio +async def test_no_budget_limits_passes(): + """Keys with empty budget_limits should pass without raising.""" + token = _make_valid_token(budget_limits=[]) + # Should not raise + await _virtual_key_multi_budget_check(valid_token=token) + + +@pytest.mark.asyncio +async def test_under_budget_passes(): + """Key with spend under all windows should pass.""" + token = _make_valid_token( + budget_limits=[ + {"budget_duration": "24h", "max_budget": 10.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": None}, + ] + ) + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=1.0, # well under both windows + ): + await _virtual_key_multi_budget_check(valid_token=token) + + +@pytest.mark.asyncio +async def test_over_first_window_raises(): + """Key exceeding the first (daily) window should raise BudgetExceededError.""" + token = _make_valid_token( + budget_limits=[ + {"budget_duration": "24h", "max_budget": 5.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 100.0, "reset_at": None}, + ] + ) + + spend_by_window = [6.0, 6.0] # over daily, under monthly + + call_count = 0 + + async def fake_get_spend(counter_key, fallback_spend): + nonlocal call_count + val = spend_by_window[call_count] + call_count += 1 + return val + + with patch( + "litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_multi_budget_check(valid_token=token) + + err = exc_info.value + assert err.status_code == 429 + assert "24h" in str(err) + assert "Key over" in str(err) + + +@pytest.mark.asyncio +async def test_over_second_window_raises(): + """Key exceeding only the monthly window should raise BudgetExceededError referencing 30d.""" + token = _make_valid_token( + budget_limits=[ + {"budget_duration": "24h", "max_budget": 50.0, "reset_at": None}, + {"budget_duration": "30d", "max_budget": 5.0, "reset_at": None}, + ] + ) + + spend_by_window = [1.0, 10.0] # under daily, over monthly + + call_count = 0 + + async def fake_get_spend(counter_key, fallback_spend): + nonlocal call_count + val = spend_by_window[call_count] + call_count += 1 + return val + + with patch( + "litellm.proxy.proxy_server.get_current_spend", side_effect=fake_get_spend + ): + with pytest.raises(litellm.BudgetExceededError) as exc_info: + await _virtual_key_multi_budget_check(valid_token=token) + + err = exc_info.value + assert err.status_code == 429 + assert "30d" in str(err) + + +@pytest.mark.asyncio +async def test_budget_limit_entry_objects_coerced(): + """BudgetLimitEntry Pydantic objects (not dicts) must be handled without KeyError. + + While budget_limits is normally serialized as List[dict], the auth check must + tolerate BudgetLimitEntry objects in case they arrive without prior serialization. + """ + from litellm.proxy._types import BudgetLimitEntry + + token = _make_valid_token(budget_limits=[]) + # Bypass Pydantic validation to simulate BudgetLimitEntry objects reaching the check + object.__setattr__( + token, + "budget_limits", + [BudgetLimitEntry(budget_duration="24h", max_budget=10.0)], + ) + + with patch( + "litellm.proxy.proxy_server.get_current_spend", + new_callable=AsyncMock, + return_value=1.0, + ): + # Should not raise TypeError / KeyError — model_dump() coerces the object + await _virtual_key_multi_budget_check(valid_token=token) diff --git a/ui/litellm-dashboard/package-lock.json b/ui/litellm-dashboard/package-lock.json index 58b0a6b611..af2edafd83 100644 --- a/ui/litellm-dashboard/package-lock.json +++ b/ui/litellm-dashboard/package-lock.json @@ -91,7 +91,6 @@ "version": "5.2.0", "resolved": "https://registry.npmjs.org/@alloc/quick-lru/-/quick-lru-5.2.0.tgz", "integrity": "sha512-UrcABB+4bUrFABwbluTIBErXwvbsU/V7TZWfmbgJfbkwiBuziS9gxdODUyuiecfdGQ85jglMW6juS3+z5TsKLw==", - "dev": true, "license": "MIT", "engines": { "node": ">=10" @@ -1773,7 +1772,6 @@ "version": "0.3.13", "resolved": "https://registry.npmjs.org/@jridgewell/gen-mapping/-/gen-mapping-0.3.13.tgz", "integrity": "sha512-2kkt/7niJ6MgEPxF0bYdQ6etZaA+fQvDcLKckhy1yIQOzaoKjBBjSj63/aLVjYE3qhRt5dvM+uUyfCg6UKCBbA==", - "dev": true, "license": "MIT", "dependencies": { "@jridgewell/sourcemap-codec": "^1.5.0", @@ -1784,7 +1782,6 @@ "version": "3.1.2", "resolved": "https://registry.npmjs.org/@jridgewell/resolve-uri/-/resolve-uri-3.1.2.tgz", "integrity": "sha512-bRISgCIjP20/tbWSPWMEi54QVPRZExkuD9lJL+UIxUKtwVJA8wW1Trb1jMs1RFXo1CBTNZ/5hpC9QvmKWdopKw==", - "dev": true, "license": "MIT", "engines": { "node": ">=6.0.0" @@ -1794,14 +1791,12 @@ "version": "1.5.5", "resolved": "https://registry.npmjs.org/@jridgewell/sourcemap-codec/-/sourcemap-codec-1.5.5.tgz", "integrity": "sha512-cYQ9310grqxueWbl+WuIUIaiUaDcj7WOq5fVhEljNVgRfOUhY9fy2zTvfoqWsnebh8Sl70VScFbICvJnLKB0Og==", - "dev": true, "license": "MIT" }, "node_modules/@jridgewell/trace-mapping": { "version": "0.3.31", "resolved": "https://registry.npmjs.org/@jridgewell/trace-mapping/-/trace-mapping-0.3.31.tgz", "integrity": "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw==", - "dev": true, "license": "MIT", "dependencies": { "@jridgewell/resolve-uri": "^3.1.0", @@ -1969,7 +1964,6 @@ "version": "2.1.5", "resolved": "https://registry.npmjs.org/@nodelib/fs.scandir/-/fs.scandir-2.1.5.tgz", "integrity": "sha512-vq24Bq3ym5HEQm2NKCr3yXDwjc7vTsEThRDnkp2DK9p1uqLR+DHurm/NOTo0KG7HYHU7eppKZj3MyqYuMBf62g==", - "dev": true, "license": "MIT", "dependencies": { "@nodelib/fs.stat": "2.0.5", @@ -1983,7 +1977,6 @@ "version": "2.0.5", "resolved": "https://registry.npmjs.org/@nodelib/fs.stat/-/fs.stat-2.0.5.tgz", "integrity": "sha512-RkhPPp2zrqDAQA/2jNhnztcPAlv64XdhIp7a7454A5ovI7Bukxgt7MX7udwAu3zg1DcpPU0rz3VV1SeaqvY4+A==", - "dev": true, "license": "MIT", "engines": { "node": ">= 8" @@ -1993,7 +1986,6 @@ "version": "1.2.8", "resolved": "https://registry.npmjs.org/@nodelib/fs.walk/-/fs.walk-1.2.8.tgz", "integrity": "sha512-oGB+UxlgWcgQkgwo8GcEGwemoTFt3FIO9ababBmaGwXIoBKZ+GTy0pP185beGg7Llih/NSHSV2XAs1lnznocSg==", - "dev": true, "license": "MIT", "dependencies": { "@nodelib/fs.scandir": "2.1.5", @@ -2317,7 +2309,7 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/@playwright/test/-/test-1.58.1.tgz", "integrity": "sha512-6LdVIUERWxQMmUSSQi0I53GgCBYgM2RpGngCPY7hSeju+VrKjq3lvs7HpJoPbDiY5QM5EYRtRX5fvrinnMAz3w==", - "dev": true, + "devOptional": true, "license": "Apache-2.0", "dependencies": { "playwright": "1.58.1" @@ -3422,14 +3414,12 @@ "version": "15.7.15", "resolved": "https://registry.npmjs.org/@types/prop-types/-/prop-types-15.7.15.tgz", "integrity": "sha512-F6bEyamV9jKGAFBEmlQnesRPGOQqS2+Uwi0Em15xenOxHaf2hv6L8YCVn3rPdPJOiJfPiCnLIRyvwVaqMY3MIw==", - "dev": true, "license": "MIT" }, "node_modules/@types/react": { "version": "18.2.48", "resolved": "https://registry.npmjs.org/@types/react/-/react-18.2.48.tgz", "integrity": "sha512-qboRCl6Ie70DQQG9hhNREz81jqC1cs9EVNcjQ1AU+jH6NFfSAhVVbrrY/+nSF+Bsk4AOwm9Qa61InvMCyV+H3w==", - "dev": true, "license": "MIT", "dependencies": { "@types/prop-types": "*", @@ -3471,7 +3461,6 @@ "version": "0.26.0", "resolved": "https://registry.npmjs.org/@types/scheduler/-/scheduler-0.26.0.tgz", "integrity": "sha512-WFHp9YUJQ6CKshqoC37iOlHnQSmxNc795UhB26CyBBttrN9svdIrUjl/NjnNmfcwtncN0h/0PPAFWv9ovP8mLA==", - "dev": true, "license": "MIT" }, "node_modules/@types/unist": { @@ -4332,14 +4321,12 @@ "version": "1.3.0", "resolved": "https://registry.npmjs.org/any-promise/-/any-promise-1.3.0.tgz", "integrity": "sha512-7UvmKalWRt1wgjL1RrGxoSJW/0QZFIegpeGvZG9kjp8vrRu55XTHbwnqq2GpXm9uLbcuhxm3IqX9OB4MZR1b2A==", - "dev": true, "license": "MIT" }, "node_modules/anymatch": { "version": "3.1.3", "resolved": "https://registry.npmjs.org/anymatch/-/anymatch-3.1.3.tgz", "integrity": "sha512-KMReFUr0B4t+D+OBkjR3KYqvocp2XaSzO55UcB6mgQMd3KbcE+mWTyvVV7D/zsdEbNnV6acZUutkiHQXvTr1Rw==", - "dev": true, "license": "ISC", "dependencies": { "normalize-path": "^3.0.0", @@ -4353,7 +4340,6 @@ "version": "2.3.1", "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-2.3.1.tgz", "integrity": "sha512-JU3teHTNjmE2VCGFzuY8EXzCDVwEqB2a8fsIvwaStHhAWJEeVd1o1QD80CU6+ZdEXXSLbSsuLwJjkCBWqRQUVA==", - "dev": true, "license": "MIT", "engines": { "node": ">=8.6" @@ -4366,7 +4352,6 @@ "version": "5.0.2", "resolved": "https://registry.npmjs.org/arg/-/arg-5.0.2.tgz", "integrity": "sha512-PYjyFOLKQ9y57JvQ6QLo8dAgNqswh8M1RMJYdQduT6xbWSgK36P/Z/v+p888pM69jMMfS8Xd8F6I1kQ/I9HUGg==", - "dev": true, "license": "MIT" }, "node_modules/argparse": { @@ -4726,7 +4711,6 @@ "version": "2.3.0", "resolved": "https://registry.npmjs.org/binary-extensions/-/binary-extensions-2.3.0.tgz", "integrity": "sha512-Ceh+7ox5qe7LJuLHoY0feh3pHuUDHAcRUeyL2VYghZwfpkNIy/+8Ocg0a3UuSoYzavmylwuLWQOf3hl0jjMMIw==", - "dev": true, "license": "MIT", "engines": { "node": ">=8" @@ -4752,7 +4736,6 @@ "version": "3.0.3", "resolved": "https://registry.npmjs.org/braces/-/braces-3.0.3.tgz", "integrity": "sha512-yQbXgO/OSZVD2IsiLlro+7Hf6Q18EJrKSEsdoMzKePKXct3gvD8oLcOQdIzGupr5Fj+EDe8gO/lxc1BzfMpxvA==", - "dev": true, "license": "MIT", "dependencies": { "fill-range": "^7.1.1" @@ -4868,7 +4851,6 @@ "version": "2.0.1", "resolved": "https://registry.npmjs.org/camelcase-css/-/camelcase-css-2.0.1.tgz", "integrity": "sha512-QOSvevhslijgYwRx6Rv7zKdMF8lbRmx+uQGx2+vDc+KI/eBnsy9kit5aj23AgGu3pa4t9AgwbnXWqS+iOY+2aA==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -4992,7 +4974,6 @@ "version": "3.6.0", "resolved": "https://registry.npmjs.org/chokidar/-/chokidar-3.6.0.tgz", "integrity": "sha512-7VT13fmjotKpGipCW9JEQAusEPE+Ei8nl6/g4FBAmIm0GOOLMua9NDDo/DWp0ZAxCr3cPq5ZpBqmPAQgDda2Pw==", - "dev": true, "license": "MIT", "dependencies": { "anymatch": "~3.1.2", @@ -5017,7 +4998,6 @@ "version": "5.1.2", "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-5.1.2.tgz", "integrity": "sha512-AOIgSQCepiJYwP3ARnGx+5VnTu2HBYdzbGP45eLw1vr3zB3vZLeyed1sC9hnbcOc9/SrMyM5RPQrkGz4aS9Zow==", - "dev": true, "license": "ISC", "dependencies": { "is-glob": "^4.0.1" @@ -5093,7 +5073,6 @@ "version": "4.1.1", "resolved": "https://registry.npmjs.org/commander/-/commander-4.1.1.tgz", "integrity": "sha512-NOKm8xhkzAjzFx8B2v5OAHT+u5pRQc2UCa2Vq9jYL/31o2wi9mxBA7LIFs3sV5VSC49z6pEhfbMULvShKj26WA==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -5154,7 +5133,6 @@ "version": "3.0.0", "resolved": "https://registry.npmjs.org/cssesc/-/cssesc-3.0.0.tgz", "integrity": "sha512-/Tb/JcjK111nNScGob5MNtsntNM1aCNUDipB/TkwZFhyDrrE47SOx/18wF2bbjgc3ZzCSKW1T5nt5EbFoAz/Vg==", - "dev": true, "license": "MIT", "bin": { "cssesc": "bin/cssesc" @@ -5568,14 +5546,12 @@ "version": "1.2.2", "resolved": "https://registry.npmjs.org/didyoumean/-/didyoumean-1.2.2.tgz", "integrity": "sha512-gxtyfqMg7GKyhQmb056K7M3xszy/myH8w+B4RT+QXBQsvAOdc3XymqDDPHx1BgPgsdAA5SIifona89YtRATDzw==", - "dev": true, "license": "Apache-2.0" }, "node_modules/dlv": { "version": "1.1.3", "resolved": "https://registry.npmjs.org/dlv/-/dlv-1.1.3.tgz", "integrity": "sha512-+HlytyjlPKnIG8XuRG8WvmBP8xs8P71y+SKKS6ZXWoEgLuePxtDoUEiH7WkdePWrQ5JBpE6aoVqfZfJUQkjXwA==", - "dev": true, "license": "MIT" }, "node_modules/doctrine": { @@ -6489,7 +6465,6 @@ "version": "1.20.1", "resolved": "https://registry.npmjs.org/fastq/-/fastq-1.20.1.tgz", "integrity": "sha512-GGToxJ/w1x32s/D2EKND7kTil4n8OVk/9mycTc4VDza13lOvpUZTGX3mFSCtV9ksdGBVzvsyAVLM6mHFThxXxw==", - "dev": true, "license": "ISC", "dependencies": { "reusify": "^1.0.4" @@ -6522,7 +6497,6 @@ "version": "6.5.0", "resolved": "https://registry.npmjs.org/fdir/-/fdir-6.5.0.tgz", "integrity": "sha512-tIbYtZbucOs0BRGqPJkshJUYdL+SDH7dVM8gjy+ERp3WAUjLEFJE+02kanyHtwjWOnwrKYBiwAmM0p4kLJAnXg==", - "dev": true, "license": "MIT", "engines": { "node": ">=12.0.0" @@ -6560,7 +6534,6 @@ "version": "7.1.1", "resolved": "https://registry.npmjs.org/fill-range/-/fill-range-7.1.1.tgz", "integrity": "sha512-YsGpe3WHLK8ZYi4tWDg2Jy3ebRz2rXowDxnld4bkQB00cc/1Zw9AWnC0i9ztDJitivtQvaI9KaLyKrc+hBW0yg==", - "dev": true, "license": "MIT", "dependencies": { "to-regex-range": "^5.0.1" @@ -6700,7 +6673,6 @@ "version": "2.3.2", "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.2.tgz", "integrity": "sha512-xiqMQR4xAeHTuB9uWm+fFRcIOgKBMiOBP+eXiyT7jsgVCq1bkVygt00oASowB7EdtpOHaaPgKt812P9ab+DDKA==", - "dev": true, "hasInstallScript": true, "license": "MIT", "optional": true, @@ -6851,7 +6823,6 @@ "version": "6.0.2", "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-6.0.2.tgz", "integrity": "sha512-XxwI8EOhVQgWp6iDL+3b0r86f4d6AX6zSU55HfB4ydCEuXLXc5FcYeOu+nnGftS4TEju/11rt4KJPTMgbfmv4A==", - "dev": true, "license": "ISC", "dependencies": { "is-glob": "^4.0.3" @@ -7349,7 +7320,6 @@ "version": "2.1.0", "resolved": "https://registry.npmjs.org/is-binary-path/-/is-binary-path-2.1.0.tgz", "integrity": "sha512-ZMERYes6pDydyuGidse7OsHxtbI7WVeUEozgR/g7rd0xUimYNlvZRE/K2MgZTjWy725IfelLeVcEM97mmtRGXw==", - "dev": true, "license": "MIT", "dependencies": { "binary-extensions": "^2.0.0" @@ -7402,7 +7372,6 @@ "version": "2.16.1", "resolved": "https://registry.npmjs.org/is-core-module/-/is-core-module-2.16.1.tgz", "integrity": "sha512-UfoeMA6fIJ8wTYFEUjelnaGI67v6+N7qXJEvQuIGa99l4xsCruSYOVSQ0uPANn4dAzm8lkYPaKLrrijLq7x23w==", - "dev": true, "license": "MIT", "dependencies": { "hasown": "^2.0.2" @@ -7463,7 +7432,6 @@ "version": "2.1.1", "resolved": "https://registry.npmjs.org/is-extglob/-/is-extglob-2.1.1.tgz", "integrity": "sha512-SbKbANkN603Vi4jEZv49LeVJMn4yGwsbzZworEoyEiutsN3nJYdbO36zfhGJ6QEDpOZIFkDtnq5JRxmvl3jsoQ==", - "dev": true, "license": "MIT", "engines": { "node": ">=0.10.0" @@ -7509,7 +7477,6 @@ "version": "4.0.3", "resolved": "https://registry.npmjs.org/is-glob/-/is-glob-4.0.3.tgz", "integrity": "sha512-xelSayHH36ZgE7ZWhli7pW34hNbNl8Ojv5KVmkJD4hBdD3th8Tfk9vYasLM+mXWOZhFkgZfxhLSnrwRr4elSSg==", - "dev": true, "license": "MIT", "dependencies": { "is-extglob": "^2.1.1" @@ -7558,7 +7525,6 @@ "version": "7.0.0", "resolved": "https://registry.npmjs.org/is-number/-/is-number-7.0.0.tgz", "integrity": "sha512-41Cifkg6e8TylSpdtTpeLVMqvSBEVzTttHvERD741+pnZ8ANv0004MRL43QKPDlK9cGvNp6NZWZUBlbGXYxxng==", - "dev": true, "license": "MIT", "engines": { "node": ">=0.12.0" @@ -7835,7 +7801,6 @@ "version": "1.21.7", "resolved": "https://registry.npmjs.org/jiti/-/jiti-1.21.7.tgz", "integrity": "sha512-/imKNG4EbWNrVjoNC/1H5/9GFy+tqjGBHCaSsN+P2RnPqjsLmv6UD3Ej+Kj8nBWaRAwyk7kK5ZUc+OEatnTR3A==", - "dev": true, "license": "MIT", "bin": { "jiti": "bin/jiti.js" @@ -8121,7 +8086,6 @@ "version": "3.1.3", "resolved": "https://registry.npmjs.org/lilconfig/-/lilconfig-3.1.3.tgz", "integrity": "sha512-/vlFKAoH5Cgt3Ie+JLhRbwOsCQePABiU3tJ1egGvyQ+33R/vcwM2Zl2QR/LzjsBeItPt3oSVXapn+m4nQDvpzw==", - "dev": true, "license": "MIT", "engines": { "node": ">=14" @@ -8134,7 +8098,6 @@ "version": "1.2.4", "resolved": "https://registry.npmjs.org/lines-and-columns/-/lines-and-columns-1.2.4.tgz", "integrity": "sha512-7ylylesZQ/PV29jhEDl3Ufjo6ZX7gCqJr5F7PKrqc93v7fzSymt1BpwEU8nAUXs8qzzvqhbjhK5QZg6Mt/HkBg==", - "dev": true, "license": "MIT" }, "node_modules/locate-path": { @@ -8588,7 +8551,6 @@ "version": "1.4.1", "resolved": "https://registry.npmjs.org/merge2/-/merge2-1.4.1.tgz", "integrity": "sha512-8q7VEgMJW4J8tcfVPy8g09NcQwZdbwFEqhe/WZkoIzjn/3TGDwtOCYtXGxA3O8tPzpczCCDgv+P2P5y00ZJOOg==", - "dev": true, "license": "MIT", "engines": { "node": ">= 8" @@ -9161,7 +9123,6 @@ "version": "4.0.8", "resolved": "https://registry.npmjs.org/micromatch/-/micromatch-4.0.8.tgz", "integrity": "sha512-PXwfBhYu0hBCPw8Dn0E+WDYb7af3dSLVWKi3HGv84IdF4TyFoC0ysxFd0Goxw7nSv4T/PzEJQxsYsEiFCKo2BA==", - "dev": true, "license": "MIT", "dependencies": { "braces": "^3.0.3", @@ -9175,7 +9136,6 @@ "version": "2.3.1", "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-2.3.1.tgz", "integrity": "sha512-JU3teHTNjmE2VCGFzuY8EXzCDVwEqB2a8fsIvwaStHhAWJEeVd1o1QD80CU6+ZdEXXSLbSsuLwJjkCBWqRQUVA==", - "dev": true, "license": "MIT", "engines": { "node": ">=8.6" @@ -9290,7 +9250,6 @@ "version": "2.7.0", "resolved": "https://registry.npmjs.org/mz/-/mz-2.7.0.tgz", "integrity": "sha512-z81GNO7nnYMEhrGh9LeymoE4+Yr0Wn5McHIZMK5cfQCl+NDX08sCZgUc9/6MHni9IWuFLm1Z3HTCXu2z9fN62Q==", - "dev": true, "license": "MIT", "dependencies": { "any-promise": "^1.0.0", @@ -9502,7 +9461,6 @@ "version": "3.0.0", "resolved": "https://registry.npmjs.org/normalize-path/-/normalize-path-3.0.0.tgz", "integrity": "sha512-6eZs5Ls3WtCisHWp9S2GUy8dqkpGi4BVSz3GaqiE6ezub0512ESztXUwUB6C6IKbQkY2Pnb/mD4WYojCRwcwLA==", - "dev": true, "license": "MIT", "engines": { "node": ">=0.10.0" @@ -9521,7 +9479,6 @@ "version": "3.0.0", "resolved": "https://registry.npmjs.org/object-hash/-/object-hash-3.0.0.tgz", "integrity": "sha512-RSn9F68PjH9HqtltsSnqYC1XXoWe9Bju5+213R98cNGttag9q9yAOTzdbsqvIa7aNm5WffBZFpWYr2aWrklWAw==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -9866,7 +9823,6 @@ "version": "1.0.7", "resolved": "https://registry.npmjs.org/path-parse/-/path-parse-1.0.7.tgz", "integrity": "sha512-LDJzPVEEEPR+y48z93A0Ed0yXb8pAByGWo/k5YYdYgpY2/2EsOsksJrq7lOHxryrVOn1ejG6oAp8ahvOIQD8sw==", - "dev": true, "license": "MIT" }, "node_modules/path-scurry": { @@ -9913,7 +9869,6 @@ "version": "4.0.3", "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-4.0.3.tgz", "integrity": "sha512-5gTmgEY/sqK6gFXLIsQNH19lWb4ebPDLA4SdLP7dsWkIXHWlG66oPuVvXSGFPppYZz8ZDZq0dYYrbHfBCVUb1Q==", - "dev": true, "license": "MIT", "engines": { "node": ">=12" @@ -9926,7 +9881,6 @@ "version": "2.3.0", "resolved": "https://registry.npmjs.org/pify/-/pify-2.3.0.tgz", "integrity": "sha512-udgsAY+fTnvv7kI7aaxbqwWNb0AHiB0qBO89PZKPkoTmGOgdbrHDKD+0B2X4uTfJ/FT1R09r9gTsjUjNJotuog==", - "dev": true, "license": "MIT", "engines": { "node": ">=0.10.0" @@ -9936,7 +9890,6 @@ "version": "4.0.7", "resolved": "https://registry.npmjs.org/pirates/-/pirates-4.0.7.tgz", "integrity": "sha512-TfySrs/5nm8fQJDcBDuUng3VOUKsd7S+zqvbOTiGXHfxX4wK31ard+hoNuvkicM/2YFzlpDgABOevKSsB4G/FA==", - "dev": true, "license": "MIT", "engines": { "node": ">= 6" @@ -9946,7 +9899,7 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright/-/playwright-1.58.1.tgz", "integrity": "sha512-+2uTZHxSCcxjvGc5C891LrS1/NlxglGxzrC4seZiVjcYVQfUa87wBL6rTDqzGjuoWNjnBzRqKmF6zRYGMvQUaQ==", - "dev": true, + "devOptional": true, "license": "Apache-2.0", "dependencies": { "playwright-core": "1.58.1" @@ -9965,7 +9918,7 @@ "version": "1.58.1", "resolved": "https://registry.npmjs.org/playwright-core/-/playwright-core-1.58.1.tgz", "integrity": "sha512-bcWzOaTxcW+VOOGBCQgnaKToLJ65d6AqfLVKEWvexyS3AS6rbXl+xdpYRMGSRBClPvyj44njOWoxjNdL/H9UNg==", - "dev": true, + "devOptional": true, "license": "Apache-2.0", "bin": { "playwright-core": "cli.js" @@ -9988,7 +9941,6 @@ "version": "8.5.6", "resolved": "https://registry.npmjs.org/postcss/-/postcss-8.5.6.tgz", "integrity": "sha512-3Ybi1tAuwAP9s0r1UQ2J4n5Y0G05bJkpUIO0/bI9MhwmD70S5aTWbXGBwxHrelT+XM1k6dM0pk+SwNkpTRN7Pg==", - "dev": true, "funding": [ { "type": "opencollective", @@ -10017,7 +9969,6 @@ "version": "15.1.0", "resolved": "https://registry.npmjs.org/postcss-import/-/postcss-import-15.1.0.tgz", "integrity": "sha512-hpr+J05B2FVYUAXHeK1YyI267J/dDDhMU6B6civm8hSY1jYJnBXxzKDKDswzJmtLHryrjhnDjqqp/49t8FALew==", - "dev": true, "license": "MIT", "dependencies": { "postcss-value-parser": "^4.0.0", @@ -10035,7 +9986,6 @@ "version": "4.1.0", "resolved": "https://registry.npmjs.org/postcss-js/-/postcss-js-4.1.0.tgz", "integrity": "sha512-oIAOTqgIo7q2EOwbhb8UalYePMvYoIeRY2YKntdpFQXNosSu3vLrniGgmH9OKs/qAkfoj5oB3le/7mINW1LCfw==", - "dev": true, "funding": [ { "type": "opencollective", @@ -10061,7 +10011,6 @@ "version": "6.0.1", "resolved": "https://registry.npmjs.org/postcss-load-config/-/postcss-load-config-6.0.1.tgz", "integrity": "sha512-oPtTM4oerL+UXmx+93ytZVN82RrlY/wPUV8IeDxFrzIjXOLF1pN+EmKPLbubvKHT2HC20xXsCAH2Z+CKV6Oz/g==", - "dev": true, "funding": [ { "type": "opencollective", @@ -10104,7 +10053,6 @@ "version": "6.2.0", "resolved": "https://registry.npmjs.org/postcss-nested/-/postcss-nested-6.2.0.tgz", "integrity": "sha512-HQbt28KulC5AJzG+cZtj9kvKB93CFCdLvog1WFLf1D+xmMvPGlBstkpTEZfK5+AN9hfJocyBFCNiqyS48bpgzQ==", - "dev": true, "funding": [ { "type": "opencollective", @@ -10130,7 +10078,6 @@ "version": "6.1.2", "resolved": "https://registry.npmjs.org/postcss-selector-parser/-/postcss-selector-parser-6.1.2.tgz", "integrity": "sha512-Q8qQfPiZ+THO/3ZrOrO0cJJKfpYCagtMUkXbnEfmgUjwXg6z/WBeOyS9APBBPCTSiDV+s4SwQGu8yFsiMRIudg==", - "dev": true, "license": "MIT", "dependencies": { "cssesc": "^3.0.0", @@ -10144,7 +10091,6 @@ "version": "4.2.0", "resolved": "https://registry.npmjs.org/postcss-value-parser/-/postcss-value-parser-4.2.0.tgz", "integrity": "sha512-1NNCs6uurfkVbeXG4S8JFT9t19m45ICnif8zWLd5oPSZ50QnwMfK+H3jv408d4jw/7Bttv5axS5IiHoLaVNHeQ==", - "dev": true, "license": "MIT" }, "node_modules/prelude-ls": { @@ -10251,7 +10197,6 @@ "version": "1.2.3", "resolved": "https://registry.npmjs.org/queue-microtask/-/queue-microtask-1.2.3.tgz", "integrity": "sha512-NuaNSa6flKT5JaSYQzJok04JzTL1CA6aGhv5rfLW3PgqA+M2ChpZQnAC8h8i4ZFkBS8X5RqkDBHA7r4hej3K9A==", - "dev": true, "funding": [ { "type": "github", @@ -11040,7 +10985,6 @@ "version": "1.0.0", "resolved": "https://registry.npmjs.org/read-cache/-/read-cache-1.0.0.tgz", "integrity": "sha512-Owdv/Ft7IjOgm/i0xvNDZ1LrRANRfew4b2prF3OWMQLxLfu3bS8FVhCsrSCMK4lR56Y9ya+AThoTpDCTxCmpRA==", - "dev": true, "license": "MIT", "dependencies": { "pify": "^2.3.0" @@ -11050,7 +10994,6 @@ "version": "3.6.0", "resolved": "https://registry.npmjs.org/readdirp/-/readdirp-3.6.0.tgz", "integrity": "sha512-hOS089on8RduqdbhvQ5Z37A0ESjsqz6qnRcffsMU3495FuTdqSm+7bhJ29JvIOsBDEEnan5DPu9t3To9VRlMzA==", - "dev": true, "license": "MIT", "dependencies": { "picomatch": "^2.2.1" @@ -11063,7 +11006,6 @@ "version": "2.3.1", "resolved": "https://registry.npmjs.org/picomatch/-/picomatch-2.3.1.tgz", "integrity": "sha512-JU3teHTNjmE2VCGFzuY8EXzCDVwEqB2a8fsIvwaStHhAWJEeVd1o1QD80CU6+ZdEXXSLbSsuLwJjkCBWqRQUVA==", - "dev": true, "license": "MIT", "engines": { "node": ">=8.6" @@ -11361,7 +11303,6 @@ "version": "1.22.11", "resolved": "https://registry.npmjs.org/resolve/-/resolve-1.22.11.tgz", "integrity": "sha512-RfqAvLnMl313r7c9oclB1HhUEAezcpLjz95wFH4LVuhk9JF/r22qmVP9AMmOU4vMX7Q8pN8jwNg/CSpdFnMjTQ==", - "dev": true, "license": "MIT", "dependencies": { "is-core-module": "^2.16.1", @@ -11402,7 +11343,6 @@ "version": "1.1.0", "resolved": "https://registry.npmjs.org/reusify/-/reusify-1.1.0.tgz", "integrity": "sha512-g6QUff04oZpHs0eG5p83rFLhHeV00ug/Yf9nZM6fLeUrPguBTkTQOdpAWWspMh55TZfVQDPaN3NQJfbVRAxdIw==", - "dev": true, "license": "MIT", "engines": { "iojs": ">=1.0.0", @@ -11458,7 +11398,6 @@ "version": "1.2.0", "resolved": "https://registry.npmjs.org/run-parallel/-/run-parallel-1.2.0.tgz", "integrity": "sha512-5l4VyZR86LZ/lDxZTR6jqL8AFE2S0IFLMP26AbjsLVADxHdhB/c0GUsH+y39UfCi3dzz8OlQuPmnaJOMoDHQBA==", - "dev": true, "funding": [ { "type": "github", @@ -12099,7 +12038,6 @@ "version": "3.35.1", "resolved": "https://registry.npmjs.org/sucrase/-/sucrase-3.35.1.tgz", "integrity": "sha512-DhuTmvZWux4H1UOnWMB3sk0sbaCVOoQZjv8u1rDoTV0HTdGem9hkAZtl4JZy8P2z4Bg0nT+YMeOFyVr4zcG5Tw==", - "dev": true, "license": "MIT", "dependencies": { "@jridgewell/gen-mapping": "^0.3.2", @@ -12135,7 +12073,6 @@ "version": "1.0.0", "resolved": "https://registry.npmjs.org/supports-preserve-symlinks-flag/-/supports-preserve-symlinks-flag-1.0.0.tgz", "integrity": "sha512-ot0WnXS9fgdkgIcePe6RHNk1WA8+muPa6cSjeR3V8K27q9BB1rTE3R1p7Hv0z1ZyAc8s6Vvv8DIyWf681MAt0w==", - "dev": true, "license": "MIT", "engines": { "node": ">= 0.4" @@ -12171,7 +12108,6 @@ "version": "3.4.19", "resolved": "https://registry.npmjs.org/tailwindcss/-/tailwindcss-3.4.19.tgz", "integrity": "sha512-3ofp+LL8E+pK/JuPLPggVAIaEuhvIz4qNcf3nA1Xn2o/7fb7s/TYpHhwGDv1ZU3PkBluUVaF8PyCHcm48cKLWQ==", - "dev": true, "license": "MIT", "dependencies": { "@alloc/quick-lru": "^5.2.0", @@ -12209,7 +12145,6 @@ "version": "3.3.3", "resolved": "https://registry.npmjs.org/fast-glob/-/fast-glob-3.3.3.tgz", "integrity": "sha512-7MptL8U0cqcFdzIzwOTHoilX9x5BrNqye7Z/LuC7kCMRio1EMSyqRK3BEAUD7sXRq4iT4AzTVuZdhgQ2TCvYLg==", - "dev": true, "license": "MIT", "dependencies": { "@nodelib/fs.stat": "^2.0.2", @@ -12226,7 +12161,6 @@ "version": "5.1.2", "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-5.1.2.tgz", "integrity": "sha512-AOIgSQCepiJYwP3ARnGx+5VnTu2HBYdzbGP45eLw1vr3zB3vZLeyed1sC9hnbcOc9/SrMyM5RPQrkGz4aS9Zow==", - "dev": true, "license": "ISC", "dependencies": { "is-glob": "^4.0.1" @@ -12254,7 +12188,6 @@ "version": "3.3.1", "resolved": "https://registry.npmjs.org/thenify/-/thenify-3.3.1.tgz", "integrity": "sha512-RVZSIV5IG10Hk3enotrhvz0T9em6cyHBLkH/YAZuKqd8hRkKhSfCGIcP2KUY0EPxndzANBmNllzWPwak+bheSw==", - "dev": true, "license": "MIT", "dependencies": { "any-promise": "^1.0.0" @@ -12264,7 +12197,6 @@ "version": "1.6.0", "resolved": "https://registry.npmjs.org/thenify-all/-/thenify-all-1.6.0.tgz", "integrity": "sha512-RNxQH/qI8/t3thXJDwcstUO4zeqo64+Uy/+sNVRBx4Xn2OX+OZ9oP+iJnNFqplFra2ZUVeKCSa2oVWi3T4uVmA==", - "dev": true, "license": "MIT", "dependencies": { "thenify": ">= 3.1.0 < 4" @@ -12306,7 +12238,6 @@ "version": "0.2.15", "resolved": "https://registry.npmjs.org/tinyglobby/-/tinyglobby-0.2.15.tgz", "integrity": "sha512-j2Zq4NyQYG5XMST4cbs02Ak8iJUdxRM0XI5QyxXuZOzKOINmWurp3smXu3y5wDcJrptwpSjgXHzIQxR0omXljQ==", - "dev": true, "license": "MIT", "dependencies": { "fdir": "^6.5.0", @@ -12373,7 +12304,6 @@ "version": "5.0.1", "resolved": "https://registry.npmjs.org/to-regex-range/-/to-regex-range-5.0.1.tgz", "integrity": "sha512-65P7iz6X5yEr1cwcgvQxbbIw7Uk3gOy5dIdtZ4rDveLqhrdJP+Li/Hx6tyK0NEb+2GCyneCMJiGqrADCSNk8sQ==", - "dev": true, "license": "MIT", "dependencies": { "is-number": "^7.0.0" @@ -12461,7 +12391,6 @@ "version": "0.1.13", "resolved": "https://registry.npmjs.org/ts-interface-checker/-/ts-interface-checker-0.1.13.tgz", "integrity": "sha512-Y/arvbn+rrz3JCKl9C4kVNfTfSm2/mEp5FSz5EsZSANGPSlQrpRI5M4PKF+mJnE52jOO90PnPSc3Ur3bTQw0gA==", - "dev": true, "license": "Apache-2.0" }, "node_modules/tsconfig-paths": { @@ -12578,7 +12507,7 @@ "version": "5.9.3", "resolved": "https://registry.npmjs.org/typescript/-/typescript-5.9.3.tgz", "integrity": "sha512-jl1vZzPDinLr9eUt3J/t7V6FgNEw9QjvBPdysz9KfQDD41fQrC2Y4vKQdiaUpFT4bXlb1RHhLpp8wtm6M5TgSw==", - "dev": true, + "devOptional": true, "license": "Apache-2.0", "bin": { "tsc": "bin/tsc", @@ -12780,7 +12709,6 @@ "version": "1.0.2", "resolved": "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz", "integrity": "sha512-EPD5q1uXyFxJpCrLnCc1nHnq3gOa6DZBocAIiI2TaSCA7VCJ1UJDMagCzIkXNsUYfD1daK//LTEQ8xiIbrHtcw==", - "dev": true, "license": "MIT" }, "node_modules/uuid": { @@ -13234,7 +13162,7 @@ "version": "8.19.0", "resolved": "https://registry.npmjs.org/ws/-/ws-8.19.0.tgz", "integrity": "sha512-blAT2mjOEIi0ZzruJfIhb3nps74PRWTCz1IjglWEEpQl5XS/UNama6u2/rjFkDDouqr4L67ry+1aGIALViWjDg==", - "dev": true, + "devOptional": true, "license": "MIT", "engines": { "node": ">=10.0.0" diff --git a/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx new file mode 100644 index 0000000000..47d2331bed --- /dev/null +++ b/ui/litellm-dashboard/src/app/(dashboard)/skills/page.tsx @@ -0,0 +1,17 @@ +"use client"; + +import ClaudeCodePluginsPanel from "@/components/claude_code_plugins"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; + +const SkillsPage = () => { + const { accessToken, userRole } = useAuthorized(); + + return ( + + ); +}; + +export default SkillsPage; diff --git a/ui/litellm-dashboard/src/app/page.tsx b/ui/litellm-dashboard/src/app/page.tsx index 44df1b5bd4..746847057f 100644 --- a/ui/litellm-dashboard/src/app/page.tsx +++ b/ui/litellm-dashboard/src/app/page.tsx @@ -624,7 +624,7 @@ function CreateKeyPageContent() { ) : page == "tag-management" ? ( - ) : page == "claude-code-plugins" ? ( + ) : page == "skills" || page == "claude-code-plugins" ? ( ) : page == "access-groups" ? ( diff --git a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx index 537aa001b7..0be6919c1f 100644 --- a/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx +++ b/ui/litellm-dashboard/src/components/AIHub/ModelHubTable.tsx @@ -5,7 +5,10 @@ import MakeModelPublicForm from "@/components/AIHub/forms/MakeModelPublicForm"; import { mcpHubColumns, MCPServerData } from "@/components/mcp_hub_table_columns"; import { modelHubColumns } from "@/components/model_hub_table_columns"; import UsefulLinksManagement from "@/components/AIHub/UsefulLinksManagement"; -import ClaudeCodeMarketplaceTab from "@/components/AIHub/ClaudeCodeMarketplaceTab"; +import { getClaudeCodePluginsList } from "@/components/networking"; +import { Plugin } from "@/components/claude_code_plugins/types"; +import SkillHubDashboard from "@/components/AIHub/SkillHubDashboard"; +import MakeSkillPublicForm from "@/components/claude_code_plugins/MakeSkillPublicForm"; import { ModelDataTable } from "@/components/model_dashboard/table"; import ModelFilters from "@/components/model_filters"; import NotificationsManager from "@/components/molecules/notifications_manager"; @@ -78,6 +81,10 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, const [selectedMcpServer, setSelectedMcpServer] = useState(null); const [isMcpModalVisible, setIsMcpModalVisible] = useState(false); const [isMakeMcpPublicModalVisible, setIsMakeMcpPublicModalVisible] = useState(false); + // Skill Hub state + const [skillHubData, setSkillHubData] = useState([]); + const [skillLoading, setSkillLoading] = useState(false); + const [isMakeSkillPublicModalVisible, setIsMakeSkillPublicModalVisible] = useState(false); const router = useRouter(); const { data: uiSettings, isLoading: isUISettingsLoading } = useUISettings(); @@ -208,6 +215,25 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, } }, [publicPage, accessToken]); + // Fetch Skill Hub data — all skills for admins, enabled-only for public page + useEffect(() => { + const fetchSkillData = async () => { + if (!accessToken) return; + try { + setSkillLoading(true); + const enabledOnly = publicPage === true; + const response = await getClaudeCodePluginsList(accessToken, enabledOnly); + setSkillHubData(response.plugins); + } catch (error) { + console.error("Error fetching skill hub data", error); + } finally { + setSkillLoading(false); + } + }; + + fetchSkillData(); + }, [accessToken, publicPage]); + const showModal = (model: ModelGroupInfo) => { setSelectedModel(model); setIsModalVisible(true); @@ -406,7 +432,7 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage, Model Hub Agent Hub MCP Hub - Claude Code Plugin Marketplace + Skill Hub @@ -492,9 +518,26 @@ const ModelHubTable: React.FC = ({ accessToken, publicPage,
- {/* Plugin Marketplace Tab */} + {/* Skill Hub Tab */} - + {publicPage == false && isAdminRole(userRole || "") && ( +
+ +
+ )} + { + const response = await getClaudeCodePluginsList(accessToken || "", publicPage === true); + setSkillHubData(response.plugins); + }} + /> @@ -1055,6 +1098,18 @@ if __name__ == "__main__": mcpHubData={mcpHubData || []} onSuccess={handleMakeMcpPublicSuccess} /> + + {/* Make Skill Public Form */} + setIsMakeSkillPublicModalVisible(false)} + accessToken={accessToken || ""} + skillsList={skillHubData} + onSuccess={async () => { + const response = await getClaudeCodePluginsList(accessToken || "", publicPage === true); + setSkillHubData(response.plugins); + }} + /> ); }; diff --git a/ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx b/ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx new file mode 100644 index 0000000000..af6668dcbe --- /dev/null +++ b/ui/litellm-dashboard/src/components/AIHub/SkillHubDashboard.tsx @@ -0,0 +1,139 @@ +import React, { useMemo, useState } from "react"; +import { Text } from "@tremor/react"; +import { SearchOutlined } from "@ant-design/icons"; +import { Input, Select } from "antd"; +import { Plugin } from "@/components/claude_code_plugins/types"; +import { ModelDataTable } from "@/components/model_dashboard/table"; +import { skillHubColumns } from "@/components/skill_hub_table_columns"; +import SkillDetail from "@/components/claude_code_plugins/skill_detail"; + +interface SkillHubDashboardProps { + skills: Plugin[]; + isLoading: boolean; + isAdmin?: boolean; + accessToken?: string | null; + publicPage?: boolean; + onPublishSuccess?: () => void; +} + +const SkillHubDashboard: React.FC = ({ + skills, + isLoading, + isAdmin, + accessToken, + publicPage = false, + onPublishSuccess, +}) => { + const [search, setSearch] = useState(""); + const [domainFilter, setDomainFilter] = useState(undefined); + const [selectedSkill, setSelectedSkill] = useState(null); + + const copyToClipboard = (text: string) => { + navigator.clipboard.writeText(text); + }; + + // Derived stats + const totalSkills = skills.length; + const domains = useMemo(() => [...new Set(skills.map((s) => s.domain).filter(Boolean))], [skills]); + const namespaces = useMemo(() => [...new Set(skills.map((s) => s.namespace).filter(Boolean))], [skills]); + + // Filtered table data + const filteredSkills = useMemo(() => { + let result = skills; + if (domainFilter) { + result = result.filter((s) => (s.domain || "General") === domainFilter); + } + if (search.trim()) { + const q = search.toLowerCase(); + result = result.filter( + (s) => + s.name.toLowerCase().includes(q) || + s.description?.toLowerCase().includes(q) || + s.domain?.toLowerCase().includes(q) || + s.namespace?.toLowerCase().includes(q) || + s.keywords?.some((k) => k.toLowerCase().includes(q)) + ); + } + return result; + }, [skills, search, domainFilter]); + + if (selectedSkill) { + return ( + setSelectedSkill(null)} + isAdmin={isAdmin} + accessToken={accessToken} + onPublishClick={onPublishSuccess} + /> + ); + } + + if (isLoading) { + return
Loading skills...
; + } + + return ( +
+ {/* Stats row */} +
+
+
Total Skills
+
{totalSkills}
+
+
+
Namespaces
+
{namespaces.length}
+
+
+
Domains
+
{domains.length}
+
+
+ + {/* Search + filters + table */} +
+
+

+ All {publicPage ? "Public " : ""}Skills +

+
+ } + placeholder="Search by name, namespace, or tag…" + value={search} + onChange={(e) => setSearch(e.target.value)} + style={{ width: 280 }} + allowClear + /> +
+
+ setSelectedSkill(skill), + copyToClipboard, + publicPage + )} + data={filteredSkills} + isLoading={false} + defaultSorting={[{ id: "name", desc: false }]} + /> +
+ + Showing {filteredSkills.length} of {totalSkills} skill{totalSkills !== 1 ? "s" : ""} + +
+
+
+ ); +}; + +export default SkillHubDashboard; diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins.tsx index b2e66349a2..67d5f6e0ad 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins.tsx @@ -7,8 +7,8 @@ import { } from "./networking"; import AddPluginForm from "./claude_code_plugins/add_plugin_form"; import PluginTable from "./claude_code_plugins/plugin_table"; +import SkillDetail from "./claude_code_plugins/skill_detail"; import { isAdminRole } from "@/utils/roles"; -import PluginInfoView from "./claude_code_plugins/plugin_info"; import NotificationsManager from "./molecules/notifications_manager"; import { Plugin, ListPluginsResponse } from "./claude_code_plugins/types"; @@ -29,27 +29,22 @@ const ClaudeCodePluginsPanel: React.FC = ({ name: string; displayName: string; } | null>(null); - const [selectedPluginId, setSelectedPluginId] = useState( - null - ); + const [selectedSkill, setSelectedSkill] = useState(null); const isAdmin = userRole ? isAdminRole(userRole) : false; const fetchPlugins = async () => { - if (!accessToken) { - return; - } + if (!accessToken) return; setIsLoading(true); try { const response: ListPluginsResponse = await getClaudeCodePluginsList( accessToken, - false // Get all plugins (enabled and disabled) + false ); - console.log(`Claude Code plugins: ${JSON.stringify(response)}`); setPluginsList(response.plugins); } catch (error) { - console.error("Error fetching Claude Code plugins:", error); + console.error("Error fetching skills:", error); } finally { setIsLoading(false); } @@ -59,21 +54,6 @@ const ClaudeCodePluginsPanel: React.FC = ({ fetchPlugins(); }, [accessToken]); - const handleAddPlugin = () => { - if (selectedPluginId) { - setSelectedPluginId(null); - } - setIsAddModalVisible(true); - }; - - const handleCloseModal = () => { - setIsAddModalVisible(false); - }; - - const handleSuccess = () => { - fetchPlugins(); - }; - const handleDeleteClick = (pluginName: string, displayName: string) => { setPluginToDelete({ name: pluginName, displayName }); }; @@ -84,79 +64,76 @@ const ClaudeCodePluginsPanel: React.FC = ({ setIsDeleting(true); try { await deleteClaudeCodePlugin(accessToken, pluginToDelete.name); - NotificationsManager.success( - `Plugin "${pluginToDelete.displayName}" deleted successfully` - ); + NotificationsManager.success(`Skill "${pluginToDelete.displayName}" deleted successfully`); fetchPlugins(); } catch (error) { - console.error("Error deleting plugin:", error); - NotificationsManager.error("Failed to delete plugin"); + console.error("Error deleting skill:", error); + NotificationsManager.error("Failed to delete skill"); } finally { setIsDeleting(false); setPluginToDelete(null); } }; - const handleDeleteCancel = () => { - setPluginToDelete(null); - }; - return (
-
-

Claude Code Plugins

-

- Manage Claude Code marketplace plugins. Add, enable, disable, or - delete plugins that will be available in your marketplace catalog. - Enabled plugins will appear in the public marketplace at{" "} - /claude-code/marketplace.json. -

-
- -
-
- - {selectedPluginId ? ( - setSelectedPluginId(null)} - accessToken={accessToken} + {selectedSkill ? ( + setSelectedSkill(null)} isAdmin={isAdmin} - onPluginUpdated={fetchPlugins} + accessToken={accessToken} + onPublishClick={fetchPlugins} /> ) : ( - setSelectedPluginId(id)} - /> + <> +
+

Skills

+

+ Register Claude Code skills. Published skills appear in the Skill Hub for all users and + are served via{" "} + /claude-code/marketplace.json. +

+
+ +
+
+ + { + const skill = pluginsList.find((p) => p.id === id); + if (skill) setSelectedSkill(skill); + }} + /> + )} setIsAddModalVisible(false)} accessToken={accessToken} - onSuccess={handleSuccess} + onSuccess={fetchPlugins} /> {pluginToDelete && ( setPluginToDelete(null)} confirmLoading={isDeleting} okText="Delete" okButtonProps={{ danger: true }} >

- Are you sure you want to delete plugin:{" "} + Are you sure you want to delete skill:{" "} {pluginToDelete.displayName}?

This action cannot be undone.

diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx new file mode 100644 index 0000000000..351502d78b --- /dev/null +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/MakeSkillPublicForm.tsx @@ -0,0 +1,249 @@ +import React, { useState, useEffect } from "react"; +import { Modal, Form, Steps, Button, Checkbox } from "antd"; +import { Text, Title, Badge } from "@tremor/react"; +import { enableClaudeCodePlugin, disableClaudeCodePlugin } from "../networking"; +import NotificationsManager from "../molecules/notifications_manager"; +import { Plugin } from "./types"; + +const { Step } = Steps; + +interface MakeSkillPublicFormProps { + visible: boolean; + onClose: () => void; + accessToken: string; + skillsList: Plugin[]; + onSuccess: () => void; +} + +const MakeSkillPublicForm: React.FC = ({ + visible, + onClose, + accessToken, + skillsList, + onSuccess, +}) => { + const [currentStep, setCurrentStep] = useState(0); + const [selectedSkills, setSelectedSkills] = useState>(new Set()); + const [loading, setLoading] = useState(false); + const [form] = Form.useForm(); + + const handleClose = () => { + setCurrentStep(0); + setSelectedSkills(new Set()); + form.resetFields(); + onClose(); + }; + + const handleNext = () => { + if (selectedSkills.size === 0) { + NotificationsManager.fromBackend("Please select at least one skill"); + return; + } + setCurrentStep(1); + }; + + const handleSkillSelection = (name: string, checked: boolean) => { + const next = new Set(selectedSkills); + if (checked) { + next.add(name); + } else { + next.delete(name); + } + setSelectedSkills(next); + }; + + const handleSelectAll = (checked: boolean) => { + if (checked) { + setSelectedSkills(new Set(skillsList.map((s) => s.name))); + } else { + setSelectedSkills(new Set()); + } + }; + + // Pre-check already-published skills when modal opens + useEffect(() => { + if (visible && skillsList.length > 0) { + setSelectedSkills(new Set(skillsList.filter((s) => s.enabled).map((s) => s.name))); + } + }, [visible, skillsList]); + + const handleSubmit = async () => { + if (selectedSkills.size === 0) { + NotificationsManager.fromBackend("Please select at least one skill"); + return; + } + + setLoading(true); + try { + const selectedSet = selectedSkills; + await Promise.all( + skillsList.map((skill) => { + const shouldBePublic = selectedSet.has(skill.name); + if (shouldBePublic && !skill.enabled) { + return enableClaudeCodePlugin(accessToken, skill.name); + } + if (!shouldBePublic && skill.enabled) { + return disableClaudeCodePlugin(accessToken, skill.name); + } + return Promise.resolve(); + }) + ); + + NotificationsManager.success(`Skill Hub updated — ${selectedSkills.size} skill(s) published`); + handleClose(); + onSuccess(); + } catch (error) { + console.error("Error publishing skills:", error); + NotificationsManager.fromBackend("Failed to update skills. Please try again."); + } finally { + setLoading(false); + } + }; + + const allSelected = + skillsList.length > 0 && skillsList.every((s) => selectedSkills.has(s.name)); + const isIndeterminate = selectedSkills.size > 0 && !allSelected; + + const renderStep1 = () => ( +
+
+ Select Skills to Publish + handleSelectAll(e.target.checked)} + disabled={skillsList.length === 0} + > + Select All ({skillsList.length}) + +
+ + + Selected skills will be visible to all users in the Skill Hub. + Deselected skills will be unpublished. + + +
+
+ {skillsList.length === 0 ? ( +
+ No skills registered yet. +
+ ) : ( + skillsList.map((skill) => ( +
+ handleSkillSelection(skill.name, e.target.checked)} + /> +
+
+ {skill.name} + {skill.enabled && ( + Public + )} +
+ {skill.description && ( + + {skill.description} + + )} +
+ {skill.domain && ( + {skill.domain} + )} +
+ )) + )} +
+
+ + {selectedSkills.size > 0 && ( +
+ + {selectedSkills.size} skill{selectedSkills.size !== 1 ? "s" : ""} will be published + +
+ )} +
+ ); + + const renderStep2 = () => ( +
+ Confirm Publish to Skill Hub + +
+ + Note: Published skills will be visible to all users in the Skill Hub tab. + Skills not in the list below will be unpublished. + +
+ +
+ Skills to be published: +
+
+ {Array.from(selectedSkills).map((name) => { + const skill = skillsList.find((s) => s.name === name); + return ( +
+ {name} + {skill?.domain && {skill.domain}} +
+ ); + })} +
+
+
+ +
+ + Total: {selectedSkills.size} skill{selectedSkills.size !== 1 ? "s" : ""} will be published + +
+
+ ); + + return ( + +
+ + + + + + {currentStep === 0 ? renderStep1() : renderStep2()} + +
+ +
+ {currentStep === 0 && ( + + )} + {currentStep === 1 && ( + + )} +
+
+
+ + ); +}; + +export default MakeSkillPublicForm; diff --git a/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx b/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx index d6a3f0f4fd..fdbef3a766 100644 --- a/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx +++ b/ui/litellm-dashboard/src/components/claude_code_plugins/add_plugin_form.tsx @@ -32,6 +32,80 @@ const PREDEFINED_CATEGORIES = [ "Documentation", ]; +interface ParsedSource { + source: "github" | "url" | "git-subdir"; + repo?: string; + url?: string; + path?: string; +} + +interface ParsePreview { + parsed: ParsedSource; + label: string; + suggestedName: string; +} + +function parseGitHubUrl(raw: string): ParsePreview | null { + // Strip protocol and trailing slashes/spaces + let s = raw.trim().replace(/^https?:\/\//, "").replace(/\/+$/, ""); + + if (!s.startsWith("github.com/")) return null; + + // Remove "github.com/" + const rest = s.slice("github.com/".length); + const parts = rest.split("/"); + + if (parts.length < 2) return null; + + const org = parts[0]; + const repo = parts[1]; + const repoBase = repo.replace(/\.git$/, ""); + + // github.com/org/repo (exactly 2 parts, or ends with .git) + if (parts.length === 2 || (parts.length === 2 && repoBase)) { + return { + parsed: { source: "github", repo: `${org}/${repoBase}` }, + label: `GitHub repo — ${org}/${repoBase}`, + suggestedName: repoBase, + }; + } + + // github.com/org/repo/tree/branch/folder or /blob/branch/folder/FILE.md + if ( + parts.length >= 5 && + (parts[2] === "tree" || parts[2] === "blob") + ) { + // parts[3] = branch, parts[4..] = path segments + const pathParts = parts.slice(4); + // If last segment looks like a file (has extension), drop it + const lastPart = pathParts[pathParts.length - 1]; + if (lastPart && lastPart.includes(".")) { + pathParts.pop(); + } + if (pathParts.length === 0) { + // Path resolved to repo root — treat as plain github source + return { + parsed: { source: "github", repo: `${org}/${repoBase}` }, + label: `GitHub repo — ${org}/${repoBase}`, + suggestedName: repoBase, + }; + } + const subPath = pathParts.join("/"); + const suggestedName = pathParts[pathParts.length - 1]; + return { + parsed: { + source: "git-subdir", + url: `https://github.com/${org}/${repoBase}`, + path: subPath, + }, + label: `GitHub subdir — ${org}/${repoBase} @ ${subPath}`, + suggestedName, + }; + } + + return null; +} + const AddPluginForm: React.FC = ({ visible, onClose, @@ -40,7 +114,20 @@ const AddPluginForm: React.FC = ({ }) => { const [form] = Form.useForm(); const [isSubmitting, setIsSubmitting] = useState(false); - const [sourceType, setSourceType] = useState<"github" | "url" | "git-subdir">("github"); + const [urlPreview, setUrlPreview] = useState(null); + + const handleUrlChange = (e: React.ChangeEvent) => { + const val = e.target.value; + const preview = parseGitHubUrl(val); + setUrlPreview(preview); + if (preview) { + // Auto-fill name only if it's currently empty + const currentName = form.getFieldValue("name"); + if (!currentName) { + form.setFieldsValue({ name: preview.suggestedName }); + } + } + }; const handleSubmit = async (values: any) => { if (!accessToken) { @@ -48,98 +135,62 @@ const AddPluginForm: React.FC = ({ return; } - // Validate plugin name + if (!urlPreview) { + MessageManager.error("Please enter a valid GitHub URL"); + return; + } + if (!validatePluginName(values.name)) { MessageManager.error( - "Plugin name must be kebab-case (lowercase letters, numbers, and hyphens only)" + "Skill name must be kebab-case (lowercase letters, numbers, and hyphens only)" ); return; } - // Validate semantic version if provided if (values.version && !isValidSemanticVersion(values.version)) { - MessageManager.error( - "Version must be in semantic versioning format (e.g., 1.0.0)" - ); + MessageManager.error("Version must be in semantic versioning format (e.g., 1.0.0)"); return; } - // Validate email if provided if (values.authorEmail && !isValidEmail(values.authorEmail)) { MessageManager.error("Invalid email format"); return; } - // Validate homepage URL if provided if (values.homepage && !isValidUrl(values.homepage)) { MessageManager.error("Invalid homepage URL format"); return; } - // Validate git URL for url/git-subdir source types - if ((sourceType === "url" || sourceType === "git-subdir") && values.url && !isValidUrl(values.url)) { - MessageManager.error("Invalid git URL format"); - return; - } - setIsSubmitting(true); try { - // Build plugin data const pluginData: any = { name: values.name.trim(), - source: - sourceType === "github" - ? { - source: "github", - repo: values.repo.trim(), - } - : sourceType === "git-subdir" - ? { - source: "git-subdir", - url: values.url.trim(), - path: values.path.trim(), - } - : { - source: "url", - url: values.url.trim(), - }, + source: urlPreview.parsed, }; - // Add optional fields - if (values.version) { - pluginData.version = values.version.trim(); - } - if (values.description) { - pluginData.description = values.description.trim(); - } + if (values.version) pluginData.version = values.version.trim(); + if (values.description) pluginData.description = values.description.trim(); if (values.authorName || values.authorEmail) { pluginData.author = {}; - if (values.authorName) { - pluginData.author.name = values.authorName.trim(); - } - if (values.authorEmail) { - pluginData.author.email = values.authorEmail.trim(); - } - } - if (values.homepage) { - pluginData.homepage = values.homepage.trim(); - } - if (values.category) { - pluginData.category = values.category; - } - if (values.keywords) { - pluginData.keywords = parseKeywords(values.keywords); + if (values.authorName) pluginData.author.name = values.authorName.trim(); + if (values.authorEmail) pluginData.author.email = values.authorEmail.trim(); } + if (values.homepage) pluginData.homepage = values.homepage.trim(); + if (values.category) pluginData.category = values.category; + if (values.keywords) pluginData.keywords = parseKeywords(values.keywords); + if (values.domain) pluginData.domain = values.domain.trim(); + if (values.namespace) pluginData.namespace = values.namespace.trim(); await registerClaudeCodePlugin(accessToken, pluginData); - MessageManager.success("Plugin registered successfully"); + MessageManager.success("Skill registered successfully"); form.resetFields(); - setSourceType("github"); + setUrlPreview(null); onSuccess(); onClose(); } catch (error) { - console.error("Error registering plugin:", error); - MessageManager.error("Failed to register plugin"); + console.error("Error registering skill:", error); + MessageManager.error("Failed to register skill"); } finally { setIsSubmitting(false); } @@ -147,19 +198,13 @@ const AddPluginForm: React.FC = ({ const handleCancel = () => { form.resetFields(); - setSourceType("github"); + setUrlPreview(null); onClose(); }; - const handleSourceTypeChange = (value: "github" | "url" | "git-subdir") => { - setSourceType(value); - // Clear repo/url/path fields when switching - form.setFieldsValue({ repo: undefined, url: undefined, path: undefined }); - }; - return ( = ({ onFinish={handleSubmit} className="mt-4" > - {/* Plugin Name */} + {/* Smart URL Input */} + + + + {/* Parsed preview */} + {urlPreview && ( +
+ Detected: {urlPreview.label} +
+ )} + + {/* Skill Name */} + - + - {/* Source Type */} - - - - - {/* GitHub Repository */} - {sourceType === "github" && ( + {/* Domain and Namespace — side by side */} +
- + - )} - - {/* Git URL */} - {(sourceType === "url" || sourceType === "git-subdir") && ( - + - )} - - {/* Git Subdir Path */} - {sourceType === "git-subdir" && ( - - - - )} - - {/* Version */} - - - +
{/* Description */}