partial variables (#1308)

tests/unit_tests/prompts/test_few_shot.py

@@ -85,3 +85,92 @@ def test_few_shot_functionality() -> None:
         "Now you try to talk about party."
     )
     assert output == expected_output
+
+
+def test_partial_init_string() -> None:
+    """Test prompt can be initialized with partial variables."""
+    prefix = "This is a test about {content}."
+    suffix = "Now you try to talk about {new_content}."
+    examples = [
+        {"question": "foo", "answer": "bar"},
+        {"question": "baz", "answer": "foo"},
+    ]
+    prompt = FewShotPromptTemplate(
+        suffix=suffix,
+        prefix=prefix,
+        input_variables=["new_content"],
+        partial_variables={"content": "animals"},
+        examples=examples,
+        example_prompt=EXAMPLE_PROMPT,
+        example_separator="\n",
+    )
+    output = prompt.format(new_content="party")
+    expected_output = (
+        "This is a test about animals.\n"
+        "foo: bar\n"
+        "baz: foo\n"
+        "Now you try to talk about party."
+    )
+    assert output == expected_output
+
+
+def test_partial_init_func() -> None:
+    """Test prompt can be initialized with partial variables."""
+    prefix = "This is a test about {content}."
+    suffix = "Now you try to talk about {new_content}."
+    examples = [
+        {"question": "foo", "answer": "bar"},
+        {"question": "baz", "answer": "foo"},
+    ]
+    prompt = FewShotPromptTemplate(
+        suffix=suffix,
+        prefix=prefix,
+        input_variables=["new_content"],
+        partial_variables={"content": lambda: "animals"},
+        examples=examples,
+        example_prompt=EXAMPLE_PROMPT,
+        example_separator="\n",
+    )
+    output = prompt.format(new_content="party")
+    expected_output = (
+        "This is a test about animals.\n"
+        "foo: bar\n"
+        "baz: foo\n"
+        "Now you try to talk about party."
+    )
+    assert output == expected_output
+
+
+def test_partial() -> None:
+    """Test prompt can be partialed."""
+    prefix = "This is a test about {content}."
+    suffix = "Now you try to talk about {new_content}."
+    examples = [
+        {"question": "foo", "answer": "bar"},
+        {"question": "baz", "answer": "foo"},
+    ]
+    prompt = FewShotPromptTemplate(
+        suffix=suffix,
+        prefix=prefix,
+        input_variables=["content", "new_content"],
+        examples=examples,
+        example_prompt=EXAMPLE_PROMPT,
+        example_separator="\n",
+    )
+    new_prompt = prompt.partial(content="foo")
+    new_output = new_prompt.format(new_content="party")
+    expected_output = (
+        "This is a test about foo.\n"
+        "foo: bar\n"
+        "baz: foo\n"
+        "Now you try to talk about party."
+    )
+    assert new_output == expected_output
+    output = prompt.format(new_content="party", content="bar")
+    expected_output = (
+        "This is a test about bar.\n"
+        "foo: bar\n"
+        "baz: foo\n"
+        "Now you try to talk about party."
+    )
+    assert output == expected_output

tests/unit_tests/prompts/test_prompt.py

@@ -108,3 +108,40 @@ def test_prompt_from_file() -> None:
     input_variables = ["question"]
     prompt = PromptTemplate.from_file(template_file, input_variables)
     assert prompt.template == "Question: {question}\nAnswer:"
+
+
+def test_partial_init_string() -> None:
+    """Test prompt can be initialized with partial variables."""
+    template = "This is a {foo} test."
+    prompt = PromptTemplate(
+        input_variables=[], template=template, partial_variables={"foo": 1}
+    )
+    assert prompt.template == template
+    assert prompt.input_variables == []
+    result = prompt.format()
+    assert result == "This is a 1 test."
+
+
+def test_partial_init_func() -> None:
+    """Test prompt can be initialized with partial variables."""
+    template = "This is a {foo} test."
+    prompt = PromptTemplate(
+        input_variables=[], template=template, partial_variables={"foo": lambda: 2}
+    )
+    assert prompt.template == template
+    assert prompt.input_variables == []
+    result = prompt.format()
+    assert result == "This is a 2 test."
+
+
+def test_partial() -> None:
+    """Test prompt can be partialed."""
+    template = "This is a {foo} test."
+    prompt = PromptTemplate(input_variables=["foo"], template=template)
+    assert prompt.template == template
+    assert prompt.input_variables == ["foo"]
+    new_prompt = prompt.partial(foo="3")
+    new_result = new_prompt.format()
+    assert new_result == "This is a 3 test."
+    result = prompt.format(foo="foo")
+    assert result == "This is a foo test."