From 80679ab906b8b1c68abbea48edbc0cf09d232a06 Mon Sep 17 00:00:00 2001
From: kaijietti <43436010+kaijietti@users.noreply.github.com>
Date: Thu, 18 Apr 2024 04:40:24 +0800
Subject: [PATCH] zep[patch]: implement add_messages and aadd_messages (#20099)

This PR implement `add_messages` and `aadd_messages` to avoid
unnecessary round-trips.
---
 .../chat_message_histories/zep.py             | 46 ++++++++++++++++++-
 1 file changed, 45 insertions(+), 1 deletion(-)

diff --git a/libs/community/langchain_community/chat_message_histories/zep.py b/libs/community/langchain_community/chat_message_histories/zep.py
index b38b1e9b49..c3a2a820d8 100644
--- a/libs/community/langchain_community/chat_message_histories/zep.py
+++ b/libs/community/langchain_community/chat_message_histories/zep.py
@@ -2,7 +2,7 @@ from __future__ import annotations
 
 import logging
 from enum import Enum
-from typing import TYPE_CHECKING, Any, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence
 
 from langchain_core.chat_history import BaseChatMessageHistory
 from langchain_core.messages import (
@@ -184,6 +184,38 @@ class ZepChatMessageHistory(BaseChatMessageHistory):
 
         self.zep_client.memory.add_memory(self.session_id, zep_memory)
 
+    def add_messages(self, messages: Sequence[BaseMessage]) -> None:
+        """Append the messages to the Zep memory history"""
+        from zep_python import Memory, Message
+
+        zep_messages = [
+            Message(
+                content=message.content,
+                role=message.type,
+                metadata=message.additional_kwargs.get("metadata", None),
+            )
+            for message in messages
+        ]
+        zep_memory = Memory(messages=zep_messages)
+
+        self.zep_client.memory.add_memory(self.session_id, zep_memory)
+
+    async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:
+        """Append the messages to the Zep memory history asynchronously"""
+        from zep_python import Memory, Message
+
+        zep_messages = [
+            Message(
+                content=message.content,
+                role=message.type,
+                metadata=message.additional_kwargs.get("metadata", None),
+            )
+            for message in messages
+        ]
+        zep_memory = Memory(messages=zep_messages)
+
+        await self.zep_client.memory.aadd_memory(self.session_id, zep_memory)
+
     def search(
         self,
         query: str,
@@ -218,3 +250,15 @@ class ZepChatMessageHistory(BaseChatMessageHistory):
             logger.warning(
                 f"Session {self.session_id} not found in Zep. Skipping delete."
             )
+
+    async def aclear(self) -> None:
+        """Clear session memory from Zep asynchronously.
+        Note that Zep is long-term storage for memory and this is not advised
+        unless you have specific data retention requirements.
+        """
+        try:
+            await self.zep_client.memory.adelete_memory(self.session_id)
+        except NotFoundError:
+            logger.warning(
+                f"Session {self.session_id} not found in Zep. Skipping delete."
+            )