agno 2.2.13__py3-none-any.whl

agno/tools/zep.py

@@ -0,0 +1,454 @@
+import uuid
+from os import getenv
+from textwrap import dedent
+from typing import Any, List, Optional
+
+from agno.tools import Toolkit
+from agno.utils.log import log_debug, log_error, log_warning
+
+try:
+    from zep_cloud import (
+        BadRequestError,
+        NotFoundError,
+    )
+    from zep_cloud import (
+        Message as ZepMessage,
+    )
+    from zep_cloud.client import AsyncZep, Zep
+except ImportError:
+    raise ImportError("`zep-cloud` package not found. Please install it with `pip install zep-cloud`")
+
+DEFAULT_INSTRUCTIONS = dedent(
+    """\
+    You have access to the users memories stored in Zep. You can interact with them using the following tools:
+    - `add_zep_message`: Add a message to the Zep session memory. Use this to add messages to the Zep session memory.
+    - `get_zep_memory`: Get the memory for the current Zep session. Use this to get the memory for the current Zep session.
+    - `search_zep_memory`: Search the Zep user graph for relevant facts. Use this to search the Zep user graph for relevant facts.
+
+    Guidelines:
+    - Use `add_zep_message` tool to add relevant messages to the users memories. You can use this tool multiple times to add multiple messages.
+    - Use `get_zep_memory` tool to get the memory for the current Zep session for additional context. This will give you a entire context of the user's memories with relevant facts.
+    - Use `search_zep_memory` tool to search the Zep user memories for relevant facts. This will give you a list of relevant facts.
+    """
+)
+
+
+class ZepTools(Toolkit):
+    def __init__(
+        self,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        api_key: Optional[str] = None,
+        ignore_assistant_messages: bool = False,
+        enable_add_zep_message: bool = True,
+        enable_get_zep_memory: bool = True,
+        enable_search_zep_memory: bool = True,
+        instructions: Optional[str] = None,
+        add_instructions: bool = False,
+        all: bool = False,
+        **kwargs,
+    ):
+        self._api_key = api_key or getenv("ZEP_API_KEY")
+        if not self._api_key:
+            raise ValueError(
+                "No Zep API key provided. Please set the ZEP_API_KEY environment variable or pass it to the ZepTools constructor."
+            )
+
+        if instructions is None:
+            self.instructions = "<Memory Instructions>\n" + DEFAULT_INSTRUCTIONS + "\n</Memory Instructions>"
+        else:
+            self.instructions = instructions
+
+        self.zep_client: Optional[Zep] = None
+        self._initialized = False
+
+        self.session_id_provided = session_id
+        self.user_id_provided = user_id
+        self.ignore_assistant_messages = ignore_assistant_messages
+
+        self.session_id: Optional[str] = None
+        self.user_id: Optional[str] = None
+
+        self.initialize()
+
+        tools: List[Any] = []
+        if enable_add_zep_message or all:
+            tools.append(self.add_zep_message)
+        if enable_get_zep_memory or all:
+            tools.append(self.get_zep_memory)
+        if enable_search_zep_memory or all:
+            tools.append(self.search_zep_memory)
+
+        super().__init__(
+            name="zep_tools", instructions=self.instructions, add_instructions=add_instructions, tools=tools, **kwargs
+        )
+
+    def initialize(self) -> bool:
+        """
+        Initialize the Zep client and ensure session/user setup.
+        """
+        if self._initialized:
+            return True
+
+        try:
+            self.zep_client = Zep(api_key=self._api_key)
+
+            # Handle session_id generation/validation
+            self.session_id = self.session_id_provided
+            if not self.session_id:
+                self.session_id = f"{uuid.uuid4()}"
+                log_debug(f"Generated new session ID: {self.session_id}")
+
+            # Handle user_id generation/validation and Zep user check/creation
+            self.user_id = self.user_id_provided
+            if not self.user_id:
+                self.user_id = f"user-{uuid.uuid4()}"
+                log_debug(f"Creating new default Zep user: {self.user_id}")
+                self.zep_client.user.add(user_id=self.user_id)  # type: ignore
+            else:
+                try:
+                    self.zep_client.user.get(self.user_id)  # type: ignore
+                    log_debug(f"Confirmed provided Zep user exists: {self.user_id}")
+                except NotFoundError:
+                    try:
+                        self.zep_client.user.add(user_id=self.user_id)  # type: ignore
+                    except BadRequestError as add_err:
+                        log_error(f"Failed to create provided user {self.user_id}: {add_err}")
+                        self.zep_client = None  # Reset client on failure
+                        return False  # Initialization failed
+
+            # Create session associated with the user
+            try:
+                self.zep_client.thread.create(thread_id=self.session_id, user_id=self.user_id)  # type: ignore
+                log_debug(f"Created session {self.session_id} for user {self.user_id}")
+            except Exception as e:
+                log_debug(f"Session may already exist: {e}")
+
+            self._initialized = True
+            return True
+
+        except Exception as e:
+            log_error(f"Failed to initialize ZepTools: {e}")
+            self.zep_client = None
+            self._initialized = False
+            return False
+
+    def add_zep_message(self, role: str, content: str) -> str:
+        """
+        Adds a message to the current Zep session memory.
+        Args:
+            role (str): The role of the message sender (e.g., 'user', 'assistant', 'system').
+            content (str): The text content of the message.
+
+        Returns:
+            A confirmation message or an error string.
+        """
+        if not self.zep_client or not self.session_id:
+            log_error("Zep client or session ID not initialized. Cannot add message.")
+            return "Error: Zep client/session not initialized."
+
+        try:
+            zep_message = ZepMessage(
+                role=role,
+                content=content,
+                role_type=role,
+            )
+
+            # Prepare ignore_roles if needed
+            ignore_roles_list = ["assistant"] if self.ignore_assistant_messages else None
+
+            # Add message to Zep memory
+            self.zep_client.thread.add_messages(  # type: ignore
+                thread_id=self.session_id,
+                messages=[zep_message],
+                ignore_roles=ignore_roles_list,
+            )
+            return f"Message from '{role}' added successfully to session {self.session_id}."
+        except Exception as e:
+            error_msg = f"Failed to add message to Zep session {self.session_id}: {e}"
+            log_error(error_msg)
+            return f"Error adding message: {e}"
+
+    def get_zep_memory(self, memory_type: str = "context") -> str:
+        """
+        Retrieves the memory for the current Zep session.
+        Args:
+            memory_type: The type of memory to retrieve ('context', 'messages').
+        Returns:
+            The requested memory content as a string, or an error string.
+        """
+        if not self.zep_client or not self.session_id:
+            log_error("Zep client or session ID not initialized. Cannot get memory.")
+            return "Error: Zep client/session not initialized."
+
+        try:
+            log_debug(f"Getting Zep memory for session {self.session_id}")
+
+            if memory_type == "context":
+                # Ensure context is a string
+                user_context = self.zep_client.thread.get_user_context(thread_id=self.session_id, mode="basic")  # type: ignore
+                log_debug(f"Memory data: {user_context}")
+                return user_context.context or "No context available."
+            elif memory_type == "messages":
+                messages_list = self.zep_client.thread.get(thread_id=self.session_id)  # type: ignore
+                # Ensure messages string representation is returned
+                return str(messages_list.messages) if messages_list.messages else "No messages available."
+            else:
+                warning_msg = f"Unsupported memory_type requested: {memory_type}. Returning empty string."
+                log_warning(warning_msg)
+                return warning_msg
+
+        except Exception as e:
+            log_error(f"Failed to get Zep memory for session {self.session_id}: {e}")
+            return f"Error getting memory for session {self.session_id}"
+
+    def search_zep_memory(self, query: str, search_scope: str = "edges") -> str:
+        """
+        Searches the Zep knowledge graph for relevant facts or nodes.
+        Args:
+            query: The search term to find relevant facts or nodes.
+            search_scope: The scope of the search to perform. Can be "edges" (for facts) or "nodes".
+        Returns:
+            A string of the search result
+        """
+        # Graph search is built on user_id not on session_id
+        if not self.zep_client or not self.user_id:
+            log_error("Zep client or user ID not initialized. Cannot search graph.")
+            return "Error: Zep client/user not initialized."
+
+        try:
+            search_response = self.zep_client.graph.search(
+                query=query,
+                user_id=self.user_id,
+                scope=search_scope,  # Can be "edges" or "nodes"
+            )
+
+            if search_scope == "edges" and search_response.edges:
+                # Return facts from edges
+                facts_str = "\n".join([f"- {edge.fact}" for edge in search_response.edges])
+                return f"Found {len(search_response.edges)} facts:\n{facts_str}"
+            elif search_scope == "nodes" and search_response.nodes:
+                # Return node summaries
+                nodes_str = "\n".join([f"- {node.name}: {node.summary}" for node in search_response.nodes])
+                return f"Found {len(search_response.nodes)} nodes:\n{nodes_str}"
+            else:
+                return f"No {search_scope} found for query: {query}"
+
+        except Exception as e:
+            log_error(f"Failed to search Zep graph for user {self.user_id}: {e}")
+            return f"Error searching graph: {e}"
+
+
+class ZepAsyncTools(Toolkit):
+    def __init__(
+        self,
+        session_id: Optional[str] = None,
+        user_id: Optional[str] = None,
+        api_key: Optional[str] = None,
+        ignore_assistant_messages: bool = False,
+        add_zep_message: bool = True,
+        get_zep_memory: bool = True,
+        search_zep_memory: bool = True,
+        instructions: Optional[str] = None,
+        add_instructions: bool = False,
+        **kwargs,
+    ):
+        self._api_key = api_key or getenv("ZEP_API_KEY")
+        if not self._api_key:
+            raise ValueError(
+                "No Zep API key provided. Please set the ZEP_API_KEY environment variable or pass it to the ZepTools constructor."
+            )
+
+        if instructions is None:
+            self.instructions = "<Memory Instructions>\n" + DEFAULT_INSTRUCTIONS + "\n</Memory Instructions>"
+        else:
+            self.instructions = instructions
+
+        self.zep_client: Optional[AsyncZep] = None
+        self._initialized = False
+
+        self.session_id_provided = session_id
+        self.user_id_provided = user_id
+        self.ignore_assistant_messages = ignore_assistant_messages
+
+        self.session_id: Optional[str] = None
+        self.user_id: Optional[str] = None
+
+        self._initialized = False
+
+        # Register methods as tools conditionally
+        tools = []
+        if add_zep_message:
+            tools.append(self.add_zep_message)
+        if get_zep_memory:
+            tools.append(self.get_zep_memory)  # type: ignore
+        if search_zep_memory:
+            tools.append(self.search_zep_memory)  # type: ignore
+
+        super().__init__(
+            name="zep_tools", instructions=self.instructions, add_instructions=add_instructions, tools=tools, **kwargs
+        )
+
+    async def initialize(self) -> bool:
+        """
+        Initialize the AsyncZep client and ensure session/user setup.
+        """
+        if self._initialized:
+            return True
+
+        try:
+            self.zep_client = AsyncZep(api_key=self._api_key)
+
+            # Handle session_id generation/validation
+            self.session_id = self.session_id_provided
+            if not self.session_id:
+                self.session_id = f"{uuid.uuid4()}"
+                log_debug(f"Generated new session ID: {self.session_id}")
+
+            # Handle user_id generation/validation and Zep user check/creation
+            self.user_id = self.user_id_provided
+            if not self.user_id:
+                self.user_id = f"user-{uuid.uuid4()}"
+                log_debug(f"Creating new default Zep user: {self.user_id}")
+                await self.zep_client.user.add(user_id=self.user_id)  # type: ignore
+            else:
+                try:
+                    await self.zep_client.user.get(self.user_id)  # type: ignore
+                    log_debug(f"Confirmed provided Zep user exists: {self.user_id}")
+                except NotFoundError:
+                    try:
+                        await self.zep_client.user.add(user_id=self.user_id)  # type: ignore
+                    except BadRequestError as add_err:
+                        log_error(f"Failed to create provided user {self.user_id}: {add_err}")
+                        self.zep_client = None  # Reset client on failure
+                        return False  # Initialization failed
+
+            # Create session associated with the user
+            try:
+                await self.zep_client.thread.create(thread_id=self.session_id, user_id=self.user_id)  # type: ignore
+                log_debug(f"Created session {self.session_id} for user {self.user_id}")
+            except Exception as e:
+                log_debug(f"Session may already exist: {e}")
+
+            self._initialized = True
+            return True
+
+        except Exception as e:
+            log_error(f"Failed to initialize ZepTools: {e}")
+            self.zep_client = None
+            self._initialized = False
+            return False
+
+    async def add_zep_message(self, role: str, content: str) -> str:
+        """
+        Adds a message to the current Zep session memory.
+        Args:
+            role (str): The role of the message sender (e.g., 'user', 'assistant', 'system').
+            content (str): The text content of the message.
+
+        Returns:
+            A confirmation message or an error string.
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        if not self.zep_client or not self.session_id:
+            log_error("Zep client or session ID not initialized. Cannot add message.")
+            return "Error: Zep client/session not initialized."
+
+        try:
+            zep_message = ZepMessage(
+                role=role,
+                content=content,
+                role_type=role,
+            )
+
+            # Prepare ignore_roles if needed
+            ignore_roles_list = ["assistant"] if self.ignore_assistant_messages else None
+
+            # Add message to Zep memory
+            await self.zep_client.thread.add_messages(  # type: ignore
+                thread_id=self.session_id,
+                messages=[zep_message],
+                ignore_roles=ignore_roles_list,
+            )
+            return f"Message from '{role}' added successfully to session {self.session_id}."
+        except Exception as e:
+            error_msg = f"Failed to add message to Zep session {self.session_id}: {e}"
+            log_error(error_msg)
+            return f"Error adding message: {e}"
+
+    async def get_zep_memory(self, memory_type: str = "context") -> str:
+        """
+        Retrieves the memory for the current Zep session.
+        Args:
+            memory_type: The type of memory to retrieve ('context', 'messages').
+        Returns:
+            The requested memory content as a string, or an error string.
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        if not self.zep_client or not self.session_id:
+            log_error("Zep client or session ID not initialized. Cannot get memory.")
+            return "Error: Zep client/session not initialized."
+
+        try:
+            if memory_type == "context":
+                # Ensure context is a string
+                user_context = await self.zep_client.thread.get_user_context(thread_id=self.session_id, mode="basic")  # type: ignore
+                log_debug(f"Memory data: {user_context}")
+                return user_context.context or "No context available."
+            elif memory_type == "messages":
+                # Ensure messages string representation is returned
+                messages_list = await self.zep_client.thread.get(thread_id=self.session_id)  # type: ignore
+                return str(messages_list.messages) if messages_list.messages else "No messages available."
+            else:
+                warning_msg = f"Unsupported memory_type requested: {memory_type}. Returning context."
+                log_warning(warning_msg)
+                return "No context available."
+
+        except Exception as e:
+            error_msg = f"Failed to get Zep memory for session {self.session_id}: {e}"
+            log_error(error_msg)
+            return f"Error getting memory: {e}"
+
+    async def search_zep_memory(self, query: str, scope: str = "edges", limit: int = 5) -> str:
+        """
+        Searches the Zep knowledge graph for relevant facts or nodes.
+        Args:
+            query: The search term to find relevant facts or nodes.
+            scope: The scope of the search to perform. Can be "edges" (for facts) or "nodes".
+            limit: The maximum number of results to return.
+        Returns:
+            A string of the search result
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        if not self.zep_client or not self.user_id:
+            log_error("Zep client or user ID not initialized. Cannot search graph.")
+            return "Error: Zep client/user not initialized."
+
+        try:
+            search_response = await self.zep_client.graph.search(  # type: ignore
+                query=query,
+                user_id=self.user_id,
+                scope=scope,  # Can be "edges" or "nodes"
+                limit=limit,
+            )
+
+            if scope == "edges" and search_response.edges:
+                # Return facts from edges
+                facts_str = "\n".join([f"- {edge.fact}" for edge in search_response.edges])
+                return f"Found {len(search_response.edges)} facts:\n{facts_str}"
+            elif scope == "nodes" and search_response.nodes:
+                # Return node summaries
+                nodes_str = "\n".join([f"- {node.name}: {node.summary}" for node in search_response.nodes])
+                return f"Found {len(search_response.nodes)} nodes:\n{nodes_str}"
+            else:
+                return f"No {scope} found for query: {query}"
+
+        except Exception as e:
+            log_error(f"Failed to search Zep graph for user {self.user_id}: {e}")
+            return f"Error searching graph: {e}"