core[minor]: Add aformat to FewShotPromptTemplate (#19652)

libs/core/tests/unit_tests/prompts/test_few_shot.py

@@ -96,7 +96,7 @@ def test_prompt_missing_input_variables() -> None:
     ).input_variables == ["foo"]
 
 
-def test_few_shot_functionality() -> None:
+async def test_few_shot_functionality() -> None:
     """Test that few shot works with examples."""
     prefix = "This is a test about {content}."
     suffix = "Now you try to talk about {new_content}."
@@ -112,13 +112,15 @@ def test_few_shot_functionality() -> None:
         example_prompt=EXAMPLE_PROMPT,
         example_separator="\n",
     )
-    output = prompt.format(content="animals", 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."
     )
+    output = prompt.format(content="animals", new_content="party")
+    assert output == expected_output
+    output = await prompt.aformat(content="animals", new_content="party")
     assert output == expected_output
 
 
@@ -365,6 +367,24 @@ class AsIsSelector(BaseExampleSelector):
         return list(self.examples)
 
 
+def test_few_shot_prompt_template_with_selector() -> None:
+    """Tests for few shot chat message template with an example selector."""
+    examples = [
+        {"question": "foo", "answer": "bar"},
+        {"question": "baz", "answer": "foo"},
+    ]
+    example_selector = AsIsSelector(examples)
+
+    few_shot_prompt = FewShotPromptTemplate(
+        input_variables=["foo"],
+        suffix="This is a {foo} test.",
+        example_prompt=EXAMPLE_PROMPT,
+        example_selector=example_selector,
+    )
+    messages = few_shot_prompt.format(foo="bar")
+    assert messages == "foo: bar\n\nbaz: foo\n\nThis is a bar test."
+
+
 def test_few_shot_chat_message_prompt_template_with_selector() -> None:
     """Tests for few shot chat message template with an example selector."""
     examples = [
@@ -421,6 +441,24 @@ class AsyncAsIsSelector(BaseExampleSelector):
         return list(self.examples)
 
 
+async def test_few_shot_prompt_template_with_selector_async() -> None:
+    """Tests for few shot chat message template with an example selector."""
+    examples = [
+        {"question": "foo", "answer": "bar"},
+        {"question": "baz", "answer": "foo"},
+    ]
+    example_selector = AsyncAsIsSelector(examples)
+
+    few_shot_prompt = FewShotPromptTemplate(
+        input_variables=["foo"],
+        suffix="This is a {foo} test.",
+        example_prompt=EXAMPLE_PROMPT,
+        example_selector=example_selector,
+    )
+    messages = await few_shot_prompt.aformat(foo="bar")
+    assert messages == "foo: bar\n\nbaz: foo\n\nThis is a bar test."
+
+
 async def test_few_shot_chat_message_prompt_template_with_selector_async() -> None:
     """Tests for few shot chat message template with an async example selector."""
     examples = [