Add Mem0 integration - support for Mem0 platform

.github/workflows/tests.yml

@@ -24,12 +24,20 @@ on:
       - "pyproject.toml"
       - ".github/workflows/tests.yml"
 
+permissions:
+  id-token: write
+  contents: read
+
 env:
   PYTHON_VERSION: "3.10"
   HATCH_VERSION: "1.14.2"
   PYTHONUNBUFFERED: "1"
   FORCE_COLOR: "1"
   OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+  COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
+  MEM0_API_KEY: ${{ secrets.MEM0_API_KEY }}
+  GOOGLE_API_KEY: ${{ secrets.GOOGLE_API_KEY }}
+  AWS_REGION: "us-east-1"
 jobs:
   linting:
     runs-on: ubuntu-latest
@@ -129,5 +137,16 @@ jobs:
 
       - name: Install Hatch
         run: pip install hatch==${{ env.HATCH_VERSION }}
+
+      # Do not authenticate on PRs from forks and on PRs created by dependabot
+      - name: AWS authentication
+        id: aws-auth
+        if: github.event_name == 'schedule' || (github.event.pull_request.head.repo.full_name == github.repository && !startsWith(github.event.pull_request.head.ref, 'dependabot/'))
+        uses: aws-actions/configure-aws-credentials@61815dcd50bd041e203e49132bacad1fd04d2708
+        with:
+          aws-region: ${{ env.AWS_REGION }}
+          role-to-assume: ${{ secrets.AWS_CI_ROLE_ARN }}
+
       - name: Run
+        if: success() && steps.aws-auth.outcome == 'success'
         run: hatch run test:integration-retry

README.md

@@ -48,6 +48,7 @@ that includes it. Once it reaches the end of its lifespan, the experiment will b
 | [`Agent`][17]; [Confirmation Policies][18]; [ConfirmationUIs][19]; [ConfirmationStrategies][20]; [`ConfirmationUIResult` and `ToolExecutionDecision`][21] [HITLBreakpointException][22] | Human in the Loop                     | December 2025     | rich         | None                                                                                                                                                                                                                                                       | [Discuss][23] |
 | [`LLMSummarizer`][24]                                                                                                                                                                   | Document Summarizer                   | January 2025      | None         | None                                                                                                                                                                                                                                                       | [Discuss][25] |
 | [`InMemoryChatMessageStore`][1]; [`ChatMessageRetriever`][2]; [`ChatMessageWriter`][3]                                                                                                  | Chat Message Store, Retriever, Writer | February 2025     | None         | <a href="https://colab.research.google.com/github/deepset-ai/haystack-cookbook/blob/main/notebooks/conversational_rag_using_memory.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>       | [Discuss][4]  |
+| [`Mem0MemoryStore`][26]                                                                                                                                                                   | MemoryStore                   | February 2025      | mem0ai         | None                                                                                                                                                                                                                                                       | -- |
 
 [1]: https://github.com/deepset-ai/haystack-experimental/blob/main/haystack_experimental/chat_message_stores/in_memory.py
 [2]: https://github.com/deepset-ai/haystack-experimental/blob/main/haystack_experimental/components/retrievers/chat_message_retriever.py
@@ -66,6 +67,7 @@ that includes it. Once it reaches the end of its lifespan, the experiment will b
 [23]: https://github.com/deepset-ai/haystack-experimental/discussions/381
 [24]: https://github.com/deepset-ai/haystack-experimental/blob/main/haystack_experimental/components/sumarizers/llm_summarizer.py
 [25]: https://github.com/deepset-ai/haystack-experimental/discussions/382
+[26]: https://github.com/deepset-ai/haystack-experimental/blob/main/haystack_experimental/memory_stores/mem0/memory_store.py
 
 ### Adopted experiments
 | Name                                                                                                                                                                                                                                       | Type                                     | Final release |

haystack_experimental/components/agents/agent.py

@@ -57,6 +57,7 @@
 )
 from haystack_experimental.components.retrievers import ChatMessageRetriever
 from haystack_experimental.components.writers import ChatMessageWriter
