diff --git a/tests/otel_tests/test_guardrails.py b/tests/otel_tests/test_guardrails.py index 979f11149..b82bc57cc 100644 --- a/tests/otel_tests/test_guardrails.py +++ b/tests/otel_tests/test_guardrails.py @@ -6,18 +6,33 @@ from typing import Optional, List, Union import uuid -async def chat_completion(session, key, messages, model: Union[str, List] = "gpt-4"): +async def chat_completion( + session, + key, + messages, + model: Union[str, List] = "gpt-4", + guardrails: Optional[List] = None, +): url = "http://0.0.0.0:4000/chat/completions" headers = { "Authorization": f"Bearer {key}", "Content-Type": "application/json", } + data = { "model": model, "messages": messages, - "guardrails": ["aporia-post-guard", "aporia-pre-guard"], + "guardrails": [ + "aporia-post-guard", + "aporia-pre-guard", + ], # default guardrails for all tests } + if guardrails is not None: + data["guardrails"] = guardrails + + print("data=", data) + async with session.post(url, headers=headers, json=data) as response: status = response.status response_text = await response.text() @@ -36,7 +51,7 @@ async def chat_completion(session, key, messages, model: Union[str, List] = "gpt @pytest.mark.asyncio -async def test_no_llm_guard_triggered(): +async def test_llm_guard_triggered_safe_request(): """ - Tests a request where no content mod is triggered - Assert that the guardrails applied are returned in the response headers @@ -80,3 +95,24 @@ async def test_llm_guard_triggered(): except Exception as e: print(e) assert "Aporia detected and blocked PII" in str(e) + + +@pytest.mark.asyncio +async def test_no_llm_guard_triggered(): + """ + - Tests a request where no content mod is triggered + - Assert that the guardrails applied are returned in the response headers + """ + async with aiohttp.ClientSession() as session: + response, headers = await chat_completion( + session, + "sk-1234", + model="fake-openai-endpoint", + messages=[{"role": "user", "content": f"Hello what's the weather"}], + guardrails=[], + ) + await asyncio.sleep(3) + + print("response=", response, "response headers", headers) + + assert "x-litellm-applied-guardrails" not in headers