zep-crewai 0.1.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

zep_crewai/__init__.py

@@ -1,28 +1,34 @@
 """
 Zep CrewAI Integration.
 
-This package provides memory integration between Zep and CrewAI agents,
-enabling persistent conversation memory and context retrieval.
+This package provides comprehensive memory integration between Zep and CrewAI agents,
+including graph storage, user storage, and tools for searching and adding data.
 
 Installation:
     pip install zep-crewai
 
 Usage:
-    from zep_crewai import ZepStorage
-    from zep_cloud.client import AsyncZep
+    from zep_crewai import ZepGraphStorage, ZepUserStorage, create_search_tool
+    from zep_cloud.client import Zep
     from crewai.memory.external.external_memory import ExternalMemory
-    from crewai import Crew
-
-    # Initialize Zep client and storage
-    zep_client = AsyncZep(api_key="your-api-key")
-    zep_storage = ZepStorage(client=zep_client, user_id="user123", thread_id="thread123")
-    external_memory = ExternalMemory(storage=zep_storage)
-
-    # Create crew with Zep memory
-    crew = Crew(
-        agents=[...],
-        tasks=[...],
-        external_memory=external_memory
+    from crewai import Agent, Crew
+
+    # Initialize Zep client
+    zep_client = Zep(api_key="your-api-key")
+
+    # For user-specific storage
+    user_storage = ZepUserStorage(client=zep_client, user_id="user123", thread_id="thread123")
+
+    # For generic knowledge graphs
+    graph_storage = ZepGraphStorage(client=zep_client, graph_id="knowledge_base")
+
+    # Create tools
+    search_tool = create_search_tool(zep_client, user_id="user123")
+
+    # Create agent with tools
+    agent = Agent(
+        role="Assistant",
+        tools=[search_tool]
     )
 """
 
@@ -35,11 +41,29 @@ from .exceptions import ZepDependencyError
 try:
     # Check for required CrewAI dependencies - just test import
     import crewai.memory.storage.interface  # noqa: F401
+    import crewai.tools  # noqa: F401
 
-    # Import our integration
+    # Import our integration components
+    from .graph_storage import ZepGraphStorage
     from .memory import ZepStorage
+    from .tools import (
+        ZepAddDataTool,
+        ZepSearchTool,
+        create_add_data_tool,
+        create_search_tool,
+    )
+    from .user_storage import ZepUserStorage
 
-    __all__ = ["ZepStorage", "ZepDependencyError"]
+    __all__ = [
+        "ZepStorage",
+        "ZepGraphStorage",
+        "ZepUserStorage",
+        "ZepSearchTool",
+        "ZepAddDataTool",
+        "create_search_tool",
+        "create_add_data_tool",
+        "ZepDependencyError",
+    ]
 
 except ImportError as e:
     raise ZepDependencyError(framework="CrewAI", install_command="pip install zep-crewai") from e

zep_crewai/graph_storage.py

