diff --git a/litellm/proxy/guardrails/guardrail_hooks/presidio.py b/litellm/proxy/guardrails/guardrail_hooks/presidio.py index adfe69315c6..2b9cbef282a 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/presidio.py +++ b/litellm/proxy/guardrails/guardrail_hooks/presidio.py @@ -552,7 +552,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): entity_type = item.get("entity_type") str_entity_type = str( - entity_type.value if hasattr(entity_type, "value") else entity_type + entity_type.value if hasattr(entity_type, "value") else entity_type # type: ignore[union-attr] ) if entity_type and str_entity_type in deny_list_strings: continue @@ -1135,7 +1135,7 @@ class _OPTIONAL_PresidioPIIMasking(CustomGuardrail): ) and remaining_chunks: last_chunk = remaining_chunks[-1] if hasattr(last_chunk, "usage") and last_chunk.usage: - assembled_model_response.usage = last_chunk.usage + assembled_model_response.usage = last_chunk.usage # type: ignore[attr-defined] # Apply PII unmasking to assembled content (unmasking tokens back to original text) await self._process_response_for_pii( diff --git a/litellm/proxy/health_endpoints/_health_endpoints.py b/litellm/proxy/health_endpoints/_health_endpoints.py index 5a4a8eae385..86a700d9c71 100644 --- a/litellm/proxy/health_endpoints/_health_endpoints.py +++ b/litellm/proxy/health_endpoints/_health_endpoints.py @@ -300,7 +300,7 @@ async def health_services_endpoint( # noqa: PLR0915 datadog_metrics_logger = DatadogMetricsLogger( start_periodic_flush=False ) - response = await datadog_metrics_logger.async_health_check() + response = await datadog_metrics_logger.async_health_check() # type: ignore[attr-defined] return { "status": response["status"], "message": ( diff --git a/tests/guardrails_tests/test_presidio_pii.py b/tests/guardrails_tests/test_presidio_pii.py index a1f6b7bfbbd..47346965ca9 100644 --- a/tests/guardrails_tests/test_presidio_pii.py +++ b/tests/guardrails_tests/test_presidio_pii.py @@ -286,7 +286,8 @@ async def test_output_parsing(): } ] - pii_masking.pii_tokens = {"": "Jane Doe", "": "034453334"} + pii_tokens = {"": "Jane Doe", "": "034453334"} + pii_masking.pii_tokens = pii_tokens response = mock_completion( model="gpt-3.5-turbo", @@ -296,7 +297,8 @@ async def test_output_parsing(): new_response = await pii_masking.async_post_call_success_hook( user_api_key_dict=UserAPIKeyAuth(), data={ - "messages": [{"role": "system", "content": "You are an helpfull assistant"}] + "messages": [{"role": "system", "content": "You are an helpfull assistant"}], + "pii_tokens": pii_tokens, }, response=response, ) diff --git a/tests/test_litellm/proxy/auth/test_handle_jwt.py b/tests/test_litellm/proxy/auth/test_handle_jwt.py index 8418dde5e9c..3c190974277 100644 --- a/tests/test_litellm/proxy/auth/test_handle_jwt.py +++ b/tests/test_litellm/proxy/auth/test_handle_jwt.py @@ -1559,6 +1559,7 @@ async def test_resolve_jwks_url_caches_resolved_jwks_uri(): jwks_url = "https://login.microsoftonline.com/tenant/discovery/keys" mock_response = MagicMock() + mock_response.status_code = 200 mock_response.json.return_value = {"jwks_uri": jwks_url} with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response) as mock_get: @@ -1587,6 +1588,7 @@ async def test_resolve_jwks_url_raises_if_no_jwks_uri_in_discovery_doc(): discovery_url = "https://example.com/.well-known/openid-configuration" mock_response = MagicMock() + mock_response.status_code = 200 mock_response.json.return_value = {"issuer": "https://example.com"} # no jwks_uri with patch.object(handler.http_handler, "get", new_callable=AsyncMock, return_value=mock_response):