375 lines
12 KiB
Python
375 lines
12 KiB
Python
import pytest
|
|
from pydantic import BaseModel
|
|
|
|
from guardrails import AsyncGuard, Validator, register_validator
|
|
from guardrails_ai.types import PassResult
|
|
from guardrails.utils.validator_utils import verify_metadata_requirements
|
|
from guardrails.types import OnFailAction
|
|
from tests.integration_tests.test_assets.custom_llm import mock_async_llm
|
|
from tests.integration_tests.test_assets.validators import (
|
|
EndsWith,
|
|
LowerCase,
|
|
OneLine,
|
|
TwoWords,
|
|
UpperCase,
|
|
ValidLength,
|
|
)
|
|
|
|
|
|
@register_validator("myrequiringvalidator", data_type="string")
|
|
class RequiringValidator(Validator):
|
|
required_metadata_keys = ["required_key"]
|
|
|
|
def validate(self, value, metadata):
|
|
return PassResult()
|
|
|
|
|
|
@register_validator("myrequiringvalidator2", data_type="string")
|
|
class RequiringValidator2(Validator):
|
|
required_metadata_keys = ["required_key2"]
|
|
|
|
def validate(self, value, metadata):
|
|
return PassResult()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"spec,metadata,error_message",
|
|
[
|
|
(
|
|
"""
|
|
<rail version="0.1">
|
|
<output>
|
|
<string name="string_name" validators="myrequiringvalidator" />
|
|
</output>
|
|
</rail>
|
|
""",
|
|
{"required_key": "a"},
|
|
"Missing required metadata keys: required_key",
|
|
),
|
|
(
|
|
"""
|
|
<rail version="0.1">
|
|
<output>
|
|
<object name="temp_name">
|
|
<string name="string_name" validators="myrequiringvalidator" />
|
|
</object>
|
|
<list name="list_name">
|
|
<string name="string_name" validators="myrequiringvalidator2" />
|
|
</list>
|
|
</output>
|
|
</rail>
|
|
""",
|
|
{"required_key": "a", "required_key2": "b"},
|
|
"Missing required metadata keys: required_key, required_key2",
|
|
),
|
|
(
|
|
"""
|
|
<rail version="0.1">
|
|
<output>
|
|
<object name="temp_name">
|
|
<list name="list_name">
|
|
<choice name="choice_name" discriminator="hi">
|
|
<case name="hello">
|
|
<string name="string_name" />
|
|
</case>
|
|
<case name="hiya">
|
|
<string name="string_name" validators="myrequiringvalidator" />
|
|
</case>
|
|
</choice>
|
|
</list>
|
|
</object>
|
|
</output>
|
|
</rail>
|
|
""",
|
|
{"required_key": "a"},
|
|
"Missing required metadata keys: required_key",
|
|
),
|
|
],
|
|
)
|
|
@pytest.mark.asyncio
|
|
async def test_required_metadata(spec, metadata, error_message):
|
|
guard: AsyncGuard = AsyncGuard.for_rail_string(spec)
|
|
|
|
missing_keys = verify_metadata_requirements({}, guard._validators)
|
|
assert set(missing_keys) == set(metadata)
|
|
|
|
not_missing_keys = verify_metadata_requirements(metadata, guard._validators)
|
|
assert not_missing_keys == []
|
|
|
|
# test async guard
|
|
with pytest.raises(ValueError) as excinfo:
|
|
await guard.parse("{}")
|
|
await guard.parse("{}", llm_api=mock_async_llm, num_reasks=0)
|
|
assert str(excinfo.value) == error_message
|
|
|
|
response = await guard.parse(
|
|
"{}", metadata=metadata, llm_api=mock_async_llm, num_reasks=0
|
|
)
|
|
assert response.error is None
|
|
|
|
|
|
empty_rail_string = """<rail version="0.1">
|
|
<output
|
|
type="string"
|
|
description="empty railspec"
|
|
/>
|
|
</rail>"""
|
|
|
|
|
|
class EmptyModel(BaseModel):
|
|
empty_field: str
|
|
|
|
|
|
r_guard_none = AsyncGuard.for_rail("tests/unit_tests/test_assets/empty.rail")
|
|
r_guard_two = AsyncGuard.for_rail("tests/unit_tests/test_assets/empty.rail")
|
|
r_guard_two.configure(num_reasks=2)
|
|
rs_guard_none = AsyncGuard.for_rail_string(empty_rail_string)
|
|
rs_guard_two = AsyncGuard.for_rail_string(empty_rail_string)
|
|
rs_guard_two.configure(num_reasks=2)
|
|
py_guard_none = AsyncGuard.for_pydantic(output_class=EmptyModel)
|
|
py_guard_two = AsyncGuard.for_pydantic(output_class=EmptyModel)
|
|
py_guard_two.configure(num_reasks=2)
|
|
s_guard_none = AsyncGuard.for_string(validators=[], description="empty railspec")
|
|
s_guard_two = AsyncGuard.for_string(validators=[], description="empty railspec")
|
|
s_guard_two.configure(num_reasks=2)
|
|
|
|
|
|
def guard_init_for_rail():
|
|
guard = AsyncGuard.for_rail("tests/unit_tests/test_assets/simple.rail")
|
|
assert (
|
|
guard.instructions.format().source.strip()
|
|
== "You are a helpful bot, who answers only with valid JSON"
|
|
)
|
|
assert guard.prompt.format().source.strip() == "Extract a string from the text"
|
|
|
|
|
|
def test_use():
|
|
guard: AsyncGuard = AsyncGuard().use(
|
|
EndsWith("a"),
|
|
OneLine(),
|
|
LowerCase(),
|
|
TwoWords(on_fail=OnFailAction.REASK),
|
|
ValidLength(0, 12, on_fail=OnFailAction.REFRAIN),
|
|
)
|
|
|
|
# print(guard.__stringify__())
|
|
assert len(guard._validators) == 5
|
|
|
|
assert isinstance(guard._validators[0], EndsWith)
|
|
assert guard._validators[0]._kwargs["end"] == "a"
|
|
assert (
|
|
guard._validators[0].on_fail_descriptor == OnFailAction.FIX
|
|
) # bc this is the default
|
|
|
|
assert isinstance(guard._validators[1], OneLine)
|
|
assert (
|
|
guard._validators[1].on_fail_descriptor == OnFailAction.EXCEPTION
|
|
) # bc this is the default
|
|
|
|
assert isinstance(guard._validators[2], LowerCase)
|
|
assert (
|
|
guard._validators[2].on_fail_descriptor == OnFailAction.EXCEPTION
|
|
) # bc this is the default
|
|
|
|
assert isinstance(guard._validators[3], TwoWords)
|
|
assert guard._validators[3].on_fail_descriptor == OnFailAction.REASK # bc we set it
|
|
|
|
assert isinstance(guard._validators[4], ValidLength)
|
|
assert guard._validators[4]._min == 0
|
|
assert guard._validators[4]._kwargs["min"] == 0
|
|
assert guard._validators[4]._max == 12
|
|
assert guard._validators[4]._kwargs["max"] == 12
|
|
assert (
|
|
guard._validators[4].on_fail_descriptor == OnFailAction.REFRAIN
|
|
) # bc we set it
|
|
|
|
# No longer a constraint
|
|
# Raises error when trying to `use` a validator on a non-string
|
|
# with pytest.raises(RuntimeError):
|
|
|
|
class TestClass(BaseModel):
|
|
another_field: str
|
|
|
|
py_guard = AsyncGuard.for_pydantic(output_class=TestClass)
|
|
py_guard.use(EndsWith("a"))
|
|
assert py_guard._validator_map.get("$") == [EndsWith("a")]
|
|
|
|
# Use a combination of prompt, instructions, messages and output validators
|
|
# Should only have the output validators in the guard,
|
|
# everything else is in the schema
|
|
guard: AsyncGuard = (
|
|
AsyncGuard()
|
|
.use(LowerCase(), OneLine(), on="prompt")
|
|
.use(UpperCase(), on="instructions")
|
|
.use(LowerCase(), on="messages")
|
|
.use(
|
|
EndsWith(end="a"), TwoWords(on_fail=OnFailAction.REASK), on="output"
|
|
) # default on="output", still explicitly set
|
|
)
|
|
|
|
# Check schemas for prompt, instructions and messages validators
|
|
prompt_validators = guard._validator_map.get("prompt")
|
|
assert len(prompt_validators) == 2
|
|
assert prompt_validators[0].__class__.__name__ == "LowerCase"
|
|
assert prompt_validators[1].__class__.__name__ == "OneLine"
|
|
|
|
instructions_validators = guard._validator_map.get("instructions")
|
|
assert len(instructions_validators) == 1
|
|
assert instructions_validators[0].__class__.__name__ == "UpperCase"
|
|
|
|
messages_validators = guard._validator_map.get("messages")
|
|
assert len(messages_validators) == 1
|
|
assert messages_validators[0].__class__.__name__ == "LowerCase"
|
|
|
|
# Check guard for output validators
|
|
assert len(guard._validators) == 6 # 2 + 1 + 1 + 2 = 6
|
|
|
|
assert isinstance(guard._validators[4], EndsWith)
|
|
assert guard._validators[4]._kwargs["end"] == "a"
|
|
assert (
|
|
guard._validators[4].on_fail_descriptor == OnFailAction.FIX
|
|
) # bc this is the default
|
|
|
|
assert isinstance(guard._validators[5], TwoWords)
|
|
assert guard._validators[5].on_fail_descriptor == OnFailAction.REASK # bc we set it
|
|
|
|
# Test with an unrecognized "on" parameter, should warn with a UserWarning
|
|
with pytest.warns(UserWarning):
|
|
guard: AsyncGuard = (
|
|
AsyncGuard()
|
|
.use(EndsWith("a"), on="response") # invalid on parameter
|
|
.use(OneLine(), on="prompt") # valid on parameter
|
|
)
|
|
|
|
|
|
# TODO: Move to integration tests; these are not unit tests...
|
|
class TestValidate:
|
|
@pytest.mark.asyncio
|
|
async def test_output_only_success(self):
|
|
guard: AsyncGuard = AsyncGuard().use(
|
|
OneLine(),
|
|
LowerCase(on_fail=OnFailAction.FIX),
|
|
TwoWords(),
|
|
ValidLength(0, 12, on_fail=OnFailAction.REFRAIN),
|
|
)
|
|
|
|
llm_output: str = "Oh Canada" # bc it meets our criteria
|
|
|
|
response = await guard.validate(llm_output)
|
|
|
|
assert response.validation_passed is True
|
|
assert response.validated_output == llm_output.lower()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_only_failure(self):
|
|
guard: AsyncGuard = AsyncGuard().use(
|
|
OneLine(on_fail=OnFailAction.NOOP),
|
|
LowerCase(on_fail=OnFailAction.FIX),
|
|
TwoWords(on_fail=OnFailAction.NOOP),
|
|
ValidLength(0, 12, on_fail=OnFailAction.REFRAIN),
|
|
)
|
|
|
|
llm_output = "Star Spangled Banner" # to stick with the theme
|
|
|
|
response = await guard.validate(llm_output)
|
|
|
|
assert response.validation_passed is False
|
|
assert response.validated_output is None
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_on_many_success(self):
|
|
# Test with a combination of prompt, output,
|
|
# instructions and messages validators
|
|
# Should still only use the output validators to validate the output
|
|
guard: AsyncGuard = (
|
|
AsyncGuard()
|
|
.use(OneLine(), LowerCase(), UpperCase(), on="messages")
|
|
.use(
|
|
LowerCase(on_fail=OnFailAction.FIX),
|
|
TwoWords(),
|
|
ValidLength(0, 12, on_fail=OnFailAction.REFRAIN),
|
|
)
|
|
)
|
|
|
|
llm_output: str = "Oh Canada" # bc it meets our criteria
|
|
|
|
response = await guard.validate(llm_output)
|
|
|
|
assert response.validation_passed is True
|
|
assert response.validated_output == llm_output.lower()
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_on_many_failure(self):
|
|
guard: AsyncGuard = (
|
|
AsyncGuard()
|
|
.use(
|
|
OneLine(on_fail=OnFailAction.NOOP),
|
|
LowerCase(on_fail=OnFailAction.NOOP),
|
|
UpperCase(on_fail=OnFailAction.NOOP),
|
|
on="messages",
|
|
)
|
|
.use(
|
|
LowerCase(on_fail=OnFailAction.FIX),
|
|
TwoWords(on_fail=OnFailAction.NOOP),
|
|
ValidLength(0, 12, on_fail=OnFailAction.REFRAIN),
|
|
on="output",
|
|
)
|
|
)
|
|
|
|
llm_output = "Star Spangled Banner" # to stick with the theme
|
|
|
|
response = await guard.validate(llm_output)
|
|
|
|
assert response.validation_passed is False
|
|
assert response.validated_output is None
|
|
|
|
|
|
def test_multi_use():
|
|
guard: AsyncGuard = (
|
|
AsyncGuard()
|
|
.use(UpperCase(), on="messages")
|
|
.use(
|
|
TwoWords(on_fail=OnFailAction.REASK),
|
|
ValidLength(0, 12, on_fail=OnFailAction.REFRAIN),
|
|
on="output",
|
|
)
|
|
)
|
|
|
|
# Check schemas for prompt, instructions and messages validators
|
|
output_validators = guard._validator_map.get("$")
|
|
|
|
assert len(output_validators) == 2
|
|
assert output_validators[0].__class__.__name__ == "TwoWords"
|
|
assert output_validators[1].__class__.__name__ == "ValidLength"
|
|
|
|
messages_validators = guard._validator_map.get("messages")
|
|
assert len(messages_validators) == 1
|
|
assert messages_validators[0].__class__.__name__ == "UpperCase"
|
|
|
|
# Check guard for output validators
|
|
assert len(guard._validators) == 3 # 2 + 1 + 1 + 2 = 6
|
|
|
|
assert isinstance(guard._validators[1], TwoWords)
|
|
assert guard._validators[1].on_fail_descriptor == OnFailAction.REASK # bc we set it
|
|
|
|
assert isinstance(guard._validators[2], ValidLength)
|
|
assert guard._validators[2]._min == 0
|
|
assert guard._validators[2]._kwargs["min"] == 0
|
|
assert guard._validators[2]._max == 12
|
|
assert guard._validators[2]._kwargs["max"] == 12
|
|
assert (
|
|
guard._validators[2].on_fail_descriptor == OnFailAction.REFRAIN
|
|
) # bc we set it
|
|
|
|
# Test with an unrecognized "on" parameter, should warn with a UserWarning
|
|
with pytest.warns(UserWarning):
|
|
guard: AsyncGuard = (
|
|
AsyncGuard()
|
|
.use(LowerCase(), on="messages")
|
|
.use(
|
|
TwoWords(on_fail=OnFailAction.REASK),
|
|
ValidLength(0, 12, on_fail=OnFailAction.REFRAIN),
|
|
on="response", # invalid "on" parameter
|
|
)
|
|
)
|