mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
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:
committed by
GitHub
Unverified
parent
89c8418705
commit
57d09afe04
@@ -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
|
||||
@@ -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."
|
||||
)
|
||||
@@ -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
|
||||
@@ -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")
|
||||
Reference in New Issue
Block a user