Python: Context providers abstraction and Mem0 implementation (#631)

* Added context provider abstractions

* Added mem0 implementation

* Example and small fixes

* Added unit tests for agent

* Added unit tests for mem0 provider

* Updated README

* Small doc updates

* Update python/packages/mem0/agent_framework_mem0/_provider.py

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>

* Small fixes in tests

* Renaming based on PR feedback

* Small fixes

* Added tests for AggregateContextProvider

* Small improvements

* More improvements based on PR feedback

* Small constant update

* Added more examples

* Added README for Mem0 examples

* Small updates to API

* Updated initialization logic

* Updates for context manager

* Updated Context class

* Dependency update

* Revert changes

* Fixed tests

---------

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: Chris <66376200+crickman@users.noreply.github.com>
This commit is contained in:
Dmytro Struk
2025-09-10 14:11:42 -07:00
committed by GitHub
Unverified
parent 89c8418705
commit 57d09afe04
25 changed files with 4166 additions and 1915 deletions
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) Microsoft Corporation.
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE
+21
View File
@@ -0,0 +1,21 @@
# Get Started with Microsoft Agent Framework Mem0
Please install this package as the extra for `agent-framework`:
```bash
pip install agent-framework[mem0]
```
## Memory Context Provider
The Mem0 context provider enables persistent memory capabilities for your agents, allowing them to remember user preferences and conversation context across different sessions and threads.
### Basic Usage Example
See the [Mem0 basic example](https://github.com/microsoft/agent-framework/tree/main/python/samples/getting_started/context_providers/mem0/mem0_basic.py) which demonstrates:
- Setting up an agent with Mem0 context provider
- Teaching the agent user preferences
- Retrieving information using remembered context across new threads
- Persistent memory
@@ -0,0 +1,15 @@
# Copyright (c) Microsoft. All rights reserved.
import importlib.metadata
from ._provider import Mem0Provider
try:
__version__ = importlib.metadata.version(__name__)
except importlib.metadata.PackageNotFoundError:
__version__ = "0.0.0" # Fallback for development mode
__all__ = [
"Mem0Provider",
"__version__",
]
@@ -0,0 +1,180 @@
# Copyright (c) Microsoft. All rights reserved.
import sys
from collections.abc import MutableSequence, Sequence
from typing import Any, Final
from agent_framework import ChatMessage, Context, ContextProvider, TextContent
from agent_framework.exceptions import ServiceInitializationError
from pydantic import PrivateAttr
if sys.version_info >= (3, 11):
from typing import Self # pragma: no cover
else:
from typing_extensions import Self # pragma: no cover
DEFAULT_CONTEXT_PROMPT: Final[str] = "## Memories\nConsider the following memories when answering user questions:"
class Mem0Provider(ContextProvider):
api_key: str | None = None
application_id: str | None = None
agent_id: str | None = None
thread_id: str | None = None
user_id: str | None = None
scope_to_per_operation_thread_id: bool = False
context_prompt: str = DEFAULT_CONTEXT_PROMPT
# Use Any to avoid forward reference issues with AsyncMemoryClient
mem0_client: Any = None
_should_close_client: bool = PrivateAttr(default=False) # Track whether we should close client connection
def __init__(
self,
api_key: str | None = None,
application_id: str | None = None,
agent_id: str | None = None,
thread_id: str | None = None,
user_id: str | None = None,
scope_to_per_operation_thread_id: bool = False,
context_prompt: str = DEFAULT_CONTEXT_PROMPT,
mem0_client: Any = None,
) -> None:
"""Initializes a new instance of the Mem0Provider class.
Args:
api_key: The API key for authenticating with the Mem0 API. If not
provided, it will attempt to use the MEM0_API_KEY environment variable.
application_id: The application ID for scoping memories or None.
agent_id: The agent ID for scoping memories or None.
thread_id: The thread ID for scoping memories or None.
user_id: The user ID for scoping memories or None.
scope_to_per_operation_thread_id: Whether to scope memories to per-operation thread ID.
context_prompt: The prompt to prepend to retrieved memories.
mem0_client: A pre-created Mem0 MemoryClient or None to create a default client.
"""
should_close_client = False
if mem0_client is None:
from mem0 import AsyncMemoryClient
mem0_client = AsyncMemoryClient(api_key=api_key)
should_close_client = True
super().__init__(
api_key=api_key, # type: ignore[reportCallIssue]
application_id=application_id, # type: ignore[reportCallIssue]
agent_id=agent_id, # type: ignore[reportCallIssue]
thread_id=thread_id, # type: ignore[reportCallIssue]
user_id=user_id, # type: ignore[reportCallIssue]
scope_to_per_operation_thread_id=scope_to_per_operation_thread_id, # type: ignore[reportCallIssue]
context_prompt=context_prompt, # type: ignore[reportCallIssue]
mem0_client=mem0_client, # type: ignore[reportCallIssue]
)
self._per_operation_thread_id: str | None = None
self._should_close_client = should_close_client
async def __aenter__(self) -> "Self":
"""Async context manager entry."""
if self.mem0_client:
await self.mem0_client.__aenter__()
return self
async def __aexit__(self, exc_type: type[BaseException] | None, exc_val: BaseException | None, exc_tb: Any) -> None:
"""Async context manager exit."""
if self._should_close_client and self.mem0_client:
await self.mem0_client.__aexit__(exc_type, exc_val, exc_tb)
async def thread_created(self, thread_id: str | None = None) -> None:
"""Called when a new thread is created.
Args:
thread_id: The ID of the thread or None.
"""
self._validate_per_operation_thread_id(thread_id)
self._per_operation_thread_id = self._per_operation_thread_id or thread_id
async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
"""Called when a new message is being added to the thread.
Args:
thread_id: The ID of the thread or None.
new_messages: New messages to add.
"""
self._validate_filters()
self._validate_per_operation_thread_id(thread_id)
self._per_operation_thread_id = self._per_operation_thread_id or thread_id
messages_list = [new_messages] if isinstance(new_messages, ChatMessage) else list(new_messages)
messages: list[dict[str, str]] = [
{"role": message.role.value, "content": message.text}
for message in messages_list
if message.role.value in {"user", "assistant", "system"} and message.text and message.text.strip()
]
if messages:
await self.mem0_client.add( # type: ignore[misc]
messages=messages,
user_id=self.user_id,
agent_id=self.agent_id,
run_id=self._per_operation_thread_id if self.scope_to_per_operation_thread_id else self.thread_id,
metadata={"application_id": self.application_id},
)
async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
"""Called before invoking the AI model to provide context.
Args:
messages: List of new messages in the thread.
Returns:
Context: Context object containing instructions with memories.
"""
self._validate_filters()
messages_list = [messages] if isinstance(messages, ChatMessage) else list(messages)
input_text = "\n".join(msg.text for msg in messages_list if msg and msg.text and msg.text.strip())
memories = await self.mem0_client.search( # type: ignore[misc]
query=input_text,
user_id=self.user_id,
agent_id=self.agent_id,
run_id=self._per_operation_thread_id if self.scope_to_per_operation_thread_id else self.thread_id,
)
line_separated_memories = "\n".join(memory.get("memory", "") for memory in memories)
content = TextContent(f"{self.context_prompt}\n{line_separated_memories}") if line_separated_memories else None
return Context(contents=[content] if content else None)
def _validate_filters(self) -> None:
"""Validates that at least one filter is provided.
Raises:
ServiceInitializationError: If no filters are provided.
"""
if not self.agent_id and not self.user_id and not self.application_id and not self.thread_id:
raise ServiceInitializationError(
"At least one of the filters: agent_id, user_id, application_id, or thread_id is required."
)
def _validate_per_operation_thread_id(self, thread_id: str | None) -> None:
"""Validates that a new thread ID doesn't conflict with an existing one when scoped.
Args:
thread_id: The new thread ID or None.
Raises:
ValueError: If a new thread ID is provided when one already exists.
"""
if (
self.scope_to_per_operation_thread_id
and thread_id
and self._per_operation_thread_id
and thread_id != self._per_operation_thread_id
):
raise ValueError(
"Mem0Provider can only be used with one thread at a time when scope_to_per_operation_thread_id is True."
)
+94
View File
@@ -0,0 +1,94 @@
[project]
name = "agent-framework-mem0"
description = "Mem0 integration for Microsoft Agent Framework."
authors = [{ name = "Microsoft", email = "SK-Support@microsoft.com"}]
readme = "README.md"
requires-python = ">=3.10"
version = "0.1.0b1"
license-files = ["LICENSE"]
urls.homepage = "https://learn.microsoft.com/en-us/semantic-kernel/overview/"
urls.source = "https://github.com/microsoft/agent-framework/tree/main/python"
urls.release_notes = "https://github.com/microsoft/agent-framework/releases?q=tag%3Apython-1&expanded=true"
urls.issues = "https://github.com/microsoft/agent-framework/issues"
classifiers = [
"License :: OSI Approved :: MIT License",
"Development Status :: 5 - Production/Stable",
"Intended Audience :: Developers",
"Programming Language :: Python :: 3",
"Programming Language :: Python :: 3.10",
"Programming Language :: Python :: 3.11",
"Programming Language :: Python :: 3.12",
"Programming Language :: Python :: 3.13",
"Framework :: Pydantic :: 2",
"Typing :: Typed",
]
dependencies = [
"agent-framework",
"mem0ai>=0.1.117",
]
[tool.uv]
prerelease = "if-necessary-or-explicit"
environments = [
"sys_platform == 'darwin'",
"sys_platform == 'linux'",
"sys_platform == 'win32'"
]
[tool.uv-dynamic-versioning]
fallback-version = "0.0.0"
[tool.pytest.ini_options]
testpaths = 'tests'
addopts = "-ra -q -r fEX"
asyncio_mode = "auto"
asyncio_default_fixture_loop_scope = "function"
filterwarnings = [
"ignore:Support for class-based `config` is deprecated:DeprecationWarning:pydantic.*"
]
timeout = 120
[tool.ruff]
extend = "../../pyproject.toml"
[tool.coverage.run]
omit = [
"**/__init__.py"
]
[tool.pyright]
extend = "../../pyproject.toml"
exclude = ['tests']
[tool.mypy]
plugins = ['pydantic.mypy']
strict = true
python_version = "3.10"
ignore_missing_imports = true
disallow_untyped_defs = true
no_implicit_optional = true
check_untyped_defs = true
warn_return_any = true
show_error_codes = true
warn_unused_ignores = false
disallow_incomplete_defs = true
disallow_untyped_decorators = true
disallow_any_unimported = true
[tool.bandit]
targets = ["agent_framework_mem0"]
exclude_dirs = ["tests"]
[tool.poe]
executor.type = "uv"
include = "../../shared_tasks.toml"
[tool.poe.tasks]
mypy = "mypy --config-file $POE_ROOT/pyproject.toml agent_framework_mem0"
test = "pytest --cov=agent_framework_mem0 --cov-report=term-missing:skip-covered tests"
[tool.uv.build-backend]
module-name = "agent_framework_mem0"
module-root = ""
[build-system]
requires = ["uv_build>=0.8.2,<0.9.0"]
build-backend = "uv_build"
@@ -0,0 +1,28 @@
# Copyright (c) Microsoft. All rights reserved.
def test_self_through_main() -> None:
try:
from agent_framework.mem0 import __version__
except ImportError:
__version__ = None
assert __version__ is not None
def test_self() -> None:
try:
from agent_framework_mem0 import __version__
except ImportError:
__version__ = None
assert __version__ is not None
def test_agent_framework() -> None:
try:
from agent_framework import __version__
except ImportError:
__version__ = None
assert __version__ is not None
+481
View File
@@ -0,0 +1,481 @@
# Copyright (c) Microsoft. All rights reserved.
# pyright: reportPrivateUsage=false
from unittest.mock import AsyncMock, patch
import pytest
from agent_framework import ChatMessage, Context, Role, TextContent
from agent_framework.exceptions import ServiceInitializationError
from agent_framework.mem0 import Mem0Provider
def test_mem0_provider_import():
"""Test that Mem0Provider can be imported."""
assert Mem0Provider is not None
@pytest.fixture
def mock_mem0_client() -> AsyncMock:
"""Create a mock Mem0 AsyncMemoryClient."""
mock_client = AsyncMock()
mock_client.add = AsyncMock()
mock_client.search = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock()
mock_client.async_client = AsyncMock()
mock_client.async_client.aclose = AsyncMock()
return mock_client
@pytest.fixture
def sample_messages() -> list[ChatMessage]:
"""Create sample chat messages for testing."""
return [
ChatMessage(role=Role.USER, text="Hello, how are you?"),
ChatMessage(role=Role.ASSISTANT, text="I'm doing well, thank you!"),
ChatMessage(role=Role.SYSTEM, text="You are a helpful assistant"),
]
class TestMem0ProviderInitialization:
"""Test initialization and configuration of Mem0Provider."""
def test_init_with_all_ids(self, mock_mem0_client: AsyncMock):
"""Test initialization with all IDs provided."""
provider = Mem0Provider(
user_id="user123",
agent_id="agent123",
application_id="app123",
thread_id="thread123",
mem0_client=mock_mem0_client,
)
assert provider.user_id == "user123"
assert provider.agent_id == "agent123"
assert provider.application_id == "app123"
assert provider.thread_id == "thread123"
def test_init_without_filters_succeeds(self, mock_mem0_client: AsyncMock):
"""Test that initialization succeeds even without filters (validation happens during invocation)."""
provider = Mem0Provider(mem0_client=mock_mem0_client)
assert provider.user_id is None
assert provider.agent_id is None
assert provider.application_id is None
assert provider.thread_id is None
def test_init_with_custom_context_prompt(self, mock_mem0_client: AsyncMock):
"""Test initialization with custom context prompt."""
custom_prompt = "## Custom Memories\nConsider these memories:"
provider = Mem0Provider(user_id="user123", context_prompt=custom_prompt, mem0_client=mock_mem0_client)
assert provider.context_prompt == custom_prompt
def test_init_with_scope_to_per_operation_thread_id(self, mock_mem0_client: AsyncMock):
"""Test initialization with scope_to_per_operation_thread_id enabled."""
provider = Mem0Provider(
user_id="user123",
scope_to_per_operation_thread_id=True,
mem0_client=mock_mem0_client,
)
assert provider.scope_to_per_operation_thread_id is True
@patch("mem0.AsyncMemoryClient")
def test_init_creates_default_client_when_none_provided(self, mock_memory_client_class: AsyncMock):
"""Test that a default client is created when none is provided."""
mock_client = AsyncMock()
mock_memory_client_class.return_value = mock_client
provider = Mem0Provider(user_id="user123", api_key="test_api_key")
mock_memory_client_class.assert_called_once_with(api_key="test_api_key")
assert provider.mem0_client == mock_client
assert provider._should_close_client is True
def test_init_with_provided_client_should_not_close(self, mock_mem0_client: AsyncMock):
"""Test that provided client should not be closed by provider."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
assert provider._should_close_client is False
class TestMem0ProviderAsyncContextManager:
"""Test async context manager behavior."""
async def test_async_context_manager_entry(self, mock_mem0_client: AsyncMock):
"""Test async context manager entry returns self."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
async with provider as ctx:
assert ctx is provider
async def test_async_context_manager_exit_closes_client_when_should_close(self):
"""Test that async context manager closes client when it should."""
mock_client = AsyncMock()
mock_client.__aenter__ = AsyncMock(return_value=mock_client)
mock_client.__aexit__ = AsyncMock()
mock_client.async_client = AsyncMock()
mock_client.async_client.aclose = AsyncMock()
with patch("mem0.AsyncMemoryClient", return_value=mock_client):
provider = Mem0Provider(user_id="user123", api_key="test_key")
assert provider._should_close_client is True
async with provider:
pass
mock_client.__aexit__.assert_called_once()
async def test_async_context_manager_exit_does_not_close_provided_client(self, mock_mem0_client: AsyncMock):
"""Test that async context manager does not close provided client."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
assert provider._should_close_client is False
async with provider:
pass
mock_mem0_client.__aexit__.assert_not_called()
class TestMem0ProviderThreadMethods:
"""Test thread lifecycle methods."""
async def test_thread_created_sets_per_operation_thread_id(self, mock_mem0_client: AsyncMock):
"""Test that thread_created sets per-operation thread ID."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
await provider.thread_created("thread123")
assert provider._per_operation_thread_id == "thread123"
async def test_thread_created_with_existing_thread_id(self, mock_mem0_client: AsyncMock):
"""Test thread_created when thread ID already exists."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
provider._per_operation_thread_id = "existing_thread"
await provider.thread_created("thread123")
# Should not overwrite existing thread ID
assert provider._per_operation_thread_id == "existing_thread"
async def test_thread_created_validation_with_scope_enabled(self, mock_mem0_client: AsyncMock):
"""Test thread_created validation when scope_to_per_operation_thread_id is enabled."""
provider = Mem0Provider(
user_id="user123",
scope_to_per_operation_thread_id=True,
mem0_client=mock_mem0_client,
)
provider._per_operation_thread_id = "existing_thread"
with pytest.raises(ValueError) as exc_info:
await provider.thread_created("different_thread")
assert "can only be used with one thread at a time" in str(exc_info.value)
async def test_messages_adding_sets_per_operation_thread_id(
self, mock_mem0_client: AsyncMock, sample_messages: list[ChatMessage]
):
"""Test that messages_adding sets per-operation thread ID."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
await provider.messages_adding("thread123", sample_messages)
assert provider._per_operation_thread_id == "thread123"
class TestMem0ProviderMessagesAdding:
"""Test messages_adding method."""
async def test_messages_adding_fails_without_filters(self, mock_mem0_client: AsyncMock):
"""Test that messages_adding fails when no filters are provided."""
provider = Mem0Provider(mem0_client=mock_mem0_client)
message = ChatMessage(role=Role.USER, text="Hello!")
with pytest.raises(ServiceInitializationError) as exc_info:
await provider.messages_adding("thread123", message)
assert "At least one of the filters" in str(exc_info.value)
async def test_messages_adding_single_message(self, mock_mem0_client: AsyncMock):
"""Test adding a single message."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
message = ChatMessage(role=Role.USER, text="Hello!")
await provider.messages_adding("thread123", message)
mock_mem0_client.add.assert_called_once()
call_args = mock_mem0_client.add.call_args
assert call_args.kwargs["messages"] == [{"role": "user", "content": "Hello!"}]
assert call_args.kwargs["user_id"] == "user123"
async def test_messages_adding_multiple_messages(
self, mock_mem0_client: AsyncMock, sample_messages: list[ChatMessage]
):
"""Test adding multiple messages."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
await provider.messages_adding("thread123", sample_messages)
mock_mem0_client.add.assert_called_once()
call_args = mock_mem0_client.add.call_args
expected_messages = [
{"role": "user", "content": "Hello, how are you?"},
{"role": "assistant", "content": "I'm doing well, thank you!"},
{"role": "system", "content": "You are a helpful assistant"},
]
assert call_args.kwargs["messages"] == expected_messages
async def test_messages_adding_with_agent_id(self, mock_mem0_client: AsyncMock, sample_messages: list[ChatMessage]):
"""Test adding messages with agent_id."""
provider = Mem0Provider(agent_id="agent123", mem0_client=mock_mem0_client)
await provider.messages_adding("thread123", sample_messages)
call_args = mock_mem0_client.add.call_args
assert call_args.kwargs["agent_id"] == "agent123"
assert call_args.kwargs["user_id"] is None
async def test_messages_adding_with_application_id(
self, mock_mem0_client: AsyncMock, sample_messages: list[ChatMessage]
):
"""Test adding messages with application_id in metadata."""
provider = Mem0Provider(user_id="user123", application_id="app123", mem0_client=mock_mem0_client)
await provider.messages_adding("thread123", sample_messages)
call_args = mock_mem0_client.add.call_args
assert call_args.kwargs["metadata"] == {"application_id": "app123"}
async def test_messages_adding_with_scope_to_per_operation_thread_id(
self, mock_mem0_client: AsyncMock, sample_messages: list[ChatMessage]
):
"""Test adding messages with scope_to_per_operation_thread_id enabled."""
provider = Mem0Provider(
user_id="user123",
thread_id="base_thread",
scope_to_per_operation_thread_id=True,
mem0_client=mock_mem0_client,
)
provider._per_operation_thread_id = "operation_thread"
await provider.messages_adding("operation_thread", sample_messages)
call_args = mock_mem0_client.add.call_args
assert call_args.kwargs["run_id"] == "operation_thread"
async def test_messages_adding_without_scope_uses_base_thread_id(
self, mock_mem0_client: AsyncMock, sample_messages: list[ChatMessage]
):
"""Test adding messages without scope uses base thread_id."""
provider = Mem0Provider(
user_id="user123",
thread_id="base_thread",
scope_to_per_operation_thread_id=False,
mem0_client=mock_mem0_client,
)
await provider.messages_adding("operation_thread", sample_messages)
call_args = mock_mem0_client.add.call_args
assert call_args.kwargs["run_id"] == "base_thread"
async def test_messages_adding_filters_empty_messages(self, mock_mem0_client: AsyncMock):
"""Test that empty or invalid messages are filtered out."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
messages = [
ChatMessage(role=Role.USER, text=""), # Empty text
ChatMessage(role=Role.USER, text=" "), # Whitespace only
ChatMessage(role=Role.USER, text="Valid message"),
]
await provider.messages_adding("thread123", messages)
call_args = mock_mem0_client.add.call_args
# Should only include the valid message
assert call_args.kwargs["messages"] == [{"role": "user", "content": "Valid message"}]
async def test_messages_adding_skips_when_no_valid_messages(self, mock_mem0_client: AsyncMock):
"""Test that mem0 client is not called when no valid messages exist."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
messages = [
ChatMessage(role=Role.USER, text=""),
ChatMessage(role=Role.USER, text=" "),
]
await provider.messages_adding("thread123", messages)
mock_mem0_client.add.assert_not_called()
class TestMem0ProviderModelInvoking:
"""Test model_invoking method."""
async def test_model_invoking_fails_without_filters(self, mock_mem0_client: AsyncMock):
"""Test that model_invoking fails when no filters are provided."""
provider = Mem0Provider(mem0_client=mock_mem0_client)
message = ChatMessage(role=Role.USER, text="What's the weather?")
with pytest.raises(ServiceInitializationError) as exc_info:
await provider.model_invoking(message)
assert "At least one of the filters" in str(exc_info.value)
async def test_model_invoking_single_message(self, mock_mem0_client: AsyncMock):
"""Test model_invoking with a single message."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
message = ChatMessage(role=Role.USER, text="What's the weather?")
# Mock search results
mock_mem0_client.search.return_value = [
{"memory": "User likes outdoor activities"},
{"memory": "User lives in Seattle"},
]
context = await provider.model_invoking(message)
mock_mem0_client.search.assert_called_once()
call_args = mock_mem0_client.search.call_args
assert call_args.kwargs["query"] == "What's the weather?"
assert call_args.kwargs["user_id"] == "user123"
assert isinstance(context, Context)
expected_instructions = (
"## Memories\nConsider the following memories when answering user questions:\n"
"User likes outdoor activities\nUser lives in Seattle"
)
assert context.contents
assert isinstance(context.contents[0], TextContent)
assert context.contents[0].text == expected_instructions
async def test_model_invoking_multiple_messages(
self, mock_mem0_client: AsyncMock, sample_messages: list[ChatMessage]
):
"""Test model_invoking with multiple messages."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
mock_mem0_client.search.return_value = [{"memory": "Previous conversation context"}]
await provider.model_invoking(sample_messages)
call_args = mock_mem0_client.search.call_args
expected_query = "Hello, how are you?\nI'm doing well, thank you!\nYou are a helpful assistant"
assert call_args.kwargs["query"] == expected_query
async def test_model_invoking_with_agent_id(self, mock_mem0_client: AsyncMock):
"""Test model_invoking with agent_id."""
provider = Mem0Provider(agent_id="agent123", mem0_client=mock_mem0_client)
message = ChatMessage(role=Role.USER, text="Hello")
mock_mem0_client.search.return_value = []
await provider.model_invoking(message)
call_args = mock_mem0_client.search.call_args
assert call_args.kwargs["agent_id"] == "agent123"
assert call_args.kwargs["user_id"] is None
async def test_model_invoking_with_scope_to_per_operation_thread_id(self, mock_mem0_client: AsyncMock):
"""Test model_invoking with scope_to_per_operation_thread_id enabled."""
provider = Mem0Provider(
user_id="user123",
thread_id="base_thread",
scope_to_per_operation_thread_id=True,
mem0_client=mock_mem0_client,
)
provider._per_operation_thread_id = "operation_thread"
message = ChatMessage(role=Role.USER, text="Hello")
mock_mem0_client.search.return_value = []
await provider.model_invoking(message)
call_args = mock_mem0_client.search.call_args
assert call_args.kwargs["run_id"] == "operation_thread"
async def test_model_invoking_no_memories_returns_none_instructions(self, mock_mem0_client: AsyncMock):
"""Test that no memories returns context with None instructions."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
message = ChatMessage(role=Role.USER, text="Hello")
mock_mem0_client.search.return_value = []
context = await provider.model_invoking(message)
assert isinstance(context, Context)
assert not context.contents
async def test_model_invoking_filters_empty_message_text(self, mock_mem0_client: AsyncMock):
"""Test that empty message text is filtered out from query."""
provider = Mem0Provider(user_id="user123", mem0_client=mock_mem0_client)
messages = [
ChatMessage(role=Role.USER, text=""),
ChatMessage(role=Role.USER, text="Valid message"),
ChatMessage(role=Role.USER, text=" "),
]
mock_mem0_client.search.return_value = []
await provider.model_invoking(messages)
call_args = mock_mem0_client.search.call_args
assert call_args.kwargs["query"] == "Valid message"
async def test_model_invoking_custom_context_prompt(self, mock_mem0_client: AsyncMock):
"""Test model_invoking with custom context prompt."""
custom_prompt = "## Custom Context\nRemember these details:"
provider = Mem0Provider(
user_id="user123",
context_prompt=custom_prompt,
mem0_client=mock_mem0_client,
)
message = ChatMessage(role=Role.USER, text="Hello")
mock_mem0_client.search.return_value = [{"memory": "Test memory"}]
context = await provider.model_invoking(message)
expected_instructions = "## Custom Context\nRemember these details:\nTest memory"
assert context.contents
assert isinstance(context.contents[0], TextContent)
assert context.contents[0].text == expected_instructions
class TestMem0ProviderValidation:
"""Test validation methods."""
def test_validate_per_operation_thread_id_success(self, mock_mem0_client: AsyncMock):
"""Test successful validation of per-operation thread ID."""
provider = Mem0Provider(
user_id="user123",
scope_to_per_operation_thread_id=True,
mem0_client=mock_mem0_client,
)
provider._per_operation_thread_id = "thread123"
# Should not raise exception for same thread ID
provider._validate_per_operation_thread_id("thread123")
# Should not raise exception for None
provider._validate_per_operation_thread_id(None)
def test_validate_per_operation_thread_id_failure(self, mock_mem0_client: AsyncMock):
"""Test validation failure for conflicting thread IDs."""
provider = Mem0Provider(
user_id="user123",
scope_to_per_operation_thread_id=True,
mem0_client=mock_mem0_client,
)
provider._per_operation_thread_id = "thread123"
with pytest.raises(ValueError) as exc_info:
provider._validate_per_operation_thread_id("different_thread")
assert "can only be used with one thread at a time" in str(exc_info.value)
def test_validate_per_operation_thread_id_disabled_scope(self, mock_mem0_client: AsyncMock):
"""Test that validation is skipped when scope is disabled."""
provider = Mem0Provider(
user_id="user123",
scope_to_per_operation_thread_id=False,
mem0_client=mock_mem0_client,
)
provider._per_operation_thread_id = "thread123"
# Should not raise exception even with different thread ID
provider._validate_per_operation_thread_id("different_thread")