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>
2026-06-16 21:04:09 +08:00 · 2025-09-10 21:11:42 +00:00
parent 89c8418705
commit 57d09afe04
25 changed files with 4166 additions and 1915 deletions
@@ -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