codekingpro/portable-devtools

1from __future__ import annotations2 3import logging4from enum import Enum5from typing import TYPE_CHECKING, Any, Dict, List, Optional, Sequence6 7from langchain_core.chat_history import BaseChatMessageHistory8from langchain_core.messages import (9    AIMessage,10    BaseMessage,11    HumanMessage,12    SystemMessage,13)14 15if TYPE_CHECKING:16    from zep_python import Memory, MemorySearchResult, Message, NotFoundError17 18logger = logging.getLogger(__name__)19 20 21class SearchScope(str, Enum):22    """Scope for the document search. Messages or Summaries?"""23 24    messages = "messages"25    """Search chat history messages."""26    summary = "summary"27    """Search chat history summaries."""28 29 30class SearchType(str, Enum):31    """Enumerator of the types of search to perform."""32 33    similarity = "similarity"34    """Similarity search."""35    mmr = "mmr"36    """Maximal Marginal Relevance reranking of similarity search."""37 38 39class ZepChatMessageHistory(BaseChatMessageHistory):40    """Chat message history that uses Zep as a backend.41 42    Recommended usage::43 44        # Set up Zep Chat History45        zep_chat_history = ZepChatMessageHistory(46            session_id=session_id,47            url=ZEP_API_URL,48            api_key=<your_api_key>,49        )50 51        # Use a standard ConversationBufferMemory to encapsulate the Zep chat history52        memory = ConversationBufferMemory(53            memory_key="chat_history", chat_memory=zep_chat_history54        )55 56 57    Zep provides long-term conversation storage for LLM apps. The server stores,58    summarizes, embeds, indexes, and enriches conversational AI chat59    histories, and exposes them via simple, low-latency APIs.60 61    For server installation instructions and more, see:62    https://docs.getzep.com/deployment/quickstart/63 64    This class is a thin wrapper around the zep-python package. Additional65    Zep functionality is exposed via the `zep_summary` and `zep_messages`66    properties.67 68    For more information on the zep-python package, see:69    https://github.com/getzep/zep-python70    """71 72    def __init__(73        self,74        session_id: str,75        url: str = "http://localhost:8000",76        api_key: Optional[str] = None,77    ) -> None:78        try:79            from zep_python import ZepClient80        except ImportError:81            raise ImportError(82                "Could not import zep-python package. "83                "Please install it with `pip install zep-python`."84            )85 86        self.zep_client = ZepClient(base_url=url, api_key=api_key)87        self.session_id = session_id88 89    @property90    def messages(self) -> List[BaseMessage]:  # type: ignore[override]91        """Retrieve messages from Zep memory"""92        zep_memory: Optional[Memory] = self._get_memory()93        if not zep_memory:94            return []95 96        messages: List[BaseMessage] = []97        # Extract summary, if present, and messages98        if zep_memory.summary:99            if len(zep_memory.summary.content) > 0:100                messages.append(SystemMessage(content=zep_memory.summary.content))101        if zep_memory.messages:102            msg: Message103            for msg in zep_memory.messages:104                metadata: Dict = {105                    "uuid": msg.uuid,106                    "created_at": msg.created_at,107                    "token_count": msg.token_count,108                    "metadata": msg.metadata,109                }110                if msg.role == "ai":111                    messages.append(112                        AIMessage(content=msg.content, additional_kwargs=metadata)113                    )114                else:115                    messages.append(116                        HumanMessage(content=msg.content, additional_kwargs=metadata)117                    )118 119        return messages120 121    @property122    def zep_messages(self) -> List[Message]:123        """Retrieve summary from Zep memory"""124        zep_memory: Optional[Memory] = self._get_memory()125        if not zep_memory:126            return []127 128        return zep_memory.messages129 130    @property131    def zep_summary(self) -> Optional[str]:132        """Retrieve summary from Zep memory"""133        zep_memory: Optional[Memory] = self._get_memory()134        if not zep_memory or not zep_memory.summary:135            return None136 137        return zep_memory.summary.content138 139    def _get_memory(self) -> Optional[Memory]:140        """Retrieve memory from Zep"""141        from zep_python import NotFoundError142 143        try:144            zep_memory: Memory = self.zep_client.memory.get_memory(self.session_id)145        except NotFoundError:146            logger.warning(147                f"Session {self.session_id} not found in Zep. Returning None"148            )149            return None150        return zep_memory151 152    def add_user_message(  # type: ignore[override]153        self, message: str, metadata: Optional[Dict[str, Any]] = None154    ) -> None:155        """Convenience method for adding a human message string to the store.156 157        Args:158            message: The string contents of a human message.159            metadata: Optional metadata to attach to the message.160        """161        self.add_message(HumanMessage(content=message), metadata=metadata)162 163    def add_ai_message(  # type: ignore[override]164        self, message: str, metadata: Optional[Dict[str, Any]] = None165    ) -> None:166        """Convenience method for adding an AI message string to the store.167 168        Args:169            message: The string contents of an AI message.170            metadata: Optional metadata to attach to the message.171        """172        self.add_message(AIMessage(content=message), metadata=metadata)173 174    def add_message(175        self, message: BaseMessage, metadata: Optional[Dict[str, Any]] = None176    ) -> None:177        """Append the message to the Zep memory history"""178        from zep_python import Memory, Message179 180        zep_message = Message(181            content=message.content, role=message.type, metadata=metadata182        )183        zep_memory = Memory(messages=[zep_message])184 185        self.zep_client.memory.add_memory(self.session_id, zep_memory)186 187    def add_messages(self, messages: Sequence[BaseMessage]) -> None:188        """Append the messages to the Zep memory history"""189        from zep_python import Memory, Message190 191        zep_messages = [192            Message(193                content=message.content,194                role=message.type,195                metadata=message.additional_kwargs.get("metadata", None),196            )197            for message in messages198        ]199        zep_memory = Memory(messages=zep_messages)200 201        self.zep_client.memory.add_memory(self.session_id, zep_memory)202 203    async def aadd_messages(self, messages: Sequence[BaseMessage]) -> None:204        """Append the messages to the Zep memory history asynchronously"""205        from zep_python import Memory, Message206 207        zep_messages = [208            Message(209                content=message.content,210                role=message.type,211                metadata=message.additional_kwargs.get("metadata", None),212            )213            for message in messages214        ]215        zep_memory = Memory(messages=zep_messages)216 217        await self.zep_client.memory.aadd_memory(self.session_id, zep_memory)218 219    def search(220        self,221        query: str,222        metadata: Optional[Dict] = None,223        search_scope: SearchScope = SearchScope.messages,224        search_type: SearchType = SearchType.similarity,225        mmr_lambda: Optional[float] = None,226        limit: Optional[int] = None,227    ) -> List[MemorySearchResult]:228        """Search Zep memory for messages matching the query"""229        from zep_python import MemorySearchPayload230 231        payload = MemorySearchPayload(232            text=query,233            metadata=metadata,234            search_scope=search_scope,235            search_type=search_type,236            mmr_lambda=mmr_lambda,237        )238 239        return self.zep_client.memory.search_memory(240            self.session_id, payload, limit=limit241        )242 243    def clear(self) -> None:244        """Clear session memory from Zep. Note that Zep is long-term storage for memory245        and this is not advised unless you have specific data retention requirements.246        """247        try:248            self.zep_client.memory.delete_memory(self.session_id)249        except NotFoundError:250            logger.warning(251                f"Session {self.session_id} not found in Zep. Skipping delete."252            )253 254    async def aclear(self) -> None:255        """Clear session memory from Zep asynchronously.256        Note that Zep is long-term storage for memory and this is not advised257        unless you have specific data retention requirements.258        """259        try:260            await self.zep_client.memory.adelete_memory(self.session_id)261        except NotFoundError:262            logger.warning(263                f"Session {self.session_id} not found in Zep. Skipping delete."264            )265