langchain/tests/unit_tests/test_dynamic_prompt.py

112 lines
3.5 KiB
Python

"""Test functionality related to dynamic prompts."""
from langchain.prompts.dynamic import DynamicPrompt
from langchain.prompts.prompt import Prompt
# FULL TEMPLATES
LONGER_TEMPLATE = """Test Prompt:
Question: who are you?
Answer: foo
Question: what are you?
Answer: bar
Question: {question}
Answer:"""
SHORTER_TEMPLATE = """Test Prompt:
Question: who are you?
Answer: foo
Question: {question}
Answer:"""
SHORTEST_TEMPLATE = """Test Prompt:
Question: {question}
Answer:"""
# DYNAMIC PROMPT COMPONENTS
PREFIX = """Test Prompt:"""
SUFFIX = """Question: {question}\nAnswer:"""
EXAMPLES = [
"""Question: who are you?\nAnswer: foo""",
"""Question: what are you?\nAnswer: bar""",
]
# INPUTS
TEST_LONG_QUESTION = """I am writing a really long question,
this probably is going to affect the example right?"""
TEST_LONGEST_QUESTION = """This question is super super super,
super super super super super super super super super super super,
super super super super long, this will affect the example right?"""
TEST_SHORT_QUESTION = "Short question?"
def test_dynamic_prompt_valid() -> None:
"""Test dynamic prompt can be successfully constructed from examples."""
input_variables = ["question"]
example_separator = "\n\n"
dynamic_prompt_cls = DynamicPrompt(
examples=EXAMPLES,
suffix=SUFFIX,
input_variables=input_variables,
example_separator=example_separator,
prefix=PREFIX,
)
prompt_cls = Prompt(input_variables=input_variables, template=LONGER_TEMPLATE)
dynamic_prompt_template = dynamic_prompt_cls.format(question="foo?")
prompt_template = prompt_cls.format(question="foo?")
assert dynamic_prompt_template == prompt_template
assert dynamic_prompt_cls.input_variables == prompt_cls.input_variables
def test_dynamic_prompt_trims_one_example() -> None:
"""Test dynamic prompt can trim one example."""
input_variables = ["question"]
example_separator = "\n\n"
dynamic_prompt_cls = DynamicPrompt(
examples=EXAMPLES,
suffix=SUFFIX,
input_variables=input_variables,
example_separator=example_separator,
prefix=PREFIX,
max_length=30,
)
dynamic_prompt = dynamic_prompt_cls.format(question=TEST_LONG_QUESTION)
shorter_prompt = SHORTER_TEMPLATE.format(question=TEST_LONG_QUESTION)
assert dynamic_prompt == shorter_prompt
def test_dynamic_prompt_trims_no_examples() -> None:
"""Test dynamic prompt can trim no examples."""
input_variables = ["question"]
example_separator = "\n\n"
dynamic_prompt_cls = DynamicPrompt(
examples=EXAMPLES,
suffix=SUFFIX,
input_variables=input_variables,
example_separator=example_separator,
prefix=PREFIX,
max_length=30,
)
dynamic_prompt = dynamic_prompt_cls.format(question=TEST_SHORT_QUESTION)
full_prompt = LONGER_TEMPLATE.format(question=TEST_SHORT_QUESTION)
assert dynamic_prompt == full_prompt
def test_dynamic_prompt_trims_all_examples() -> None:
"""Test dynamic prompt can trim all examples."""
input_variables = ["question"]
example_separator = "\n\n"
dynamic_prompt_cls = DynamicPrompt(
examples=EXAMPLES,
suffix=SUFFIX,
input_variables=input_variables,
example_separator=example_separator,
prefix=PREFIX,
max_length=30,
)
dynamic_prompt = dynamic_prompt_cls.format(question=TEST_LONGEST_QUESTION)
full_prompt = SHORTEST_TEMPLATE.format(question=TEST_LONGEST_QUESTION)
assert dynamic_prompt == full_prompt