mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-13 15:07:47 +00:00
fix: reject status code and tests.
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user