github/langchain
1
0
Fork 0
mirror of https://github.com/hwchase17/langchain.git synced 2025-09-14 14:05:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Zep: Summary Search and Example (#12686)

Browse Source 
Zep now has the ability to search over chat history summaries. This PR
adds support for doing so. More here: https://blog.getzep.com/zep-v0-17/

@baskaryan @eyurtsev
This commit is contained in:
Daniel Chalef
2023-11-02 16:31:11 -07:00
committed by GitHub
parent 526313002c
commit cc3d3920e3
2 changed files with 175 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
libs/langchain/langchain/retrievers/zep.py
View File

@@ -11,7 +11,16 @@ from langchain.pydantic_v1 import root_validator
from langchain.schema import BaseRetriever, Document
if TYPE_CHECKING:
from zep_python import MemorySearchResult
from zep_python.memory import MemorySearchResult
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):
@@ -62,6 +71,8 @@ class ZepRetriever(BaseRetriever):
"""Zep session ID."""
top_k: Optional[int]
"""Number of items to return."""
search_scope: SearchScope = SearchScope.messages
"""Which documents to search. Messages or Summaries?"""
search_type: SearchType = SearchType.similarity
"""Type of search to perform (similarity / mmr)"""
mmr_lambda: Optional[float] = None
@@ -82,7 +93,7 @@ class ZepRetriever(BaseRetriever):
)
return values
def _search_result_to_doc(
def _messages_search_result_to_doc(
self, results: List[MemorySearchResult]
) -> List[Document]:
return [
@@ -94,6 +105,23 @@ class ZepRetriever(BaseRetriever):
if r.message
]
def _summary_search_result_to_doc(
self, results: List[MemorySearchResult]
) -> List[Document]:
return [
Document(
page_content=r.summary.content,
metadata={
"score": r.dist,
"uuid": r.summary.uuid,
"created_at": r.summary.created_at,
"token_count": r.summary.token_count,
},
)
for r in results
if r.summary
]
def _get_relevant_documents(
self,
query: str,
@@ -101,14 +129,15 @@ class ZepRetriever(BaseRetriever):
run_manager: CallbackManagerForRetrieverRun,
metadata: Optional[Dict[str, Any]] = None,
) -> List[Document]:
from zep_python import MemorySearchPayload
from zep_python.memory import MemorySearchPayload
if not self.zep_client:
raise RuntimeError("Zep client not initialized.")
payload: MemorySearchPayload = MemorySearchPayload(
payload = MemorySearchPayload(
text=query,
metadata=metadata,
search_scope=self.search_scope,
search_type=self.search_type,
mmr_lambda=self.mmr_lambda,
)
@@ -117,7 +146,10 @@ class ZepRetriever(BaseRetriever):
self.session_id, payload, limit=self.top_k
)
return self._search_result_to_doc(results)
if self.search_scope == SearchScope.summary:
return self._summary_search_result_to_doc(results)
return self._messages_search_result_to_doc(results)
async def _aget_relevant_documents(
self,
@@ -126,14 +158,15 @@ class ZepRetriever(BaseRetriever):
run_manager: AsyncCallbackManagerForRetrieverRun,
metadata: Optional[Dict[str, Any]] = None,
) -> List[Document]:
from zep_python import MemorySearchPayload
from zep_python.memory import MemorySearchPayload
if not self.zep_client:
raise RuntimeError("Zep client not initialized.")
payload: MemorySearchPayload = MemorySearchPayload(
payload = MemorySearchPayload(
text=query,
metadata=metadata,
search_scope=self.search_scope,
search_type=self.search_type,
mmr_lambda=self.mmr_lambda,
)
@@ -142,4 +175,7 @@ class ZepRetriever(BaseRetriever):
self.session_id, payload, limit=self.top_k
)
return self._search_result_to_doc(results)
if self.search_scope == SearchScope.summary:
return self._summary_search_result_to_doc(results)
return self._messages_search_result_to_doc(results)