CalvinBackup/invariant-gateway
1
0
Fork 0
mirror of https://github.com/invariantlabs-ai/invariant-gateway.git synced 2026-07-19 00:17:23 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add tests for the guardrails integration for open_ai route.

Browse Source
This commit is contained in:
Hemang
2025-03-18 13:51:04 +01:00
committed by Hemang Sarkar
parent 5dc9485831
commit 22acb38131
5 changed files with 276 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.github/workflows/tests_ci.yml
+1
View File
@@ -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
run.sh
+13 -1
View File
@@ -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 $@
}
tests/integration/docker-compose.test.yml
+4
View File
@@ -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: []
tests/integration/guardrails/test_guardrails_open_ai.py
+250
View File
@@ -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"
)
tests/integration/resources/guardrails/find_capital_guardrails.py
+8
View File
@@ -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"