import langchain_visualizer  # isort:skip  # noqa: F401

from fvalues import FValue
from langchain import FewShotPromptTemplate, PromptTemplate


def test_few_shot_f():
    examples = [
        {"word": "happy", "antonym": "sad"},
        {"word": "tall", "antonym": "short"},
        # Should be able to handle extra keys that is not exists in input_variables
        {"word": "better", "antonym": "worse", "extra": "extra"},
    ]

    example_prompt = PromptTemplate(
        input_variables=["word", "antonym"],
        template="w={word},a={antonym}",
    )

    few_shot_prompt = FewShotPromptTemplate(
        examples=examples,
        example_prompt=example_prompt,
        prefix="Give the antonym of every input:",
        suffix="w={input},a=",
        input_variables=["input"],
        example_separator="  ",
    )

    s = few_shot_prompt.format(input="big")
    assert s == (
        "Give the antonym of every input:  "
        "w=happy,a=sad  w=tall,a=short  w=better,a=worse  w=big,a="
    )
    print([repr(x) for x in s.flatten().parts])
    assert s.flatten().parts == (
        "Give the antonym of every input:",
        FValue(source="self.example_separator", value="  ", formatted="  "),
        "w=",
        FValue(source="word", value="happy", formatted="happy"),
        ",a=",
        FValue(source="antonym", value="sad", formatted="sad"),
        FValue(source="self.example_separator", value="  ", formatted="  "),
        "w=",
        FValue(source="word", value="tall", formatted="tall"),
        ",a=",
        FValue(source="antonym", value="short", formatted="short"),
        FValue(source="self.example_separator", value="  ", formatted="  "),
        "w=",
        FValue(source="word", value="better", formatted="better"),
        ",a=",
        FValue(source="antonym", value="worse", formatted="worse"),
        FValue(source="self.example_separator", value="  ", formatted="  "),
        "w=",
        FValue(source="input", value="big", formatted="big"),
        ",a=",
    )

langchain.FewShotPromptTemplate langchain.PromptTemplate
[((455, 544), 'langchain.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['word', 'antonym']", 'template': '"""w={word},a={antonym}"""'}), "(input_variables=['word', 'antonym'], template=\n    'w={word},a={antonym}')\n", (469, 544), False, 'from langchain import FewShotPromptTemplate, PromptTemplate\n'), ((586, 782), 'langchain.FewShotPromptTemplate', 'FewShotPromptTemplate', ([], {'examples': 'examples', 'example_prompt': 'example_prompt', 'prefix': '"""Give the antonym of every input:"""', 'suffix': '"""w={input},a="""', 'input_variables': "['input']", 'example_separator': '"""  """'}), "(examples=examples, example_prompt=example_prompt,\n    prefix='Give the antonym of every input:', suffix='w={input},a=',\n    input_variables=['input'], example_separator='  ')\n", (607, 782), False, 'from langchain import FewShotPromptTemplate, PromptTemplate\n'), ((1146, 1213), 'fvalues.FValue', 'FValue', ([], {'source': '"""self.example_separator"""', 'value': '"""  """', 'formatted': '"""  """'}), "(source='self.example_separator', value='  ', formatted='  ')\n", (1152, 1213), False, 'from fvalues import FValue\n'), ((1237, 1292), 'fvalues.FValue', 'FValue', ([], {'source': '"""word"""', 'value': '"""happy"""', 'formatted': '"""happy"""'}), "(source='word', value='happy', formatted='happy')\n", (1243, 1292), False, 'from fvalues import FValue\n'), ((1317, 1371), 'fvalues.FValue', 'FValue', ([], {'source': '"""antonym"""', 'value': '"""sad"""', 'formatted': '"""sad"""'}), "(source='antonym', value='sad', formatted='sad')\n", (1323, 1371), False, 'from fvalues import FValue\n'), ((1381, 1448), 'fvalues.FValue', 'FValue', ([], {'source': '"""self.example_separator"""', 'value': '"""  """', 'formatted': '"""  """'}), "(source='self.example_separator', value='  ', formatted='  ')\n", (1387, 1448), False, 'from fvalues import FValue\n'), ((1472, 1525), 'fvalues.FValue', 'FValue', ([], {'source': '"""word"""', 'value': '"""tall"""', 'formatted': '"""tall"""'}), "(source='word', value='tall', formatted='tall')\n", (1478, 1525), False, 'from fvalues import FValue\n'), ((1550, 1608), 'fvalues.FValue', 'FValue', ([], {'source': '"""antonym"""', 'value': '"""short"""', 'formatted': '"""short"""'}), "(source='antonym', value='short', formatted='short')\n", (1556, 1608), False, 'from fvalues import FValue\n'), ((1618, 1685), 'fvalues.FValue', 'FValue', ([], {'source': '"""self.example_separator"""', 'value': '"""  """', 'formatted': '"""  """'}), "(source='self.example_separator', value='  ', formatted='  ')\n", (1624, 1685), False, 'from fvalues import FValue\n'), ((1709, 1766), 'fvalues.FValue', 'FValue', ([], {'source': '"""word"""', 'value': '"""better"""', 'formatted': '"""better"""'}), "(source='word', value='better', formatted='better')\n", (1715, 1766), False, 'from fvalues import FValue\n'), ((1791, 1849), 'fvalues.FValue', 'FValue', ([], {'source': '"""antonym"""', 'value': '"""worse"""', 'formatted': '"""worse"""'}), "(source='antonym', value='worse', formatted='worse')\n", (1797, 1849), False, 'from fvalues import FValue\n'), ((1859, 1926), 'fvalues.FValue', 'FValue', ([], {'source': '"""self.example_separator"""', 'value': '"""  """', 'formatted': '"""  """'}), "(source='self.example_separator', value='  ', formatted='  ')\n", (1865, 1926), False, 'from fvalues import FValue\n'), ((1950, 2002), 'fvalues.FValue', 'FValue', ([], {'source': '"""input"""', 'value': '"""big"""', 'formatted': '"""big"""'}), "(source='input', value='big', formatted='big')\n", (1956, 2002), False, 'from fvalues import FValue\n')]