Merge branch 'main' into guardrails-from-header

tests/integration/anthropic/test_anthropic_header_with_invariant_key.py

@@ -27,12 +27,10 @@ async def test_gateway_with_invariant_key_in_anthropic_key_header(
     """Test the Anthropic gateway with Invariant key in the Anthropic key"""
     anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
     dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
+    invariant_key_suffix = f";invariant-auth={os.getenv('INVARIANT_API_KEY')}"
     with patch.dict(
         os.environ,
-        {
-            "ANTHROPIC_API_KEY": anthropic_api_key
-            + ";invariant-auth=<not needed for test>"
-        },
+        {"ANTHROPIC_API_KEY": anthropic_api_key + invariant_key_suffix},
     ):
         client = anthropic.Anthropic(
             http_client=Client(),

tests/integration/anthropic/test_anthropic_with_tool_call.py

@@ -12,10 +12,11 @@ from typing import Dict, List
 # 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
+
 import anthropic
 import pytest
 import requests
-from httpx import Client
 
 # Pytest plugins
 pytest_plugins = ("pytest_asyncio",)
@@ -26,14 +27,8 @@ class WeatherAgent:
 
     def __init__(self, gateway_url, push_to_explorer):
         self.dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-        invariant_api_key = os.environ.get("INVARIANT_API_KEY", "None")
-        self.client = anthropic.Anthropic(
-            http_client=Client(
-                headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
-            ),
-            base_url=f"{gateway_url}/api/v1/gateway/{self.dataset_name}/anthropic"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/anthropic",
+        self.client = get_anthropic_client(
+            gateway_url, push_to_explorer, self.dataset_name
         )
         self.get_weather_function = {
             "name": "get_weather",

tests/integration/anthropic/test_anthropic_without_tool_call.py

@@ -8,10 +8,10 @@ import uuid
 # Add integration folder (parent) to sys.path
 sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-import anthropic
+from utils import get_anthropic_client
+
 import pytest
 import requests
-from httpx import Client
 
 # Pytest plugins
 pytest_plugins = ("pytest_asyncio",)
@@ -26,15 +26,10 @@ async def test_response_without_tool_call(
 ):
     """Test the Anthropic gateway without tool calling."""
     dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-    invariant_api_key = os.environ.get("INVARIANT_API_KEY", "None")
-
-    client = anthropic.Anthropic(
-        http_client=Client(
-            headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/anthropic",
+    client = get_anthropic_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
     )
 
     cities = ["zurich", "new york", "london"]
@@ -91,16 +86,7 @@ async def test_streaming_response_without_tool_call(
 ):
     """Test the Anthropic gateway without tool calling."""
     dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-    invariant_api_key = os.environ.get("INVARIANT_API_KEY", "None")
-
-    client = anthropic.Anthropic(
-        http_client=Client(
-            headers={"Invariant-Authorization": f"Bearer {invariant_api_key}"},
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/anthropic",
-    )
+    client = get_anthropic_client(gateway_url, push_to_explorer, dataset_name)
 
     cities = ["zurich", "new york", "london"]
     queries = [

tests/integration/docker-compose.test.yml

@@ -60,6 +60,7 @@ services:
   app-api:
     container_name: invariant-gateway-test-explorer-app-api
     image: ghcr.io/invariantlabs-ai/explorer/app-api:latest
+    pull_policy: always
     platform: linux/amd64
     depends_on:
       database:

tests/integration/gemini/test_generate_content_with_tool_calls.py

@@ -8,9 +8,10 @@ import uuid
 # 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
+
 import pytest
 import requests
-from google import genai
 from google.genai import types
 
 # Pytest plugins
@@ -143,18 +144,7 @@ async def test_generate_content_with_tool_call(
     without streaming.
     """
     dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
-
-    client = genai.Client(
-        api_key=os.getenv("GEMINI_API_KEY"),
-        http_options={
-            "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/gemini",
-            "headers": {
-                "invariant-authorization": "Bearer <some-key>"
-            },  # This key is not used for local tests
-        },
-    )
+    client = get_gemini_client(gateway_url, push_to_explorer, dataset_name)
 
     request = {
         "model": "gemini-2.0-flash",

tests/integration/gemini/test_generate_content_without_tool_calls.py

@@ -10,6 +10,8 @@ from unittest.mock import patch
 # 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
+
 import pytest
 import PIL.Image
 import requests
@@ -29,17 +31,8 @@ async def test_generate_content(
 ):
     """Test the generate content gateway calls without tool calling."""
     dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
-    client = genai.Client(
-        api_key=os.getenv("GEMINI_API_KEY"),
-        http_options={
-            "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/gemini",
-            "headers": {
-                "invariant-authorization": "Bearer <some-key>"
-            },  # This key is not used for local tests
-        },
-    )
+    client = get_gemini_client(gateway_url, push_to_explorer, dataset_name)
+
     request = {
         "model": "gemini-2.0-flash",
         "contents": "What is the capital of France?",
@@ -115,18 +108,8 @@ async def test_generate_content_with_image(
 ):
     """Test that generate content gateway calls work with image."""
     dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
+    client = get_gemini_client(gateway_url, push_to_explorer, dataset_name)
 
-    client = genai.Client(
-        api_key=os.getenv("GEMINI_API_KEY"),
-        http_options={
-            "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/gemini",
-            "headers": {
-                "invariant-authorization": "Bearer <some-key>"
-            },  # This key is not used for local tests
-        },
-    )
 
     image_path = Path(__file__).parent.parent / "resources" / "images" / "two-cats.png"
     image = PIL.Image.open(image_path)
@@ -181,9 +164,10 @@ async def test_generate_content_with_invariant_key_in_gemini_key_header(
     """Test the generate content gateway calls with the Invariant API Key in the Gemini Key header."""
     dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
     gemini_api_key = os.getenv("GEMINI_API_KEY")
+    invariant_key_suffix = f";invariant-auth={os.getenv('INVARIANT_API_KEY')}"
     with patch.dict(
         os.environ,
-        {"GEMINI_API_KEY": gemini_api_key + ";invariant-auth=<not needed for test>"},
+        {"GEMINI_API_KEY": gemini_api_key + invariant_key_suffix},
     ):
         client = genai.Client(
             api_key=os.getenv("GEMINI_API_KEY"),
@@ -194,14 +178,14 @@ async def test_generate_content_with_invariant_key_in_gemini_key_header(
 
         chat_response = client.models.generate_content(
             model="gemini-2.0-flash",
-            contents="What is the capital of Spain?",
+            contents="What is the capital of Denmark?",
             config={
                 "maxOutputTokens": 100,
             },
         )
 
         # Verify the chat response
-        assert "MADRID" in chat_response.candidates[0].content.parts[0].text.upper()
+        assert "COPENHAGEN" in chat_response.candidates[0].content.parts[0].text.upper()
         expected_assistant_message = chat_response.candidates[0].content.parts[0].text
 
         # Wait for the trace to be saved
@@ -228,7 +212,7 @@ async def test_generate_content_with_invariant_key_in_gemini_key_header(
         assert trace["messages"] == [
             {
                 "role": "user",
-                "content": [{"text": "What is the capital of Spain?", "type": "text"}],
+                "content": [{"text": "What is the capital of Denmark?", "type": "text"}],
             },
             {
                 "role": "assistant",

tests/integration/guardrails/test_guardrails_anthropic.py

@@ -8,10 +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, 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",)
@@ -32,16 +33,10 @@ async def test_message_content_guardrail_from_file(
         pytest.fail("No INVARIANT_API_KEY set, failing")
 
     dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-
-    client = Anthropic(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/anthropic",
+    client = get_anthropic_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
     )
 
     request = {
@@ -161,16 +156,10 @@ async def test_tool_call_guardrail_from_file(
     }
 
     dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-
-    client = Anthropic(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/anthropic",
+    client = get_anthropic_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
     )
 
     if not do_stream:
@@ -255,16 +244,10 @@ async def test_input_from_guardrail_from_file(
         pytest.fail("No INVARIANT_API_KEY set, failing")
 
     dataset_name = f"test-dataset-anthropic-{uuid.uuid4()}"
-
-    client = Anthropic(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/anthropic",
+    client = get_anthropic_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
     )
 
     request = {
@@ -332,3 +315,276 @@ 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"
+    )
+
+
+@pytest.mark.skipif(
+    not os.getenv("ANTHROPIC_API_KEY"), reason="No ANTHROPIC_API_KEY set"
+)
+@pytest.mark.parametrize(
+    "do_stream, is_block_action",
+    [(True, True), (True, False), (False, True), (False, False)],
+)
+async def test_preguardrailing_with_guardrails_from_explorer(
+    explorer_api_url, gateway_url, do_stream, is_block_action
+):
+    """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 "pun detected in user message" if:\n   (msg: Message)\n   "pun" in msg.content and msg.role == "user"',
+        action="block" if is_block_action else "log",
+        invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"),
+    )
+
+    user_prompt = "Tell me a one sentence pun."
+    request = {
+        "model": "claude-3-5-sonnet-20241022",
+        "messages": [
+            {
+                "role": "user",
+                "content": user_prompt,
+            }
+        ],
+        "max_tokens": 100,
+    }
+    if is_block_action:
+        if do_stream:
+            with pytest.raises(APIStatusError) as exc_info:
+                chat_response = client.messages.create(
+                    **request,
+                    stream=True,
+                )
+                for _ in chat_response:
+                    pass
+
+            assert "[Invariant] The request did not pass the guardrails" in str(
+                exc_info.value
+            )
+        else:
+            with pytest.raises(BadRequestError) as exc_info:
+                chat_response = client.messages.create(
+                    **request,
+                    stream=False,
+                )
+
+            assert exc_info.value.status_code == 400
+            assert "[Invariant] The request did not pass the guardrails" in str(
+                exc_info.value
+            )
+            assert "pun detected in user message" in str(exc_info.value)
+
+    else:
+        if do_stream:
+            _ = client.messages.create(
+                **request,
+                stream=True,
+            )
+        else:
+            _ = client.messages.create(
+                **request,
+                stream=False,
+            )
+
+    # 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) == 1
+    trace_id = traces[0]["id"]
+
+    # Fetch the 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 if not is_block_action else 1
+    assert trace["messages"][0] == {
+        "role": "user",
+        "content": user_prompt,
+    }
+    if not is_block_action:
+        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) == 1
+    assert (
+        annotations[0]["content"] == "pun detected in user message"
+        and annotations[0]["extra_metadata"]["source"] == "guardrails-error"
+        and annotations[0]["extra_metadata"]["guardrail-action"] == "block"
+        if is_block_action
+        else "log"
+    )

tests/integration/guardrails/test_guardrails_gemini.py

@@ -8,9 +8,10 @@ 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, create_dataset, add_guardrail_to_dataset
+
 import pytest
 import requests
-from httpx import Client
 from google import genai
 
 # Pytest plugins
@@ -30,17 +31,10 @@ async def test_message_content_guardrail_from_file(
         pytest.fail("No INVARIANT_API_KEY set, failing")
 
     dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
-
-    client = genai.Client(
-        api_key=os.getenv("GEMINI_API_KEY"),
-        http_options={
-            "headers": {
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-            "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/gemini",
-        },
+    client = get_gemini_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
     )
 
     request = {
@@ -141,17 +135,10 @@ async def test_tool_call_guardrail_from_file(
     )
 
     dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
-
-    client = genai.Client(
-        api_key=os.getenv("GEMINI_API_KEY"),
-        http_options={
-            "headers": {
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-            "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/gemini",
-        },
+    client = get_gemini_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
     )
 
     request = {
@@ -244,17 +231,10 @@ async def test_input_from_guardrail_from_file(
         pytest.fail("No INVARIANT_API_KEY set, failing")
 
     dataset_name = f"test-dataset-gemini-{uuid.uuid4()}"
-
-    client = genai.Client(
-        api_key=os.getenv("GEMINI_API_KEY"),
-        http_options={
-            "headers": {
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-            "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
-            if push_to_explorer
-            else f"{gateway_url}/api/v1/gateway/gemini",
-        },
+    client = get_gemini_client(
+        gateway_url,
+        push_to_explorer,
+        dataset_name,
     )
 
     request = {
@@ -323,6 +303,259 @@ 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"
+    )
+
+
+@pytest.mark.skipif(not os.getenv("GEMINI_API_KEY"), reason="No GEMINI_API_KEY set")
+@pytest.mark.parametrize(
+    "do_stream, is_block_action",
+    [(True, True), (True, False), (False, True), (False, False)],
+)
+async def test_preguardrailing_with_guardrails_from_explorer(
+    explorer_api_url, gateway_url, do_stream, is_block_action
+):
+    """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 "pun detected in user message" if:\n   (msg: Message)\n   "pun" in msg.content and msg.role == "user"',
+        action="block" if is_block_action else "log",
+        invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"),
+    )
+
+    user_prompt = "Tell me a one sentence pun."
+    request = {
+        "model": "gemini-2.0-flash",
+        "contents": user_prompt,
+        "config": {
+            "maxOutputTokens": 100,
+        },
+    }
+    if is_block_action:
+        if do_stream:
+            chat_response = client.models.generate_content_stream(**request)
+
+            assert_is_streamed_refusal(
+                chat_response,
+                [
+                    "[Invariant] The request did not pass the guardrails",
+                    "pun detected in user message",
+                ],
+            )
+        else:
+            with pytest.raises(genai.errors.ClientError) as exc_info:
+                chat_response = client.models.generate_content(**request)
+            assert "[Invariant] The request did not pass the guardrails" in str(
+                exc_info.value
+            )
+            assert "pun detected in user message" in str(exc_info.value)
+    else:
+        if do_stream:
+            response = client.models.generate_content_stream(**request)
+            for _ in response:
+                pass
+        else:
+            _ = client.models.generate_content(**request)
+
+    # 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) == 1
+    trace_id = traces[0]["id"]
+
+    # Fetch the 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 if not is_block_action else 1
+    assert trace["messages"][0] == {
+        "role": "user",
+        "content": [
+            {
+                "type": "text",
+                "text": user_prompt,
+            }
+        ],
+    }
+    if not is_block_action:
+        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) == 1
+    assert (
+        annotations[0]["content"] == "pun detected in user message"
+        and annotations[0]["extra_metadata"]["source"] == "guardrails-error"
+        and annotations[0]["extra_metadata"]["guardrail-action"] == "block"
+        if is_block_action
+        else "log"
+    )
+
+
 def is_refusal(chunk):
     return (
         len(chunk.candidates) == 1

tests/integration/guardrails/test_guardrails_open_ai.py

@@ -8,10 +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, 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",)
@@ -30,17 +31,7 @@ async def test_message_content_guardrail_from_file(
         pytest.fail("No INVARIANT_API_KEY set, failing")
 
     dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)
 
     request = {
         "model": "gpt-4o",
@@ -161,17 +152,7 @@ async def test_tool_call_guardrail_from_file(
     }
 
     dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)
 
     if not do_stream:
         with pytest.raises(BadRequestError) as exc_info:
@@ -259,17 +240,7 @@ async def test_input_from_guardrail_from_file(
         pytest.fail("No INVARIANT_API_KEY set, failing")
 
     dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
-            },
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)
 
     request = {
         "model": "gpt-4o",
@@ -349,3 +320,268 @@ 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
+        )
+
+    # 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"
+    )
+
+
+@pytest.mark.skipif(not os.getenv("OPENAI_API_KEY"), reason="No OPENAI_API_KEY set")
+@pytest.mark.parametrize(
+    "do_stream, is_block_action",
+    [(True, True), (True, False), (False, True), (False, False)],
+)
+async def test_preguardrailing_with_guardrails_from_explorer(
+    explorer_api_url, gateway_url, do_stream, is_block_action
+):
+    """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 "pun detected in user message" if:\n   (msg: Message)\n   "pun" in msg.content and msg.role == "user"',
+        action="block" if is_block_action else "log",
+        invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"),
+    )
+
+    user_prompt = "Tell me a one sentence pun."
+    request = {
+        "model": "gpt-4o",
+        "messages": [
+            {
+                "role": "user",
+                "content": user_prompt,
+            }
+        ],
+        "max_tokens": 100,
+    }
+    if is_block_action:
+        if do_stream:
+            with pytest.raises(APIError) as exc_info:
+                chat_response = client.chat.completions.create(
+                    **request,
+                    stream=True,
+                )
+                for _ in chat_response:
+                    pass
+
+            assert "[Invariant] The request did not pass the guardrails" in str(
+                exc_info.value
+            )
+        else:
+            with pytest.raises(BadRequestError) as exc_info:
+                chat_response = client.chat.completions.create(
+                    **request,
+                    stream=False,
+                )
+
+            assert exc_info.value.status_code == 400
+            assert "[Invariant] The request did not pass the guardrails" in str(
+                exc_info.value
+            )
+            assert "pun detected in user message" in str(exc_info.value)
+    else:
+        if do_stream:
+            _ = client.chat.completions.create(
+                **request,
+                stream=True,
+            )
+        else:
+            _ = client.chat.completions.create(
+                **request,
+                stream=False,
+            )
+
+    # 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) == 1
+    trace_id = traces[0]["id"]
+
+    # Fetch the trace
+    trace_response = requests.get(
+        f"{explorer_api_url}/api/v1/trace/{trace_id}",
+        timeout=5,
+    )
+    trace = trace_response.json()
+
+    assert len(trace["messages"]) == 1 if is_block_action else 2
+    assert trace["messages"][0] == {
+        "role": "user",
+        "content": user_prompt,
+    }
+    if not is_block_action:
+        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) == 1
+    assert (
+        annotations[0]["content"] == "pun detected in user message"
+        and annotations[0]["extra_metadata"]["source"] == "guardrails-error"
+        and annotations[0]["extra_metadata"]["guardrail-action"] == "block"
+        if is_block_action
+        else "log"
+    )

tests/integration/guardrails/test_header_guardrails.py

@@ -1,4 +1,4 @@
-"""Test the guardrails from file with the OpenAI route."""
+"""Test the guardrails from header with the OpenAI route."""
 
 import os
 import sys
@@ -136,9 +136,7 @@ raise "Users must not mention the magic phrase 'Abracadabra'" if:
     "do_stream, push_to_explorer",
     [(True, True), (True, False), (False, True), (False, False)],
 )
-async def test_invalid_guardrail_in_header(
-    explorer_api_url, gateway_url, do_stream, push_to_explorer
-):
+async def test_invalid_guardrail_in_header(gateway_url, do_stream, push_to_explorer):
     """Test the message content guardrail."""
     if not os.getenv("INVARIANT_API_KEY"):
         pytest.fail("No INVARIANT_API_KEY set, failing")
@@ -178,7 +176,8 @@ raise "Users must not mention the magic phrase 'Abracadabra'" if:
                 stream=False,
             )
 
-        assert "Gateway: Guardrails check failed" in str(
+        print(exc_info.value.message, flush=True)
+        assert "Failed to create policy from policy source." in str(
             exc_info.value
         ), "guardrails check fails because of an invalid guardrailing rule"
         assert "illegal statement" in str(

tests/integration/open_ai/test_chat_with_tool_call.py

@@ -9,10 +9,10 @@ import uuid
 # 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
+
 import pytest
 import requests
-from httpx import Client
-from openai import OpenAI
 
 # Pytest plugins
 pytest_plugins = ("pytest_asyncio",)
@@ -28,17 +28,7 @@ async def test_chat_completion_with_tool_call_without_streaming(
     without streaming.
     """
     dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": "Bearer <some-key>"
-            },  # This key is not used for local tests
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)
 
     chat_response = client.chat.completions.create(
         model="gpt-4o",
@@ -146,17 +136,7 @@ async def test_chat_completion_with_tool_call_with_streaming(
     while streaming.
     """
     dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": "Bearer <some-key>"
-            },  # This key is not used for local tests
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)
 
     chat_response = client.chat.completions.create(
         model="gpt-4o",

tests/integration/open_ai/test_chat_without_tool_calls.py

@@ -11,6 +11,8 @@ from unittest.mock import patch
 # 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
+
 import pytest
 import requests
 from httpx import Client
@@ -30,17 +32,7 @@ async def test_chat_completion(
 ):
     """Test the chat completions gateway calls without tool calling."""
     dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
-
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": "Bearer <some-key>"
-            },  # This key is not used for local tests
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)
 
     chat_response = client.chat.completions.create(
         model="gpt-4o",
@@ -103,17 +95,8 @@ async def test_chat_completion_with_image(
 ):
     """Test the chat completions gateway works with image."""
     dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
+    client = get_open_ai_client(gateway_url, push_to_explorer, dataset_name)
 
-    client = OpenAI(
-        http_client=Client(
-            headers={
-                "Invariant-Authorization": "Bearer <some-key>"
-            },  # This key is not used for local tests
-        ),
-        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
-        if push_to_explorer
-        else f"{gateway_url}/api/v1/gateway/openai",
-    )
     image_path = Path(__file__).parent.parent / "resources" / "images" / "two-cats.png"
     with image_path.open("rb") as image_file:
         base64_image = base64.b64encode(image_file.read()).decode("utf-8")
@@ -189,9 +172,10 @@ async def test_chat_completion_with_invariant_key_in_openai_key_header(
     """Test the chat completions gateway calls with the Invariant API Key in the OpenAI Key header."""
     dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}"
     openai_api_key = os.getenv("OPENAI_API_KEY")
+    invariant_key_suffix = f";invariant-auth={os.getenv('INVARIANT_API_KEY')}"
     with patch.dict(
         os.environ,
-        {"OPENAI_API_KEY": openai_api_key + ";invariant-auth=<not needed for test>"},
+        {"OPENAI_API_KEY": openai_api_key + invariant_key_suffix},
     ):
         client = OpenAI(
             http_client=Client(),
@@ -252,7 +236,7 @@ async def test_chat_completion_with_openai_exception(gateway_url, do_stream):
     client = OpenAI(
         http_client=Client(
             headers={
-                "Invariant-Authorization": "Bearer <some-key>"
+                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
             },  # This key is not used for local tests
         ),
         base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai",

tests/integration/resources/guardrails/integration_test_guardrails_via_file.py

@@ -1,15 +1,15 @@
-# For Anthropic and OpenAI tests:
+# For Anthropic and OpenAI tests
 raise "Madrid detected in the response" if:
     (msg: Message)
    "Madrid" in msg.content
 
-# For Gemini tests:
+# For Gemini tests
 raise "Dublin detected in the response" if:
     (msg: Message)
     (chunk: str) in text(msg.content)
     "Dublin" in chunk
 
-# For all provider tests:
+# For all provider tests
 raise "get_capital is called with Germany as argument" if:
     (call: ToolCall)
     call is tool:get_capital

tests/integration/utils.py

@@ -0,0 +1,105 @@
+"""Common utilities for integration tests."""
+
+import os
+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
+
+
+def get_open_ai_client(
+    gateway_url: str, push_to_explorer: bool, dataset_name: str
+) -> OpenAI:
+    """Create an OpenAI client for integration tests."""
+    return OpenAI(
+        http_client=Client(
+            headers={
+                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
+            },
+        ),
+        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/openai"
+        if push_to_explorer
+        else f"{gateway_url}/api/v1/gateway/openai",
+    )
+
+
+def get_anthropic_client(
+    gateway_url: str, push_to_explorer: bool, dataset_name: str
+) -> Anthropic:
+    """Create an Anthropic client for integration tests."""
+    return Anthropic(
+        http_client=Client(
+            headers={
+                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
+            },
+        ),
+        base_url=f"{gateway_url}/api/v1/gateway/{dataset_name}/anthropic"
+        if push_to_explorer
+        else f"{gateway_url}/api/v1/gateway/anthropic",
+    )
+
+
+def get_gemini_client(
+    gateway_url: str, push_to_explorer: bool, dataset_name: str
+) -> genai.Client:
+    """Create a Gemini client for integration tests."""
+    return genai.Client(
+        api_key=os.getenv("GEMINI_API_KEY"),
+        http_options={
+            "base_url": f"{gateway_url}/api/v1/gateway/{dataset_name}/gemini"
+            if push_to_explorer
+            else f"{gateway_url}/api/v1/gateway/gemini",
+            "headers": {
+                "Invariant-Authorization": f"Bearer {os.getenv('INVARIANT_API_KEY')}"
+            },
+        },
+    )
+
+
+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()