+from haystack_experimental.memory_stores.types import MemoryStore
 
 logger = logging.getLogger(__name__)
 
@@ -146,6 +147,7 @@ def __init__(
         confirmation_strategies: dict[str, ConfirmationStrategy] | None = None,
         tool_invoker_kwargs: dict[str, Any] | None = None,
         chat_message_store: ChatMessageStore | None = None,
+        memory_store: MemoryStore | None = None,
     ) -> None:
         """
         Initialize the agent component.
@@ -164,6 +166,9 @@ def __init__(
         :param raise_on_tool_invocation_failure: Should the agent raise an exception when a tool invocation fails?
             If set to False, the exception will be turned into a chat message and passed to the LLM.
         :param tool_invoker_kwargs: Additional keyword arguments to pass to the ToolInvoker.
+        :param chat_message_store: The ChatMessageStore that the agent can use to store
+            and retrieve chat messages history.
+        :param memory_store: The memory store that the agent can use to store and retrieve memories.
         :raises TypeError: If the chat_generator does not support tools parameter in its run method.
         :raises ValueError: If the exit_conditions are not valid.
         """
@@ -186,6 +191,7 @@ def __init__(
         self._chat_message_writer = (
             ChatMessageWriter(chat_message_store=chat_message_store) if chat_message_store else None
         )
+        self._memory_store = memory_store
 
     def _initialize_fresh_execution(
         self,
@@ -198,6 +204,7 @@ def _initialize_fresh_execution(
         tools: ToolsType | list[str] | None = None,
         confirmation_strategy_context: dict[str, Any] | None = None,
         chat_message_store_kwargs: dict[str, Any] | None = None,
+        memory_store_kwargs: dict[str, Any] | None = None,
         **kwargs: dict[str, Any],
     ) -> _ExecutionContext:
         """
@@ -209,29 +216,62 @@ def _initialize_fresh_execution(
         :param system_prompt: System prompt for the agent. If provided, it overrides the default system prompt.
         :param tools: Optional list of Tool objects, a Toolset, or list of tool names to use for this run.
             When passing tool names, tools are selected from the Agent's originally configured tools.
+
+        :param memory_store_kwargs: Optional dictionary of keyword arguments to pass to the MemoryStore.
+            For example, it can include the `user_id`, `run_id`, and `agent_id` parameters
+            for storing and retrieving memories.
         :param confirmation_strategy_context: Optional dictionary for passing request-scoped resources
             to confirmation strategies.
         :param chat_message_store_kwargs: Optional dictionary of keyword arguments to pass to the ChatMessageStore.
+            For example, it can include the `chat_history_id` and `last_k` parameters for retrieving chat history.
         :param kwargs: Additional data to pass to the State used by the Agent.
         """
         system_prompt = system_prompt or self.system_prompt
-        if system_prompt is not None:
-            messages = [ChatMessage.from_system(system_prompt)] + messages
+        retrieved_memory = None
+        updated_system_prompt = system_prompt
+
+        # Retrieve memories from the memory store
+        if self._memory_store:
+            retrieved_memories = self._memory_store.search_memories(query=messages[-1].text, **memory_store_kwargs)  # type: ignore[arg-type]
+
+            # we combine the memories into a single string
+            combined_memory = "\n".join(
+                f"- MEMORY #{idx + 1}: {memory.text}" for idx, memory in enumerate(retrieved_memories)
+            )
+            retrieved_memory = ChatMessage.from_system(text=combined_memory)
+
+        if retrieved_memory:
+            memory_instruction = (
+                "\n\nWhen messages start with `[MEMORY]`, treat them as long-term "
+                "context and use them to guide the response if relevant."
+            )
+            updated_system_prompt = f"{system_prompt}{memory_instruction}"
+
+            memory_text = f"Here are the relevant memories for the user's query: {retrieved_memory.text}"
+            print(memory_text)
+            updated_memory = ChatMessage.from_system(text=memory_text)
+        else:
+            updated_memory = None
+
+        combined_messages = messages + [updated_memory] if updated_memory else messages
+        if updated_system_prompt is not None:
+            combined_messages = [ChatMessage.from_system(updated_system_prompt)] + combined_messages
 
         # NOTE: difference with parent method to add chat message retrieval
         if self._chat_message_retriever:
             retriever_kwargs = _select_kwargs(self._chat_message_retriever, chat_message_store_kwargs or {})
             if "chat_history_id" in retriever_kwargs:
                 messages = self._chat_message_retriever.run(
-                    current_messages=messages,
+                    current_messages=combined_messages,
                     **retriever_kwargs,
                 )["messages"]
+                combined_messages = messages
 
-        if all(m.is_from(ChatRole.SYSTEM) for m in messages):
+        if all(m.is_from(ChatRole.SYSTEM) for m in combined_messages):
             logger.warning("All messages provided to the Agent component are system messages. This is not recommended.")
 
         state = State(schema=self.state_schema, data=kwargs)
-        state.set("messages", messages)
+        state.set("messages", combined_messages)
 
         streaming_callback = select_streaming_callback(  # type: ignore[call-overload]
             init_callback=self.streaming_callback, runtime_callback=streaming_callback, requires_async=requires_async
@@ -329,6 +369,7 @@ def run(  # type: ignore[override]  # noqa: PLR0915 PLR0912
         tools: ToolsType | list[str] | None = None,
         confirmation_strategy_context: dict[str, Any] | None = None,
         chat_message_store_kwargs: dict[str, Any] | None = None,
+        memory_store_kwargs: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> dict[str, Any]:
         """
@@ -352,6 +393,19 @@ def run(  # type: ignore[override]  # noqa: PLR0915 PLR0912
             can use for non-blocking user interaction.
         :param chat_message_store_kwargs: Optional dictionary of keyword arguments to pass to the ChatMessageStore.
             For example, it can include the `chat_history_id` and `last_k` parameters for retrieving chat history.
+        :param memory_store_kwargs: Optional dictionary of keyword arguments to pass to the MemoryStore.
+            It can include:
+            - `user_id`: The user ID to search and add memories from.
+            - `run_id`: The run ID to search and add memories from.
+            - `agent_id`: The agent ID to search and add memories from.
+            - `search_criteria`: A dictionary of containing kwargs for the `search_memories` method.
+                This can include:
+                - `filters`: A dictionary of filters to search for memories.
+                - `query`: The query to search for memories.
+                    Note: If you pass this, the user query passed to the agent will be
+                    ignored for memory retrieval.
+                - `top_k`: The number of memories to return.
+                - `include_memory_metadata`: Whether to include the memory metadata in the ChatMessage.
         :param kwargs: Additional data to pass to the State schema used by the Agent.
             The keys must match the schema defined in the Agent's `state_schema`.
         :returns:
@@ -362,6 +416,8 @@ def run(  # type: ignore[override]  # noqa: PLR0915 PLR0912
         :raises RuntimeError: If the Agent component wasn't warmed up before calling `run()`.
         :raises BreakpointException: If an agent breakpoint is triggered.
         """
+        memory_store_kwargs = memory_store_kwargs or {}
+
         agent_inputs = {
             "messages": messages,
             "streaming_callback": streaming_callback,
@@ -392,6 +448,7 @@ def run(  # type: ignore[override]  # noqa: PLR0915 PLR0912
                 tools=tools,
                 confirmation_strategy_context=confirmation_strategy_context,
                 chat_message_store_kwargs=chat_message_store_kwargs,
+                memory_store_kwargs=memory_store_kwargs,
                 **kwargs,
             )
 
@@ -547,6 +604,11 @@ def run(  # type: ignore[override]  # noqa: PLR0915 PLR0912
         if msgs := result.get("messages"):
             result["last_message"] = msgs[-1]
 
+            # Add the new conversation as memories to the memory store
+            if self._memory_store:
+                new_memories = [message for message in msgs if message.role.value != "system"]
+                self._memory_store.add_memories(messages=new_memories, **memory_store_kwargs)
+
         # Write messages to ChatMessageStore if configured
         if self._chat_message_writer:
             writer_kwargs = _select_kwargs(self._chat_message_writer, chat_message_store_kwargs or {})
@@ -567,6 +629,7 @@ async def run_async(  # type: ignore[override] # noqa: PLR0915
         tools: ToolsType | list[str] | None = None,
         confirmation_strategy_context: dict[str, Any] | None = None,
         chat_message_store_kwargs: dict[str, Any] | None = None,
+        memory_store_kwargs: dict[str, Any] | None = None,
         **kwargs: Any,
     ) -> dict[str, Any]:
         """
@@ -593,6 +656,20 @@ async def run_async(  # type: ignore[override] # noqa: PLR0915
             can use for non-blocking user interaction.
         :param chat_message_store_kwargs: Optional dictionary of keyword arguments to pass to the ChatMessageStore.
             For example, it can include the `chat_history_id` and `last_k` parameters for retrieving chat history.
+        :param kwargs: Additional data to pass to the State schema used by the Agent.
+        :param memory_store_kwargs: Optional dictionary of keyword arguments to pass to the MemoryStore.
+            It can include:
+            - `user_id`: The user ID to search and add memories from.
+            - `run_id`: The run ID to search and add memories from.
+            - `agent_id`: The agent ID to search and add memories from.
+            - `search_criteria`: A dictionary of containing kwargs for the `search_memories` method.
+                This can include:
+                - `filters`: A dictionary of filters to search for memories.
+                - `query`: The query to search for memories.
+                    Note: If you pass this, the user query passed to the agent will be
+                    ignored for memory retrieval.
+                - `top_k`: The number of memories to return.
+                - `include_memory_metadata`: Whether to include the memory metadata in the ChatMessage.
         :param kwargs: Additional data to pass to the State schema used by the Agent.
             The keys must match the schema defined in the Agent's `state_schema`.
         :returns:
@@ -603,6 +680,8 @@ async def run_async(  # type: ignore[override] # noqa: PLR0915
         :raises RuntimeError: If the Agent component wasn't warmed up before calling `run_async()`.
         :raises BreakpointException: If an agent breakpoint is triggered.
         """
+        memory_store_kwargs = memory_store_kwargs or {}
+
         agent_inputs = {
             "messages": messages,
             "streaming_callback": streaming_callback,
@@ -631,6 +710,7 @@ async def run_async(  # type: ignore[override] # noqa: PLR0915
                 tools=tools,
                 confirmation_strategy_context=confirmation_strategy_context,
                 chat_message_store_kwargs=chat_message_store_kwargs,
+                memory_store_kwargs=memory_store_kwargs,
                 **kwargs,
             )
 
@@ -773,6 +853,11 @@ async def run_async(  # type: ignore[override] # noqa: PLR0915
         if msgs := result.get("messages"):
             result["last_message"] = msgs[-1]
 
+            # Add the new conversation as memories to the memory store
+            if self._memory_store:
+                new_memories = [message for message in msgs if message.role.value != "system"]
+                self._memory_store.add_memories(messages=new_memories, **memory_store_kwargs)
+
         # Write messages to ChatMessageStore if configured
         if self._chat_message_writer:
             writer_kwargs = _select_kwargs(self._chat_message_writer, chat_message_store_kwargs or {})

haystack_experimental/memory_stores/__init__.py

@@ -0,0 +1,7 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from .types import MemoryStore
+
+__all__ = ["MemoryStore"]

haystack_experimental/memory_stores/mem0/__init__.py

@@ -0,0 +1,16 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import sys
+from typing import TYPE_CHECKING
+
+from lazy_imports import LazyImporter
+
+_import_structure = {"memory_store": ["Mem0MemoryStore"]}
+
+if TYPE_CHECKING:
+    from .memory_store import Mem0MemoryStore as Mem0MemoryStore
+
+else:
+    sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure)