diff --git a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py index 1ae87e99c9e..737a2d21c99 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aim/aim.py @@ -58,6 +58,7 @@ class AimGuardrail(CustomGuardrail): self.api_base = ( api_base or os.environ.get("AIM_API_BASE") or "https://api.aim.security" ) + self.api_base = self.api_base.rstrip("/") self.ws_api_base = self.api_base.replace("http://", "ws://").replace( "https://", "wss://" ) diff --git a/tests/guardrails_tests/test_aim_guardrail_base_url.py b/tests/guardrails_tests/test_aim_guardrail_base_url.py new file mode 100644 index 00000000000..3e230bb32cc --- /dev/null +++ b/tests/guardrails_tests/test_aim_guardrail_base_url.py @@ -0,0 +1,21 @@ +import pytest +import os +from litellm.proxy.guardrails.guardrail_hooks.aim.aim import CustomGuardrail + +def test_aim_base_url_trailing_slash(): + # Test with trailing slash + guardrail = CustomGuardrail(api_base="https://api.aim.security/") + assert guardrail.api_base == "https://api.aim.security" + assert guardrail.ws_api_base.startswith("wss://api.aim.security") + + # Test without trailing slash + guardrail2 = CustomGuardrail(api_base="https://api.aim.security") + assert guardrail2.api_base == "https://api.aim.security" + assert guardrail2.ws_api_base.startswith("wss://api.aim.security") + + # Test with environment variable + os.environ["AIM_API_BASE"] = "https://api.aim.security/" + guardrail3 = CustomGuardrail(api_base=None) + assert guardrail3.api_base == "https://api.aim.security" + assert guardrail3.ws_api_base.startswith("wss://api.aim.security") + del os.environ["AIM_API_BASE"]