test_few_shot_chat_message_prompt_template_with_selector() — langchain Function Reference

Source Code

libs/core/tests/unit_tests/prompts/test_few_shot.py lines 409–442

def test_few_shot_chat_message_prompt_template_with_selector() -> None:
    """Tests for few shot chat message template with an example selector."""
    examples = [
        {"input": "2+2", "output": "4"},
        {"input": "2+3", "output": "5"},
    ]
    example_selector = AsIsSelector(examples)
    example_prompt = ChatPromptTemplate.from_messages(
        [
            HumanMessagePromptTemplate.from_template("{input}"),
            AIMessagePromptTemplate.from_template("{output}"),
        ]
    )

    few_shot_prompt = FewShotChatMessagePromptTemplate(
        input_variables=["input"],
        example_prompt=example_prompt,
        example_selector=example_selector,
    )
    final_prompt: ChatPromptTemplate = (
        SystemMessagePromptTemplate.from_template("You are a helpful AI Assistant")
        + few_shot_prompt
        + HumanMessagePromptTemplate.from_template("{input}")
    )
    expected = [
        SystemMessage(content="You are a helpful AI Assistant", additional_kwargs={}),
        HumanMessage(content="2+2", additional_kwargs={}),
        AIMessage(content="4", additional_kwargs={}),
        HumanMessage(content="2+3", additional_kwargs={}),
        AIMessage(content="5", additional_kwargs={}),
        HumanMessage(content="100 + 1", additional_kwargs={}),
    ]
    messages = final_prompt.format_messages(input="100 + 1")
    assert messages == expected