Python: support mem0ai 2.x

python/packages/mem0/agent_framework_mem0/_context_provider.py

@@ -65,7 +65,10 @@ def __init__(
             source_id: Unique identifier for this provider instance.
             mem0_client: A pre-created Mem0 MemoryClient or None to create a default client.
             api_key: The API key for authenticating with the Mem0 API.
-            application_id: The application ID for scoping memories.
+            application_id: The application ID for scoping memories. Platform-only:
+                the OSS ``AsyncMemory`` client does not recognize an application
+                scope (it scopes only by user_id/agent_id in this provider), so
+                application_id cannot be used with an OSS client.
             agent_id: The agent ID for scoping memories.
             user_id: The user ID for scoping memories.
             context_prompt: The prompt to prepend to retrieved memories.
@@ -125,13 +128,9 @@ async def before_run(
             agent_kwargs = self._build_search_kwargs(input_text, "agent_id", self.agent_id)
             search_tasks.append(self.mem0_client.search(**agent_kwargs))  # type: ignore[reportUnknownMemberType, reportUnknownArgumentType]
 
-        # Fall back to an app-scoped search when only application_id is configured
+        # Fall back to an app-scoped search when only application_id is configured.
         if not search_tasks and self.application_id:
-            app_kwargs: dict[str, Any] = {"query": input_text}
-            if isinstance(self.mem0_client, AsyncMemory):
-                app_kwargs["app_id"] = self.application_id
-            else:
-                app_kwargs["filters"] = {"app_id": self.application_id}
+            app_kwargs: dict[str, Any] = {"query": input_text, "filters": self._build_filters()}
             search_tasks.append(self.mem0_client.search(**app_kwargs))  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
         if not search_tasks:
             return
@@ -219,43 +218,54 @@ def get_role_value(role: Any) -> str:
         if messages:
             add_kwargs: dict[str, Any] = {
                 "messages": messages,
-                "user_id": self.user_id,
-                "agent_id": self.agent_id,
             }
 
-            # Inject the application scope using the matching signature format for each SDK variant
             if isinstance(self.mem0_client, AsyncMemory):
-                if self.application_id:
-                    add_kwargs["app_id"] = self.application_id
+                add_kwargs["user_id"] = self.user_id
+                add_kwargs["agent_id"] = self.agent_id
             else:
-                if self.application_id:
-                    add_kwargs["filters"] = {"app_id": self.application_id}
+                add_kwargs["filters"] = self._build_filters()
 
             await self.mem0_client.add(**add_kwargs)  # type: ignore[misc, call-arg]
 
     # -- Internal methods ------------------------------------------------------
 
     def _validate_filters(self) -> None:
-        """Validates that at least one filter is provided."""
+        """Validates that at least one usable filter is provided for the configured client."""
         if not self.agent_id and not self.user_id and not self.application_id:
             raise ValueError("At least one of the filters: agent_id, user_id, or application_id is required.")
+        if isinstance(self.mem0_client, AsyncMemory) and self.application_id:
+            raise ValueError(
+                "application_id is not supported by the OSS AsyncMemory client, which scopes "
+                "memories only by user_id/agent_id. Remove application_id or use AsyncMemoryClient."
+            )
 
     def _build_search_kwargs(self, input_text: str, entity_key: str, entity_value: str) -> dict[str, Any]:
         """Build search keyword arguments formatted for OSS vs Platform clients."""
         filters: dict[str, Any] = {"query": input_text}
 
-        if isinstance(self.mem0_client, AsyncMemory):
-            # AsyncMemory (OSS) expects direct kwargs
-            filters[entity_key] = entity_value
-            if self.application_id:
-                filters["app_id"] = self.application_id
-        else:
-            # AsyncMemoryClient (Platform) expects a filters dict
-            filters["filters"] = {entity_key: entity_value}
-            if self.application_id:
-                filters["filters"]["app_id"] = self.application_id
+        if self.application_id and isinstance(self.mem0_client, AsyncMemory):
+            raise ValueError(
+                "application_id is not supported by the OSS AsyncMemory client, which scopes "
+                "memories only by user_id/agent_id. Remove application_id or use AsyncMemoryClient."
+            )
+
+        filters["filters"] = {entity_key: entity_value}
+        if self.application_id and not isinstance(self.mem0_client, AsyncMemory):
+            filters["filters"]["app_id"] = self.application_id
 
         return filters
 
+    def _build_filters(self) -> dict[str, Any]:
+        """Build identity filters from initialization parameters."""
+        filters: dict[str, Any] = {}
+        if self.user_id:
+            filters["user_id"] = self.user_id
+        if self.agent_id:
+            filters["agent_id"] = self.agent_id
+        if self.application_id:
+            filters["app_id"] = self.application_id
+        return filters
+
 
 __all__ = ["Mem0ContextProvider"]

python/packages/mem0/pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
 ]
 dependencies = [
     "agent-framework-core>=1.8.1,<2",
-    "mem0ai>=1.0.0,<2",
+    "mem0ai>=2.0.0,<3",
 ]
 
 [tool.uv]

python/packages/mem0/tests/test_mem0_context_provider.py

@@ -198,8 +198,8 @@ async def test_search_query_combines_input_messages(self, mock_mem0_client: Asyn
         call_kwargs = mock_mem0_client.search.call_args.kwargs
         assert call_kwargs["query"] == "Hello\nWorld"
 
-    async def test_oss_client_passes_direct_kwargs(self, mock_oss_mem0_client: AsyncMock) -> None:
-        """OSS AsyncMemory client should receive user_id as direct kwarg, not in filters."""
+    async def test_oss_client_passes_filters_dict(self, mock_oss_mem0_client: AsyncMock) -> None:
+        """OSS AsyncMemory client should receive entity IDs in a filters dict (mem0 >=2.0)."""
         mock_oss_mem0_client.search.return_value = [{"memory": "User likes Python"}]
         provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_oss_mem0_client, user_id="u1")
         session = AgentSession(session_id="test-session")
@@ -214,32 +214,50 @@ async def test_oss_client_passes_direct_kwargs(self, mock_oss_mem0_client: Async
 
         call_kwargs = mock_oss_mem0_client.search.call_args.kwargs
         assert call_kwargs["query"] == "Hello"
-        assert call_kwargs["user_id"] == "u1"
-        assert "filters" not in call_kwargs
+        assert call_kwargs["filters"] == {"user_id": "u1"}
+        assert "user_id" not in call_kwargs
 
-    @pytest.mark.asyncio
-    async def test_oss_client_all_scoping_params_except_app_id(self, mock_oss_mem0_client: AsyncMock) -> None:
-        """OSS client with all scoping parameters passes them as isolated concurrent kwargs."""
+    async def test_oss_client_rejects_application_id_with_user_or_agent(self, mock_oss_mem0_client: AsyncMock) -> None:
+        """OSS client rejects application_id even when user_id/agent_id are provided."""
         mock_oss_mem0_client.search.return_value = []
 
-        provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_oss_mem0_client, user_id="u1", agent_id="a1")
+        provider = Mem0ContextProvider(
+            source_id="mem0",
+            mem0_client=mock_oss_mem0_client,
+            user_id="u1",
+            agent_id="a1",
+            application_id="app1",
+        )
 
         mock_context = MagicMock(spec=SessionContext)
         mock_msg = MagicMock()
         mock_msg.text = "hello"
         mock_context.input_messages = [mock_msg]
         mock_context.response = None
 
-        await provider.before_run(
-            agent=MagicMock(), session=MagicMock(spec=AgentSession), context=mock_context, state={}
-        )
+        with pytest.raises(ValueError, match="application_id is not supported"):
+            await provider.before_run(
+                agent=MagicMock(), session=MagicMock(spec=AgentSession), context=mock_context, state={}
+            )
+
+        mock_oss_mem0_client.search.assert_not_awaited()
+
+    async def test_oss_client_rejects_application_id_only(self, mock_oss_mem0_client: AsyncMock) -> None:
+        """OSS client with only application_id set raises and never searches."""
+        provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_oss_mem0_client, application_id="app1")
+        session = AgentSession(session_id="test-session")
+        ctx = SessionContext(input_messages=[Message(role="user", contents=["Hello"])], session_id="s1")
 
