Add tests for guardrails integration with explorer.

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

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)

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)

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)

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"
+    )

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

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"
+    )

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()