diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 501185b207..1ca45f907e 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -119,11 +119,8 @@ class CustomGuardrail(CustomLogger): """ if "guardrails" in data: return data["guardrails"] - metadata = data.get("metadata") or {} - requested_guardrails = metadata.get("guardrails") or [] - if requested_guardrails: - return requested_guardrails - return requested_guardrails + metadata = data.get("litellm_metadata") or data.get("metadata", {}) + return metadata.get("guardrails") or [] def _guardrail_is_in_requested_guardrails( self, diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index a10a39a6a5..2de5ec1ee1 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -90,6 +90,17 @@ async def anthropic_response( # noqa: PLR0915 user_api_key_dict=user_api_key_dict, data=data, call_type="text_completion" ) + tasks = [] + tasks.append( + proxy_logging_obj.during_call_hook( + data=data, + user_api_key_dict=user_api_key_dict, + call_type=ProxyBaseLLMRequestProcessing._get_pre_call_type( + route_type="anthropic_messages" # type: ignore + ), + ) + ) + ### ROUTE THE REQUESTs ### router_model_names = llm_router.model_names if llm_router is not None else [] @@ -97,23 +108,21 @@ async def anthropic_response( # noqa: PLR0915 if ( llm_router is not None and data["model"] in router_model_names ): # model in router model list - llm_response = asyncio.create_task(llm_router.aanthropic_messages(**data)) + llm_coro = llm_router.aanthropic_messages(**data) elif ( llm_router is not None and llm_router.model_group_alias is not None and data["model"] in llm_router.model_group_alias ): # model set in model_group_alias - llm_response = asyncio.create_task(llm_router.aanthropic_messages(**data)) + llm_coro = llm_router.aanthropic_messages(**data) elif ( llm_router is not None and data["model"] in llm_router.deployment_names ): # model in router deployments, calling a specific deployment on the router - llm_response = asyncio.create_task( - llm_router.aanthropic_messages(**data, specific_deployment=True) - ) + llm_coro = llm_router.aanthropic_messages(**data, specific_deployment=True) elif ( llm_router is not None and data["model"] in llm_router.get_model_ids() ): # model in router model list - llm_response = asyncio.create_task(llm_router.aanthropic_messages(**data)) + llm_coro = llm_router.aanthropic_messages(**data) elif ( llm_router is not None and data["model"] not in router_model_names @@ -122,9 +131,9 @@ async def anthropic_response( # noqa: PLR0915 or len(llm_router.pattern_router.patterns) > 0 ) ): # model in router deployments, calling a specific deployment on the router - llm_response = asyncio.create_task(llm_router.aanthropic_messages(**data)) + llm_coro = llm_router.aanthropic_messages(**data) elif user_model is not None: # `litellm --model ` - llm_response = asyncio.create_task(litellm.anthropic_messages(**data)) + llm_coro = litellm.anthropic_messages(**data) else: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -134,8 +143,16 @@ async def anthropic_response( # noqa: PLR0915 }, ) - # Await the llm_response task - response = await llm_response + tasks.append(llm_coro) + + # wait for call to end + llm_responses = asyncio.gather( + *tasks + ) # run the moderation check in parallel to the actual llm api call + + responses = await llm_responses + + response = responses[1] hidden_params = getattr(response, "_hidden_params", {}) or {} model_id = hidden_params.get("model_id", None) or "" @@ -183,6 +200,11 @@ async def anthropic_response( # noqa: PLR0915 headers=dict(fastapi_response.headers), ) + ### CALL HOOKS ### - modify outgoing data + response = await proxy_logging_obj.post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response # type: ignore + ) + verbose_proxy_logger.info("\nResponse from Litellm:\n{}".format(response)) return response except Exception as e: diff --git a/tests/test_litellm/integrations/test_custom_guardrail.py b/tests/test_litellm/integrations/test_custom_guardrail.py index e71b68ab93..182e013492 100644 --- a/tests/test_litellm/integrations/test_custom_guardrail.py +++ b/tests/test_litellm/integrations/test_custom_guardrail.py @@ -1,4 +1,4 @@ -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import pytest @@ -82,3 +82,101 @@ class TestCustomGuardrailDeploymentHook: # Verify messages were updated in result assert result["messages"] == mock_result["messages"] assert result["messages"] != original_messages + + +class TestCustomGuardrailShouldRunGuardrail: + + def test_should_run_guardrail_with_litellm_metadata(self): + """Test that should_run_guardrail works with litellm_metadata pattern""" + from litellm.types.guardrails import GuardrailEventHooks + + custom_guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=False, + event_hook=GuardrailEventHooks.pre_call + ) + + # Test with guardrails in litellm_metadata + data = { + "model": "gpt-3.5-turbo", + "litellm_metadata": { + "guardrails": ["test_guardrail"] + } + } + + result = custom_guardrail.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + + assert result is True + + def test_should_run_guardrail_with_metadata(self): + """Test that should_run_guardrail works with metadata pattern""" + from litellm.types.guardrails import GuardrailEventHooks + + custom_guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=False, + event_hook=GuardrailEventHooks.pre_call + ) + + # Test with guardrails in metadata + data = { + "model": "gpt-3.5-turbo", + "metadata": { + "guardrails": ["test_guardrail"] + } + } + + result = custom_guardrail.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + + assert result is True + + def test_should_run_guardrail_with_root_level_guardrails(self): + """Test that should_run_guardrail works with root level guardrails""" + from litellm.types.guardrails import GuardrailEventHooks + + custom_guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=False, + event_hook=GuardrailEventHooks.pre_call + ) + + # Test with guardrails at root level + data = { + "model": "gpt-3.5-turbo", + "guardrails": ["test_guardrail"] + } + + result = custom_guardrail.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + + assert result is True + + + def test_should_run_guardrail_no_matching_guardrail(self): + """Test that should_run_guardrail returns False when guardrail name doesn't match""" + from litellm.types.guardrails import GuardrailEventHooks + + custom_guardrail = CustomGuardrail( + guardrail_name="test_guardrail", + default_on=False, + event_hook=GuardrailEventHooks.pre_call + ) + + # Test with different guardrail name + data = { + "model": "gpt-3.5-turbo", + "litellm_metadata": { + "guardrails": ["different_guardrail"] + } + } + + result = custom_guardrail.should_run_guardrail( + data=data, event_type=GuardrailEventHooks.pre_call + ) + + assert result is False