Add support for memory distinguished by custom key (resolves #2584)

Co-Authored-By: Joe Moura <joao@crewai.com>

tests/memory/custom_key_memory_test.py

@@ -0,0 +1,57 @@
+import pytest
+from unittest.mock import patch, MagicMock
+
+from crewai.memory.short_term.short_term_memory import ShortTermMemory
+from crewai.memory.short_term.short_term_memory_item import ShortTermMemoryItem
+from crewai.agent import Agent
+from crewai.crew import Crew
+from crewai.task import Task
+
+
+@pytest.fixture
+def short_term_memory():
+    """Fixture to create a ShortTermMemory instance"""
+    agent = Agent(
+        role="Researcher",
+        goal="Search relevant data and provide results",
+        backstory="You are a researcher at a leading tech think tank.",
+        tools=[],
+        verbose=True,
+    )
+
+    task = Task(
+        description="Perform a search on specific topics.",
+        expected_output="A list of relevant URLs based on the search query.",
+        agent=agent,
+    )
+    return ShortTermMemory(crew=Crew(agents=[agent], tasks=[task]))
+
+
+def test_save_with_custom_key(short_term_memory):
+    """Test that save method correctly passes custom_key to storage"""
+    with patch.object(short_term_memory.storage, 'save') as mock_save:
+        short_term_memory.save(
+            value="Test data",
+            metadata={"task": "test_task"},
+            agent="test_agent",
+            custom_key="user123",
+        )
+        
+        called_args = mock_save.call_args[0]
+        called_kwargs = mock_save.call_args[1]
+        
+        assert "custom_key" in called_args[1]
+        assert called_args[1]["custom_key"] == "user123"
+
+
+def test_search_with_custom_key(short_term_memory):
+    """Test that search method correctly passes custom_key to storage"""
+    expected_results = [{"context": "Test data", "metadata": {"custom_key": "user123"}, "score": 0.95}]
+    
+    with patch.object(short_term_memory.storage, 'search', return_value=expected_results) as mock_search:
+        results = short_term_memory.search("test query", custom_key="user123")
+        
+        mock_search.assert_called_once()
+        filter_arg = mock_search.call_args[1].get('filter')
+        assert filter_arg == {"custom_key": {"$eq": "user123"}}
+        assert results == expected_results