@@ -0,0 +1,147 @@
+"""
+Zep Graph Storage for CrewAI.
+
+This module provides graph-based storage that integrates Zep's knowledge graph
+capabilities with CrewAI's memory system.
+"""
+
+import logging
+from typing import Any
+
+from crewai.memory.storage.interface import Storage
+from zep_cloud.client import Zep
+from zep_cloud.types import SearchFilters
+
+from .utils import search_graph_and_compose_context
+
+
+class ZepGraphStorage(Storage):
+    """
+    Storage implementation for Zep's generic knowledge graphs.
+
+    This class provides persistent storage and retrieval of structured knowledge
+    using Zep's graph capabilities for non-user-specific data.
+    """
+
+    def __init__(
+        self,
+        client: Zep,
+        graph_id: str,
+        search_filters: SearchFilters | None = None,
+        facts_limit: int = 20,
+        entity_limit: int = 5,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Initialize ZepGraphStorage with a Zep client instance.
+
+        Args:
+            client: An initialized Zep instance (sync client)
+            graph_id: Identifier for the knowledge graph
+            search_filters: Optional filters for search operations
+            facts_limit: Maximum number of facts (edges) to retrieve for context
+            entity_limit: Maximum number of entities (nodes) to retrieve for context
+            **kwargs: Additional configuration options
+        """
+        if not isinstance(client, Zep):
+            raise TypeError("client must be an instance of Zep")
+
+        if not graph_id:
+            raise ValueError("graph_id is required")
+
+        self._client = client
+        self._graph_id = graph_id
+        self._search_filters = search_filters
+        self._facts_limit = facts_limit
+        self._entity_limit = entity_limit
+        self._config = kwargs
+
+        self._logger = logging.getLogger(__name__)
+
+    def save(self, value: Any, metadata: dict[str, Any] | None = None) -> None:
+        """
+        Save data to the Zep knowledge graph.
+
+        Routes storage based on metadata.type:
+        - "json": Store as JSON data
+        - "text": Store as text data (default)
+        - "message": Store as message data
+
+        Args:
+            value: The content to store
+            metadata: Metadata including type information
+        """
+        metadata = metadata or {}
+        content_str = str(value)
+        content_type = metadata.get("type", "text")
+
+        # Validate content type
+        if content_type not in ["json", "text", "message"]:
+            content_type = "text"
+
+        try:
+            # Add data to the graph
+            self._client.graph.add(
+                graph_id=self._graph_id,
+                data=content_str,
+                type=content_type,
+            )
+
+            self._logger.debug(
+                f"Saved {content_type} data to graph {self._graph_id}: {content_str[:100]}..."
+            )
+
+        except Exception as e:
+            self._logger.error(f"Error saving to Zep graph: {e}")
+            raise
+
+    def search(
+        self, query: str, limit: int = 10, score_threshold: float = 0.0
+    ) -> dict[str, Any] | list[Any]:
+        """
+        Search the Zep knowledge graph and return composed context.
+
+        Performs parallel searches across edges, nodes, and episodes,
+        then returns a composed context string.
+
+        Args:
+            query: Search query string from the agent
+            limit: Maximum number of results per scope
+            score_threshold: Minimum relevance score (not used in Zep, kept for interface compatibility)
+
+        Returns:
+            List with a single dict containing the composed context string
+        """
+        try:
+            # Use the shared utility function for graph search and context composition
+            context = search_graph_and_compose_context(
+                client=self._client,
+                query=query,
+                graph_id=self._graph_id,
+                facts_limit=self._facts_limit,
+                entity_limit=self._entity_limit,
+                episodes_limit=limit,
+                search_filters=self._search_filters,
+            )
+
+            if context:
+                self._logger.info(f"Composed context for query: {query}")
+                return [
+                    {"memory": context, "type": "graph_context", "source": "graph", "query": query}
+                ]
+
+            self._logger.info(f"No results found for query: {query}")
+            return []
+
+        except Exception as e:
+            self._logger.error(f"Error searching graph: {e}")
+            return []
+
+    def reset(self) -> None:
+        """Reset is not implemented for graph storage as graphs should persist."""
+        pass
+
+    @property
+    def graph_id(self) -> str:
+        """Get the graph ID."""
+        return self._graph_id

zep_crewai/tools.py

@@ -0,0 +1,318 @@
+"""
+Zep CrewAI Tools.
+
+This module provides CrewAI tools for interacting with Zep memory storage,
+including graph and user memory operations.
+"""
+
+import logging
+from typing import Any
+
+from crewai.tools import BaseTool
+from pydantic import BaseModel, Field
+from zep_cloud.client import Zep
+
+logger = logging.getLogger(__name__)
+
+
+class SearchMemoryInput(BaseModel):
+    """Input schema for memory search tool."""
+
+    query: str = Field(..., description="The search query to find relevant memories")
+    limit: int = Field(default=10, description="Maximum number of results to return")
+    scope: str = Field(
+        default="edges",
+        description="Scope of search: 'edges' (facts), 'nodes' (entities), 'episodes', or 'all'",
+    )
+
+
+class AddGraphDataInput(BaseModel):
+    """Input schema for adding data to graph."""
+
+    data: str = Field(..., description="The data/information to store in the graph")
+    data_type: str = Field(default="text", description="Type of data: 'text', 'json', or 'message'")
+
+
+class ZepSearchTool(BaseTool):
+    """
+    Tool for searching Zep memory storage.
+
+    Can search either graph memory or user memory depending on initialization.
+    """
+
+    name: str = "Zep Memory Search"
+    description: str = "Search Zep memory storage for relevant information"
+    args_schema: type[BaseModel] = SearchMemoryInput
+
+    def __init__(
+        self, client: Zep, graph_id: str | None = None, user_id: str | None = None, **kwargs: Any
+    ):
+        """
+        Initialize search tool bound to either a graph or user.
+
+        Args:
+            client: Zep client instance
+            graph_id: Graph ID for generic knowledge graph search
+            user_id: User ID for user-specific graph search
+            **kwargs: Additional configuration
+        """
+        if not graph_id and not user_id:
+            raise ValueError("Either graph_id or user_id must be provided")
+
+        if graph_id and user_id:
+            raise ValueError("Only one of graph_id or user_id should be provided")
+
+        # Update description based on target
+        if graph_id:
+            kwargs["description"] = f"Search Zep graph '{graph_id}' for relevant information"
+        else:
+            kwargs["description"] = f"Search user '{user_id}' memories for relevant information"
+
+        super().__init__(**kwargs)
+
+        # Store as private attributes to avoid Pydantic validation
+        self._client = client
+        self._graph_id = graph_id
+        self._user_id = user_id
+
+    @property
+    def client(self) -> Zep:
+        """Get the Zep client."""
+        return self._client
+
+    @property
+    def graph_id(self) -> str | None:
+        """Get the graph ID."""
+        return self._graph_id
+
+    @property
+    def user_id(self) -> str | None:
+        """Get the user ID."""
+        return self._user_id
+
+    def _run(self, query: str, limit: int = 10, scope: str = "edges") -> str:
+        """
+        Execute the search operation.
+
+        Args:
+            query: Search query
+            limit: Maximum results
+            scope: Search scope
+
+        Returns:
+            Formatted search results
+        """
+        try:
+            results = []
+
+            if scope == "all":
+                # Search all scopes
+                scopes = ["edges", "nodes", "episodes"]
+            else:
+                scopes = [scope]
+
+            for search_scope in scopes:
+                if self._graph_id:
+                    # Search graph memory
+                    search_results = self._client.graph.search(
+                        graph_id=self._graph_id, query=query, limit=limit, scope=search_scope
+                    )
+                else:
+                    # Search user memory
+                    search_results = self._client.graph.search(
+                        user_id=self._user_id, query=query, limit=limit, scope=search_scope
+                    )
+
+                # Process results based on scope
+                if search_scope == "edges" and search_results.edges:
+                    for edge in search_results.edges:
+                        results.append(
+                            {
+                                "type": "fact",
+                                "content": edge.fact,
+                                "name": edge.name,
+                                "created_at": str(edge.created_at) if edge.created_at else None,
+                            }
+                        )
+
+                elif search_scope == "nodes" and search_results.nodes:
+                    for node in search_results.nodes:
+                        results.append(
+                            {
+                                "type": "entity",
+                                "content": f"{node.name}: {node.summary}",
+                                "name": node.name,
+                                "created_at": str(node.created_at) if node.created_at else None,
+                            }
+                        )
+
+                elif search_scope == "episodes" and search_results.episodes:
+                    for episode in search_results.episodes:
+                        results.append(
+                            {
+                                "type": "episode",
+                                "content": episode.content,
+                                "source": episode.source,
+                                "role": episode.role,
+                                "created_at": str(episode.created_at)
+                                if episode.created_at
+                                else None,
+                            }
+                        )
+
+            if not results:
+                return f"No results found for query: '{query}'"
+
+            # Format results for agent consumption
+            formatted = f"Found {len(results)} relevant memories:\n\n"
+            for i, result in enumerate(results, 1):
+                result_type = result.get("type", "unknown")
+                formatted += f"{i}. [{result_type.upper() if result_type else 'UNKNOWN'}] {result['content']}\n"
+                if result.get("created_at"):
+                    formatted += f"   (Created: {result['created_at']})\n"
+                formatted += "\n"
+
+            logger.info(f"Found {len(results)} memories for query: {query}")
+            return formatted
+
+        except Exception as e:
+            error_msg = f"Error searching Zep memory: {str(e)}"
+            logger.error(error_msg)
+            return error_msg
+
+
+class ZepAddDataTool(BaseTool):
+    """
+    Tool for adding data to Zep memory storage.
+
+    Can add data to either graph memory or user memory depending on initialization.
+    """
+
+    name: str = "Zep Add Data"
+    description: str = "Add data to Zep memory storage"
+    args_schema: type[BaseModel] = AddGraphDataInput
+
+    def __init__(
+        self, client: Zep, graph_id: str | None = None, user_id: str | None = None, **kwargs: Any
+    ):
+        """
+        Initialize add data tool bound to either a graph or user.
+
+        Args:
+            client: Zep client instance
+            graph_id: Graph ID for generic knowledge graph
+            user_id: User ID for user-specific graph
+            **kwargs: Additional configuration
+        """
+        if not graph_id and not user_id:
+            raise ValueError("Either graph_id or user_id must be provided")
+
+        if graph_id and user_id:
+            raise ValueError("Only one of graph_id or user_id should be provided")
+
+        # Update description based on target
+        if graph_id:
+            kwargs["description"] = f"Add data to Zep graph '{graph_id}'"
+        else:
+            kwargs["description"] = f"Add data to user '{user_id}' memory"
+
+        super().__init__(**kwargs)
+
+        # Store as private attributes to avoid Pydantic validation
+        self._client = client
+        self._graph_id = graph_id
+        self._user_id = user_id
+
+    @property
+    def client(self) -> Zep:
+        """Get the Zep client."""
+        return self._client
+
+    @property
+    def graph_id(self) -> str | None:
+        """Get the graph ID."""
+        return self._graph_id
+
+    @property
+    def user_id(self) -> str | None:
+        """Get the user ID."""
+        return self._user_id
+
+    def _run(self, data: str, data_type: str = "text") -> str:
+        """
+        Execute the add data operation.
+
+        Args:
+            data: Data to store
+            data_type: Type of data
+
+        Returns:
+            Success or error message
+        """
+        try:
+            # Validate data type
+            if data_type not in ["text", "json", "message"]:
+                data_type = "text"
+
+            if self._graph_id:
+                # Add to graph memory
+                self._client.graph.add(graph_id=self._graph_id, type=data_type, data=data)
+
+                success_msg = f"Successfully added {data_type} data to graph '{self._graph_id}'"
+                logger.debug(f"Added data to graph {self._graph_id}: {data[:100]}...")
+
+            else:
+                # Add to user graph memory
+                self._client.graph.add(user_id=self._user_id, type=data_type, data=data)
+
+                success_msg = (
+                    f"Successfully added {data_type} data to user '{self._user_id}' memory"
+                )
+                logger.debug(f"Added data to user {self._user_id}: {data[:100]}...")
+
+            return success_msg
+
+        except Exception as e:
+            error_msg = f"Error adding data to Zep: {str(e)}"
+            logger.error(error_msg)
+            return error_msg
+
+
+def create_search_tool(
+    client: Zep, graph_id: str | None = None, user_id: str | None = None
+) -> ZepSearchTool:
+    """
+    Create a search tool bound to a Zep client.
+
+    Args:
+        client: Zep client instance
+        graph_id: Optional graph ID for generic knowledge graph
+        user_id: Optional user ID for user-specific graph
+
+    Returns:
+        ZepSearchTool instance
+
+    Raises:
+        ValueError: If neither or both IDs are provided
+    """
+    return ZepSearchTool(client=client, graph_id=graph_id, user_id=user_id)
+
+
+def create_add_data_tool(
+    client: Zep, graph_id: str | None = None, user_id: str | None = None
+) -> ZepAddDataTool:
+    """
+    Create an add data tool bound to a Zep client.
+
+    Args:
+        client: Zep client instance
+        graph_id: Optional graph ID for generic knowledge graph
+        user_id: Optional user ID for user-specific graph
+
+    Returns:
+        ZepAddDataTool instance
+
+    Raises:
+        ValueError: If neither or both IDs are provided
+    """
+    return ZepAddDataTool(client=client, graph_id=graph_id, user_id=user_id)