kork/tests/test_prompt_adapter.py

from langchain.prompts import PromptTemplate

from kork.prompt_adapter import FewShotPromptValue, FewShotTemplate


def test_few_shot_template() -> None:
    """Test few shot template."""
    prompt_template = PromptTemplate(
        template="meow\n\n",
        input_variables=[],
    )
    few_shot_template = FewShotTemplate(
        instruction_template=prompt_template,
        examples=[("foo", "bar")],
        input_variables=["query"],
    )

    prompt_value = few_shot_template.format_prompt(query="query")

    assert isinstance(prompt_value, FewShotPromptValue)
    # Test to_string()
    assert (
        prompt_value.to_string()
        == "meow\n\nInput: foo\n\nOutput: bar\n\nInput: query\n\nOutput:"
    )

    # Test to_messages()
    messages = prompt_value.to_messages()
    assert len(messages) == 4
    assert messages[0].content == "meow"
    assert messages[1].content == "foo"
    assert messages[2].content == "bar"
    assert messages[3].content == "query"


def test_few_shot_template_with_no_examples() -> None:
    """Test few shot template without any examples"""
    prompt_template = PromptTemplate(
        template="meow\n\n",
        input_variables=[],
    )
    few_shot_template = FewShotTemplate(
        instruction_template=prompt_template,
        examples=[],
        input_variables=["query"],
    )

    prompt_value = few_shot_template.format_prompt(query="query")

    assert isinstance(prompt_value, FewShotPromptValue)
    # Test to_string()
    assert prompt_value.to_string() == "meow\n\nInput: query\n\nOutput:"

    # Test to_messages()
    messages = prompt_value.to_messages()
    assert len(messages) == 2
    assert messages[0].content == "meow"
    assert messages[1].content == "query"