From 927793d0881bf0924935873ab530864cea18951b Mon Sep 17 00:00:00 2001
From: Christophe Bornet <cbornet@hotmail.com>
Date: Fri, 5 Apr 2024 16:25:27 +0200
Subject: [PATCH] Merge pull request #20038

* Implement aformat_messages for ChatMessagePromptTemplate
---
 libs/core/langchain_core/prompts/chat.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/libs/core/langchain_core/prompts/chat.py b/libs/core/langchain_core/prompts/chat.py
index cd27e0ef5c2..0df0a4bb146 100644
--- a/libs/core/langchain_core/prompts/chat.py
+++ b/libs/core/langchain_core/prompts/chat.py
@@ -297,6 +297,17 @@ class BaseStringMessagePromptTemplate(BaseMessagePromptTemplate, ABC):
             Formatted message.
         """
 
+    async def aformat(self, **kwargs: Any) -> BaseMessage:
+        """Format the prompt template.
+
+        Args:
+            **kwargs: Keyword arguments to use for formatting.
+
+        Returns:
+            Formatted message.
+        """
+        return self.format(**kwargs)
+
     def format_messages(self, **kwargs: Any) -> List[BaseMessage]:
         """Format messages from kwargs.
 
@@ -308,6 +319,9 @@ class BaseStringMessagePromptTemplate(BaseMessagePromptTemplate, ABC):
         """
         return [self.format(**kwargs)]
 
+    async def aformat_messages(self, **kwargs: Any) -> List[BaseMessage]:
+        return [await self.aformat(**kwargs)]
+
     @property
     def input_variables(self) -> List[str]:
         """
@@ -350,6 +364,12 @@ class ChatMessagePromptTemplate(BaseStringMessagePromptTemplate):
             content=text, role=self.role, additional_kwargs=self.additional_kwargs
         )
 
+    async def aformat(self, **kwargs: Any) -> BaseMessage:
+        text = await self.prompt.aformat(**kwargs)
+        return ChatMessage(
+            content=text, role=self.role, additional_kwargs=self.additional_kwargs
+        )
+
 
 _StringImageMessagePromptTemplateT = TypeVar(
     "_StringImageMessagePromptTemplateT", bound="_StringImageMessagePromptTemplate"