You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
langchain/tests/unit_tests/prompts/test_length_based_example_s...

49 lines
1.8 KiB
Python

"""Test functionality related to dynamic prompts."""
import pytest
from langchain.prompts.example_selector.length_based import LengthBasedExampleSelector
from langchain.prompts.prompt import PromptTemplate
EXAMPLES = [
{"question": "Question: who are you?\nAnswer: foo"},
{"question": "Question: who are you?\nAnswer: foo"},
]
@pytest.fixture
def selector() -> LengthBasedExampleSelector:
"""Get length based selector to use in tests."""
prompts = PromptTemplate(input_variables=["question"], template="{question}")
selector = LengthBasedExampleSelector(
examples=EXAMPLES,
example_prompt=prompts,
max_length=25,
)
return selector
def test_dynamic_prompt_valid(selector: LengthBasedExampleSelector) -> None:
"""Test dynamic prompt can be successfully constructed from examples."""
short_question = "Short question?"
output = selector.select_examples({"question": short_question})
assert output == EXAMPLES
def test_dynamic_prompt_trims_one_example(selector: LengthBasedExampleSelector) -> None:
"""Test dynamic prompt can trim one example."""
long_question = """I am writing a really long question,
this probably is going to affect the example right?"""
output = selector.select_examples({"question": long_question})
assert output == EXAMPLES[:1]
def test_dynamic_prompt_trims_all_examples(
selector: LengthBasedExampleSelector,
) -> None:
"""Test dynamic prompt can trim all examples."""
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?"""
output = selector.select_examples({"question": longest_question})
assert output == []