diff --git a/.github/workflows/tests_ci.yml b/.github/workflows/tests_ci.yml index b81a8fb..71dfe5e 100644 --- a/.github/workflows/tests_ci.yml +++ b/.github/workflows/tests_ci.yml @@ -36,4 +36,5 @@ jobs: OPENAI_API_KEY: ${{ secrets.INVARIANT_TESTING_OPENAI_KEY }} ANTHROPIC_API_KEY: ${{ secrets.INVARIANT_TESTING_ANTHROPIC_KEY}} GEMINI_API_KEY: ${{ secrets.INVARIANT_TESTING_GEMINI_KEY }} + GUARDRAILS_API_KEY: ${{ secrets.INVARIANT_TESTING_GUARDRAILS_KEY }} run: ./run.sh integration-tests -s -vv diff --git a/run.sh b/run.sh index db3f273..6cc9931 100755 --- a/run.sh +++ b/run.sh @@ -85,10 +85,21 @@ integration_tests() { fi echo "File successfully downloaded: $FILE" + TEST_GUARDRAILS_FILE_PATH="tests/integration/resources/guardrails/find_capital_guardrails.py" + if [[ -n "$TEST_GUARDRAILS_FILE_PATH" ]]; then + if [[ -f "$TEST_GUARDRAILS_FILE_PATH" ]]; then + TEST_GUARDRAILS_FILE_PATH=$(realpath "$TEST_GUARDRAILS_FILE_PATH") + else + echo "Error: Specified test guardrails file does not exist: $TEST_GUARDRAILS_FILE_PATH" + exit 1 + fi + fi + + # Start containers GATEWAY_PATH=$(pwd)/gateway docker compose -f tests/integration/docker-compose.test.yml down GATEWAY_PATH=$(pwd)/gateway docker compose -f tests/integration/docker-compose.test.yml build - GATEWAY_PATH=$(pwd)/gateway docker compose -f tests/integration/docker-compose.test.yml up -d + GUARDRAILS_FILE_PATH="$TEST_GUARDRAILS_FILE_PATH" GATEWAY_PATH=$(pwd)/gateway docker compose -f tests/integration/docker-compose.test.yml up -d until [ "$(docker inspect -f '{{.State.Health.Status}}' invariant-gateway-test-explorer-app-api)" = "healthy" ]; do echo "Explorer backend app-api instance container starting..." @@ -121,6 +132,7 @@ integration_tests() { -e OPENAI_API_KEY="$OPENAI_API_KEY" \ -e ANTHROPIC_API_KEY="$ANTHROPIC_API_KEY"\ -e GEMINI_API_KEY="$GEMINI_API_KEY" \ + -e GUARDRAILS_API_KEY="$GUARDRAILS_API_KEY" \ --env-file ./tests/integration/.env.test \ invariant-gateway-tests $@ } diff --git a/tests/integration/docker-compose.test.yml b/tests/integration/docker-compose.test.yml index 5d58b20..26ba1c6 100644 --- a/tests/integration/docker-compose.test.yml +++ b/tests/integration/docker-compose.test.yml @@ -34,10 +34,14 @@ services: - .env.test environment: - DEV_MODE=true + - GUARDRAILS_FILE_PATH=${GUARDRAILS_FILE_PATH:+/srv/resources/guardrails.py} volumes: - type: bind source: ${GATEWAY_PATH} target: /srv/gateway + - type: bind + source: ${GUARDRAILS_FILE_PATH:-/dev/null} + target: /srv/resources/guardrails.py networks: - invariant-gateway-web-test ports: [] diff --git a/tests/integration/guardrails/test_guardrails_open_ai.py b/tests/integration/guardrails/test_guardrails_open_ai.py new file mode 100644 index 0000000..44af91c --- /dev/null +++ b/tests/integration/guardrails/test_guardrails_open_ai.py @@ -0,0 +1,250 @@ +"""Test the guardrails from file.""" + +import os +import sys +import uuid +import time + +# Add integration folder (parent) to sys.path +sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + +import pytest +import requests +from httpx import Client +from openai import OpenAI, BadRequestError, APIError + +# Pytest plugins +pytest_plugins = ("pytest_asyncio",) + + +@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_message_content_guardrail_from_file( + explorer_api_url, gateway_url, do_stream, push_to_explorer +): + """Test the message content guardrail.""" + if not os.getenv("GUARDRAILS_API_KEY"): + pytest.fail("No GUARDRAILS_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('GUARDRAILS_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": "What is the capital of Spain?"}], + } + + 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 response did not pass the guardrails" in str( + exc_info.value + ) + assert "Madrid detected in the response" 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 response did not pass the guardrails" + in exc_info.value.message + ) + assert "Madrid detected in the response" 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() + + assert len(trace["messages"]) == 2 + assert trace["messages"][0] == { + "role": "user", + "content": "What is the capital of Spain?", + } + + # 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"] == "Madrid detected in the response" + and annotations[0]["extra_metadata"]["source"] == "guardrails-error" + ) + assert ( + annotations[1]["content"] == "Madrid detected in the response" + and annotations[1]["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_tool_call_guardrail_from_file( + explorer_api_url, gateway_url, do_stream, push_to_explorer +): + """Test the message content guardrail.""" + if not os.getenv("GUARDRAILS_API_KEY"): + pytest.fail("No GUARDRAILS_API_KEY set, failing") + + tools = [ + { + "type": "function", + "function": { + "name": "get_capital", + "description": "Get the capital for a given country.", + "parameters": { + "type": "object", + "properties": { + "country_name": { + "type": "string", + "description": "The country name example Germany or France.", + } + }, + "required": ["country_name"], + }, + }, + } + ] + system_message = { + "role": "system", + "content": "Use the get_capital tool call to get the capital of a country. If the user input doesn't contain a country name, fail the request with a pretty message. If the get_capital tool call returns 'not_found' then fail the request with a pretty message. Do not return the capital if the get_capital tool call returns 'not_found'.", + } + request = { + "model": "gpt-4o", + "messages": [ + system_message, + {"role": "user", "content": "What is the capital of Germany?"}, + ], + "tools": tools, + } + + dataset_name = f"test-dataset-open-ai-{uuid.uuid4()}" + + client = OpenAI( + http_client=Client( + headers={ + "Invariant-Authorization": f"Bearer {os.getenv('GUARDRAILS_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", + ) + + 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 response did not pass the guardrails" in str( + exc_info.value + ) + assert "get_capital is called with Germany as argument" 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 response did not pass the guardrails" + in exc_info.value.message + ) + assert "get_capital is called with Germany as argument" 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() + + assert len(trace["messages"]) == 3 + assert trace["messages"][0] == system_message + assert trace["messages"][1] == { + "role": "user", + "content": "What is the capital of Germany?", + } + + # 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"] + == "get_capital is called with Germany as argument" + and annotations[0]["extra_metadata"]["source"] == "guardrails-error" + ) diff --git a/tests/integration/resources/guardrails/find_capital_guardrails.py b/tests/integration/resources/guardrails/find_capital_guardrails.py new file mode 100644 index 0000000..bfa6418 --- /dev/null +++ b/tests/integration/resources/guardrails/find_capital_guardrails.py @@ -0,0 +1,8 @@ +raise "Madrid detected in the response" if: + (msg: Message) + "Madrid" in msg.content + +raise "get_capital is called with Germany as argument" if: + (call: ToolCall) + call is tool:get_capital + call.function.arguments["country_name"] == "Germany" \ No newline at end of file