diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py index 09b32b1..22a104d 100644 --- a/gateway/integrations/explorer.py +++ b/gateway/integrations/explorer.py @@ -14,7 +14,7 @@ DEFAULT_API_URL = "https://explorer.invariantlabs.ai" def create_annotations_from_guardrails_errors( - guardrails_errors: List[dict], + guardrails_errors: List[dict], action: str = "block" ) -> List[AnnotationCreate]: """Create Explorer annotations from the guardrails errors.""" annotations = [] @@ -48,7 +48,10 @@ def create_annotations_from_guardrails_errors( AnnotationCreate( content=content, address=r, - extra_metadata={"source": "guardrails-error"}, + extra_metadata={ + "source": "guardrails-error", + "guardrail-action": action, + }, ) ) return annotations diff --git a/gateway/routes/anthropic.py b/gateway/routes/anthropic.py index 09ce85b..24cf097 100644 --- a/gateway/routes/anthropic.py +++ b/gateway/routes/anthropic.py @@ -167,7 +167,7 @@ async def push_to_explorer( """Pushes the full trace to the Invariant Explorer""" guardrails_execution_result = guardrails_execution_result or {} annotations = create_annotations_from_guardrails_errors( - guardrails_execution_result.get("errors", []) + guardrails_execution_result.get("errors", []), action="block" ) # Execute the logging guardrails before pushing to Explorer @@ -177,7 +177,7 @@ async def push_to_explorer( response_json=merged_response, ) logging_annotations = create_annotations_from_guardrails_errors( - logging_guardrails_execution_result.get("errors", []) + logging_guardrails_execution_result.get("errors", []), action="log" ) # Update the annotations with the logging guardrails annotations.extend(logging_annotations) diff --git a/gateway/routes/gemini.py b/gateway/routes/gemini.py index b390461..6d4a409 100644 --- a/gateway/routes/gemini.py +++ b/gateway/routes/gemini.py @@ -404,7 +404,7 @@ async def push_to_explorer( """Pushes the full trace to the Invariant Explorer""" guardrails_execution_result = guardrails_execution_result or {} annotations = create_annotations_from_guardrails_errors( - guardrails_execution_result.get("errors", []) + guardrails_execution_result.get("errors", []), action="block" ) # Execute the logging guardrails before pushing to Explorer @@ -414,7 +414,7 @@ async def push_to_explorer( response_json=response_json, ) logging_annotations = create_annotations_from_guardrails_errors( - logging_guardrails_execution_result.get("errors", []) + logging_guardrails_execution_result.get("errors", []), action="log" ) # Update the annotations with the logging guardrails annotations.extend(logging_annotations) diff --git a/gateway/routes/open_ai.py b/gateway/routes/open_ai.py index f4a20f4..f929a2c 100644 --- a/gateway/routes/open_ai.py +++ b/gateway/routes/open_ai.py @@ -221,7 +221,7 @@ class InstrumentedOpenAIStreamResponse(InstrumentedStreamingResponse): # push will happen in on_end async def on_end(self): - """Sends full merged response to the exploree.""" + """Sends full merged response to the explorer.""" # don't block on the response from explorer (.create_task) if self.context.dataset_name: asyncio.create_task( @@ -437,7 +437,9 @@ async def push_to_explorer( and merged_response["choices"][0].get("finish_reason") not in FINISH_REASON_TO_PUSH_TRACE ): - annotations = create_annotations_from_guardrails_errors(guardrails_errors) + annotations = create_annotations_from_guardrails_errors( + guardrails_errors, action="block" + ) # Execute the logging guardrails before pushing to Explorer logging_guardrails_execution_result = await get_guardrails_check_result( @@ -446,7 +448,7 @@ async def push_to_explorer( response_json=merged_response, ) logging_annotations = create_annotations_from_guardrails_errors( - logging_guardrails_execution_result.get("errors", []) + logging_guardrails_execution_result.get("errors", []), action="log" ) # Update the annotations with the logging guardrails annotations.extend(logging_annotations) diff --git a/tests/integration/guardrails/test_guardrails_anthropic.py b/tests/integration/guardrails/test_guardrails_anthropic.py index aaa349a..f61d9e6 100644 --- a/tests/integration/guardrails/test_guardrails_anthropic.py +++ b/tests/integration/guardrails/test_guardrails_anthropic.py @@ -8,12 +8,11 @@ import time # Add integration folder (parent) to sys.path sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from utils import get_anthropic_client +from utils import get_anthropic_client, create_dataset, add_guardrail_to_dataset import pytest import requests -from httpx import Client -from anthropic import Anthropic, APIStatusError, BadRequestError +from anthropic import APIStatusError, BadRequestError # Pytest plugins pytest_plugins = ("pytest_asyncio",) @@ -316,3 +315,149 @@ async def test_input_from_guardrail_from_file( == "Users must not mention the magic phrase 'Fight Club'" and annotations[0]["extra_metadata"]["source"] == "guardrails-error" ) + +@pytest.mark.skipif(not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set") +@pytest.mark.parametrize("do_stream", [True, False]) +async def test_with_guardrails_from_explorer(explorer_api_url, gateway_url, do_stream): + """Test that the guardrails from the explorer work.""" + dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}" + client = get_anthropic_client( + gateway_url, push_to_explorer=True, dataset_name=dataset_name + ) + + dataset_creation_response = await create_dataset( + explorer_api_url, + invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), + dataset_name=dataset_name, + ) + dataset_id = dataset_creation_response["id"] + _ = await add_guardrail_to_dataset( + explorer_api_url, + dataset_id=dataset_id, + policy='raise "ogre detected in response" if:\n (msg: Message)\n "ogre" in msg.content and msg.role == "assistant"', + action="block", + invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), + ) + _ = await add_guardrail_to_dataset( + explorer_api_url, + dataset_id=dataset_id, + policy='raise "Fiona detected in response" if:\n (msg: Message)\n "Fiona" in msg.content', + action="log", + invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), + ) + + # Ask about the capital of Spain + # This should not be blocked by the guardrails from the explorer when we push to explorer + # because the file based guardrails are overridden by the explorer guardrails + spain_request = { + "model": "claude-3-5-sonnet-20241022", + "messages": [{"role": "user", "content": "What is the capital of Spain?"}], + "max_tokens": 100, + } + if not do_stream: + chat_response = client.messages.create( + **spain_request, + stream=False, + ) + + assert "Madrid" in chat_response.content[0].text + else: + chat_response = client.messages.create( + **spain_request, + stream=True, + ) + + merged_content = "" + for chunk in chat_response: + if chunk.type == "content_block_delta": + merged_content += chunk.delta.text + assert "Madrid" in merged_content + + # Ask about Shrek + # This should be blocked by the guardrails from the explorer + user_prompt = "What kind of a creature is Shrek? What is his Shrek's wife's name? Only answer these questions with single sentences, don't add any extra details." + shrek_request = { + "model": "claude-3-5-sonnet-20241022", + "messages": [ + { + "role": "user", + "content": user_prompt, + } + ], + "max_tokens": 100, + } + if not do_stream: + with pytest.raises(BadRequestError) as exc_info: + chat_response = client.messages.create( + **shrek_request, + stream=False, + ) + + assert exc_info.value.status_code == 400 + assert "[Invariant] The response did not pass the guardrails" in str( + exc_info.value + ) + # Only the block guardrail should be triggered here + assert "ogre detected in response" in str(exc_info.value) + assert "Fiona detected in response" not in str(exc_info.value) + else: + with pytest.raises(APIStatusError) as exc_info: + chat_response = client.messages.create( + **shrek_request, + stream=True, + ) + + for _ in chat_response: + pass + assert "[Invariant] The response did not pass the guardrails" in str( + exc_info.value + ) + # Only the block guardrail should be triggered here + assert "ogre detected in response" in str(exc_info.value) + assert "Fiona detected in response" not in str(exc_info.value) + + # Wait for the trace to be saved + # This is needed because the trace is saved asynchronously + time.sleep(2) + + # Fetch the trace ids for the dataset + traces_response = requests.get( + f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces", + timeout=5, + ) + traces = traces_response.json() + assert len(traces) == 2 + trace_id = traces[1]["id"] + + # Fetch the second trace + trace_response = requests.get( + f"{explorer_api_url}/api/v1/trace/{trace_id}", + timeout=5, + ) + trace = trace_response.json() + + assert len(trace["messages"]) == 2 + assert trace["messages"][0] == { + "role": "user", + "content": user_prompt, + } + assert trace["messages"][1].get("role") == "assistant" + + # Fetch annotations + annotations_response = requests.get( + f"{explorer_api_url}/api/v1/trace/{trace_id}/annotations", + timeout=5, + ) + annotations = annotations_response.json() + + assert len(annotations) == 2 + assert ( + annotations[0]["content"] == "ogre detected in response" + and annotations[0]["extra_metadata"]["source"] == "guardrails-error" + and annotations[0]["extra_metadata"]["guardrail-action"] == "block" + ) + assert ( + annotations[1]["content"] == "Fiona detected in response" + and annotations[1]["extra_metadata"]["source"] == "guardrails-error" + and annotations[1]["extra_metadata"]["guardrail-action"] == "log" + ) diff --git a/tests/integration/guardrails/test_guardrails_gemini.py b/tests/integration/guardrails/test_guardrails_gemini.py index e452284..6fc0945 100644 --- a/tests/integration/guardrails/test_guardrails_gemini.py +++ b/tests/integration/guardrails/test_guardrails_gemini.py @@ -8,7 +8,7 @@ import time # Add integration folder (parent) to sys.path sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from utils import get_gemini_client +from utils import get_gemini_client, create_dataset, add_guardrail_to_dataset import pytest import requests @@ -303,6 +303,147 @@ async def test_input_from_guardrail_from_file( ) +@pytest.mark.skipif(not os.getenv("GEMINI_API_KEY"), reason="No GEMINI_API_KEY set") +@pytest.mark.parametrize("do_stream", [True, False]) +async def test_with_guardrails_from_explorer(explorer_api_url, gateway_url, do_stream): + """Test that the guardrails from the explorer work.""" + dataset_name = f"test-dataset-gemini-{uuid.uuid4()}" + client = get_gemini_client( + gateway_url, push_to_explorer=True, dataset_name=dataset_name + ) + + dataset_creation_response = await create_dataset( + explorer_api_url, + invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), + dataset_name=dataset_name, + ) + dataset_id = dataset_creation_response["id"] + _ = await add_guardrail_to_dataset( + explorer_api_url, + dataset_id=dataset_id, + policy='raise "ogre detected in response" if:\n (msg: Message)\n "ogre" in msg.content and msg.role == "assistant"', + action="block", + invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), + ) + _ = await add_guardrail_to_dataset( + explorer_api_url, + dataset_id=dataset_id, + policy='raise "Fiona detected in response" if:\n (msg: Message)\n "Fiona" in msg.content', + action="log", + invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), + ) + + # Ask about the capital of Spain + # This should not be blocked by the guardrails from the explorer when we push to explorer + # because the file based guardrails are overridden by the explorer guardrails + spain_request = { + "model": "gemini-2.0-flash", + "contents": "What is the capital of Spain?", + "config": { + "maxOutputTokens": 100, + }, + } + if not do_stream: + chat_response = client.models.generate_content(**spain_request) + + assert "Madrid" in chat_response.candidates[0].content.parts[0].text + else: + chat_response = client.models.generate_content_stream(**spain_request) + + merged_content = "" + for chunk in chat_response: + if ( + chunk.candidates + and chunk.candidates[0].content + and chunk.candidates[0].content.parts + ): + for text_part in chunk.candidates[0].content.parts: + merged_content += text_part.text + assert "Madrid" in merged_content + + # Ask about Shrek + # This should be blocked by the guardrails from the explorer + user_prompt = "What kind of a creature is Shrek? What is his Shrek's wife's name? Only answer these questions with single sentences, don't add any extra details." + shrek_request = { + "model": "gemini-2.0-flash", + "contents": user_prompt, + "config": { + "maxOutputTokens": 100, + }, + } + if not do_stream: + with pytest.raises(genai.errors.ClientError) as exc_info: + client.models.generate_content(**shrek_request) + + assert "[Invariant] The response did not pass the guardrails" in str( + exc_info.value + ) + # Only the block guardrail should be triggered here + assert "ogre detected in response" in str(exc_info.value) + assert "Fiona detected in response" not in str(exc_info.value) + else: + response = client.models.generate_content_stream(**shrek_request) + + assert_is_streamed_refusal( + response, + [ + "[Invariant] The response did not pass the guardrails", + "ogre detected in response", + ], + ) + + # Wait for the trace to be saved + # This is needed because the trace is saved asynchronously + time.sleep(2) + + # Fetch the trace ids for the dataset + traces_response = requests.get( + f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces", + timeout=5, + ) + traces = traces_response.json() + assert len(traces) == 2 + trace_id = traces[1]["id"] + + # Fetch the second trace + trace_response = requests.get( + f"{explorer_api_url}/api/v1/trace/{trace_id}", + timeout=5, + ) + trace = trace_response.json() + + assert len(trace["messages"]) == 2 + assert trace["messages"][0] == { + "role": "user", + "content": [ + { + "type": "text", + "text": user_prompt, + } + ], + } + assert trace["messages"][1].get("role") == "assistant" + + # Fetch annotations + annotations_response = requests.get( + f"{explorer_api_url}/api/v1/trace/{trace_id}/annotations", + timeout=5, + ) + annotations = annotations_response.json() + + assert len(annotations) == 2 + assert ( + annotations[0]["content"] == "ogre detected in response" + and annotations[0]["extra_metadata"]["source"] == "guardrails-error" + and annotations[0]["extra_metadata"]["guardrail-action"] == "block" + ) + assert ( + annotations[1]["content"] == "Fiona detected in response" + and annotations[1]["extra_metadata"]["source"] == "guardrails-error" + and annotations[1]["extra_metadata"]["guardrail-action"] == "log" + ) + + def is_refusal(chunk): return ( len(chunk.candidates) == 1 diff --git a/tests/integration/guardrails/test_guardrails_open_ai.py b/tests/integration/guardrails/test_guardrails_open_ai.py index c15989a..b0c6b24 100644 --- a/tests/integration/guardrails/test_guardrails_open_ai.py +++ b/tests/integration/guardrails/test_guardrails_open_ai.py @@ -8,12 +8,11 @@ import time # Add integration folder (parent) to sys.path sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from utils import get_open_ai_client +from utils import get_open_ai_client, create_dataset, add_guardrail_to_dataset import pytest import requests -from httpx import Client -from openai import OpenAI, BadRequestError, APIError +from openai import BadRequestError, APIError # Pytest plugins pytest_plugins = ("pytest_asyncio",) @@ -321,3 +320,150 @@ async def test_input_from_guardrail_from_file( == "Users must not mention the magic phrase 'Fight Club'" and annotations[0]["extra_metadata"]["source"] == "guardrails-error" ) + + +@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set") +@pytest.mark.parametrize("do_stream", [True, False]) +async def test_with_guardrails_from_explorer(explorer_api_url, gateway_url, do_stream): + """Test that the guardrails from the explorer work.""" + dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}" + client = get_open_ai_client( + gateway_url, push_to_explorer=True, dataset_name=dataset_name + ) + + dataset_creation_response = await create_dataset( + explorer_api_url, + invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), + dataset_name=dataset_name, + ) + dataset_id = dataset_creation_response["id"] + _ = await add_guardrail_to_dataset( + explorer_api_url, + dataset_id=dataset_id, + policy='raise "ogre detected in response" if:\n (msg: Message)\n "ogre" in msg.content and msg.role == "assistant"', + action="block", + invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), + ) + _ = await add_guardrail_to_dataset( + explorer_api_url, + dataset_id=dataset_id, + policy='raise "Fiona detected in response" if:\n (msg: Message)\n "Fiona" in msg.content', + action="log", + invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), + ) + + # Ask about the capital of Spain + # This should not be blocked by the guardrails from the explorer when we push to explorer + # because the file based guardrails are overridden by the explorer guardrails + spain_request = { + "model": "gpt-4o", + "messages": [{"role": "user", "content": "What is the capital of Spain?"}], + "max_tokens": 100, + } + if not do_stream: + chat_response = client.chat.completions.create( + **spain_request, + stream=False, + ) + + assert "Madrid" in chat_response.choices[0].message.content + else: + chat_response = client.chat.completions.create( + **spain_request, + stream=True, + ) + + merged_content = "" + for chunk in chat_response: + if chunk.choices[0].delta.content: + merged_content += chunk.choices[0].delta.content + assert "Madrid" in merged_content + + # Ask about Shrek + # This should be blocked by the guardrails from the explorer + user_prompt = "What kind of a creature is Shrek? What is his Shrek's wife's name? Only answer these questions with single sentences, don't add any extra details." + shrek_request = { + "model": "gpt-4o", + "messages": [ + { + "role": "user", + "content": user_prompt, + } + ], + "max_tokens": 100, + } + if not do_stream: + with pytest.raises(BadRequestError) as exc_info: + chat_response = client.chat.completions.create( + **shrek_request, + stream=False, + ) + + assert exc_info.value.status_code == 400 + assert "[Invariant] The response did not pass the guardrails" in str( + exc_info.value + ) + # Only the block guardrail should be triggered here + assert "ogre detected in response" in str(exc_info.value) + assert "Fiona detected in response" not in str(exc_info.value) + else: + with pytest.raises(APIError) as exc_info: + chat_response = client.chat.completions.create( + **shrek_request, + stream=True, + ) + + for _ in chat_response: + pass + assert "[Invariant] The response did not pass the guardrails" in str( + exc_info.value + ) + # Only the block guardrail should be triggered here + assert "ogre detected in response" in str(exc_info.value) + assert "Fiona detected in response" not in str(exc_info.value) + + # Wait for the trace to be saved + # This is needed because the trace is saved asynchronously + time.sleep(2) + + # Fetch the trace ids for the dataset + traces_response = requests.get( + f"{explorer_api_url}/api/v1/dataset/byuser/developer/{dataset_name}/traces", + timeout=5, + ) + traces = traces_response.json() + assert len(traces) == 2 + trace_id = traces[1]["id"] + + # Fetch the second trace + trace_response = requests.get( + f"{explorer_api_url}/api/v1/trace/{trace_id}", + timeout=5, + ) + trace = trace_response.json() + + assert len(trace["messages"]) == 2 + assert trace["messages"][0] == { + "role": "user", + "content": user_prompt, + } + assert trace["messages"][1].get("role") == "assistant" + + # Fetch annotations + annotations_response = requests.get( + f"{explorer_api_url}/api/v1/trace/{trace_id}/annotations", + timeout=5, + ) + annotations = annotations_response.json() + + assert len(annotations) == 2 + assert ( + annotations[0]["content"] == "ogre detected in response" + and annotations[0]["extra_metadata"]["source"] == "guardrails-error" + and annotations[0]["extra_metadata"]["guardrail-action"] == "block" + ) + assert ( + annotations[1]["content"] == "Fiona detected in response" + and annotations[1]["extra_metadata"]["source"] == "guardrails-error" + and annotations[1]["extra_metadata"]["guardrail-action"] == "log" + ) diff --git a/tests/integration/utils.py b/tests/integration/utils.py index 39c82b6..6df9ce9 100644 --- a/tests/integration/utils.py +++ b/tests/integration/utils.py @@ -1,7 +1,10 @@ """Common utilities for integration tests.""" import os -from httpx import Client +import uuid +from typing import Any, Dict, Literal, Optional + +from httpx import AsyncClient, Client from openai import OpenAI from google import genai from anthropic import Anthropic @@ -54,3 +57,49 @@ def get_gemini_client( }, }, ) + + +async def create_dataset( + explorer_api_url: str, + invariant_authorization: str, + dataset_name: Optional[str] = None, +) -> Dict[str, Any]: + """Create a dataset in the Explorer API.""" + client = Client(base_url=explorer_api_url) + response = client.post( + "/api/v1/dataset/create", + json={"name": dataset_name if dataset_name else f"test-dataset-{uuid.uuid4()}"}, + headers={"Authorization": invariant_authorization}, + timeout=5, + ) + if response.status_code != 200: + raise ValueError( + f"Failed to create dataset: {response.status_code}, {response.text}" + ) + return response.json() + + +async def add_guardrail_to_dataset( + explorer_api_url: str, + dataset_id: str, + policy: str, + action: Literal["block", "log"], + invariant_authorization: str, +) -> Dict[str, Any]: + """Add a guardrail to a dataset.""" + client = Client(base_url=explorer_api_url) + response = client.post( + f"/api/v1/dataset/{dataset_id}/policy", + json={ + "action": action, + "policy": policy, + "name": f"test-guardrail-{uuid.uuid4()}", + }, + headers={"Authorization": invariant_authorization}, + timeout=5, + ) + if response.status_code != 200: + raise ValueError( + f"Failed to add guardrail: {response.status_code}, {response.text}" + ) + return response.json()