Pipelined Guardrails (#32)

* initial draft: pipelined guardrails

* documentation on stream instrumentation

* more comments

* fix: return earlier

* non-streaming case

* handle non-streaming case

* fix more cases

* simplify request instrumentation

* improve comments

* fix import issues

* extend tests for input guardrailing

* anthropic integration of pipelined and pre-guardrailing

* fix gemini streamed refusal

tests/integration/guardrails/test_guardrails_anthropic.py

@@ -238,3 +238,97 @@ async def test_tool_call_guardrail_from_file(
             == "get_capital is called with Germany as argument"
             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, push_to_explorer",
+    [(True, True), (True, False), (False, True), (False, False)],
+)
+async def test_input_from_guardrail_from_file(
+    explorer_api_url, gateway_url, do_stream, push_to_explorer
+):
+    """Test input guardrail enforcement with Anthropic."""
+    if not os.getenv("INVARIANT_API_KEY"):
+        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",
+    )
+
+    request = {
+        "model": "claude-3-5-sonnet-20241022",
+        "max_tokens": 100,
+        "messages": [{"role": "user", "content": "Tell me more about Fight Club."}],
+    }
+
+    if not do_stream:
+        with pytest.raises(BadRequestError) as exc_info:
+            _ = 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 "Users must not mention the magic phrase 'Fight Club'" in str(
+            exc_info.value
+        )
+
+    else:
+        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 exc_info.value.message
+        )
+        assert "Users must not mention the magic phrase 'Fight Club'" in str(
+            exc_info.value.body
+        )
+
+    if push_to_explorer:
+        time.sleep(2)
+        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"]
+
+        trace_response = requests.get(
+            f"{explorer_api_url}/api/v1/trace/{trace_id}",
+            timeout=5,
+        )
+        # in case of input guardrailing, the pushed trace will not contain a response
+        trace = trace_response.json()
+        assert len(trace["messages"]) == 1, "Only the user message should be present"
+        assert trace["messages"][0] == {
+            "role": "user",
+            "content": "Tell me more about Fight Club.",
+        }
+
+        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"]
+            == "Users must not mention the magic phrase 'Fight Club'"
+            and annotations[0]["extra_metadata"]["source"] == "guardrails-error"
+        )

tests/integration/guardrails/test_guardrails_gemini.py

