Skip to content

feat: multiple context injection modes mem0 #6850

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: multiple context injection modes mem0 #6850

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
67 changes: 59 additions & 8 deletions python/docs/src/user-guide/agentchat-user-guide/memory.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -542,9 +542,9 @@
"source": [
"## Mem0Memory Example\n",
"\n",
"`autogen_ext.memory.mem0.Mem0Memory` provides integration with `Mem0.ai`'s memory system. It supports both cloud-based and local backends, offering advanced memory capabilities for agents. The implementation handles proper retrieval and context updating, making it suitable for production environments.\n",
"`autogen_ext.memory.mem0.Mem0Memory` provides integration with `Mem0.ai`'s memory system. It supports both cloud-based and local backends, offering advanced memory capabilities for agents. The implementation handles proper retrieval and context updating with configurable injection modes, making it suitable for production environments.\n",
"\n",
"In the following example, we'll demonstrate how to use `Mem0Memory` to maintain persistent memories across conversations:"
"In the following example, we'll demonstrate how to use `Mem0Memory` to maintain persistent memories across conversations:"
]
},
{
Expand All @@ -556,14 +556,15 @@
"from autogen_agentchat.agents import AssistantAgent\n",
"from autogen_agentchat.ui import Console\n",
"from autogen_core.memory import MemoryContent, MemoryMimeType\n",
"from autogen_ext.memory.mem0 import Mem0Memory\n",
"from autogen_ext.memory.mem0 import Mem0Memory, ContextInjectionMode\n",
"from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
"\n",
"# Initialize Mem0 cloud memory (requires API key)\n",
"# For local deployment, use is_cloud=False with appropriate config\n",
"mem0_memory = Mem0Memory(\n",
" is_cloud=True,\n",
" limit=5, # Maximum number of memories to retrieve\n",
" context_injection_mode=ContextInjectionMode.SYSTEM_MESSAGE, # Default mode\n",
")\n",
"\n",
"# Add user preferences to memory\n",
Expand Down Expand Up @@ -595,12 +596,12 @@
"\n",
"# Ask about the weather\n",
"stream = assistant_agent.run_stream(task=\"What are my dietary preferences?\")\n",
"await Console(stream)"
"await Console(stream)\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"cell_type": "markdown",
"source": [
"The example above demonstrates how Mem0Memory can be used with an assistant agent. The memory integration ensures that:\n",
"\n",
Expand All @@ -618,20 +619,70 @@
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"# Serialize the memory configuration\n",
"config_json = mem0_memory.dump_component().model_dump_json()\n",
"print(f\"Memory config JSON: {config_json[:100]}...\")"
]
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"### Context Injection Modes\n",
"\n",
"Mem0Memory supports two different modes for injecting retrieved memories into the model context:"
]
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"#### System Message Mode (Default)\n",
"\n",
"This is the traditional approach where memories are added as system messages:"
]
},
{
"metadata": {},
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"# Default behavior - memories added as system messages\n",
"mem0_memory = Mem0Memory(\n",
" is_cloud=False,\n",
" config={\"path\": \":memory:\"},\n",
" context_injection_mode=ContextInjectionMode.SYSTEM_MESSAGE, # Default\n",
")"
]
},
{
"metadata": {},
"cell_type": "markdown",
"source": [
"#### Function Call Mode\n",
"\n",
"This mode injects memories as function call results, this is suitable for models that do not support multiple system messages:"
]
},
{
"metadata": {},
"source": []
"cell_type": "code",
"outputs": [],
"execution_count": null,
"source": [
"# Function call mode - memories added as function execution results\n",
"mem0_memory = Mem0Memory(\n",
" is_cloud=False,\n",
" config={\"path\": \":memory:\"},\n",
" context_injection_mode=ContextInjectionMode.FUNCTION_CALL,\n",
")"
]
}
],
"metadata": {
Expand Down
3 changes: 2 additions & 1 deletion python/packages/autogen-ext/src/autogen_ext/memory/mem0/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from ._mem0 import Mem0Memory, Mem0MemoryConfig
from ._mem0 import ContextInjectionMode, Mem0Memory, Mem0MemoryConfig

__all__ = [
"Mem0Memory",
"Mem0MemoryConfig",
"ContextInjectionMode",
]
75 changes: 66 additions & 9 deletions python/packages/autogen-ext/src/autogen_ext/memory/mem0/_mem0.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,13 @@
import uuid
from contextlib import redirect_stderr, redirect_stdout
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional, TypedDict, cast

from autogen_core import CancellationToken, Component, ComponentBase
from autogen_core import CancellationToken, Component, ComponentBase, FunctionCall
from autogen_core.memory import Memory, MemoryContent, MemoryQueryResult, UpdateContextResult
from autogen_core.model_context import ChatCompletionContext
from autogen_core.models import SystemMessage
from autogen_core.models import AssistantMessage, FunctionExecutionResult, FunctionExecutionResultMessage, SystemMessage
from mem0 import Memory as Memory0
from mem0 import MemoryClient
from pydantic import BaseModel, Field
Expand All @@ -18,6 +19,13 @@
logging.getLogger("chromadb").setLevel(logging.ERROR)


class ContextInjectionMode(Enum):
"""Enum for context injection modes."""

SYSTEM_MESSAGE = "system_message"
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we also add USER_MESSAGE? Because it is naturally a choice if the model doesn't support system message in the middle.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

My concern with adding a USER_MESSAGE mode is that I think we'd then require a more opinionated message format to tell the LLM to use the memory retrieved but also don't reply as if the user actually sent them in a message.

FUNCTION_CALL = "function_call"
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Add description for each mode enum. Also, note that using the function call mode may cause problem with models that requires the tool definition to be available when calling the model.



class Mem0MemoryConfig(BaseModel):
"""Configuration for Mem0Memory component."""

Expand All @@ -32,6 +40,10 @@ class Mem0MemoryConfig(BaseModel):
config: Optional[Dict[str, Any]] = Field(
default=None, description="Configuration dictionary for local Mem0 client. Required if is_cloud=False."
)
context_injection_mode: ContextInjectionMode = Field(
default=ContextInjectionMode.SYSTEM_MESSAGE,
description="Mode for injecting memories into context: 'system_message' or 'function_call'.",
)


class MemoryResult(TypedDict, total=False):
Expand Down Expand Up @@ -68,15 +80,16 @@ class Mem0Memory(Memory, Component[Mem0MemoryConfig], ComponentBase[Mem0MemoryCo
.. code-block:: python

import asyncio
from autogen_ext.memory.mem0 import Mem0Memory
from autogen_ext.memory.mem0 import Mem0Memory, ContextInjectionMode
from autogen_core.memory import MemoryContent


async def main() -> None:
# Create a local Mem0Memory (no API key required)
# Create a local Mem0Memory with function call injection mode
memory = Mem0Memory(
is_cloud=False,
config={"path": ":memory:"}, # Use in-memory storage for testing
context_injection_mode=ContextInjectionMode.FUNCTION_CALL,
)
print("Memory initialized successfully!")

Expand Down Expand Up @@ -111,19 +124,20 @@ async def main() -> None:
import asyncio
from autogen_agentchat.agents import AssistantAgent
from autogen_core.memory import MemoryContent
from autogen_ext.memory.mem0 import Mem0Memory
from autogen_ext.memory.mem0 import Mem0Memory, ContextInjectionMode
from autogen_ext.models.openai import OpenAIChatCompletionClient


async def main() -> None:
# Create a model client
model_client = OpenAIChatCompletionClient(model="gpt-4.1")

# Create a Mem0 memory instance
# Create a Mem0 memory instance with system message injection (default)
memory = Mem0Memory(
user_id="user123",
is_cloud=False,
config={"path": ":memory:"}, # Use in-memory storage for testing
context_injection_mode=ContextInjectionMode.SYSTEM_MESSAGE,
)

# Add something to memory
Expand Down Expand Up @@ -157,6 +171,7 @@ async def main() -> None:
is_cloud: Whether to use cloud Mem0 client (True) or local client (False).
api_key: API key for cloud Mem0 client. It will read from the environment MEM0_API_KEY if not provided.
config: Configuration dictionary for local Mem0 client. Required if is_cloud=False.
context_injection_mode: Mode for injecting memories into context ('system_message' or 'function_call').
"""

component_type = "memory"
Expand All @@ -170,6 +185,7 @@ def __init__(
is_cloud: bool = True,
api_key: Optional[str] = None,
config: Optional[Dict[str, Any]] = None,
context_injection_mode: ContextInjectionMode = ContextInjectionMode.SYSTEM_MESSAGE,
) -> None:
# Validate parameters
if not is_cloud and config is None:
Expand All @@ -181,6 +197,7 @@ def __init__(
self._is_cloud = is_cloud
self._api_key = api_key
self._config = config
self._context_injection_mode = context_injection_mode

# Initialize client
if self._is_cloud:
Expand Down Expand Up @@ -210,6 +227,11 @@ def config(self) -> Optional[Dict[str, Any]]:
"""Get the configuration for the Mem0 client."""
return self._config

@property
def context_injection_mode(self) -> ContextInjectionMode:
"""Get the context injection mode."""
return self._context_injection_mode

async def add(
self,
content: MemoryContent,
Expand Down Expand Up @@ -366,7 +388,8 @@ async def update_context(

This method retrieves the conversation history from the model context,
uses the last message as a query to find relevant memories, and then
adds those memories to the context as a system message.
adds those memories to the context either as a system message or as
function call messages based on the configured injection mode.

Args:
model_context: The model context to update.
Expand All @@ -392,8 +415,40 @@ async def update_context(
memory_strings = [f"{i}. {str(memory.content)}" for i, memory in enumerate(query_results.results, 1)]
memory_context = "\nRelevant memories:\n" + "\n".join(memory_strings)

# Add as system message
await model_context.add_message(SystemMessage(content=memory_context))
if self._context_injection_mode == ContextInjectionMode.SYSTEM_MESSAGE:
# Add as system message (original behavior)
await model_context.add_message(SystemMessage(content=memory_context))

elif self._context_injection_mode == ContextInjectionMode.FUNCTION_CALL:
# Add as function call result messages
# Generate a unique call ID
call_id = f"call_{uuid.uuid4().hex[:20]}"

# Create the function call
function_call = FunctionCall(
id=call_id,
name="retrieve_mem0memory",
arguments="{}", # No parameters as specified
)

# Create AssistantMessage with the function call
assistant_message = AssistantMessage(
content=[function_call], source="memory_system", type="AssistantMessage"
)

# Create the function execution result
function_result = FunctionExecutionResult(
content=memory_context, name="retrieve_mem0memory", call_id=call_id, is_error=False
)

# Create FunctionExecutionResultMessage
result_message = FunctionExecutionResultMessage(
content=[function_result], type="FunctionExecutionResultMessage"
)

# Add both messages to the context
await model_context.add_message(assistant_message)
await model_context.add_message(result_message)

return UpdateContextResult(memories=query_results)

Expand Down Expand Up @@ -432,6 +487,7 @@ def _from_config(cls, config: Mem0MemoryConfig) -> Self:
is_cloud=config.is_cloud,
api_key=config.api_key,
config=config.config,
context_injection_mode=config.context_injection_mode,
)

def _to_config(self) -> Mem0MemoryConfig:
Expand All @@ -446,4 +502,5 @@ def _to_config(self) -> Mem0MemoryConfig:
is_cloud=self._is_cloud,
api_key=self._api_key,
config=self._config,
context_injection_mode=self._context_injection_mode,
)
Loading
Loading