zep: Memory Retriever MMR Support & Docs Updates (#11954)

- Update Zep Memory and Retriever docstrings
- Zep Memory Retriever: Add support for native MMR
- Add MMR example to existing ZepRetriever Notebook

@baskaryan

libs/langchain/langchain/memory/zep_memory.py

@@ -7,7 +7,7 @@ from langchain.memory.chat_message_histories import ZepChatMessageHistory
 
 
 class ZepMemory(ConversationBufferMemory):
-    """Persist your chain history to the Zep Memory Server.
+    """Persist your chain history to the Zep MemoryStore.
 
     The number of messages returned by Zep and when the Zep server summarizes chat
     histories is configurable. See the Zep documentation for more details.
@@ -36,11 +36,11 @@ class ZepMemory(ConversationBufferMemory):
     in the call to `self.memory.save_context`.
 
 
-    About Zep
+    Zep - Fast, scalable building blocks for LLM Apps
     =========
-    Zep provides long-term conversation storage for LLM apps. The server stores,
-    summarizes, embeds, indexes, and enriches conversational AI chat
-    histories, and exposes them via simple, low-latency APIs.
+    Zep is an open source platform for productionizing LLM apps. Go from a prototype
+    built in LangChain or LlamaIndex, or a custom app, to production in minutes without
+    rewriting code.
 
     For server installation instructions and more, see:
     https://docs.getzep.com/deployment/quickstart/

libs/langchain/langchain/retrievers/zep.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from enum import Enum
 from typing import TYPE_CHECKING, Any, Dict, List, Optional
 
 from langchain.callbacks.manager import (
@@ -13,28 +14,58 @@ if TYPE_CHECKING:
     from zep_python import MemorySearchResult
 
 
+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 ZepRetriever(BaseRetriever):
-    """`Zep` long-term memory store retriever.
+    """`Zep` MemoryStore Retriever.
 
     Search your user's long-term chat history with Zep.
 
+    Zep offers both simple semantic search and Maximal Marginal Relevance (MMR)
+    reranking of search results.
+
     Note: You will need to provide the user's `session_id` to use this retriever.
 
-    More on Zep:
-    Zep provides long-term conversation storage for LLM apps. The server stores,
-    summarizes, embeds, indexes, and enriches conversational AI chat
-    histories, and exposes them via simple, low-latency APIs.
+    Args:
+        url: URL of your Zep server (required)
+        api_key: Your Zep API key (optional)
+        session_id: Identifies your user or a user's session (required)
+        top_k: Number of documents to return (default: 3, optional)
+        search_type: Type of search to perform (similarity / mmr) (default: similarity,
+                                                                    optional)
+        mmr_lambda: Lambda value for MMR search. Defaults to 0.5 (optional)
+
+    Zep - Fast, scalable building blocks for LLM Apps
+    =========
+    Zep is an open source platform for productionizing LLM apps. Go from a prototype
+    built in LangChain or LlamaIndex, or a custom app, to production in minutes without
+    rewriting code.
 
     For server installation instructions, see:
     https://docs.getzep.com/deployment/quickstart/
     """
 
-    zep_client: Any
+    zep_client: Optional[Any] = None
     """Zep client."""
+    url: str
+    """URL of your Zep server."""
+    api_key: Optional[str] = None
+    """Your Zep API key."""
     session_id: str
     """Zep session ID."""
     top_k: Optional[int]
-    """Number of documents to return."""
+    """Number of items to return."""
+    search_type: SearchType = SearchType.similarity
+    """Type of search to perform (similarity / mmr)"""
+    mmr_lambda: Optional[float] = None
+    """Lambda value for MMR search."""
 
     @root_validator(pre=True)
     def create_client(cls, values: dict) -> dict:
@@ -68,12 +99,18 @@ class ZepRetriever(BaseRetriever):
         query: str,
         *,
         run_manager: CallbackManagerForRetrieverRun,
-        metadata: Optional[Dict] = None,
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> List[Document]:
         from zep_python import MemorySearchPayload
 
+        if not self.zep_client:
+            raise RuntimeError("Zep client not initialized.")
+
         payload: MemorySearchPayload = MemorySearchPayload(
-            text=query, metadata=metadata
+            text=query,
+            metadata=metadata,
+            search_type=self.search_type,
+            mmr_lambda=self.mmr_lambda,
         )
 
         results: List[MemorySearchResult] = self.zep_client.memory.search_memory(
@@ -87,12 +124,18 @@ class ZepRetriever(BaseRetriever):
         query: str,
         *,
         run_manager: AsyncCallbackManagerForRetrieverRun,
-        metadata: Optional[Dict] = None,
+        metadata: Optional[Dict[str, Any]] = None,
     ) -> List[Document]:
         from zep_python import MemorySearchPayload
 
+        if not self.zep_client:
+            raise RuntimeError("Zep client not initialized.")
+
         payload: MemorySearchPayload = MemorySearchPayload(
-            text=query, metadata=metadata
+            text=query,
+            metadata=metadata,
+            search_type=self.search_type,
+            mmr_lambda=self.mmr_lambda,
         )
 
         results: List[MemorySearchResult] = await self.zep_client.memory.asearch_memory(