361 lines
12 KiB
Python
361 lines
12 KiB
Python
import pytest
|
|
|
|
from guardrails import AsyncGuard
|
|
from guardrails.utils import docs_utils
|
|
from guardrails.classes.llm.llm_response import LLMResponse
|
|
from tests.integration_tests.test_assets.custom_llm import mock_async_llm
|
|
from tests.integration_tests.test_assets.fixtures import ( # noqa
|
|
fixture_llm_output,
|
|
fixture_rail_spec,
|
|
fixture_validated_output,
|
|
)
|
|
|
|
from .mock_llm_outputs import entity_extraction
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_entity_extraction_with_reask(mocker):
|
|
"""Test that the entity extraction works with re-asking."""
|
|
mock_invoke_llm = mocker.patch(
|
|
"guardrails.llm_providers.AsyncArbitraryCallable.invoke_llm",
|
|
)
|
|
mock_invoke_llm.side_effect = [
|
|
LLMResponse(
|
|
output=entity_extraction.LLM_OUTPUT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
),
|
|
LLMResponse(
|
|
# TODO: Re-enable once field level reasking is supported
|
|
# output=entity_extraction.LLM_OUTPUT_REASK,
|
|
output=entity_extraction.LLM_OUTPUT_FULL_REASK,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
),
|
|
]
|
|
|
|
content = docs_utils.read_pdf("docs/examples/data/chase_card_agreement.pdf")
|
|
guard = AsyncGuard.for_rail_string(entity_extraction.RAIL_SPEC_WITH_REASK)
|
|
|
|
final_output = await guard(
|
|
llm_api=mock_async_llm,
|
|
prompt_params={"document": content[:6000]},
|
|
num_reasks=1,
|
|
)
|
|
|
|
# Assertions are made on the guard state object.
|
|
assert final_output.validation_passed is True
|
|
assert final_output.validated_output == entity_extraction.VALIDATED_OUTPUT_REASK_2
|
|
|
|
guard_history = guard.history
|
|
call = guard_history.first
|
|
|
|
# Check that the guard was only called once and
|
|
# has the correct number of re-asks.
|
|
assert guard_history.length == 1
|
|
assert call.iterations.length == 2
|
|
|
|
# For orginal prompt and output
|
|
first = call.iterations.first
|
|
assert (
|
|
first.inputs.messages[0]["content"]._source
|
|
== entity_extraction.NON_OPENAI_COMPILED_PROMPT
|
|
)
|
|
# Same as above
|
|
assert (
|
|
call.compiled_messages[0]["content"]
|
|
== entity_extraction.NON_OPENAI_COMPILED_PROMPT
|
|
)
|
|
assert first.prompt_tokens_consumed == 123
|
|
assert first.completion_tokens_consumed == 1234
|
|
assert first.raw_output == entity_extraction.LLM_OUTPUT
|
|
assert first.validation_response == entity_extraction.VALIDATED_OUTPUT_REASK_1
|
|
|
|
# For re-asked prompt and output
|
|
final = call.iterations.last
|
|
assert (
|
|
final.inputs.messages[1]["content"]._source
|
|
== entity_extraction.NON_OPENAI_COMPILED_PROMPT_REASK
|
|
)
|
|
# Same as above
|
|
assert (
|
|
call.reask_messages[0][1]["content"]
|
|
== entity_extraction.NON_OPENAI_COMPILED_PROMPT_REASK
|
|
)
|
|
|
|
# TODO: Re-enable once field level reasking is supported
|
|
# assert final.raw_output == entity_extraction.LLM_OUTPUT_REASK
|
|
assert final.raw_output == entity_extraction.LLM_OUTPUT_FULL_REASK
|
|
assert call.guarded_output == entity_extraction.VALIDATED_OUTPUT_REASK_2
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_entity_extraction_with_noop(mocker):
|
|
mock_invoke_llm = mocker.patch(
|
|
"guardrails.llm_providers.AsyncArbitraryCallable.invoke_llm",
|
|
)
|
|
mock_invoke_llm.side_effect = [
|
|
LLMResponse(
|
|
output=entity_extraction.LLM_OUTPUT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
)
|
|
]
|
|
content = docs_utils.read_pdf("docs/examples/data/chase_card_agreement.pdf")
|
|
guard = AsyncGuard.for_rail_string(entity_extraction.RAIL_SPEC_WITH_NOOP)
|
|
final_output = await guard(
|
|
llm_api=mock_async_llm,
|
|
prompt_params={"document": content[:6000]},
|
|
num_reasks=1,
|
|
)
|
|
|
|
# Assertions are made on the guard state object.
|
|
|
|
# Old assertion which is wrong
|
|
# This should not pass validation and therefore will not have a validated output
|
|
# assert final_output.validated_output == entity_extraction.VALIDATED_OUTPUT_NOOP
|
|
|
|
assert final_output.validation_passed is False
|
|
assert final_output.validated_output is not None
|
|
assert final_output.validated_output["fees"]
|
|
assert final_output.validated_output["interest_rates"]
|
|
|
|
call = guard.history.first
|
|
|
|
# Check that the guard was called once
|
|
# and did not have to reask
|
|
assert guard.history.length == 1
|
|
assert call.iterations.length == 1
|
|
|
|
# For orginal prompt and output
|
|
assert (
|
|
call.compiled_messages[0]["content"]
|
|
== entity_extraction.NON_OPENAI_COMPILED_PROMPT
|
|
)
|
|
assert call.raw_outputs.last == entity_extraction.LLM_OUTPUT
|
|
assert call.validation_response == entity_extraction.VALIDATED_OUTPUT_NOOP
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_entity_extraction_with_noop_pydantic(mocker):
|
|
mock_invoke_llm = mocker.patch(
|
|
"guardrails.llm_providers.AsyncArbitraryCallable.invoke_llm",
|
|
)
|
|
mock_invoke_llm.side_effect = [
|
|
LLMResponse(
|
|
output=entity_extraction.LLM_OUTPUT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
)
|
|
]
|
|
content = docs_utils.read_pdf("docs/examples/data/chase_card_agreement.pdf")
|
|
guard = AsyncGuard.for_pydantic(
|
|
entity_extraction.PYDANTIC_RAIL_WITH_NOOP,
|
|
messages=[
|
|
{
|
|
"role": "user",
|
|
"content": entity_extraction.PYDANTIC_PROMPT,
|
|
}
|
|
],
|
|
)
|
|
final_output = await guard(
|
|
llm_api=mock_async_llm,
|
|
prompt_params={"document": content[:6000]},
|
|
num_reasks=1,
|
|
)
|
|
|
|
# Assertions are made on the guard state object.
|
|
assert final_output.validation_passed is False
|
|
assert final_output.validated_output is not None
|
|
assert final_output.validated_output["fees"]
|
|
assert final_output.validated_output["interest_rates"]
|
|
|
|
call = guard.history.first
|
|
|
|
# Check that the guard was called once
|
|
# and did not have toreask
|
|
assert guard.history.length == 1
|
|
assert call.iterations.length == 1
|
|
|
|
# For orginal prompt and output
|
|
assert (
|
|
call.compiled_messages[0]["content"]
|
|
== entity_extraction.NON_OPENAI_COMPILED_PROMPT
|
|
)
|
|
assert call.raw_outputs.last == entity_extraction.LLM_OUTPUT
|
|
assert call.validation_response == entity_extraction.VALIDATED_OUTPUT_NOOP
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_entity_extraction_with_filter(mocker):
|
|
"""Test that the entity extraction works with re-asking."""
|
|
mock_invoke_llm = mocker.patch(
|
|
"guardrails.llm_providers.AsyncArbitraryCallable.invoke_llm",
|
|
)
|
|
mock_invoke_llm.side_effect = [
|
|
LLMResponse(
|
|
output=entity_extraction.LLM_OUTPUT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
)
|
|
]
|
|
|
|
content = docs_utils.read_pdf("docs/examples/data/chase_card_agreement.pdf")
|
|
guard = AsyncGuard.for_rail_string(entity_extraction.RAIL_SPEC_WITH_FILTER)
|
|
final_output = await guard(
|
|
llm_api=mock_async_llm,
|
|
prompt_params={"document": content[:6000]},
|
|
num_reasks=1,
|
|
)
|
|
|
|
# Assertions are made on the guard state object.
|
|
assert final_output.validation_passed is False
|
|
assert final_output.validated_output is None
|
|
|
|
call = guard.history.first
|
|
|
|
# Check that the guard state object has the correct number of re-asks.
|
|
assert guard.history.length == 1
|
|
assert call.iterations.length == 1
|
|
|
|
# For orginal prompt and output
|
|
assert (
|
|
call.compiled_messages[0]["content"]
|
|
== entity_extraction.NON_OPENAI_COMPILED_PROMPT
|
|
)
|
|
assert call.raw_outputs.last == entity_extraction.LLM_OUTPUT
|
|
assert call.validation_response == entity_extraction.VALIDATED_OUTPUT_FILTER
|
|
assert call.guarded_output is None
|
|
assert call.status == "fail"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_entity_extraction_with_fix(mocker):
|
|
"""Test that the entity extraction works with re-asking."""
|
|
mock_invoke_llm = mocker.patch(
|
|
"guardrails.llm_providers.AsyncArbitraryCallable.invoke_llm",
|
|
)
|
|
mock_invoke_llm.side_effect = [
|
|
LLMResponse(
|
|
output=entity_extraction.LLM_OUTPUT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
)
|
|
]
|
|
|
|
content = docs_utils.read_pdf("docs/examples/data/chase_card_agreement.pdf")
|
|
guard = AsyncGuard.for_rail_string(entity_extraction.RAIL_SPEC_WITH_FIX)
|
|
final_output = await guard(
|
|
llm_api=mock_async_llm,
|
|
prompt_params={"document": content[:6000]},
|
|
num_reasks=1,
|
|
)
|
|
|
|
# Assertions are made on the guard state object.
|
|
assert final_output.validation_passed is True
|
|
assert final_output.validated_output == entity_extraction.VALIDATED_OUTPUT_FIX
|
|
|
|
call = guard.history.first
|
|
|
|
# Check that the guard state object has the correct number of re-asks.
|
|
assert guard.history.length == 1
|
|
|
|
# For orginal prompt and output
|
|
assert (
|
|
call.compiled_messages[0]["content"]
|
|
== entity_extraction.NON_OPENAI_COMPILED_PROMPT
|
|
)
|
|
assert call.raw_outputs.last == entity_extraction.LLM_OUTPUT
|
|
assert call.guarded_output == entity_extraction.VALIDATED_OUTPUT_FIX
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_entity_extraction_with_refrain(mocker):
|
|
"""Test that the entity extraction works with re-asking."""
|
|
mock_invoke_llm = mocker.patch(
|
|
"guardrails.llm_providers.AsyncArbitraryCallable.invoke_llm",
|
|
)
|
|
mock_invoke_llm.side_effect = [
|
|
LLMResponse(
|
|
output=entity_extraction.LLM_OUTPUT,
|
|
prompt_token_count=123,
|
|
response_token_count=1234,
|
|
)
|
|
]
|
|
|
|
content = docs_utils.read_pdf("docs/examples/data/chase_card_agreement.pdf")
|
|
guard = AsyncGuard.for_rail_string(entity_extraction.RAIL_SPEC_WITH_REFRAIN)
|
|
final_output = await guard(
|
|
llm_api=mock_async_llm,
|
|
prompt_params={"document": content[:6000]},
|
|
num_reasks=1,
|
|
)
|
|
# Assertions are made on the guard state object.
|
|
|
|
assert final_output.validation_passed is False
|
|
assert final_output.validated_output == entity_extraction.VALIDATED_OUTPUT_REFRAIN
|
|
|
|
call = guard.history.first
|
|
|
|
# Check that the guard state object has the correct number of re-asks.
|
|
assert guard.history.length == 1
|
|
|
|
# For orginal prompt and output
|
|
assert (
|
|
call.compiled_messages[0]["content"]
|
|
== entity_extraction.NON_OPENAI_COMPILED_PROMPT
|
|
)
|
|
assert call.raw_outputs.last == entity_extraction.LLM_OUTPUT
|
|
assert call.guarded_output == entity_extraction.VALIDATED_OUTPUT_REFRAIN
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_rail_spec_output_parse(rail_spec, llm_output, validated_output):
|
|
"""Test that the rail_spec fixture is working."""
|
|
guard = AsyncGuard.for_rail_string(rail_spec)
|
|
output = await guard.parse(
|
|
llm_output,
|
|
llm_api=mock_async_llm,
|
|
)
|
|
assert output.validated_output == validated_output
|
|
|
|
|
|
@pytest.fixture
|
|
def string_rail_spec():
|
|
return """
|
|
<rail version="0.1">
|
|
<output
|
|
type="string"
|
|
validators="two-words"
|
|
on-fail-two-words="fix"
|
|
/>
|
|
<prompt>
|
|
Hi please make me a string
|
|
</prompt>
|
|
</rail>
|
|
"""
|
|
|
|
|
|
@pytest.fixture
|
|
def string_llm_output():
|
|
return "string output yes"
|
|
|
|
|
|
@pytest.fixture
|
|
def validated_string_output():
|
|
return "string output"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_string_rail_spec_output_parse(
|
|
string_rail_spec, string_llm_output, validated_string_output
|
|
):
|
|
"""Test that the string_rail_spec fixture is working."""
|
|
guard: AsyncGuard = AsyncGuard.for_rail_string(string_rail_spec)
|
|
output = await guard.parse(
|
|
string_llm_output,
|
|
llm_api=mock_async_llm,
|
|
num_reasks=0,
|
|
)
|
|
assert output.validated_output == validated_string_output
|