Files
invariant-gateway/tests/test-client.py
Luca Beurer-Kellner c4dd3f3b19 Allow to specify different API keys for the guardrailing service (#36)
* minor refactor for getting invariant api keys for guardrailing

* allow different guardrailing api key

* tests

* fix comment + import

* improved unauthorized handling
2025-04-03 12:15:30 +02:00

34 lines
840 B
Python

"""
Simple (non-streaming) test client for the Gateway (uses OpenAI integration).
"""
from openai import OpenAI
from httpx import Client
import os
# unicode escape everything
guardrails = """
raise "Rule 1: Do not talk about Fight Club" if:
(msg: Message)
"fight club" in msg.content
""".encode("unicode_escape")
openai_client = OpenAI(
default_headers={
"Invariant-Authorization": "Bearer " + os.getenv("INVARIANT_API_KEY"),
"Invariant-Guardrails": guardrails,
},
base_url="http://localhost:8005/api/v1/gateway/non-streaming/openai",
)
response = openai_client.chat.completions.create(
model="gpt-4",
messages=[
{
"role": "user",
"content": "What can you tell me about fight club?",
}
],
)
print("Response: ", response.choices[0].message.content)