@@ -63,8 +63,13 @@ async def test_message_content_guardrail_from_file(
 
     else:
         response = client.models.generate_content_stream(**request)
-        for chunk in response:
-            assert "Dublin" not in str(chunk)
+        assert_is_streamed_refusal(
+            response,
+            [
+                "[Invariant] The response did not pass the guardrails",
+                "Dublin detected in the response",
+            ],
+        )
 
     if push_to_explorer:
         # Wait for the trace to be saved
@@ -172,8 +177,13 @@ async def test_tool_call_guardrail_from_file(
             **request,
         )
 
-        for chunk in response:
-            assert "Madrid" not in str(chunk)
+        assert_is_streamed_refusal(
+            response,
+            [
+                "[Invariant] The response did not pass the guardrails",
+                "get_capital is called with Germany as argument",
+            ],
+        )
 
     if push_to_explorer:
         # Wait for the trace to be saved
@@ -219,3 +229,122 @@ async def test_tool_call_guardrail_from_file(
             == "get_capital is called with Germany as argument"
             and annotations[0]["extra_metadata"]["source"] == "guardrails-error"
         )
+
+
+@pytest.mark.skipif(not os.getenv("GEMINI_API_KEY"), reason="No GEMINI_API_KEY set")
+@pytest.mark.parametrize(
+    "do_stream, push_to_explorer",
+    [(True, True), (True, False), (False, True), (False, False)],
+)
+async def test_input_from_guardrail_from_file(
+    explorer_api_url, gateway_url, do_stream, push_to_explorer
+):
+    """Test input guardrail enforcement with Gemini."""
+    if not os.getenv("INVARIANT_API_KEY"):
+        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",
+        },
+    )
+
+    request = {
+        "model": "gemini-2.0-flash",
+        "contents": "Tell me more about Fight Club.",
+        "config": {
+            "maxOutputTokens": 200,
+        },
+    }
+
+    if not do_stream:
+        with pytest.raises(genai.errors.ClientError) as exc_info:
+            client.models.generate_content(**request)
+
+        assert "[Invariant] The request did not pass the guardrails" in str(
+            exc_info.value
+        )
+        assert "Users must not mention the magic phrase 'Fight Club'" in str(
+            exc_info.value
+        )
+
+    else:
+        response = client.models.generate_content_stream(**request)
+
+        assert_is_streamed_refusal(
+            response,
+            [
+                "[Invariant] The request did not pass the guardrails",
+                "Users must not mention the magic phrase 'Fight Club'",
+            ],
+        )
+
+    if push_to_explorer:
+        time.sleep(2)
+        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"]
+
+        trace_response = requests.get(
+            f"{explorer_api_url}/api/v1/trace/{trace_id}",
+            timeout=5,
+        )
+        trace = trace_response.json()
+
+        assert len(trace["messages"]) == 1
+        assert trace["messages"][0] == {
+            "role": "user",
+            "content": [{"type": "text", "text": "Tell me more about Fight Club."}],
+        }
+
+        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"]
+            == "Users must not mention the magic phrase 'Fight Club'"
+            and annotations[0]["extra_metadata"]["source"] == "guardrails-error"
+        )
+
+
+def is_refusal(chunk):
+    return (
+        len(chunk.candidates) == 1
+        and chunk.candidates[0].content.parts[0].text.startswith("[Invariant]")
+        and chunk.prompt_feedback is not None
+        and "BlockedReason.SAFETY" in str(chunk.prompt_feedback)
+    )
+
+
+def assert_is_streamed_refusal(response, expected_message_components: list[str]):
+    """
+    Validates that the streamed response contains a refusal at the end (or as only message).
+    """
+    num_chunks = 0
+    for c in response:
+        num_chunks += 1
+
+    assert num_chunks >= 1, "Expected at least one chunk"
+    # last chunk must be a refusal
+    assert is_refusal(c)
+
+    for emc in expected_message_components:
+        assert (
+            emc in c.model_dump_json()
+        ), f"Expected message component {emc} not found in refusal message: {c.model_dump_json()}"

tests/integration/guardrails/test_guardrails_open_ai.py

@@ -244,3 +244,108 @@ async def test_tool_call_guardrail_from_file(
             == "get_capital is called with Germany as argument"
             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, push_to_explorer",
+    [(True, True), (True, False), (False, True), (False, False)],
+)
+async def test_input_from_guardrail_from_file(
+    explorer_api_url, 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")
+
+    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",
+    )
+
+    request = {
+        "model": "gpt-4o",
+        "messages": [{"role": "user", "content": "Tell me more about Fight Club."}],
+    }
+
+    if not do_stream:
+        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 "Users must not mention the magic phrase 'Fight Club'" in str(
+            exc_info.value
+        )
+
+    else:
+        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 exc_info.value.message
+        )
+        assert "Users must not mention the magic phrase 'Fight Club'" in str(
+            exc_info.value.body
+        )
+
+    if push_to_explorer:
+        # 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()
+
+        # in case of input guardrailing, the pushed trace will not contain a response
+        assert len(trace["messages"]) == 1
+        assert trace["messages"][0] == {
+            "role": "user",
+            "content": "Tell me more about Fight Club.",
+        }
+
+        # 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"]
+            == "Users must not mention the magic phrase 'Fight Club'"
+            and annotations[0]["extra_metadata"]["source"] == "guardrails-error"
+        )

tests/integration/resources/guardrails/find_capital_guardrails.py

@@ -13,4 +13,10 @@ raise "Dublin detected in the response" if:
 raise "get_capital is called with Germany as argument" if:
     (call: ToolCall)
     call is tool:get_capital
-    call.function.arguments["country_name"] == "Germany"
+    call.function.arguments["country_name"] == "Germany"
+
+# For input guardrailing specifically
+raise "Users must not mention the magic phrase 'Fight Club'" if:
+    (msg: Message)
+    msg.role == "user"
+    "Fight Club" in msg.content