wip

src/crewai/memory/storage/mem0_storage.py

@@ -70,7 +70,7 @@ class Mem0Storage(Storage):
         """
         Returns:
             dict: A filter dictionary containing AND conditions for querying data.
-                - Includes user_id if memory_type is 'external'.
+                - Includes user_id, if memory_type is 'external'.
                 - Includes run_id if memory_type is 'short_term' and mem0_run_id is present.
         """
         filter = {
@@ -79,7 +79,14 @@ class Mem0Storage(Storage):
 
         # Add user_id condition if the memory type is external
         if self.memory_type == "external":
-            filter["AND"].append({"user_id": self.config.get("user_id", "")})
+            user_id = self.config.get("user_id", "")
+            agent_id = self.config.get("agent_id", "")
+
+            if user_id:
+                filter["AND"].append({"user_id": user_id})
+
+            if agent_id:
+                filter["AND"].append({"agent_id": agent_id})
 
         # Add run_id condition if the memory type is short_term and a run ID is set
         if self.memory_type == "short_term" and self.mem0_run_id:
@@ -89,6 +96,7 @@ class Mem0Storage(Storage):
 
     def save(self, value: Any, metadata: Dict[str, Any]) -> None:
         user_id = self.config.get("user_id", "")
+        agent_id = self.config.get("agent_id", "")
         assistant_message = [{"role" : "assistant","content" : value}]
 
         base_metadata = {
@@ -106,7 +114,7 @@ class Mem0Storage(Storage):
 
         if self.memory_type == "external":
             params["user_id"] = user_id
-
+            params["agent_id"] = agent_id
 
         if params:
             # MemoryClient-specific overrides

tests/storage/test_mem0_storage.py

@@ -270,3 +270,18 @@ def test_search_method_with_memory_client(mem0_storage_with_memory_client_using_
 
     assert len(results) == 2
     assert results[0]["content"] == "Result 1"
+
+def test_save_memory_using_agent_entity():
+    config = {
+        "agent_id": "agent-123",
+    }
+    mem0_storage = Mem0Storage(type="short_term", config=config)
+
+    mem0_storage.save("test memory", {"key": "value"})
+
+    mem0_storage.memory.add.assert_called_once_with(
+        [{'role': 'assistant' , 'content': 'test memory'}],
+        infer=True,
+        metadata={"type": "short_term", "key": "value"},
+        agent_id="agent-123",
+    )