anthropic integration of pipelined and pre-guardrailing

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_open_ai.py

@@ -330,7 +330,7 @@ async def test_input_from_guardrail_from_file(
         trace = trace_response.json()
 
         # in case of input guardrailing, the pushed trace will not contain a response
-        assert len(trace["messages"]) == 1
+        assert len(trace["messages"]) == 1, "Trace should only contain the user message"
         assert trace["messages"][0] == {
             "role": "user",
             "content": "Tell me more about Fight Club.",