diff --git a/libs/core/langchain_core/chat_history.py b/libs/core/langchain_core/chat_history.py index 211c7c8466..f8d9a4dabf 100644 --- a/libs/core/langchain_core/chat_history.py +++ b/libs/core/langchain_core/chat_history.py @@ -1,7 +1,7 @@ from __future__ import annotations from abc import ABC, abstractmethod -from typing import List +from typing import List, Union from langchain_core.messages import ( AIMessage, @@ -42,21 +42,27 @@ class BaseChatMessageHistory(ABC): messages: List[BaseMessage] """A list of Messages stored in-memory.""" - def add_user_message(self, message: str) -> None: + def add_user_message(self, message: Union[HumanMessage, str]) -> None: """Convenience method for adding a human message string to the store. Args: - message: The string contents of a human message. + message: The human message to add """ - self.add_message(HumanMessage(content=message)) + if isinstance(message, HumanMessage): + self.add_message(message) + else: + self.add_message(HumanMessage(content=message)) - def add_ai_message(self, message: str) -> None: + def add_ai_message(self, message: Union[AIMessage, str]) -> None: """Convenience method for adding an AI message string to the store. Args: - message: The string contents of an AI message. + message: The AI message to add. """ - self.add_message(AIMessage(content=message)) + if isinstance(message, AIMessage): + self.add_message(message) + else: + self.add_message(AIMessage(content=message)) @abstractmethod def add_message(self, message: BaseMessage) -> None: