mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
243ed4393d
A test that asserts on the error inside its own except block passes when the call stops raising, because nothing runs the handler. That is the exact case the test exists to catch, so the regression lands green. Rewrites all 111 such blocks into pytest.raises, which fails when the call succeeds, and selects PT017 in ruff-tests.toml so no new one lands.
357 lines
12 KiB
Python
357 lines
12 KiB
Python
import pytest
|
|
import asyncio
|
|
import aiohttp, openai
|
|
from openai import OpenAI, AsyncOpenAI
|
|
from typing import Optional, List, Union
|
|
import json
|
|
from litellm._uuid import uuid
|
|
|
|
|
|
async def chat_completion(
|
|
session,
|
|
key,
|
|
messages,
|
|
model: Union[str, List] = "gpt-5.5",
|
|
guardrails: Optional[List] = None,
|
|
):
|
|
url = "http://0.0.0.0:4000/chat/completions"
|
|
headers = {
|
|
"Authorization": f"Bearer {key}",
|
|
"Content-Type": "application/json",
|
|
}
|
|
|
|
data = {
|
|
"model": model,
|
|
"messages": messages,
|
|
}
|
|
|
|
if guardrails is not None:
|
|
data["guardrails"] = guardrails
|
|
|
|
print("data=", data)
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
print(response_text)
|
|
print()
|
|
|
|
if status != 200:
|
|
raise Exception(response_text)
|
|
|
|
# response headers
|
|
response_headers = dict(response.headers)
|
|
print("response headers=", response_headers)
|
|
|
|
return await response.json(), response_headers
|
|
|
|
|
|
async def generate_key(
|
|
session, guardrails: Optional[List] = None, team_id: Optional[str] = None
|
|
):
|
|
url = "http://0.0.0.0:4000/key/generate"
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
data = {}
|
|
if guardrails:
|
|
data["guardrails"] = guardrails
|
|
if team_id:
|
|
data["team_id"] = team_id
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
print(response_text)
|
|
print()
|
|
|
|
if status != 200:
|
|
raise Exception(f"Request did not return a 200 status code: {status}")
|
|
|
|
return await response.json()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skip(reason="Aporia account disabled")
|
|
async def test_llm_guard_triggered_safe_request():
|
|
"""
|
|
- Tests a request where no content mod is triggered
|
|
- Assert that the guardrails applied are returned in the response headers
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
response, headers = await chat_completion(
|
|
session,
|
|
"sk-1234",
|
|
model="fake-openai-endpoint",
|
|
messages=[{"role": "user", "content": f"Hello what's the weather"}],
|
|
guardrails=[
|
|
"aporia-post-guard",
|
|
"aporia-pre-guard",
|
|
],
|
|
)
|
|
await asyncio.sleep(3)
|
|
|
|
print("response=", response, "response headers", headers)
|
|
|
|
assert "x-litellm-applied-guardrails" in headers
|
|
|
|
assert (
|
|
headers["x-litellm-applied-guardrails"]
|
|
== "aporia-pre-guard,aporia-post-guard"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.skip(reason="Aporia account disabled")
|
|
async def test_llm_guard_triggered():
|
|
"""
|
|
- Tests a request where no content mod is triggered
|
|
- Assert that the guardrails applied are returned in the response headers
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
with pytest.raises(Exception, match="Aporia detected and blocked PII") as exc_info:
|
|
response, headers = await chat_completion(
|
|
session,
|
|
"sk-1234",
|
|
model="fake-openai-endpoint",
|
|
messages=[
|
|
{"role": "user", "content": f"Hello my name is ishaan@berri.ai"}
|
|
],
|
|
guardrails=[
|
|
"aporia-post-guard",
|
|
"aporia-pre-guard",
|
|
],
|
|
)
|
|
e = exc_info.value
|
|
print(e)
|
|
assert "Aporia detected and blocked PII" in str(e)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_no_llm_guard_triggered():
|
|
"""
|
|
- Tests a request where no content mod is triggered
|
|
- Assert that the guardrails applied are returned in the response headers
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
response, headers = await chat_completion(
|
|
session,
|
|
"sk-1234",
|
|
model="fake-openai-endpoint",
|
|
messages=[{"role": "user", "content": f"Hello what's the weather"}],
|
|
guardrails=[],
|
|
)
|
|
await asyncio.sleep(3)
|
|
|
|
print("response=", response, "response headers", headers)
|
|
|
|
assert "x-litellm-applied-guardrails" not in headers
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_guardrails_with_api_key_controls():
|
|
"""
|
|
- Make two API Keys
|
|
- Key 1 with no guardrails
|
|
- Key 2 with guardrails
|
|
- Request to Key 1 -> should be success with no guardrails
|
|
- Request to Key 2 -> should be error since guardrails are triggered
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
key_with_guardrails = await generate_key(
|
|
session=session,
|
|
guardrails=[
|
|
"bedrock-pre-guard",
|
|
],
|
|
)
|
|
|
|
key_with_guardrails = key_with_guardrails["key"]
|
|
|
|
key_without_guardrails = await generate_key(session=session, guardrails=None)
|
|
|
|
key_without_guardrails = key_without_guardrails["key"]
|
|
|
|
# test no guardrails triggered for key without guardrails
|
|
response, headers = await chat_completion(
|
|
session,
|
|
key_without_guardrails,
|
|
model="fake-openai-endpoint",
|
|
messages=[{"role": "user", "content": f"Hello what's the weather"}],
|
|
)
|
|
await asyncio.sleep(3)
|
|
|
|
print("response=", response, "response headers", headers)
|
|
assert "x-litellm-applied-guardrails" not in headers
|
|
|
|
# test guardrails triggered for key with guardrails
|
|
response, headers = await chat_completion(
|
|
session,
|
|
key_with_guardrails,
|
|
model="fake-openai-endpoint",
|
|
messages=[{"role": "user", "content": f"Hello my name is ishaan@berri.ai"}],
|
|
)
|
|
|
|
assert "x-litellm-applied-guardrails" in headers
|
|
assert headers["x-litellm-applied-guardrails"] == "bedrock-pre-guard"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_bedrock_guardrail_triggered():
|
|
"""
|
|
- Tests a request where our bedrock guardrail should be triggered
|
|
- Assert that the guardrails applied are returned in the response headers
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
with pytest.raises(Exception, match="Violated guardrail policy") as exc_info:
|
|
response, headers = await chat_completion(
|
|
session,
|
|
"sk-1234",
|
|
model="fake-openai-endpoint",
|
|
messages=[{"role": "user", "content": "Hello do you like coffee?"}],
|
|
guardrails=["bedrock-pre-guard"],
|
|
)
|
|
e = exc_info.value
|
|
print(e)
|
|
assert "Violated guardrail policy" in str(e)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_custom_guardrail_during_call_triggered():
|
|
"""
|
|
- Tests a request where our bedrock guardrail should be triggered
|
|
- Assert that the guardrails applied are returned in the response headers
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
with pytest.raises(Exception, match="Guardrail failed words - `litellm` detected") as exc_info:
|
|
response, headers = await chat_completion(
|
|
session,
|
|
"sk-1234",
|
|
model="fake-openai-endpoint",
|
|
messages=[{"role": "user", "content": f"Hello do you like litellm?"}],
|
|
guardrails=["custom-during-guard"],
|
|
)
|
|
e = exc_info.value
|
|
print(e)
|
|
assert "Guardrail failed words - `litellm` detected" in str(e)
|
|
|
|
|
|
async def create_team(session, guardrails: Optional[List] = None):
|
|
url = "http://0.0.0.0:4000/team/new"
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
data = {"guardrails": guardrails}
|
|
|
|
print("request data=", data)
|
|
|
|
async with session.post(url, headers=headers, json=data) as response:
|
|
status = response.status
|
|
response_text = await response.text()
|
|
|
|
print(response_text)
|
|
print()
|
|
|
|
if status != 200:
|
|
raise Exception(f"Request did not return a 200 status code: {status}")
|
|
|
|
return await response.json()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_guardrails_with_team_controls():
|
|
"""
|
|
- Create a team with guardrails
|
|
- Make two API Keys
|
|
- Key 1 not associated with team
|
|
- Key 2 associated with team (inherits team guardrails)
|
|
- Request with Key 1 -> should be success with no guardrails
|
|
- Request with Key 2 -> should error since team guardrails are triggered
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
|
|
# Create team with guardrails
|
|
team = await create_team(
|
|
session=session,
|
|
guardrails=[
|
|
"bedrock-pre-guard",
|
|
],
|
|
)
|
|
|
|
print("team=", team)
|
|
|
|
team_id = team["team_id"]
|
|
|
|
# Create key with team association
|
|
key_with_team = await generate_key(session=session, team_id=team_id)
|
|
key_with_team = key_with_team["key"]
|
|
|
|
# Create key without team
|
|
key_without_team = await generate_key(
|
|
session=session,
|
|
)
|
|
key_without_team = key_without_team["key"]
|
|
|
|
# Test no guardrails triggered for key without a team
|
|
response, headers = await chat_completion(
|
|
session,
|
|
key_without_team,
|
|
model="fake-openai-endpoint",
|
|
messages=[{"role": "user", "content": "Hello my name is ishaan@berri.ai"}],
|
|
)
|
|
await asyncio.sleep(3)
|
|
|
|
print("response=", response, "response headers", headers)
|
|
assert "x-litellm-applied-guardrails" not in headers
|
|
|
|
response, headers = await chat_completion(
|
|
session,
|
|
key_with_team,
|
|
model="fake-openai-endpoint",
|
|
messages=[{"role": "user", "content": "Hello my name is ishaan@berri.ai"}],
|
|
)
|
|
|
|
print("response headers=", json.dumps(headers, indent=4))
|
|
|
|
assert "x-litellm-applied-guardrails" in headers
|
|
assert headers["x-litellm-applied-guardrails"] == "bedrock-pre-guard"
|
|
|
|
|
|
async def get_guardrail_lb_counts(session):
|
|
"""Get the current guardrail load balancing call counts from the proxy."""
|
|
url = "http://0.0.0.0:4000/guardrail/lb/counts"
|
|
headers = {"Authorization": "Bearer sk-1234", "Content-Type": "application/json"}
|
|
|
|
async with session.get(url, headers=headers) as response:
|
|
if response.status == 200:
|
|
return await response.json()
|
|
return None
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_guardrail_load_balancing():
|
|
"""
|
|
Test that guardrail load balancing distributes requests across multiple guardrail instances.
|
|
|
|
- Make 20 requests with the lb-test-guard guardrail
|
|
- Verify that both GuardrailForLBTestingA and GuardrailForLBTestingB are called
|
|
- Verify reasonable distribution (both should have at least some calls)
|
|
"""
|
|
async with aiohttp.ClientSession() as session:
|
|
num_requests = 20
|
|
|
|
# Make multiple requests with the load-balanced guardrail
|
|
for i in range(num_requests):
|
|
response, headers = await chat_completion(
|
|
session,
|
|
"sk-1234",
|
|
model="fake-openai-endpoint",
|
|
messages=[{"role": "user", "content": f"Hello request {i}"}],
|
|
guardrails=["lb-test-guard"],
|
|
)
|
|
|
|
# Verify guardrail was applied
|
|
assert "x-litellm-applied-guardrails" in headers
|
|
assert headers["x-litellm-applied-guardrails"] == "lb-test-guard"
|
|
|
|
# All requests should succeed - the test passes if we get here
|
|
# The actual load balancing verification is done by checking proxy logs
|
|
# which should show alternating calls to GuardrailForLBTestingA and GuardrailForLBTestingB
|
|
print(f"Successfully made {num_requests} requests with load-balanced guardrail")
|