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
This commit is contained in:
Elliot 2024-01-30 03:45:55 +08:00 committed by GitHub
parent 47bd58dc11
commit 39eb00d304
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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(