-        # Re-aligned assertion: We expect 2 separate concurrent calls instead of 1 combined call
-        assert mock_oss_mem0_client.search.call_count == 2
-        mock_oss_mem0_client.search.assert_any_call(query="hello", user_id="u1")
-        mock_oss_mem0_client.search.assert_any_call(query="hello", agent_id="a1")
+        with pytest.raises(ValueError, match="application_id is not supported"):
+            await provider.before_run(
+                agent=cast(Any, None),
+                session=session,
+                context=ctx,
+                state=session.state.setdefault(provider.source_id, {}),
+            )  # type: ignore[arg-type]
+
+        mock_oss_mem0_client.search.assert_not_awaited()
 
-    @pytest.mark.asyncio
     async def test_platform_client_passes_filters_dict_except_app_id(self, mock_mem0_client: AsyncMock) -> None:
         """Platform client passes scoping parameters concurrently inside the nested filters dictionary."""
         mock_mem0_client.search.return_value = []
@@ -266,6 +284,25 @@ async def test_platform_client_passes_filters_dict_except_app_id(self, mock_mem0
         mock_mem0_client.search.assert_any_call(query="hello", filters={"user_id": "u1"})
         mock_mem0_client.search.assert_any_call(query="hello", filters={"agent_id": "a1"})
 
+    async def test_platform_client_keeps_app_id(self, mock_mem0_client: AsyncMock) -> None:
+        """Platform client keeps app_id in filters for each entity-scoped partition."""
+        mock_mem0_client.search.return_value = []
+
+        provider = Mem0ContextProvider(
+            source_id="mem0", mem0_client=mock_mem0_client, user_id="u1", application_id="app1"
+        )
+        session = AgentSession(session_id="test-session")
+        ctx = SessionContext(input_messages=[Message(role="user", contents=["Hello"])], session_id="s1")
+
+        await provider.before_run(
+            agent=cast(Any, None),
+            session=session,
+            context=ctx,
+            state=session.state.setdefault(provider.source_id, {}),
+        )  # type: ignore[arg-type]
+
+        mock_mem0_client.search.assert_awaited_once_with(query="Hello", filters={"user_id": "u1", "app_id": "app1"})
+
 
 # -- after_run tests -----------------------------------------------------------
 
@@ -293,7 +330,9 @@ async def test_stores_input_and_response(self, mock_mem0_client: AsyncMock) -> N
             {"role": "user", "content": "question"},
             {"role": "assistant", "content": "answer"},
         ]
