Update Mem0Storage to support v2 API with run_id parameter

[devin-ai-integration]

Update Mem0Storage to support v2 API with run_id parameter

Description

This PR updates the Mem0Storage class to support Mem0's v2 API, which includes support for memories associated with specific conversation sessions using the run_id parameter.

Changes

• Added support for version parameter (defaulting to v2)
• Added support for run_id parameter
• Updated save and search methods to include version and run_id
• Added tests for the updated implementation
• Updated existing tests to work with the new implementation

Related Issue

Fixes #2776

Link to Devin run

https://app.devin.ai/sessions/e4382d66a48f4e0b9b49c323ce534c9a

Requested by

Joe Moura (joao@crewai.com)

Example Usage

crew = Crew(
    agents=[agent1, agent2],
    tasks=[task1, task2],
    memory_config={
        "provider": "mem0",
        "config": {
            "version": "v2",  # Optional, defaults to v2
            "run_id": "session-123",  # Optional, to associate memories with a specific session
            "api_key": "your-mem0-api-key"  # Optional, defaults to MEM0_API_KEY env var
        }
    }
)

[devin-ai-integration]

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

• Address comments on this PR. Add '(aside)' to your comment to have me ignore it.
• Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

• Disable automatic comment and CI monitoring

[joaomdmoura]

Disclaimer: This review was made by a crew of AI Agents.

Code Review Comment for PR #2777: Mem0Storage V2 API Support

Overview

The implementation of support for v2 API in the Mem0Storage is commendable and allows for improved session management through the introduction of the run_id parameter. This change enhances how memories are handled, especially in conversational contexts where context-switching is critical.

Code Quality Findings

1. src/crewai/memory/storage/mem0_storage.py

a) Configuration Initialization

• Current Implementation:

  self.version = config.get("version", "v2")
  self.run_id = config.get("run_id")

• Improvement Suggestion:
  It would be beneficial to implement a configuration validation method to ensure that the specified version is acceptable and to raise informative errors when necessary. This will enhance robustness.

  def _validate_config(self):
      """Validate configuration parameters."""
      if self.version not in ["v1.1", "v2"]:
          raise ValueError(f"Unsupported version: {self.version}. Please use 'v1.1' or 'v2'.")

b) Parameter Management

• Issue: Parameter handling across methods is scattered, potentially leading to maintenance challenges.

• Recommendation:
  Introduce a centralized parameter management method to streamline the process of parameter construction.

  def _build_params(self, base_params: Dict[str, Any]) -> Dict[str, Any]:
      """Centralize parameter building."""
      if isinstance(self.memory, MemoryClient):
          base_params["version"] = self.version
          if self.run_id:
              base_params["run_id"] = self.run_id
      return base_params

c) Search Method Optimization

• Current Implementation:
  The search method is somewhat fragmented, affecting readability and potential performance.

• Recommendation:
  Creating a helper function to construct and manage search parameters will enhance both efficiency and clarity.

  def search(self, query: str, limit: int = 3, score_threshold: float = 0.35) -> List[Any]:
      params = self._build_params({
          "query": query,
          "limit": limit,
          "metadata": {"type": self.type} if hasattr(self, "type") else None
      })
      results = self.memory.search(**params)
      return [r for r in results["results"] if r["score"] >= score_threshold]

2. tests/storage/test_mem0_storage.py

• Improvement: Test Organization
  Introduce pytest marks to better categorize tests, making it easier to run tests related to specific API versions.

  @pytest.mark.v2_api

3. tests/storage/test_mem0_storage_v2.py

Positive Aspects

• The coverage for the new functionality is comprehensive, and the separation of test cases is commendable. Clear naming conventions enhance the maintainability of these tests.

Suggestions for Improvement

• Edge Case Testing: Add more tests to cover various edge cases for the run_id parameter, ensuring that testing reflects different lengths and formats.

  @pytest.mark.parametrize("run_id", [None, "", "test-123", "a" * 256])
  def test_run_id_edge_cases(mem0_storage_with_v2_api, run_id):
      mem0_storage, mock_client = mem0_storage_with_v2_api
      mem0_storage.run_id = run_id
      # Test implementation for edge cases

• Error Handling Tests: Ensure that tests are included for invalid configurations which help bolster the resilience of the implementation.

  def test_invalid_version_handling():
      with pytest.raises(ValueError, match="Unsupported version"):
          Mem0Storage(type="short_term", crew=MockCrew(
              memory_config={"config": {"version": "invalid"}}
          ))

General Recommendations

• Documentation: Ensure type hints are provided consistently, and that docstrings explain v2 API-specific features and migration paths from v1.1 to v2.
• Error Handling: Implement robust validation and error responses for unsupported versions, and ensure that important operations are logged.
• Performance Considerations: Evaluate caching mechanisms for frequently accessed parameters, optimize result filtering, and consider batch operations to improve performance.
• Code Organization: Modularize version-specific code and use constants for repeated values throughout the codebase.

Security Considerations

• Validate the run_id input to prevent potential injection attacks or improper requests.
• Proper encoding of search queries and establishing rate limits on memory operations can significantly improve security.

Migration Guide Suggestion

To assist users in migrating from v1.1 to v2 API, a clear migration guide is essential.

```markdown
## Migrating to v2 API

1. Update configuration:
```python
memory_config = {
    "provider": "mem0",
    "config": {
        "version": "v2",
        "run_id": "your-session-id"  # Optional
    }
}
```

2. Update search and memory storage calls to adapt to the new response formats and inclusion of the run_id.
```

These changes significantly enhance the functionality and user experience of the Mem0Storage component, preparing it for future extensions and robustness. Thank you for your hard work on this PR!

[rusXL]

@joaomdmoura

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.
You should rather call memory.add() twice, once with run_id for short-term session memory and once without run_id for long-term memory.

Secondly, there is no ability to turn off agent memories. The agent name is always being passed as agent ID with every add API call, but in some situations, it is necessary that agents are memory-less.

Thirdly, it would be cool to have support for mem0 new features such as memory inclusion, custom categories, custom instructions, etc. - at the project (mem0 client) level and every add API call.

[rusXL]

@joaomdmoura

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. You should rather call memory.add() twice, once with run_id for short-term session memory and once without run_id for long-term memory.

Secondly, there is no ability to turn off agent memories. The agent name is always being passed as agent ID with every add API call, but in some situations, it is necessary that agents are memory-less.

Thirdly, it would be cool to have support for mem0 new features such as memory inclusion, custom categories, custom instructions, etc. - at the project (mem0 client) level and every add API call.

UPD: sorry this regards short-term and long-term memory with mem0 (not external)