Add ChatMessageHistory feature for multi-round dialogues in REST sessions (Issue #2284)

Co-Authored-By: Joe Moura <joao@crewai.com>

src/crewai/memory/__init__.py

@@ -2,5 +2,15 @@ from .entity.entity_memory import EntityMemory
 from .long_term.long_term_memory import LongTermMemory
 from .short_term.short_term_memory import ShortTermMemory
 from .user.user_memory import UserMemory
+from .chat_history.chat_message_history import ChatMessageHistory
+from .chat_history.chat_message import ChatMessage, MessageRole
 
-__all__ = ["UserMemory", "EntityMemory", "LongTermMemory", "ShortTermMemory"]
+__all__ = [
+    "UserMemory",
+    "EntityMemory",
+    "LongTermMemory",
+    "ShortTermMemory",
+    "ChatMessageHistory",
+    "ChatMessage",
+    "MessageRole",
+]

src/crewai/memory/chat_history/__init__.py

@@ -0,0 +1,4 @@
+from crewai.memory.chat_history.chat_message import ChatMessage, MessageRole
+from crewai.memory.chat_history.chat_message_history import ChatMessageHistory
+
+__all__ = ["ChatMessage", "MessageRole", "ChatMessageHistory"]

src/crewai/memory/chat_history/chat_message.py

@@ -0,0 +1,53 @@
+from datetime import datetime
+from enum import Enum
+from typing import Any, Dict, Optional
+
+
+class MessageRole(str, Enum):
+    """Enum for message roles in a chat."""
+    HUMAN = "human"
+    AI = "ai"
+    SYSTEM = "system"
+
+
+class ChatMessage:
+    """
+    Represents a single message in a chat history.
+    
+    Attributes:
+        role: The role of the message sender (human, ai, or system).
+        content: The content of the message.
+        timestamp: When the message was created.
+        metadata: Additional information about the message.
+    """
+    
+    def __init__(
+        self,
+        role: MessageRole,
+        content: str,
+        timestamp: Optional[datetime] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+    ):
+        self.role = role
+        self.content = content
+        self.timestamp = timestamp or datetime.now()
+        self.metadata = metadata or {}
+    
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert the message to a dictionary."""
+        return {
+            "role": self.role.value,
+            "content": self.content,
+            "timestamp": self.timestamp.isoformat(),
+            "metadata": self.metadata,
+        }
+    
+    @classmethod
+    def from_dict(cls, data: Dict[str, Any]) -> "ChatMessage":
+        """Create a message from a dictionary."""
+        return cls(
+            role=MessageRole(data["role"]),
+            content=data["content"],
+            timestamp=datetime.fromisoformat(data["timestamp"]),
+            metadata=data.get("metadata", {}),
+        )

src/crewai/memory/chat_history/chat_message_history.py

@@ -0,0 +1,180 @@
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from pydantic import PrivateAttr
+
+from crewai.memory.chat_history.chat_message import ChatMessage, MessageRole
+from crewai.memory.memory import Memory
+from crewai.memory.storage.rag_storage import RAGStorage
+
+
+class ChatMessageHistory(Memory):
+    """
+    ChatMessageHistory class for storing and retrieving chat messages.
+    
+    This class allows for maintaining conversation context across multiple
+    interactions within a single session, similar to Langchain's ChatMessageHistory.
+    
+    Attributes:
+        messages: A list of ChatMessage objects representing the conversation history.
+    """
+    
+    _memory_provider: Optional[str] = PrivateAttr()
+    _messages: List[ChatMessage] = PrivateAttr(default_factory=list)
+    
+    def __init__(self, crew=None, embedder_config=None, storage=None, path=None):
+        if crew and hasattr(crew, "memory_config") and crew.memory_config is not None:
+            memory_provider = crew.memory_config.get("provider")
+        else:
+            memory_provider = None
+
+        if memory_provider == "mem0":
+            try:
+                from crewai.memory.storage.mem0_storage import Mem0Storage
+            except ImportError:
+                raise ImportError(
+                    "Mem0 is not installed. Please install it with `pip install mem0ai`."
+                )
+            storage = Mem0Storage(type="chat_history", crew=crew)
+        else:
+            storage = (
+                storage
+                if storage
+                else RAGStorage(
+                    type="chat_history",
+                    embedder_config=embedder_config,
+                    crew=crew,
+                    path=path,
+                )
+            )
+        super().__init__(storage=storage)
+        self._memory_provider = memory_provider
+        self._messages = []
+    
+    def add_message(
+        self,
+        role: MessageRole,
+        content: str,
+        metadata: Optional[Dict[str, Any]] = None,
+        agent: Optional[str] = None,
+    ) -> None:
+        """
+        Add a message to the chat history.
+        
+        Args:
+            role: The role of the message sender (human, ai, or system).
+            content: The content of the message.
+            metadata: Additional information about the message.
+            agent: The agent associated with the message.
+        """
+        message = ChatMessage(role=role, content=content, metadata=metadata)
+        self._messages.append(message)
+        
+        # Save to storage for persistence and retrieval
+        metadata = metadata or {}
+        if agent:
+            metadata["agent"] = agent
+        
+        # Add role and timestamp to metadata
+        metadata["role"] = role.value
+        metadata["timestamp"] = message.timestamp.isoformat()
+        
+        super().save(value=content, metadata=metadata, agent=agent)
+    
+    def add_human_message(
+        self,
+        content: str,
+        metadata: Optional[Dict[str, Any]] = None,
+        agent: Optional[str] = None,
+    ) -> None:
+        """Add a human message to the chat history."""
+        self.add_message(MessageRole.HUMAN, content, metadata, agent)
+    
+    def add_ai_message(
+        self,
+        content: str,
+        metadata: Optional[Dict[str, Any]] = None,
+        agent: Optional[str] = None,
+    ) -> None:
+        """Add an AI message to the chat history."""
+        self.add_message(MessageRole.AI, content, metadata, agent)
+    
+    def add_system_message(
+        self,
+        content: str,
+        metadata: Optional[Dict[str, Any]] = None,
+        agent: Optional[str] = None,
+    ) -> None:
+        """Add a system message to the chat history."""
+        self.add_message(MessageRole.SYSTEM, content, metadata, agent)
+    
+    def get_messages(self) -> List[ChatMessage]:
+        """Get all messages in the chat history."""
+        return self._messages
+    
+    def get_messages_as_dict(self) -> List[Dict[str, Any]]:
+        """Get all messages in the chat history as dictionaries."""
+        return [message.to_dict() for message in self._messages]
+    
+    def clear(self) -> None:
+        """Clear all messages from the chat history."""
+        self._messages = []
+        self.reset()
+    
+    def reset(self) -> None:
+        """Reset the storage."""
+        try:
+            self.storage.reset()
+        except Exception as e:
+            raise Exception(
+                f"An error occurred while resetting the chat message history: {e}"
+            )
+    
+    def search(
+        self,
+        query: str,
+        limit: int = 5,
+        score_threshold: float = 0.35,
+    ) -> List[Dict[str, Any]]:
+        """
+        Search for messages in the chat history.
+        
+        Args:
+            query: The search query.
+            limit: The maximum number of results to return.
+            score_threshold: The minimum similarity score for results.
+            
+        Returns:
+            A list of dictionaries containing the search results.
+        """
+        results = self.storage.search(
+            query=query, limit=limit, score_threshold=score_threshold
+        )
+        
+        # Convert the search results to ChatMessage objects
+        messages = []
+        for result in results:
+            try:
+                role = result["metadata"].get("role", "ai")
+                content = result["context"]
+                timestamp = result["metadata"].get("timestamp")
+                if timestamp:
+                    timestamp = datetime.fromisoformat(timestamp)
+                else:
+                    timestamp = datetime.now()
+                
+                metadata = {k: v for k, v in result["metadata"].items() 
+                           if k not in ["role", "timestamp"]}
+                
+                message = ChatMessage(
+                    role=MessageRole(role),
+                    content=content,
+                    timestamp=timestamp,
+                    metadata=metadata,
+                )
+                messages.append(message.to_dict())
+            except Exception:
+                # Skip invalid messages
+                continue
+        
+        return messages