From 39eb00d30468769510b13737386b9a032c26a31a Mon Sep 17 00:00:00 2001
From: Elliot <elliotliunj@gmail.com>
Date: Tue, 30 Jan 2024 03:45:55 +0800
Subject: [PATCH] community[patch]: Adapt more parameters related to
 MemorySearchPayload for the search method of ZepChatMessageHistory (#15441)

- **Description:** To adapt more parameters related to
MemorySearchPayload for the search method of ZepChatMessageHistory,
  - **Issue:** None,
  - **Dependencies:** None,
  - **Twitter handle:** None
---
 .../chat_message_histories/zep.py             | 35 +++++++++++++++++--
 1 file changed, 32 insertions(+), 3 deletions(-)

diff --git a/libs/community/langchain_community/chat_message_histories/zep.py b/libs/community/langchain_community/chat_message_histories/zep.py
index 3899f89ba6e..45de39f150e 100644
--- a/libs/community/langchain_community/chat_message_histories/zep.py
+++ b/libs/community/langchain_community/chat_message_histories/zep.py
@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from enum import Enum
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from langchain_core.chat_history import BaseChatMessageHistory
@@ -17,6 +18,24 @@ if TYPE_CHECKING:
 logger = logging.getLogger(__name__)
 
 
+class SearchScope(str, Enum):
+    """Which documents to search. Messages or Summaries?"""
+
+    messages = "messages"
+    """Search chat history messages."""
+    summary = "summary"
+    """Search chat history summaries."""
+
+
+class SearchType(str, Enum):
+    """Enumerator of the types of search to perform."""
+
+    similarity = "similarity"
+    """Similarity search."""
+    mmr = "mmr"
+    """Maximal Marginal Relevance reranking of similarity search."""
+
+
 class ZepChatMessageHistory(BaseChatMessageHistory):
     """Chat message history that uses Zep as a backend.
 
@@ -166,13 +185,23 @@ class ZepChatMessageHistory(BaseChatMessageHistory):
         self.zep_client.memory.add_memory(self.session_id, zep_memory)
 
     def search(
-        self, query: str, metadata: Optional[Dict] = None, limit: Optional[int] = None
+        self,
+        query: str,
+        metadata: Optional[Dict] = None,
+        search_scope: SearchScope = SearchScope.messages,
+        search_type: SearchType = SearchType.similarity,
+        mmr_lambda: Optional[float] = None,
+        limit: Optional[int] = None,
     ) -> List[MemorySearchResult]:
         """Search Zep memory for messages matching the query"""
         from zep_python import MemorySearchPayload
 
-        payload: MemorySearchPayload = MemorySearchPayload(
-            text=query, metadata=metadata
+        payload = MemorySearchPayload(
+            text=query,
+            metadata=metadata,
+            search_scope=search_scope,
+            search_type=search_type,
+            mmr_lambda=mmr_lambda,
         )
 
         return self.zep_client.memory.search_memory(