-        assert call_kwargs["user_id"] == "u1"
+        assert call_kwargs["filters"] == {"user_id": "u1"}
+        assert "user_id" not in call_kwargs
+        assert "agent_id" not in call_kwargs
         assert "run_id" not in call_kwargs
 
     async def test_only_stores_user_assistant_system(self, mock_mem0_client: AsyncMock) -> None:
@@ -358,6 +397,7 @@ async def test_no_run_id_in_storage(self, mock_mem0_client: AsyncMock) -> None:
         )  # type: ignore[arg-type]
 
         assert "run_id" not in mock_mem0_client.add.call_args.kwargs
+        assert "run_id" not in mock_mem0_client.add.call_args.kwargs["filters"]
 
     async def test_validates_filters(self, mock_mem0_client: AsyncMock) -> None:
         """Raises ValueError when no filters."""
@@ -374,10 +414,10 @@ async def test_validates_filters(self, mock_mem0_client: AsyncMock) -> None:
                 state=session.state,
             )  # type: ignore[arg-type]
 
-    async def test_stores_with_application_id_filters(self, mock_mem0_client: AsyncMock) -> None:
-        """application_id is passed in filters."""
+    async def test_platform_stores_identity_fields_in_filters(self, mock_mem0_client: AsyncMock) -> None:
+        """Platform add receives all identity fields in filters for mem0ai 2.x."""
         provider = Mem0ContextProvider(
-            source_id="mem0", mem0_client=mock_mem0_client, user_id="u1", application_id="app1"
+            source_id="mem0", mem0_client=mock_mem0_client, user_id="u1", agent_id="a1", application_id="app1"
         )
         session = AgentSession(session_id="test-session")
         ctx = SessionContext(input_messages=[Message(role="user", contents=["hi"])], session_id="s1")
@@ -390,7 +430,48 @@ async def test_stores_with_application_id_filters(self, mock_mem0_client: AsyncM
             state=session.state.setdefault(provider.source_id, {}),
         )  # type: ignore[arg-type]
 
