diff --git a/docs/concepts/memory.mdx b/docs/concepts/memory.mdx index da228f4e92..431f87816a 100644 --- a/docs/concepts/memory.mdx +++ b/docs/concepts/memory.mdx @@ -720,7 +720,16 @@ crew = Crew( ``` ### Advanced Mem0 Configuration +When using Mem0 Client, you can customize the memory configuration further, by using parameters like 'includes', 'excludes', 'custom_categories' and 'run_id' (this is only for short-term memory). +You can find more details in the [Mem0 documentation](https://docs.mem0.ai/). ```python + +new_categories = [ + {"lifestyle_management_concerns": "Tracks daily routines, habits, hobbies and interests including cooking, time management and work-life balance"}, + {"seeking_structure": "Documents goals around creating routines, schedules, and organized systems in various life areas"}, + {"personal_information": "Basic information about the user including name, preferences, and personality traits"} +] + crew = Crew( agents=[...], tasks=[...], @@ -732,6 +741,10 @@ crew = Crew( "org_id": "my_org_id", # Optional "project_id": "my_project_id", # Optional "api_key": "custom-api-key" # Optional - overrides env var + "run_id": "my_run_id", # Optional - for short-term memory + "includes": "include1", # Optional + "excludes": "exclude1", # Optional + "custom_categories": new_categories # Optional - custom categories for user memory }, "user_memory": {} } diff --git a/src/crewai/memory/storage/mem0_storage.py b/src/crewai/memory/storage/mem0_storage.py index 4fe75e3f5b..07a108ef14 100644 --- a/src/crewai/memory/storage/mem0_storage.py +++ b/src/crewai/memory/storage/mem0_storage.py @@ -13,47 +13,67 @@ class Mem0Storage(Storage): """ Extends Storage to handle embedding and searching across entities using Mem0. """ - def __init__(self, type, crew=None, config=None): super().__init__() - supported_types = ["user", "short_term", "long_term", "entities", "external"] - if type not in supported_types: - raise ValueError( - f"Invalid type '{type}' for Mem0Storage. Must be one of: " - + ", ".join(supported_types) - ) + self._validate_type(type) self.memory_type = type self.crew = crew - self.config = config or {} + # TODO: Memory config will be removed in the future the config will be passed as a parameter - self.memory_config = self.config or getattr(crew, "memory_config", {}) or {} + self.config = config or getattr(crew, "memory_config", {}).get("config", {}) or {} + + self._validate_user_id() + self._extract_config_values() + self._initialize_memory() - # User ID is required for user memory type "user" since it's used as a unique identifier for the user. - user_id = self._get_user_id() - if type == "user" and not user_id: + def _validate_type(self, type): + supported_types = {"user", "short_term", "long_term", "entities", "external"} + if type not in supported_types: + raise ValueError( + f"Invalid type '{type}' for Mem0Storage. Must be one of: {', '.join(supported_types)}" + ) + + def _validate_user_id(self): + if self.memory_type == "user" and not self.config.get("user_id", ""): raise ValueError("User ID is required for user memory type") - # API key in memory config overrides the environment variable - config = self._get_config() - mem0_api_key = config.get("api_key") or os.getenv("MEM0_API_KEY") - mem0_org_id = config.get("org_id") - mem0_project_id = config.get("project_id") - mem0_local_config = config.get("local_mem0_config") - - # Initialize MemoryClient or Memory based on the presence of the mem0_api_key - if mem0_api_key: - if mem0_org_id and mem0_project_id: - self.memory = MemoryClient( - api_key=mem0_api_key, org_id=mem0_org_id, project_id=mem0_project_id - ) - else: - self.memory = MemoryClient(api_key=mem0_api_key) + def _extract_config_values(self): + cfg = self.config + self.mem0_run_id = cfg.get("run_id") + self.includes = cfg.get("includes") + self.excludes = cfg.get("excludes") + self.custom_categories = cfg.get("custom_categories") + + def _initialize_memory(self): + api_key = self.config.get("api_key") or os.getenv("MEM0_API_KEY") + org_id = self.config.get("org_id") + project_id = self.config.get("project_id") + local_config = self.config.get("local_mem0_config") + + if api_key: + self.memory = ( + MemoryClient(api_key=api_key, org_id=org_id, project_id=project_id) + if org_id and project_id + else MemoryClient(api_key=api_key) + ) + if self.custom_categories: + self.memory.update_project(custom_categories=self.custom_categories) else: - if mem0_local_config and len(mem0_local_config): - self.memory = Memory.from_config(mem0_local_config) - else: - self.memory = Memory() + self.memory = ( + Memory.from_config(local_config) + if local_config and len(local_config) + else Memory() + ) + + def _get_agent_name(self) -> str: + if not self.crew: + return "" + + agents = self.crew.agents + agents = [self._sanitize_role(agent.role) for agent in agents] + agents = "_".join(agents) + return sanitize_collection_name(name=agents,max_collection_length=MAX_AGENT_ID_LENGTH_MEM0) def _sanitize_role(self, role: str) -> str: """ @@ -62,86 +82,78 @@ def _sanitize_role(self, role: str) -> str: return role.replace("\n", "").replace(" ", "_").replace("/", "_") def save(self, value: Any, metadata: Dict[str, Any]) -> None: - user_id = self._get_user_id() + user_id = self.config.get("user_id", "") agent_name = self._get_agent_name() - params = None + + base_metadata = { + "short_term": "short_term", + "long_term": "long_term", + "entities": "entity", + "external": "external" + } + + # Shared base params + params: dict[str, Any] = { + "agent_id": agent_name, + "metadata": {"type": base_metadata[self.memory_type], **metadata} + } + + # Type-specific overrides if self.memory_type == "short_term": - params = { - "agent_id": agent_name, - "infer": False, - "metadata": {"type": "short_term", **metadata}, - } + params["infer"] = False elif self.memory_type == "long_term": - params = { - "agent_id": agent_name, - "infer": False, - "metadata": {"type": "long_term", **metadata}, - } + params["infer"] = False elif self.memory_type == "entities": - params = { - "agent_id": agent_name, - "infer": False, - "metadata": {"type": "entity", **metadata}, - } + params["infer"] = False elif self.memory_type == "external": - params = { - "user_id": user_id, - "agent_id": agent_name, - "metadata": {"type": "external", **metadata}, - } + params["user_id"] = user_id + if params: + # MemoryClient-specific overrides if isinstance(self.memory, MemoryClient): - params["output_format"] = "v1.1" + params["version"] = "v2" + params["includes"] = self.includes + params["excludes"] = self.excludes + + if self.memory_type == "short_term": + params["run_id"] = self.mem0_run_id + self.memory.add(value, **params) - def search( - self, - query: str, - limit: int = 3, - score_threshold: float = 0.35, - ) -> List[Any]: - params = {"query": query, "limit": limit, "output_format": "v1.1"} - if user_id := self._get_user_id(): + def search(self,query: str,limit: int = 3,score_threshold: float = 0.35) -> List[Any]: + params = { + "query": query, + "limit": limit, + "version": "v2" + } + + if user_id := self.config.get("user_id", ""): params["user_id"] = user_id agent_name = self._get_agent_name() - if self.memory_type == "short_term": - params["agent_id"] = agent_name - params["metadata"] = {"type": "short_term"} - elif self.memory_type == "long_term": - params["agent_id"] = agent_name - params["metadata"] = {"type": "long_term"} - elif self.memory_type == "entities": - params["agent_id"] = agent_name - params["metadata"] = {"type": "entity"} - elif self.memory_type == "external": - params["agent_id"] = agent_name - params["metadata"] = {"type": "external"} + params["agent_id"] = agent_name + + memory_type_map = { + "short_term": {"type": "short_term"}, + "long_term": {"type": "long_term"}, + "entities": {"type": "entity"}, + "external": {"type": "external"}, + } + + if self.memory_type in memory_type_map: + params["metadata"] = memory_type_map[self.memory_type] + if self.memory_type == "short_term": + params["run_id"] = self.mem0_run_id # Discard the filters for now since we create the filters # automatically when the crew is created. if isinstance(self.memory, Memory): - del params["metadata"], params["output_format"] + del params["metadata"], params["version"], params["run_id"] results = self.memory.search(**params) return [r for r in results["results"] if r["score"] >= score_threshold] - def _get_user_id(self) -> str: - return self._get_config().get("user_id", "") - - def _get_agent_name(self) -> str: - if not self.crew: - return "" - - agents = self.crew.agents - agents = [self._sanitize_role(agent.role) for agent in agents] - agents = "_".join(agents) - return sanitize_collection_name(name=agents,max_collection_length=MAX_AGENT_ID_LENGTH_MEM0) - - def _get_config(self) -> Dict[str, Any]: - return self.config or getattr(self, "memory_config", {}).get("config", {}) or {} - def reset(self): if self.memory: self.memory.reset() diff --git a/tests/storage/test_mem0_storage.py b/tests/storage/test_mem0_storage.py index b651521f92..3f70410252 100644 --- a/tests/storage/test_mem0_storage.py +++ b/tests/storage/test_mem0_storage.py @@ -1,14 +1,10 @@ -import os from unittest.mock import MagicMock, patch import pytest from mem0.client.main import MemoryClient from mem0.memory.main import Memory -from crewai.agent import Agent -from crewai.crew import Crew from crewai.memory.storage.mem0_storage import Mem0Storage -from crewai.task import Task # Define the class (if not already defined) @@ -59,10 +55,11 @@ def mem0_storage_with_mocked_config(mock_mem0_memory): } # Instantiate the class with memory_config + # Parameters like run_id, includes, and excludes doesn't matter in Memory OSS crew = MockCrew( memory_config={ "provider": "mem0", - "config": {"user_id": "test_user", "local_mem0_config": config}, + "config": {"user_id": "test_user", "local_mem0_config": config, "run_id": "my_run_id", "includes": "include1","excludes": "exclude1"}, } ) @@ -99,6 +96,9 @@ def mem0_storage_with_memory_client_using_config_from_crew(mock_mem0_memory_clie "api_key": "ABCDEFGH", "org_id": "my_org_id", "project_id": "my_project_id", + "run_id": "my_run_id", + "includes": "include1", + "excludes": "exclude1", }, } ) @@ -154,11 +154,37 @@ def test_mem0_storage_with_explict_config( assert ( mem0_storage_with_memory_client_using_explictly_config.config == expected_config ) - assert ( - mem0_storage_with_memory_client_using_explictly_config.memory_config - == expected_config + + +def test_mem0_storage_updates_project_with_custom_categories(mock_mem0_memory_client): + mock_mem0_memory_client.update_project = MagicMock() + + new_categories = [ + {"lifestyle_management_concerns": "Tracks daily routines, habits, hobbies and interests including cooking, time management and work-life balance"}, + ] + + crew = MockCrew( + memory_config={ + "provider": "mem0", + "config": { + "user_id": "test_user", + "api_key": "ABCDEFGH", + "org_id": "my_org_id", + "project_id": "my_project_id", + "custom_categories": new_categories, + }, + } ) + with patch.object(MemoryClient, "__new__", return_value=mock_mem0_memory_client): + _ = Mem0Storage(type="short_term", crew=crew) + + mock_mem0_memory_client.update_project.assert_called_once_with( + custom_categories=new_categories + ) + + + def test_save_method_with_memory_oss(mem0_storage_with_mocked_config): """Test save method for different memory types""" @@ -195,7 +221,10 @@ def test_save_method_with_memory_client(mem0_storage_with_memory_client_using_co agent_id="Test_Agent", infer=False, metadata={"type": "short_term", "key": "value"}, - output_format="v1.1" + version="v2", + run_id="my_run_id", + includes="include1", + excludes="exclude1", ) @@ -232,7 +261,8 @@ def test_search_method_with_memory_client(mem0_storage_with_memory_client_using_ agent_id="Test_Agent", metadata={"type": "short_term"}, user_id="test_user", - output_format='v1.1' + version='v2', + run_id="my_run_id", ) assert len(results) == 1