Repository: crewAIInc/crewAI
File: src/crewai/memory/storage/mem0_storage.py (Python)

Code before:
    """

    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.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 {}

        # 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:
            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")
        self.mem0_run_id = config.get("run_id")

        # 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)
        else:
            if mem0_local_config and len(mem0_local_config):
                self.memory = Memory.from_config(mem0_local_config)
            else:
                self.memory = Memory()

    def _sanitize_role(self, role: str) -> str:
        """
        Sanitizes agent roles to ensure valid directory names.
        """
        return role.replace("\n", "").replace(" ", "_").replace("/", "_")

    def save(self, value: Any, metadata: Dict[str, Any]) -> None:

Review comment (suggestion):
First, notice if you provide `run_id`, that means you create short-term memory for a user session and no long-term memory is created - [mem0 docs](https://docs.mem0.ai/platform/quickstart#4-1-create-memories). 
`run_id` should be only allowed to be used for short-term memory.

Thirdly, it would be cool to have support for mem0 [new features](https://docs.mem0.ai/features/platform-overview) such as (I find the most relevant once) `memory inclusion`, `custom categories`.

You can add support for `custom categories` in init:
````py
def __init__(self, type, crew=None, config=None):
...
# get new_categories from config, then
self.memory.update_project(custom_categories=new_categories)
``` 

Code after:
    """
    Extends Storage to handle embedding and searching across entities using Mem0.
    """
    def __init__(self, type, crew=None, config=None):
        super().__init__()

        self._validate_type(type)
        self.memory_type = type
        self.crew = crew

        # TODO: Memory config will be removed in the future the config will be passed as a parameter
        self.config = config or getattr(crew, "memory_config", {}).get("config", {}) or {}

        self._validate_user_id()
        self._extract_config_values()
        self._initialize_memory()

    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")

    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")
        self.infer = cfg.get("infer", False)

    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: