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>

python/packages/main/tests/main/conftest.py

@@ -151,11 +151,19 @@ class MockBaseChatClient(BaseChatClient):
         logger.debug(f"Running base chat client inner, with: {messages=}, {chat_options=}, {kwargs=}")
         if not self.run_responses:
             return ChatResponse(messages=ChatMessage(role="assistant", text=f"test response - {messages[0].text}"))
+
+        response = self.run_responses.pop(0)
+
         if chat_options.tool_choice == "none":
             return ChatResponse(
-                messages=ChatMessage(role="assistant", text="I broke out of the function invocation loop...")
+                messages=ChatMessage(
+                    role="assistant",
+                    text="I broke out of the function invocation loop...",
+                ),
+                conversation_id=response.conversation_id,
             )
-        return self.run_responses.pop(0)
+
+        return response
 
     @override
     async def _inner_get_streaming_response(

python/packages/main/tests/main/test_agents.py

@@ -1,6 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
-from collections.abc import AsyncIterable
+from collections.abc import AsyncIterable, MutableSequence, Sequence
 from uuid import uuid4
 
 from pytest import raises
@@ -15,10 +15,12 @@ from agent_framework import (
     ChatMessage,
     ChatMessageList,
     ChatResponse,
+    Contents,
     HostedCodeInterpreterTool,
     Role,
     TextContent,
 )
+from agent_framework._memory import AggregateContextProvider, Context, ContextProvider
 from agent_framework.exceptions import AgentExecutionException
 
 
@@ -100,6 +102,7 @@ async def test_chat_client_agent_prepare_thread_and_messages(chat_client: ChatCl
 
     _, result_messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
         thread=thread,
+        context=Context(),
         input_messages=[ChatMessage(role=Role.USER, text="Test")],
     )
 
@@ -184,3 +187,210 @@ async def test_chat_client_agent_author_name_is_used_from_response(chat_client_b
     result = await agent.run("Hello")
     assert result.text == "test response"
     assert result.messages[0].author_name == "TestAuthor"
+
+
+# Mock context provider for testing
+class MockContextProvider(ContextProvider):
+    context_contents: list[Contents] | None = None
+    thread_created_called: bool = False
+    messages_adding_called: bool = False
+    model_invoking_called: bool = False
+    thread_created_thread_id: str | None = None
+    messages_adding_thread_id: str | None = None
+    new_messages: list[ChatMessage] = []
+
+    def __init__(self, contents: list[Contents] | None = None) -> None:
+        super().__init__()
+        self.context_contents = contents
+        self.thread_created_called = False
+        self.messages_adding_called = False
+        self.model_invoking_called = False
+        self.thread_created_thread_id = None
+        self.messages_adding_thread_id = None
+        self.new_messages = []
+
+    async def thread_created(self, thread_id: str | None) -> None:
+        self.thread_created_called = True
+        self.thread_created_thread_id = thread_id
+
+    async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
+        self.messages_adding_called = True
+        self.messages_adding_thread_id = thread_id
+        if isinstance(new_messages, ChatMessage):
+            self.new_messages.append(new_messages)
+        else:
+            self.new_messages.extend(new_messages)
+
+    async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
+        self.model_invoking_called = True
+        return Context(contents=self.context_contents)
+
+
+async def test_chat_agent_context_providers_model_invoking(chat_client: ChatClientProtocol) -> None:
+    """Test that context providers' model_invoking is called during agent run."""
+    mock_provider = MockContextProvider(contents=[TextContent("Test context instructions")])
+    agent = ChatAgent(chat_client=chat_client, context_providers=mock_provider)
+
+    await agent.run("Hello")
+
+    assert mock_provider.model_invoking_called
+
+
+async def test_chat_agent_context_providers_thread_created(chat_client_base: ChatClientProtocol) -> None:
+    """Test that context providers' thread_created is called during agent run."""
+    mock_provider = MockContextProvider()
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")])],
+            conversation_id="test-thread-id",
+        )
+    ]
+
+    agent = ChatAgent(chat_client=chat_client_base, context_providers=mock_provider)
+
+    await agent.run("Hello")
+
+    assert mock_provider.thread_created_called
+    assert mock_provider.thread_created_thread_id == "test-thread-id"
+
+
+async def test_chat_agent_context_providers_messages_adding(chat_client: ChatClientProtocol) -> None:
+    """Test that context providers' messages_adding is called during agent run."""
+    mock_provider = MockContextProvider()
+    agent = ChatAgent(chat_client=chat_client, context_providers=mock_provider)
+
+    await agent.run("Hello")
+
+    assert mock_provider.messages_adding_called
+    # Should be called with both input and response messages
+    assert len(mock_provider.new_messages) >= 2
+
+
+async def test_chat_agent_context_instructions_in_messages(chat_client: ChatClientProtocol) -> None:
+    """Test that AI context instructions are included in messages."""
+    mock_provider = MockContextProvider(contents=[TextContent("Context-specific instructions")])
+    agent = ChatAgent(chat_client=chat_client, instructions="Agent instructions", context_providers=mock_provider)
+
+    # We need to test the _prepare_thread_and_messages method directly
+    context = Context(contents=[TextContent("Context-specific instructions")])
+    _, messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
+        thread=None, context=context, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
+    )
+
+    # Should have agent instructions, context instructions, and user message
+    assert len(messages) == 3
+    assert messages[0].role == Role.SYSTEM
+    assert messages[0].text == "Agent instructions"
+    assert messages[1].role == Role.SYSTEM
+    assert messages[1].text == "Context-specific instructions"
+    assert messages[2].role == Role.USER
+    assert messages[2].text == "Hello"
+
+
+async def test_chat_agent_context_instructions_without_agent_instructions(chat_client: ChatClientProtocol) -> None:
+    """Test that AI context instructions work when agent has no instructions."""
+    agent = ChatAgent(chat_client=chat_client)  # No instructions
+    context = Context(contents=[TextContent("Context-only instructions")])
+
+    _, messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
+        thread=None, context=context, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
+    )
+
+    # Should have context instructions and user message only
+    assert len(messages) == 2
+    assert messages[0].role == Role.SYSTEM
+    assert messages[0].text == "Context-only instructions"
+    assert messages[1].role == Role.USER
+    assert messages[1].text == "Hello"
+
+
+async def test_chat_agent_no_context_instructions(chat_client: ChatClientProtocol) -> None:
+    """Test behavior when AI context has no instructions."""
+    agent = ChatAgent(chat_client=chat_client, instructions="Agent instructions")
+    context = Context()  # No instructions
+
+    _, messages = await agent._prepare_thread_and_messages(  # type: ignore[reportPrivateUsage]
+        thread=None, context=context, input_messages=[ChatMessage(role=Role.USER, text="Hello")]
+    )
+
+    # Should have agent instructions and user message only
+    assert len(messages) == 2
+    assert messages[0].role == Role.SYSTEM
+    assert messages[0].text == "Agent instructions"
+    assert messages[1].role == Role.USER
+    assert messages[1].text == "Hello"
+
+
+async def test_chat_agent_run_stream_context_providers(chat_client: ChatClientProtocol) -> None:
+    """Test that context providers work with run_stream method."""
+    mock_provider = MockContextProvider(contents=[TextContent("Stream context instructions")])
+    agent = ChatAgent(chat_client=chat_client, context_providers=mock_provider)
+
+    # Collect all stream updates
+    updates: list[AgentRunResponseUpdate] = []
+    async for update in agent.run_stream("Hello"):
+        updates.append(update)
+
+    # Verify context provider was called
+    assert mock_provider.model_invoking_called
+    assert mock_provider.thread_created_called
+    assert mock_provider.messages_adding_called
+
+
+async def test_chat_agent_multiple_context_providers(chat_client: ChatClientProtocol) -> None:
+    """Test that multiple context providers work together."""
+    provider1 = MockContextProvider(contents=[TextContent("First provider instructions")])
+    provider2 = MockContextProvider(contents=[TextContent("Second provider instructions")])
+
+    agent = ChatAgent(chat_client=chat_client, context_providers=[provider1, provider2])
+
+    await agent.run("Hello")
+
+    # Both providers should be called
+    assert provider1.model_invoking_called
+    assert provider1.thread_created_called
+    assert provider1.messages_adding_called
+
+    assert provider2.model_invoking_called
+    assert provider2.thread_created_called
+    assert provider2.messages_adding_called
+
+
+async def test_chat_agent_aggregate_context_provider_combines_instructions() -> None:
+    """Test that AggregateContextProvider combines instructions from multiple providers."""
+    provider1 = MockContextProvider(contents=[TextContent("First instruction")])
+    provider2 = MockContextProvider(contents=[TextContent("Second instruction")])
+
+    aggregate = AggregateContextProvider()
+    aggregate.providers.append(provider1)
+    aggregate.providers.append(provider2)
+
+    # Test model_invoking combines instructions
+    result = await aggregate.model_invoking([ChatMessage(role=Role.USER, text="Test")])
+
+    assert result.contents
+    assert isinstance(result.contents[0], TextContent)
+    assert isinstance(result.contents[1], TextContent)
+    assert result.contents[0].text == "First instruction"
+    assert result.contents[1].text == "Second instruction"
+
+
+async def test_chat_agent_context_providers_with_thread_service_id(chat_client_base: ChatClientProtocol) -> None:
+    """Test context providers with service-managed thread."""
+    mock_provider = MockContextProvider()
+    chat_client_base.run_responses = [
+        ChatResponse(
+            messages=[ChatMessage(role=Role.ASSISTANT, contents=[TextContent("test response")])],
+            conversation_id="service-thread-123",
+        )
+    ]
+
+    agent = ChatAgent(chat_client=chat_client_base, context_providers=mock_provider)
+
+    # Use existing service-managed thread
+    thread = AgentThread(service_thread_id="existing-thread-id")
+    await agent.run("Hello", thread=thread)
+
+    # messages_adding should be called with the service thread ID from response
+    assert mock_provider.messages_adding_called
+    assert mock_provider.messages_adding_thread_id == "service-thread-123"  # Updated thread ID from response

python/packages/main/tests/main/test_memory.py

@@ -0,0 +1,302 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from collections.abc import MutableSequence, Sequence
+from unittest.mock import AsyncMock, Mock
+
+from agent_framework import ChatMessage, Contents, Role, TextContent
+from agent_framework._memory import AggregateContextProvider, Context, ContextProvider
+
+
+class MockContextProvider(ContextProvider):
+    """Mock ContextProvider for testing."""
+
+    context_contents: list[Contents] | None = None
+    thread_created_called: bool = False
+    messages_adding_called: bool = False
+    model_invoking_called: bool = False
+    thread_created_thread_id: str | None = None
+    messages_adding_thread_id: str | None = None
+    messages_adding_new_messages: ChatMessage | Sequence[ChatMessage] | None = None
+    model_invoking_messages: ChatMessage | MutableSequence[ChatMessage] | None = None
+
+    def __init__(self, context_contents: list[Contents] | None = None) -> None:
+        super().__init__()
+        self.context_contents = context_contents
+        self.thread_created_called = False
+        self.messages_adding_called = False
+        self.model_invoking_called = False
+        self.thread_created_thread_id = None
+        self.messages_adding_thread_id = None
+        self.messages_adding_new_messages = None
+        self.model_invoking_messages = None
+
+    async def thread_created(self, thread_id: str | None) -> None:
+        """Track thread_created calls."""
+        self.thread_created_called = True
+        self.thread_created_thread_id = thread_id
+
+    async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
+        """Track messages_adding calls."""
+        self.messages_adding_called = True
+        self.messages_adding_thread_id = thread_id
+        self.messages_adding_new_messages = new_messages
+
+    async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
+        """Track model_invoking calls and return context."""
+        self.model_invoking_called = True
+        self.model_invoking_messages = messages
+        context = Context()
+        context.contents = self.context_contents
+        return context
+
+
+class TestAggregateContextProvider:
+    """Tests for AggregateContextProvider class."""
+
+    def test_init_with_no_providers(self) -> None:
+        """Test initialization with no providers."""
+        aggregate = AggregateContextProvider()
+        assert aggregate.providers == []
+
+    def test_init_with_none_providers(self) -> None:
+        """Test initialization with None providers."""
+        aggregate = AggregateContextProvider(None)
+        assert aggregate.providers == []
+
+    def test_init_with_providers(self) -> None:
+        """Test initialization with providers."""
+        provider1 = MockContextProvider([TextContent("instructions1")])
+        provider2 = MockContextProvider([TextContent("instructions1")])
+        providers = [provider1, provider2]
+
+        aggregate = AggregateContextProvider(providers)
+        assert len(aggregate.providers) == 2
+        assert aggregate.providers[0] is provider1
+        assert aggregate.providers[1] is provider2
+
+    def test_add_provider(self) -> None:
+        """Test adding a provider."""
+        aggregate = AggregateContextProvider()
+        provider = MockContextProvider([TextContent("instructions")])
+
+        aggregate.add(provider)
+        assert len(aggregate.providers) == 1
+        assert aggregate.providers[0] is provider
+
+    def test_add_multiple_providers(self) -> None:
+        """Test adding multiple providers."""
+        aggregate = AggregateContextProvider()
+        provider1 = MockContextProvider([TextContent("instructions1")])
+        provider2 = MockContextProvider([TextContent("instructions2")])
+
+        aggregate.add(provider1)
+        aggregate.add(provider2)
+
+        assert len(aggregate.providers) == 2
+        assert aggregate.providers[0] is provider1
+        assert aggregate.providers[1] is provider2
+
+    async def test_thread_created_with_no_providers(self) -> None:
+        """Test thread_created with no providers."""
+        aggregate = AggregateContextProvider()
+
+        # Should not raise an exception
+        await aggregate.thread_created("thread-123")
+
+    async def test_thread_created_with_providers(self) -> None:
+        """Test thread_created calls all providers."""
+        provider1 = MockContextProvider([TextContent("instructions1")])
+        provider2 = MockContextProvider([TextContent("instructions2")])
+        aggregate = AggregateContextProvider([provider1, provider2])
+
+        thread_id = "thread-123"
+        await aggregate.thread_created(thread_id)
+
+        assert provider1.thread_created_called
+        assert provider1.thread_created_thread_id == thread_id
+        assert provider2.thread_created_called
+        assert provider2.thread_created_thread_id == thread_id
+
+    async def test_thread_created_with_none_thread_id(self) -> None:
+        """Test thread_created with None thread_id."""
+        provider = MockContextProvider([TextContent("instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        await aggregate.thread_created(None)
+
+        assert provider.thread_created_called
+        assert provider.thread_created_thread_id is None
+
+    async def test_messages_adding_with_no_providers(self) -> None:
+        """Test messages_adding with no providers."""
+        aggregate = AggregateContextProvider()
+        message = ChatMessage(text="Hello", role=Role.USER)
+
+        # Should not raise an exception
+        await aggregate.messages_adding("thread-123", message)
+
+    async def test_messages_adding_with_single_message(self) -> None:
+        """Test messages_adding with a single message."""
+        provider1 = MockContextProvider([TextContent("instructions1")])
+        provider2 = MockContextProvider([TextContent("instructions2")])
+        aggregate = AggregateContextProvider([provider1, provider2])
+
+        thread_id = "thread-123"
+        message = ChatMessage(text="Hello", role=Role.USER)
+        await aggregate.messages_adding(thread_id, message)
+
+        assert provider1.messages_adding_called
+        assert provider1.messages_adding_thread_id == thread_id
+        assert provider1.messages_adding_new_messages == message
+        assert provider2.messages_adding_called
+        assert provider2.messages_adding_thread_id == thread_id
+        assert provider2.messages_adding_new_messages == message
+
+    async def test_messages_adding_with_message_sequence(self) -> None:
+        """Test messages_adding with a sequence of messages."""
+        provider = MockContextProvider([TextContent("instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        thread_id = "thread-123"
+        messages = [
+            ChatMessage(text="Hello", role=Role.USER),
+            ChatMessage(text="Hi there", role=Role.ASSISTANT),
+        ]
+        await aggregate.messages_adding(thread_id, messages)
+
+        assert provider.messages_adding_called
+        assert provider.messages_adding_thread_id == thread_id
+        assert provider.messages_adding_new_messages == messages
+
+    async def test_model_invoking_with_no_providers(self) -> None:
+        """Test model_invoking with no providers."""
+        aggregate = AggregateContextProvider()
+        message = ChatMessage(text="Hello", role=Role.USER)
+
+        context = await aggregate.model_invoking(message)
+
+        assert isinstance(context, Context)
+        assert not context.contents
+
+    async def test_model_invoking_with_single_provider(self) -> None:
+        """Test model_invoking with a single provider."""
+        provider = MockContextProvider([TextContent("Test instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        message = ChatMessage(text="Hello", role=Role.USER)
+        context = await aggregate.model_invoking(message)
+
+        assert provider.model_invoking_called
+        assert provider.model_invoking_messages == message
+        assert isinstance(context, Context)
+
+        assert context.contents
+        assert isinstance(context.contents[0], TextContent)
+        assert context.contents[0].text == "Test instructions"
+
+    async def test_model_invoking_with_multiple_providers(self) -> None:
+        """Test model_invoking combines contexts from multiple providers."""
+        provider1 = MockContextProvider([TextContent("Instructions 1")])
+        provider2 = MockContextProvider([TextContent("Instructions 2")])
+        provider3 = MockContextProvider([TextContent("Instructions 3")])
+        aggregate = AggregateContextProvider([provider1, provider2, provider3])
+
+        messages = [ChatMessage(text="Hello", role=Role.USER)]
+        context = await aggregate.model_invoking(messages)
+
+        assert provider1.model_invoking_called
+        assert provider1.model_invoking_messages == messages
+        assert provider2.model_invoking_called
+        assert provider2.model_invoking_messages == messages
+        assert provider3.model_invoking_called
+        assert provider3.model_invoking_messages == messages
+
+        assert isinstance(context, Context)
+
+        assert context.contents
+        assert isinstance(context.contents[0], TextContent)
+        assert isinstance(context.contents[1], TextContent)
+        assert isinstance(context.contents[2], TextContent)
+        assert context.contents[0].text == "Instructions 1"
+        assert context.contents[1].text == "Instructions 2"
+        assert context.contents[2].text == "Instructions 3"
+
+    async def test_model_invoking_with_none_instructions(self) -> None:
+        """Test model_invoking filters out None instructions."""
+        provider1 = MockContextProvider([TextContent("Instructions 1")])
+        provider2 = MockContextProvider(None)  # None instructions
+        provider3 = MockContextProvider([TextContent("Instructions 3")])
+        aggregate = AggregateContextProvider([provider1, provider2, provider3])
+
+        message = ChatMessage(text="Hello", role=Role.USER)
+        context = await aggregate.model_invoking(message)
+
+        assert isinstance(context, Context)
+        assert context.contents
+        assert isinstance(context.contents[0], TextContent)
+        assert isinstance(context.contents[1], TextContent)
+        assert context.contents[0].text == "Instructions 1"
+        assert context.contents[1].text == "Instructions 3"
+
+    async def test_model_invoking_with_all_none_instructions(self) -> None:
+        """Test model_invoking when all providers return None instructions."""
+        provider1 = MockContextProvider(None)
+        provider2 = MockContextProvider(None)
+        aggregate = AggregateContextProvider([provider1, provider2])
+
+        message = ChatMessage(text="Hello", role=Role.USER)
+        context = await aggregate.model_invoking(message)
+
+        assert isinstance(context, Context)
+        assert not context.contents
+
+    async def test_model_invoking_with_mutable_sequence(self) -> None:
+        """Test model_invoking with MutableSequence of messages."""
+        provider = MockContextProvider([TextContent("Test instructions")])
+        aggregate = AggregateContextProvider([provider])
+
+        messages = [ChatMessage(text="Hello", role=Role.USER)]
+        context = await aggregate.model_invoking(messages)
+
+        assert provider.model_invoking_called
+        assert provider.model_invoking_messages == messages
+        assert isinstance(context, Context)
+        assert context.contents
+        assert isinstance(context.contents[0], TextContent)
+        assert context.contents[0].text == "Test instructions"
+
+    async def test_async_methods_concurrent_execution(self) -> None:
+        """Test that async methods execute providers concurrently."""
+        # Use AsyncMock to verify concurrent execution
+        provider1 = Mock(spec=ContextProvider)
+        provider1.thread_created = AsyncMock()
+        provider1.messages_adding = AsyncMock()
+        provider1.model_invoking = AsyncMock(return_value=Context(contents=[TextContent("Test 1")]))
+
+        provider2 = Mock(spec=ContextProvider)
+        provider2.thread_created = AsyncMock()
+        provider2.messages_adding = AsyncMock()
+        provider2.model_invoking = AsyncMock(return_value=Context(contents=[TextContent("Test 2")]))
+
+        aggregate = AggregateContextProvider([provider1, provider2])
+
+        # Test thread_created
+        await aggregate.thread_created("thread-123")
+        provider1.thread_created.assert_called_once_with("thread-123")
+        provider2.thread_created.assert_called_once_with("thread-123")
+
+        # Test messages_adding
+        message = ChatMessage(text="Hello", role=Role.USER)
+        await aggregate.messages_adding("thread-123", message)
+        provider1.messages_adding.assert_called_once_with("thread-123", message)
+        provider2.messages_adding.assert_called_once_with("thread-123", message)
+
+        # Test model_invoking
+        context = await aggregate.model_invoking(message)
+        provider1.model_invoking.assert_called_once_with(message)
+        provider2.model_invoking.assert_called_once_with(message)
+        assert context.contents
+        assert isinstance(context.contents[0], TextContent)
+        assert isinstance(context.contents[1], TextContent)
+        assert context.contents[0].text == "Test 1"
+        assert context.contents[1].text == "Test 2"