From 37e6d8a3aa98ac6bd1daf2eb24b59b00f2d61fb3 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 19 Aug 2024 18:53:12 -0700 Subject: [PATCH] test_llm_guard_triggered --- tests/otel_tests/test_guardrails.py | 33 ++++++++++++++++++++++++----- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/tests/otel_tests/test_guardrails.py b/tests/otel_tests/test_guardrails.py index db2027ce1..979f11149 100644 --- a/tests/otel_tests/test_guardrails.py +++ b/tests/otel_tests/test_guardrails.py @@ -6,7 +6,7 @@ from typing import Optional, List, Union import uuid -async def chat_completion(session, key, model: Union[str, List] = "gpt-4"): +async def chat_completion(session, key, messages, model: Union[str, List] = "gpt-4"): url = "http://0.0.0.0:4000/chat/completions" headers = { "Authorization": f"Bearer {key}", @@ -14,9 +14,7 @@ async def chat_completion(session, key, model: Union[str, List] = "gpt-4"): } data = { "model": model, - "messages": [ - {"role": "user", "content": f"Hello! {str(uuid.uuid4())}"}, - ], + "messages": messages, "guardrails": ["aporia-post-guard", "aporia-pre-guard"], } @@ -45,7 +43,10 @@ async def test_no_llm_guard_triggered(): """ async with aiohttp.ClientSession() as session: response, headers = await chat_completion( - session, "sk-1234", model="fake-openai-endpoint" + session, + "sk-1234", + model="fake-openai-endpoint", + messages=[{"role": "user", "content": f"Hello what's the weather"}], ) await asyncio.sleep(3) @@ -57,3 +58,25 @@ async def test_no_llm_guard_triggered(): headers["x-litellm-applied-guardrails"] == "aporia-pre-guard,aporia-post-guard" ) + + +@pytest.mark.asyncio +async def test_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: + try: + response, headers = await chat_completion( + session, + "sk-1234", + model="fake-openai-endpoint", + messages=[ + {"role": "user", "content": f"Hello my name is ishaan@berri.ai"} + ], + ) + pytest.fail("Should have thrown an exception") + except Exception as e: + print(e) + assert "Aporia detected and blocked PII" in str(e)