-
Notifications
You must be signed in to change notification settings - Fork 7.8k
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.
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
savy-91
wants to merge
1
commit into
microsoft:main
Choose a base branch
from
savy-91:feat/multiple_context_injection_modes_mem0
base: main
Could not load branches
Branch not found: {{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline,
and old review comments may become outdated.
Open
Changes from all commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
3 changes: 2 additions & 1 deletion
3
python/packages/autogen-ext/src/autogen_ext/memory/mem0/__init__.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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", | ||
] |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -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 | ||
|
@@ -18,6 +19,13 @@ | |
logging.getLogger("chromadb").setLevel(logging.ERROR) | ||
|
||
|
||
class ContextInjectionMode(Enum): | ||
"""Enum for context injection modes.""" | ||
|
||
SYSTEM_MESSAGE = "system_message" | ||
FUNCTION_CALL = "function_call" | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe 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.""" | ||
|
||
|
@@ -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): | ||
|
@@ -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!") | ||
|
||
|
@@ -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 | ||
|
@@ -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" | ||
|
@@ -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: | ||
|
@@ -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: | ||
|
@@ -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, | ||
|
@@ -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. | ||
|
@@ -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) | ||
|
||
|
@@ -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: | ||
|
@@ -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, | ||
) |
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
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.There was a problem hiding this comment.
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.