-        assert mock_mem0_client.add.call_args.kwargs["filters"] == {"app_id": "app1"}
+        call_kwargs = mock_mem0_client.add.call_args.kwargs
+        assert call_kwargs["filters"] == {"user_id": "u1", "agent_id": "a1", "app_id": "app1"}
+        assert "user_id" not in call_kwargs
+        assert "agent_id" not in call_kwargs
+
+    async def test_oss_stores_identity_fields_as_direct_kwargs(self, mock_oss_mem0_client: AsyncMock) -> None:
+        """OSS add keeps user_id/agent_id as direct kwargs because AsyncMemory.add uses that signature."""
+        provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_oss_mem0_client, user_id="u1", agent_id="a1")
+        session = AgentSession(session_id="test-session")
+        ctx = SessionContext(input_messages=[Message(role="user", contents=["hi"])], session_id="s1")
+        ctx._response = AgentResponse(messages=[])
+
+        await provider.after_run(
+            agent=cast(Any, None),
+            session=session,
+            context=ctx,
+            state=session.state.setdefault(provider.source_id, {}),
+        )  # type: ignore[arg-type]
+
+        call_kwargs = mock_oss_mem0_client.add.call_args.kwargs
+        assert call_kwargs["user_id"] == "u1"
+        assert call_kwargs["agent_id"] == "a1"
+        assert "filters" not in call_kwargs
+
+    async def test_oss_storage_rejects_application_id(self, mock_oss_mem0_client: AsyncMock) -> None:
+        """OSS storage rejects Platform-only application_id because AsyncMemory.add has no app_id parameter."""
+        provider = Mem0ContextProvider(
+            source_id="mem0", mem0_client=mock_oss_mem0_client, user_id="u1", application_id="app1"
+        )
+        session = AgentSession(session_id="test-session")
+        ctx = SessionContext(input_messages=[Message(role="user", contents=["hi"])], session_id="s1")
+        ctx._response = AgentResponse(messages=[])
+
+        with pytest.raises(ValueError, match="application_id is not supported"):
+            await provider.after_run(
+                agent=cast(Any, None),
+                session=session,
+                context=ctx,
+                state=session.state.setdefault(provider.source_id, {}),
+            )  # type: ignore[arg-type]
+
+        mock_oss_mem0_client.add.assert_not_awaited()
 
 
 # -- _validate_filters tests --------------------------------------------------
@@ -416,6 +497,25 @@ def test_passes_with_application_id(self, mock_mem0_client: AsyncMock) -> None:
         provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_mem0_client, application_id="app1")
         provider._validate_filters()
 
+    def test_oss_application_id_only_raises(self, mock_oss_mem0_client: AsyncMock) -> None:
+        """OSS client with only application_id is rejected because application scope is Platform-only."""
+        provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_oss_mem0_client, application_id="app1")
+        with pytest.raises(ValueError, match="application_id is not supported"):
+            provider._validate_filters()
+
+    def test_oss_application_id_with_user_id_raises(self, mock_oss_mem0_client: AsyncMock) -> None:
+        """OSS client rejects application_id even with a supported user scope."""
+        provider = Mem0ContextProvider(
+            source_id="mem0", mem0_client=mock_oss_mem0_client, user_id="u1", application_id="app1"
+        )
+        with pytest.raises(ValueError, match="application_id is not supported"):
+            provider._validate_filters()
+
+    def test_oss_passes_with_user_id(self, mock_oss_mem0_client: AsyncMock) -> None:
+        """OSS client with user_id is accepted."""
+        provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_oss_mem0_client, user_id="u1")
+        provider._validate_filters()
+
 
 # -- _build_search_kwargs tests -----------------------------------------------------
 
@@ -469,6 +569,15 @@ def test_no_run_id_in_search_filters(self, mock_mem0_client: AsyncMock) -> None:
         assert "run_id" not in result.get("filters", {})
         assert "run_id" not in result
 
+    def test_oss_search_filters_reject_app_id(self, mock_oss_mem0_client: AsyncMock) -> None:
+        """OSS search filters reject application_id because app_id is only supported by Platform."""
+        provider = Mem0ContextProvider(
+            source_id="mem0", mem0_client=mock_oss_mem0_client, user_id="u1", application_id="app1"
+        )
+
+        with pytest.raises(ValueError, match="application_id is not supported"):
+            provider._build_search_kwargs("test query", "user_id", "u1")
+
     def test_empty_when_no_params(self, mock_mem0_client: AsyncMock) -> None:
         # Validates base query payload generation
         provider = Mem0ContextProvider(source_id="mem0", mem0_client=mock_mem0_client)
@@ -477,9 +586,7 @@ def test_empty_when_no_params(self, mock_mem0_client: AsyncMock) -> None:
 
         assert result == {"query": "test query", "filters": {"custom_key": "custom_val"}}
 
-    @pytest.mark.asyncio
     async def test_before_run_application_only_fallback(self, mock_mem0_client: AsyncMock) -> None:
-
         provider = Mem0ContextProvider(
             source_id="mem0", mem0_client=mock_mem0_client, application_id="app_fallback_test"
         )
@@ -498,7 +605,10 @@ async def test_before_run_application_only_fallback(self, mock_mem0_client: Asyn
         )
 
         # Verify that an application-scoped search task executed successfully
-        assert mock_mem0_client.search.call_count == 1
+        mock_mem0_client.search.assert_awaited_once_with(
+            query="Retrieve systemic fallback memory traces",
+            filters={"app_id": "app_fallback_test"},
+        )
         mock_context.extend_messages.assert_called_once()