참고소스 수정본
This commit is contained in:
LASTA_DEV01\lasta
766
참고/guardrails-main/tests/integration_tests/test_streaming.py
Normal file
766
참고/guardrails-main/tests/integration_tests/test_streaming.py
Normal file
@@ -0,0 +1,766 @@
|
||||
# 3 tests
|
||||
# 1. Test streaming with LiteLLMCallable (mock openai.Completion.create)
|
||||
# 2. Test streaming with OpenAIChatCallable (mock openai.ChatCompletion.create)
|
||||
# 3. Test string schema streaming
|
||||
# Using the LowerCase Validator, and a custom validator to show new streaming behavior
|
||||
import json
|
||||
from typing import Any, Callable, Dict, Iterable, List, Optional, Union
|
||||
|
||||
import openai
|
||||
import pytest
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import guardrails as gd
|
||||
from guardrails.utils.casting_utils import to_int
|
||||
from guardrails.validator_base import (
|
||||
ErrorSpan,
|
||||
FailResult,
|
||||
OnFailAction,
|
||||
PassResult,
|
||||
ValidationResult,
|
||||
Validator,
|
||||
register_validator,
|
||||
)
|
||||
from tests.integration_tests.test_assets.validators import LowerCase, MockDetectPII
|
||||
|
||||
expected_raw_output = {"statement": "I am DOING well, and I HOPE you aRe too."}
|
||||
expected_fix_output = {"statement": "i am doing well, and i hope you are too."}
|
||||
expected_noop_output = {"statement": "I am DOING well, and I HOPE you aRe too."}
|
||||
expected_filter_refrain_output = {}
|
||||
|
||||
|
||||
@register_validator(name="minsentencelength", data_type=["string", "list"])
|
||||
class MinSentenceLengthValidator(Validator):
|
||||
def __init__(
|
||||
self,
|
||||
min: Optional[int] = None,
|
||||
max: Optional[int] = None,
|
||||
on_fail: Optional[Callable] = None,
|
||||
):
|
||||
super().__init__(
|
||||
on_fail=on_fail,
|
||||
min=min,
|
||||
max=max,
|
||||
)
|
||||
self._min = to_int(min)
|
||||
self._max = to_int(max)
|
||||
|
||||
def sentence_split(self, value):
|
||||
return list(map(lambda x: x + ".", value.split(".")[:-1]))
|
||||
|
||||
def validate(self, value: Union[str, List], metadata: Dict) -> ValidationResult:
|
||||
sentences = self.sentence_split(value)
|
||||
error_spans = []
|
||||
index = 0
|
||||
for sentence in sentences:
|
||||
if len(sentence) < self._min:
|
||||
error_spans.append(
|
||||
ErrorSpan(
|
||||
start=index,
|
||||
end=index + len(sentence),
|
||||
reason=f"Sentence has length less than {self._min}. "
|
||||
f"Please return a longer output, "
|
||||
f"that is shorter than {self._max} characters.",
|
||||
)
|
||||
)
|
||||
if len(sentence) > self._max:
|
||||
error_spans.append(
|
||||
ErrorSpan(
|
||||
start=index,
|
||||
end=index + len(sentence),
|
||||
reason=f"Sentence has length greater than {self._max}. "
|
||||
f"Please return a shorter output, "
|
||||
f"that is shorter than {self._max} characters.",
|
||||
)
|
||||
)
|
||||
index = index + len(sentence)
|
||||
if len(error_spans) > 0:
|
||||
return FailResult(
|
||||
validated_chunk=value,
|
||||
error_spans=error_spans,
|
||||
error_message=f"Sentence has length less than {self._min}. "
|
||||
f"Please return a longer output, "
|
||||
f"that is shorter than {self._max} characters.",
|
||||
)
|
||||
return PassResult(validated_chunk=value)
|
||||
|
||||
def validate_stream(self, chunk: Any, metadata: Dict, **kwargs) -> ValidationResult:
|
||||
return super().validate_stream(chunk, metadata, **kwargs)
|
||||
|
||||
|
||||
class Delta:
|
||||
content: str
|
||||
|
||||
def __init__(self, content):
|
||||
self.content = content
|
||||
|
||||
|
||||
class Choice:
|
||||
text: str
|
||||
finish_reason: str
|
||||
index: int
|
||||
delta: Delta
|
||||
|
||||
def __init__(self, text, delta, finish_reason, index=0):
|
||||
self.index = index
|
||||
self.delta = delta
|
||||
self.text = text
|
||||
self.finish_reason = finish_reason
|
||||
|
||||
|
||||
class MockOpenAIV1ChunkResponse:
|
||||
choices: list
|
||||
model: str
|
||||
|
||||
def __init__(self, choices, model):
|
||||
self.choices = choices
|
||||
self.model = model
|
||||
|
||||
|
||||
def mock_openai_completion_create(chunks):
|
||||
# Returns a generator
|
||||
def gen():
|
||||
for chunk in chunks:
|
||||
yield MockOpenAIV1ChunkResponse(
|
||||
choices=[
|
||||
Choice(
|
||||
text=chunk,
|
||||
delta=None,
|
||||
finish_reason=None,
|
||||
)
|
||||
],
|
||||
model="OpenAI model name",
|
||||
)
|
||||
|
||||
return gen()
|
||||
|
||||
|
||||
def mock_openai_chat_completion_create(chunks):
|
||||
# Returns a generator
|
||||
def gen():
|
||||
for chunk in chunks:
|
||||
yield MockOpenAIV1ChunkResponse(
|
||||
choices=[
|
||||
Choice(
|
||||
text=None,
|
||||
delta=Delta(content=chunk),
|
||||
finish_reason=None,
|
||||
)
|
||||
],
|
||||
model="OpenAI model name",
|
||||
)
|
||||
|
||||
return gen()
|
||||
|
||||
|
||||
class LowerCaseFix(BaseModel):
|
||||
statement: str = Field(
|
||||
description="Validates whether the text is in lower case.",
|
||||
validators=[LowerCase(on_fail=OnFailAction.FIX)],
|
||||
)
|
||||
|
||||
|
||||
class LowerCaseNoop(BaseModel):
|
||||
statement: str = Field(
|
||||
description="Validates whether the text is in lower case.",
|
||||
validators=[LowerCase(on_fail=OnFailAction.NOOP)],
|
||||
)
|
||||
|
||||
|
||||
class LowerCaseFilter(BaseModel):
|
||||
statement: str = Field(
|
||||
description="Validates whether the text is in lower case.",
|
||||
validators=[LowerCase(on_fail=OnFailAction.FILTER)],
|
||||
)
|
||||
|
||||
|
||||
class LowerCaseRefrain(BaseModel):
|
||||
statement: str = Field(
|
||||
description="Validates whether the text is in lower case.",
|
||||
validators=[LowerCase(on_fail=OnFailAction.REFRAIN)],
|
||||
)
|
||||
|
||||
|
||||
expected_minsentence_noop_output = ""
|
||||
|
||||
|
||||
class MinSentenceLengthNoOp(BaseModel):
|
||||
statement: str = Field(
|
||||
description="Validates whether the text is in lower case.",
|
||||
validators=[MinSentenceLengthValidator(on_fail=OnFailAction.NOOP)],
|
||||
)
|
||||
|
||||
|
||||
STR_PROMPT = "Say something nice to me."
|
||||
|
||||
PROMPT = """
|
||||
Say something nice to me.
|
||||
|
||||
${gr.complete_json_suffix}
|
||||
"""
|
||||
|
||||
JSON_LLM_CHUNKS = [
|
||||
'{"statement":',
|
||||
' "I am DOING',
|
||||
" well, and I",
|
||||
" HOPE you aRe",
|
||||
' too."}',
|
||||
]
|
||||
|
||||
MESSAGES = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": PROMPT,
|
||||
}
|
||||
]
|
||||
|
||||
STR_MESSAGES = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": STR_PROMPT,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"guard, expected_validated_output",
|
||||
[
|
||||
(
|
||||
gd.Guard.for_pydantic(output_class=LowerCaseNoop, messages=MESSAGES),
|
||||
expected_noop_output,
|
||||
),
|
||||
(
|
||||
gd.Guard.for_pydantic(output_class=LowerCaseFix, messages=MESSAGES),
|
||||
expected_fix_output,
|
||||
),
|
||||
(
|
||||
gd.Guard.for_pydantic(output_class=LowerCaseFilter, messages=MESSAGES),
|
||||
expected_filter_refrain_output,
|
||||
),
|
||||
(
|
||||
gd.Guard.for_pydantic(output_class=LowerCaseRefrain, messages=MESSAGES),
|
||||
expected_filter_refrain_output,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_streaming_with_openai_callable(
|
||||
mocker,
|
||||
guard,
|
||||
expected_validated_output,
|
||||
):
|
||||
"""Test streaming with LiteLLMCallable.
|
||||
|
||||
Mocks openai.Completion.create.
|
||||
"""
|
||||
mocker.patch(
|
||||
"openai.resources.Completions.create",
|
||||
return_value=mock_openai_completion_create(JSON_LLM_CHUNKS),
|
||||
)
|
||||
|
||||
method = openai.completions.create
|
||||
|
||||
method.__name__ = "mock_openai_completion_create"
|
||||
|
||||
generator = guard(
|
||||
method,
|
||||
engine="text-davinci-003",
|
||||
max_tokens=10,
|
||||
temperature=0,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert isinstance(generator, Iterable)
|
||||
|
||||
for op in generator:
|
||||
actual_output = op
|
||||
|
||||
assert actual_output.raw_llm_output == json.dumps(expected_raw_output)
|
||||
assert actual_output.validated_output == expected_validated_output
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"guard, expected_validated_output",
|
||||
[
|
||||
(
|
||||
gd.Guard.for_pydantic(output_class=LowerCaseNoop, messages=MESSAGES),
|
||||
expected_noop_output,
|
||||
),
|
||||
(
|
||||
gd.Guard.for_pydantic(output_class=LowerCaseFix, messages=MESSAGES),
|
||||
expected_fix_output,
|
||||
),
|
||||
(
|
||||
gd.Guard.for_pydantic(output_class=LowerCaseFilter, messages=MESSAGES),
|
||||
expected_filter_refrain_output,
|
||||
),
|
||||
(
|
||||
gd.Guard.for_pydantic(output_class=LowerCaseRefrain, messages=MESSAGES),
|
||||
expected_filter_refrain_output,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_streaming_with_openai_chat_callable(
|
||||
mocker,
|
||||
guard,
|
||||
expected_validated_output,
|
||||
):
|
||||
"""Test streaming with OpenAIChatCallable.
|
||||
|
||||
Mocks openai.ChatCompletion.create.
|
||||
"""
|
||||
mocker.patch(
|
||||
"openai.resources.chat.completions.Completions.create",
|
||||
return_value=mock_openai_chat_completion_create(JSON_LLM_CHUNKS),
|
||||
)
|
||||
|
||||
method = openai.chat.completions.create
|
||||
|
||||
method.__name__ = "mock_openai_chat_completion_create"
|
||||
|
||||
generator = guard(
|
||||
method,
|
||||
model="gpt-3.5-turbo",
|
||||
max_tokens=10,
|
||||
temperature=0,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert isinstance(generator, Iterable)
|
||||
|
||||
actual_output = ""
|
||||
for op in generator:
|
||||
actual_output = op
|
||||
|
||||
assert actual_output.raw_llm_output == json.dumps(expected_raw_output)
|
||||
assert actual_output.validated_output == expected_validated_output
|
||||
|
||||
|
||||
STR_LLM_CHUNKS = [
|
||||
# 38 characters
|
||||
"This sentence is simply just ",
|
||||
"too long.",
|
||||
# 25 characters long
|
||||
"This ",
|
||||
"sentence ",
|
||||
"is 2 ",
|
||||
"short."
|
||||
# 29 characters long
|
||||
"This sentence is just ",
|
||||
"right.",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"guard, expected_error_spans",
|
||||
[
|
||||
(
|
||||
gd.Guard.for_string(
|
||||
# only the middle sentence should pass
|
||||
validators=[
|
||||
MinSentenceLengthValidator(26, 30, on_fail=OnFailAction.NOOP)
|
||||
],
|
||||
messages=STR_MESSAGES,
|
||||
),
|
||||
# each value is a tuple
|
||||
# first is expected text inside span
|
||||
# second is the reason for failure
|
||||
[
|
||||
[
|
||||
"This sentence is simply just too long.",
|
||||
(
|
||||
"Sentence has length greater than 30. "
|
||||
"Please return a shorter output, "
|
||||
"that is shorter than 30 characters."
|
||||
),
|
||||
],
|
||||
[
|
||||
"This sentence is 2 short.",
|
||||
(
|
||||
"Sentence has length less than 26. "
|
||||
"Please return a longer output, "
|
||||
"that is shorter than 30 characters."
|
||||
),
|
||||
],
|
||||
],
|
||||
)
|
||||
],
|
||||
)
|
||||
def test_string_schema_streaming_with_openai_chat(mocker, guard, expected_error_spans):
|
||||
"""Test string schema streaming with OpenAIChatCallable.
|
||||
|
||||
Mocks openai.ChatCompletion.create.
|
||||
"""
|
||||
mocker.patch(
|
||||
"openai.resources.chat.completions.Completions.create",
|
||||
return_value=mock_openai_chat_completion_create(STR_LLM_CHUNKS),
|
||||
)
|
||||
|
||||
method = openai.chat.completions.create
|
||||
|
||||
method.__name__ = "mock_openai_chat_completion_create"
|
||||
generator = guard(
|
||||
method,
|
||||
model="gpt-3.5-turbo",
|
||||
max_tokens=10,
|
||||
temperature=0,
|
||||
stream=True,
|
||||
)
|
||||
|
||||
assert isinstance(generator, Iterable)
|
||||
|
||||
accumulated_output = ""
|
||||
for op in generator:
|
||||
accumulated_output += op.raw_llm_output
|
||||
error_spans = guard.error_spans_in_output()
|
||||
|
||||
assert len(error_spans) == len(expected_error_spans)
|
||||
for error_span, expected in zip(error_spans, expected_error_spans):
|
||||
assert accumulated_output[error_span.start : error_span.end] == expected[0]
|
||||
assert error_span.reason == expected[1]
|
||||
# TODO assert something about these error spans
|
||||
|
||||
|
||||
POETRY_CHUNKS = [
|
||||
'"John, under ',
|
||||
"GOLDEN bridges",
|
||||
", roams,\n",
|
||||
"SAN Francisco's ",
|
||||
"hills, his HOME.\n",
|
||||
"Dreams of",
|
||||
" FOG, and salty AIR,\n",
|
||||
"In his HEART",
|
||||
", he's always THERE.",
|
||||
]
|
||||
|
||||
|
||||
def test_noop_behavior_two_validators(mocker):
|
||||
mocker.patch(
|
||||
"openai.resources.chat.completions.Completions.create",
|
||||
return_value=mock_openai_chat_completion_create(POETRY_CHUNKS),
|
||||
)
|
||||
|
||||
guard = gd.Guard().use(
|
||||
MockDetectPII(
|
||||
on_fail=OnFailAction.NOOP,
|
||||
pii_entities="pii",
|
||||
replace_map={"John": "<PERSON>", "SAN Francisco's": "<LOCATION>"},
|
||||
),
|
||||
LowerCase(on_fail=OnFailAction.NOOP),
|
||||
)
|
||||
prompt = """Write me a 4 line poem about John in San Francisco.
|
||||
Make every third word all caps."""
|
||||
gen = guard(
|
||||
llm_api=openai.chat.completions.create,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
model="gpt-4",
|
||||
stream=True,
|
||||
)
|
||||
text = ""
|
||||
original = ""
|
||||
for res in gen:
|
||||
original = original + res.raw_llm_output
|
||||
text = text + res.validated_output
|
||||
assert (
|
||||
text
|
||||
== """"John, under GOLDEN bridges, roams,
|
||||
SAN Francisco's hills, his HOME.
|
||||
Dreams of FOG, and salty AIR,
|
||||
In his HEART, he's always THERE."""
|
||||
)
|
||||
assert (
|
||||
original
|
||||
== """"John, under GOLDEN bridges, roams,
|
||||
SAN Francisco's hills, his HOME.
|
||||
Dreams of FOG, and salty AIR,
|
||||
In his HEART, he's always THERE."""
|
||||
)
|
||||
|
||||
|
||||
def test_fix_behavior_one_validator(mocker):
|
||||
mocker.patch(
|
||||
"openai.resources.chat.completions.Completions.create",
|
||||
return_value=mock_openai_chat_completion_create(POETRY_CHUNKS),
|
||||
)
|
||||
|
||||
guard = gd.Guard().use(
|
||||
LowerCase(on_fail=OnFailAction.FIX),
|
||||
)
|
||||
prompt = """Write me a 4 line poem about John in San Francisco.
|
||||
Make every third word all caps."""
|
||||
gen = guard(
|
||||
llm_api=openai.chat.completions.create,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
model="gpt-4",
|
||||
stream=True,
|
||||
)
|
||||
text = ""
|
||||
original = ""
|
||||
for res in gen:
|
||||
original = original + res.raw_llm_output
|
||||
text = text + res.validated_output
|
||||
assert (
|
||||
text
|
||||
== """"john, under golden bridges, roams,
|
||||
san francisco's hills, his home.dreams of fog, and salty air,
|
||||
in his heart, he's always there."""
|
||||
)
|
||||
assert (
|
||||
original
|
||||
== """"John, under GOLDEN bridges, roams,
|
||||
SAN Francisco's hills, his HOME.
|
||||
Dreams of FOG, and salty AIR,
|
||||
In his HEART, he's always THERE."""
|
||||
)
|
||||
|
||||
|
||||
def test_fix_behavior_two_validators(mocker):
|
||||
mocker.patch(
|
||||
"openai.resources.chat.completions.Completions.create",
|
||||
return_value=mock_openai_chat_completion_create(POETRY_CHUNKS),
|
||||
)
|
||||
|
||||
guard = gd.Guard().use(
|
||||
MockDetectPII(
|
||||
on_fail=OnFailAction.FIX,
|
||||
pii_entities="pii",
|
||||
replace_map={"John": "<PERSON>", "SAN Francisco's": "<LOCATION>"},
|
||||
),
|
||||
LowerCase(on_fail=OnFailAction.FIX),
|
||||
)
|
||||
prompt = """Write me a 4 line poem about John in San Francisco.
|
||||
Make every third word all caps."""
|
||||
gen = guard(
|
||||
llm_api=openai.chat.completions.create,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
model="gpt-4",
|
||||
stream=True,
|
||||
)
|
||||
text = ""
|
||||
original = ""
|
||||
for res in gen:
|
||||
original = original + res.raw_llm_output
|
||||
text = text + res.validated_output
|
||||
assert (
|
||||
text
|
||||
== """"<PERSON>, under golden bridges, roams,
|
||||
<LOCATION> hills, his home.dreams of fog, and salty air,
|
||||
in his heart, he's always there."""
|
||||
)
|
||||
assert (
|
||||
original
|
||||
== """"John, under GOLDEN bridges, roams,
|
||||
SAN Francisco's hills, his HOME.
|
||||
Dreams of FOG, and salty AIR,
|
||||
In his HEART, he's always THERE."""
|
||||
)
|
||||
|
||||
|
||||
def test_fix_behavior_three_validators(mocker):
|
||||
mocker.patch(
|
||||
"openai.resources.chat.completions.Completions.create",
|
||||
return_value=mock_openai_chat_completion_create(POETRY_CHUNKS),
|
||||
)
|
||||
|
||||
guard = gd.Guard().use(
|
||||
MockDetectPII(
|
||||
on_fail=OnFailAction.FIX,
|
||||
pii_entities="pii",
|
||||
replace_map={"John": "<PERSON>", "SAN Francisco's": "<LOCATION>"},
|
||||
),
|
||||
LowerCase(on_fail=OnFailAction.FIX),
|
||||
# UpperCase(on_fail=OnFailAction.FIX),
|
||||
MockDetectPII(
|
||||
on_fail=OnFailAction.FIX,
|
||||
pii_entities="pii",
|
||||
replace_map={
|
||||
"John": "REDACTED!!",
|
||||
"SAN Francisco's": "REDACTED!!",
|
||||
"GOLDEN": "purple!!",
|
||||
},
|
||||
),
|
||||
)
|
||||
prompt = """Write me a 4 line poem about John in San Francisco.
|
||||
Make every third word all caps."""
|
||||
gen = guard(
|
||||
llm_api=openai.chat.completions.create,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
model="gpt-4",
|
||||
stream=True,
|
||||
)
|
||||
text = ""
|
||||
original = ""
|
||||
for res in gen:
|
||||
original = original + res.raw_llm_output
|
||||
text = text + res.validated_output
|
||||
print("FINAL TEXT", text)
|
||||
assert (
|
||||
text
|
||||
== """"REDACTED!!, under purple!! bridges, roams,
|
||||
<LOCATION> hills, his home.dreams of fog, and salty air,
|
||||
in his heart, he's always there."""
|
||||
)
|
||||
assert (
|
||||
original
|
||||
== """"John, under GOLDEN bridges, roams,
|
||||
SAN Francisco's hills, his HOME.
|
||||
Dreams of FOG, and salty AIR,
|
||||
In his HEART, he's always THERE."""
|
||||
)
|
||||
|
||||
|
||||
# This case does not work!
|
||||
# def test_fix_behavior_three_validators_overlap(mocker):
|
||||
# mocker.patch(
|
||||
# "openai.resources.chat.completions.Completions.create",
|
||||
# return_value=mock_openai_chat_completion_create(POETRY_CHUNKS),
|
||||
# )
|
||||
|
||||
# guard = gd.Guard().use(
|
||||
# MockDetectPII(
|
||||
# on_fail=OnFailAction.FIX,
|
||||
# pii_entities="pii",
|
||||
# replace_map={"John": "<PERSON>", "SAN Francisco's": "<LOCATION>"},
|
||||
# ),
|
||||
# LowerCase(on_fail=OnFailAction.FIX),
|
||||
# # UpperCase(on_fail=OnFailAction.FIX),
|
||||
# MockDetectPII(
|
||||
# on_fail=OnFailAction.FIX,
|
||||
# pii_entities="pii",
|
||||
# replace_map={
|
||||
# "John, under GOLDEN": "REDACTED!!",
|
||||
# "SAN Francisco's hills": "REDACTED!!",
|
||||
# "GOLDEN bridges": "gold!!!!",
|
||||
# },
|
||||
# ),
|
||||
# )
|
||||
# prompt = """Write me a 4 line poem about John in San Francisco.
|
||||
# Make every third word all caps."""
|
||||
# gen = guard(
|
||||
# llm_api=openai.chat.completions.create,
|
||||
# messages=MESSAGES,
|
||||
# model="gpt-4",
|
||||
# stream=True,
|
||||
# )
|
||||
# text = ""
|
||||
# original = ""
|
||||
# for res in gen:
|
||||
# original = original + res.raw_llm_output
|
||||
# text = text + res.validated_output
|
||||
# print("TEXT", text)
|
||||
# assert (
|
||||
# text
|
||||
# == """"REDACTED!!, under gold!!!! bridges, roams,
|
||||
# <LOCATION> hills, his home.
|
||||
# dreams of fog, and salty air,
|
||||
# in his heart, he's always there."""
|
||||
# )
|
||||
# assert (
|
||||
# original
|
||||
# == """"John, under GOLDEN bridges, roams,
|
||||
# SAN Francisco's hills, his HOME.
|
||||
# Dreams of FOG, and salty AIR,
|
||||
# In his HEART, he's always THERE."""
|
||||
# )
|
||||
|
||||
|
||||
def test_refrain_behavior(mocker):
|
||||
mocker.patch(
|
||||
"openai.resources.chat.completions.Completions.create",
|
||||
return_value=mock_openai_chat_completion_create(POETRY_CHUNKS),
|
||||
)
|
||||
|
||||
guard = gd.Guard().use(
|
||||
MockDetectPII(
|
||||
on_fail=OnFailAction.REFRAIN,
|
||||
pii_entities="pii",
|
||||
replace_map={"John": "<PERSON>", "SAN Francisco's": "<LOCATION>"},
|
||||
),
|
||||
LowerCase(on_fail=OnFailAction.FIX),
|
||||
)
|
||||
|
||||
prompt = """Write me a 4 line poem about John in San Francisco.
|
||||
Make every third word all caps."""
|
||||
gen = guard(
|
||||
llm_api=openai.chat.completions.create,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
model="gpt-4",
|
||||
stream=True,
|
||||
)
|
||||
text = ""
|
||||
original = ""
|
||||
for res in gen:
|
||||
original = original + res.raw_llm_output
|
||||
text = text + res.validated_output
|
||||
assert text == ""
|
||||
assert (
|
||||
original
|
||||
== """"John, under GOLDEN bridges, roams,
|
||||
SAN Francisco's hills, his HOME.
|
||||
"""
|
||||
)
|
||||
|
||||
|
||||
def test_filter_behavior(mocker):
|
||||
mocker.patch(
|
||||
"openai.resources.chat.completions.Completions.create",
|
||||
return_value=mock_openai_chat_completion_create(POETRY_CHUNKS),
|
||||
)
|
||||
|
||||
guard = gd.Guard().use(
|
||||
MockDetectPII(
|
||||
on_fail=OnFailAction.FIX,
|
||||
pii_entities="pii",
|
||||
replace_map={"John": "<PERSON>", "SAN Francisco's": "<LOCATION>"},
|
||||
),
|
||||
LowerCase(on_fail=OnFailAction.FILTER),
|
||||
)
|
||||
prompt = """Write me a 4 line poem about John in San Francisco.
|
||||
Make every third word all caps."""
|
||||
gen = guard(
|
||||
llm_api=openai.chat.completions.create,
|
||||
messages=[
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt,
|
||||
}
|
||||
],
|
||||
model="gpt-4",
|
||||
stream=True,
|
||||
)
|
||||
text = ""
|
||||
original = ""
|
||||
for res in gen:
|
||||
original = original + res.raw_llm_output
|
||||
text = text + res.validated_output
|
||||
assert text == ""
|
||||
assert (
|
||||
original
|
||||
== """"John, under GOLDEN bridges, roams,
|
||||
SAN Francisco's hills, his HOME.
|
||||
"""
|
||||
)
|
||||
Reference in New Issue
Block a user