153 lines
5.1 KiB
Python
153 lines
5.1 KiB
Python
from typing import Dict
|
|
import pytest
|
|
|
|
import guardrails as gd
|
|
from guardrails import register_validator
|
|
from guardrails.classes.llm.llm_response import LLMResponse
|
|
from guardrails.validator_base import OnFailAction
|
|
from guardrails_ai.types import FailResult, ValidationResult
|
|
from tests.integration_tests.test_assets.custom_llm import mock_async_llm, mock_llm
|
|
|
|
from .test_assets import pydantic, string
|
|
|
|
|
|
def test_parsing_reask(mocker):
|
|
"""Test re-asking when response is not parseable."""
|
|
mock_invoke_llm = mocker.patch(
|
|
"guardrails.llm_providers.ArbitraryCallable._invoke_llm",
|
|
)
|
|
mock_invoke_llm.side_effect = [
|
|
LLMResponse(
|
|
output=pydantic.PARSING_UNPARSEABLE_LLM_OUTPUT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
),
|
|
LLMResponse(
|
|
output=pydantic.PARSING_EXPECTED_LLM_OUTPUT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
),
|
|
]
|
|
|
|
guard = gd.Guard.for_pydantic(
|
|
output_class=pydantic.PersonalDetails,
|
|
messages=[{"role": "user", "content": pydantic.PARSING_INITIAL_PROMPT}],
|
|
)
|
|
|
|
final_output = guard(
|
|
llm_api=mock_llm,
|
|
prompt_params={"document": pydantic.PARSING_DOCUMENT},
|
|
num_reasks=1,
|
|
)
|
|
|
|
assert final_output.validated_output == pydantic.PARSING_EXPECTED_OUTPUT
|
|
|
|
call = guard.history.first
|
|
|
|
# Check that the guard state object has the correct number of re-asks.
|
|
assert call.iterations.length == 2
|
|
|
|
# For orginal prompt and output
|
|
assert call.compiled_messages[0]["content"] == pydantic.PARSING_COMPILED_PROMPT
|
|
assert call.iterations.first.raw_output == pydantic.PARSING_UNPARSEABLE_LLM_OUTPUT
|
|
assert call.iterations.first.guarded_output is None
|
|
|
|
# For re-asked prompt and output
|
|
assert call.iterations.last.inputs.messages[1]["content"] == gd.Prompt(
|
|
pydantic.PARSING_COMPILED_REASK
|
|
)
|
|
# Same as above
|
|
assert call.reask_messages[0][1]["content"] == pydantic.PARSING_COMPILED_REASK
|
|
assert call.raw_outputs.last == pydantic.PARSING_EXPECTED_LLM_OUTPUT
|
|
assert call.guarded_output == pydantic.PARSING_EXPECTED_OUTPUT
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_async_parsing_reask(mocker):
|
|
"""Test re-asking when response is not parseable during async flow."""
|
|
mock_invoke_llm = mocker.patch(
|
|
"guardrails.llm_providers.AsyncArbitraryCallable.invoke_llm",
|
|
)
|
|
mock_invoke_llm.side_effect = [
|
|
LLMResponse(
|
|
output=pydantic.PARSING_UNPARSEABLE_LLM_OUTPUT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
),
|
|
LLMResponse(
|
|
output=pydantic.PARSING_EXPECTED_LLM_OUTPUT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
),
|
|
]
|
|
|
|
guard = gd.AsyncGuard.for_pydantic(
|
|
output_class=pydantic.PersonalDetails,
|
|
messages=[{"role": "user", "content": pydantic.PARSING_INITIAL_PROMPT}],
|
|
)
|
|
|
|
final_output = await guard(
|
|
llm_api=mock_async_llm,
|
|
prompt_params={"document": pydantic.PARSING_DOCUMENT},
|
|
num_reasks=1,
|
|
)
|
|
|
|
assert final_output.validated_output == pydantic.PARSING_EXPECTED_OUTPUT
|
|
|
|
call = guard.history.first
|
|
|
|
# Check that the guard state object has the correct number of re-asks.
|
|
assert call.iterations.length == 2
|
|
|
|
# For orginal prompt and output
|
|
assert call.compiled_messages[0]["content"] == pydantic.PARSING_COMPILED_PROMPT
|
|
assert call.iterations.first.raw_output == pydantic.PARSING_UNPARSEABLE_LLM_OUTPUT
|
|
assert call.iterations.first.guarded_output is None
|
|
|
|
# For re-asked prompt and output
|
|
|
|
assert call.iterations.last.inputs.messages[1]["content"] == gd.Prompt(
|
|
pydantic.PARSING_COMPILED_REASK
|
|
)
|
|
# Same as above
|
|
assert call.reask_messages[0][1]["content"] == pydantic.PARSING_COMPILED_REASK
|
|
assert call.raw_outputs.last == pydantic.PARSING_EXPECTED_LLM_OUTPUT
|
|
assert call.guarded_output == pydantic.PARSING_EXPECTED_OUTPUT
|
|
|
|
|
|
def test_reask_prompt_instructions(mocker):
|
|
"""Test that the re-ask prompt and instructions are correct.
|
|
|
|
This is done implicitly, since if the incorrect prompt or
|
|
instructions are used, the mock LLM will raise a KeyError.
|
|
"""
|
|
|
|
mocker.patch(
|
|
"guardrails.llm_providers.LiteLLMCallable._invoke_llm",
|
|
return_value=LLMResponse(
|
|
output=string.MSG_LLM_OUTPUT_CORRECT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
),
|
|
)
|
|
|
|
@register_validator(name="always_fail", data_type="string")
|
|
def always_fail(value: str, metadata: Dict) -> ValidationResult:
|
|
return FailResult(error_message=f"Value {value} should fail.")
|
|
|
|
# We don't support tuple syntax for for_string and never have
|
|
# Once the validator function is decorated though, it becomes a Validator class
|
|
guard = gd.Guard.for_string(
|
|
validators=[always_fail(OnFailAction.REASK)],
|
|
description="Some description",
|
|
)
|
|
|
|
guard.parse(
|
|
llm_output="Tomato Cheese Pizza",
|
|
model="gpt-3.5-turbo",
|
|
messages=[
|
|
{"role": "system", "content": "Some content"},
|
|
{"role": "user", "content": "Some prompt"},
|
|
],
|
|
)
|