From 68298cdc828037ad722b725674cdd3ce8d35ef5e Mon Sep 17 00:00:00 2001
From: William FH <13333726+hinthornw@users.noreply.github.com>
Date: Wed, 20 Mar 2024 09:59:32 -0700
Subject: [PATCH] [Feat] Accept non-dict if only 1 prompt input variable
 (#19156)

For prompt templates with only 1 variable (common in e.g.,
MessageGraph), it's convenient to wrap the incoming object in the
variable before formatting.


The downside of this, of course, would be that some number of
invocations will successfully format when the user may have intended to
format it properly before
---
 libs/core/langchain_core/prompts/base.py      | 13 +++--
 libs/core/langchain_core/prompts/chat.py      | 48 +++++++++++++++++--
 .../tests/unit_tests/prompts/test_chat.py     | 13 +++++
 3 files changed, 66 insertions(+), 8 deletions(-)

diff --git a/libs/core/langchain_core/prompts/base.py b/libs/core/langchain_core/prompts/base.py
index 9f8ae5082f9..e1353937dd7 100644
--- a/libs/core/langchain_core/prompts/base.py
+++ b/libs/core/langchain_core/prompts/base.py
@@ -89,10 +89,15 @@ class BasePromptTemplate(
 
     def _format_prompt_with_error_handling(self, inner_input: Dict) -> PromptValue:
         if not isinstance(inner_input, dict):
-            raise TypeError(
-                f"Expected mapping type as input to {self.__class__.__name__}. "
-                f"Received {type(inner_input)}."
-            )
+            if len(self.input_variables) == 1:
+                var_name = self.input_variables[0]
+                inner_input = {var_name: inner_input}
+
+            else:
+                raise TypeError(
+                    f"Expected mapping type as input to {self.__class__.__name__}. "
+                    f"Received {type(inner_input)}."
+                )
         missing = set(self.input_variables).difference(inner_input)
         if missing:
             raise KeyError(
diff --git a/libs/core/langchain_core/prompts/chat.py b/libs/core/langchain_core/prompts/chat.py
index eb2ff2174d3..1c74bdd9c49 100644
--- a/libs/core/langchain_core/prompts/chat.py
+++ b/libs/core/langchain_core/prompts/chat.py
@@ -574,11 +574,51 @@ class ChatPromptTemplate(BaseChatPromptTemplate):
                 ("human", "{user_input}"),
             ])
 
-            messages = template.format_messages(
-                name="Bob",
-                user_input="What is your name?"
+            prompt_value = template.invoke(
+                {
+                    "name": "Bob",
+                    "user_input": "What is your name?"
+                }
             )
-    """
+            # Output:
+            # ChatPromptValue(
+            #    messages=[
+            #        SystemMessage(content='You are a helpful AI bot. Your name is Bob.'),
+            #        HumanMessage(content='Hello, how are you doing?'),
+            #        AIMessage(content="I'm doing well, thanks!"),
+            #        HumanMessage(content='What is your name?')
+            #    ]
+            #)
+
+    Single-variable template:
+
+        If your prompt has only a single input variable (i.e., 1 instance of "{variable_nams}"),
+        and you invoke the template with a non-dict object, the prompt template will
+        inject the provided argument into that variable location.
+
+
+        .. code-block:: python
+
+            from langchain_core.prompts import ChatPromptTemplate
+
+            template = ChatPromptTemplate.from_messages([
+                ("system", "You are a helpful AI bot. Your name is Carl."),
+                ("human", "{user_input}"),
+            ])
+
+            prompt_value = template.invoke("Hello, there!")
+            # Equivalent to
+            # prompt_value = template.invoke({"user_input": "Hello, there!"})
+
+            # Output:
+            #  ChatPromptValue(
+            #     messages=[
+            #         SystemMessage(content='You are a helpful AI bot. Your name is Carl.'),
+            #         HumanMessage(content='Hello, there!'),
+            #     ]
+            # )
+
+    """  # noqa: E501
 
     input_variables: List[str]
     """List of input variables in template messages. Used for validation."""
diff --git a/libs/core/tests/unit_tests/prompts/test_chat.py b/libs/core/tests/unit_tests/prompts/test_chat.py
index 7d1915cd60e..bcc1633cc61 100644
--- a/libs/core/tests/unit_tests/prompts/test_chat.py
+++ b/libs/core/tests/unit_tests/prompts/test_chat.py
@@ -533,3 +533,16 @@ def test_chat_prompt_message_placeholder_partial() -> None:
     assert prompt.format_messages() == []
     prompt = prompt.partial(history=[("system", "foo")])
     assert prompt.format_messages() == [SystemMessage(content="foo")]
+
+
+def test_messages_prompt_accepts_list() -> None:
+    prompt = ChatPromptTemplate.from_messages([MessagesPlaceholder("history")])
+    value = prompt.invoke([("user", "Hi there")])  # type: ignore
+    assert value.to_messages() == [HumanMessage(content="Hi there")]
+
+    # Assert still raises a nice error
+    prompt = ChatPromptTemplate.from_messages(
+        [("system", "You are a {foo}"), MessagesPlaceholder("history")]
+    )
+    with pytest.raises(TypeError):
+        prompt.invoke([("user", "Hi there")])  # type: ignore