fix: reject status code and tests.

This commit is contained in:
Abhijit L
2025-09-28 16:56:28 +05:30
parent 8169c61a48
commit 0d7256cd4e
3 changed files with 32 additions and 21 deletions
@@ -22,8 +22,8 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"
guardrail_name=guardrail.get("guardrail_name", ""),
javelin_guard_name=litellm_params.guard_name,
event_hook=litellm_params.mode,
default_on=litellm_params.default_on,
api_version=litellm_params.api_version,
default_on=litellm_params.default_on or False,
api_version=litellm_params.api_version or "v1",
config=litellm_params.config,
metadata=litellm_params.metadata,
application=litellm_params.application,
@@ -31,7 +31,6 @@ class JavelinGuardrail(CustomGuardrail):
metadata: Optional[Dict] = None,
config: Optional[Dict] = None,
application: Optional[str] = None,
event_hook: Optional[str] = None,
**kwargs,
):
f"""
@@ -73,7 +72,8 @@ class JavelinGuardrail(CustomGuardrail):
self.api_base,
self.api_version,
)
super().__init__(guardrail_name=guardrail_name, event_hook=event_hook, default_on=default_on, **kwargs)
super().__init__(guardrail_name=guardrail_name, default_on=default_on, **kwargs)
async def call_javelin_guard(
self,
@@ -105,7 +105,7 @@ class JavelinGuardrail(CustomGuardrail):
response = await self.async_handler.post(
url=url,
headers=headers,
json=request,
json=dict(request),
)
verbose_proxy_logger.debug(
"Javelin Guardrail: Javelin guard API response: %s", response.json()
@@ -178,8 +178,9 @@ class JavelinGuardrail(CustomGuardrail):
from litellm.proxy.common_utils.callback_utils import (
add_guardrail_to_applied_guardrails_header,
)
from litellm.litellm_core_utils.prompt_templates.common_utils import (get_last_user_message)
from litellm.litellm_core_utils.prompt_templates.common_utils import (
get_last_user_message,
)
verbose_proxy_logger.debug("Javelin Guardrail: pre_call_hook")
verbose_proxy_logger.debug("Javelin Guardrail: Request data: %s", data)
@@ -270,14 +271,14 @@ class JavelinGuardrail(CustomGuardrail):
# Raise HTTPException to prevent the request from going to the LLM
raise HTTPException(
status_code=400,
status_code=500,
detail={
"error": "Violated guardrail policy",
"javelin_guardrail_response": javelin_response,
"reject_prompt": reject_prompt,
},
)
add_guardrail_to_applied_guardrails_header(
request_data=data, guardrail_name=self.guardrail_name
)
@@ -67,11 +67,14 @@ async def test_javelin_guardrail_reject_prompt():
call_type="completion")
# Verify the exception details
assert exc_info.value.status_code == 400
assert exc_info.value.status_code == 500
assert "Violated guardrail policy" in str(exc_info.value.detail)
assert "javelin_guardrail_response" in exc_info.value.detail
assert "reject_prompt" in exc_info.value.detail
assert exc_info.value.detail["reject_prompt"] == "Unable to complete request, prompt injection/jailbreak detected"
detail_dict = exc_info.value.detail
assert isinstance(detail_dict, dict)
detail_dict = dict(detail_dict)
assert "javelin_guardrail_response" in detail_dict
assert "reject_prompt" in detail_dict
assert detail_dict["reject_prompt"] == "Unable to complete request, prompt injection/jailbreak detected"
#test trustsafety guardrail
@pytest.mark.asyncio
@@ -139,11 +142,14 @@ async def test_javelin_guardrail_trustsafety():
call_type="completion")
# Verify the exception details
assert exc_info.value.status_code == 400
assert exc_info.value.status_code == 500
assert "Violated guardrail policy" in str(exc_info.value.detail)
assert "javelin_guardrail_response" in exc_info.value.detail
assert "reject_prompt" in exc_info.value.detail
assert exc_info.value.detail["reject_prompt"] == "Unable to complete request, trust & safety violation detected"
detail_dict = exc_info.value.detail
assert isinstance(detail_dict, dict)
detail_dict = dict(detail_dict) # Ensure type checker knows it's a dict
assert "javelin_guardrail_response" in detail_dict
assert "reject_prompt" in detail_dict
assert detail_dict["reject_prompt"] == "Unable to complete request, trust & safety violation detected"
#test language detection guardrail
@pytest.mark.asyncio
@@ -197,11 +203,14 @@ async def test_javelin_guardrail_language_detection():
call_type="completion")
# Verify the exception details
assert exc_info.value.status_code == 400
assert exc_info.value.status_code == 500
assert "Violated guardrail policy" in str(exc_info.value.detail)
assert "javelin_guardrail_response" in exc_info.value.detail
assert "reject_prompt" in exc_info.value.detail
assert exc_info.value.detail["reject_prompt"] == "Unable to complete request, language violation detected"
detail_dict = exc_info.value.detail
assert isinstance(detail_dict, dict)
detail_dict = dict(detail_dict) # Ensure type checker knows it's a dict
assert "javelin_guardrail_response" in detail_dict
assert "reject_prompt" in detail_dict
assert detail_dict["reject_prompt"] == "Unable to complete request, language violation detected"
@pytest.mark.asyncio
@@ -237,4 +246,5 @@ async def test_javelin_guardrail_no_user_message():
# Verify the response is unchanged
assert response is not None
assert isinstance(response, dict)
assert response["messages"] == original_messages