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/agent_framework/__init__.py

@@ -12,6 +12,7 @@ from ._agents import *  # noqa: F403
 from ._clients import *  # noqa: F403
 from ._logging import *  # noqa: F403
 from ._mcp import *  # noqa: F403
+from ._memory import *  # noqa: F403
 from ._threads import *  # noqa: F403
 from ._tools import *  # noqa: F403
 from ._types import *  # noqa: F403

python/packages/main/agent_framework/_agents.py

@@ -3,7 +3,6 @@
 import sys
 from collections.abc import AsyncIterable, Callable, MutableMapping, Sequence
 from contextlib import AbstractAsyncContextManager, AsyncExitStack
-from itertools import chain
 from typing import Any, ClassVar, Literal, Protocol, TypeVar, runtime_checkable
 from uuid import uuid4
 
@@ -12,6 +11,7 @@ from pydantic import BaseModel, Field, PrivateAttr
 from ._clients import BaseChatClient, ChatClientProtocol
 from ._logging import get_logger
 from ._mcp import MCPTool
+from ._memory import AggregateContextProvider, Context, ContextProvider
 from ._pydantic import AFBaseModel
 from ._threads import AgentThread, ChatMessageStore, deserialize_thread_state, thread_on_new_messages
 from ._tools import FUNCTION_INVOKING_CHAT_CLIENT_MARKER, ToolProtocol
@@ -137,12 +137,13 @@ class BaseAgent(AFBaseModel):
        name: The name of the agent, can be None.
        description: The description of the agent.
        display_name: The display name of the agent, which is either the name or id.
-
+       context_providers: The collection of multiple context providers to include during agent invocation.
     """
 
     id: str = Field(default_factory=lambda: str(uuid4()))
     name: str | None = None
     description: str | None = None
+    context_providers: AggregateContextProvider | None = None
 
     async def _notify_thread_of_new_messages(
         self, thread: AgentThread, new_messages: ChatMessage | Sequence[ChatMessage]
@@ -214,6 +215,7 @@ class ChatAgent(BaseAgent):
         user: str | None = None,
         additional_properties: dict[str, Any] | None = None,
         chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
+        context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
         **kwargs: Any,
     ) -> None:
         """Create a ChatAgent.
@@ -248,6 +250,7 @@ class ChatAgent(BaseAgent):
             additional_properties: additional properties to include in the request.
             chat_message_store_factory: factory function to create an instance of ChatMessageStore. If not provided,
                 the default in-memory store will be used.
+            context_providers: The collection of multiple context providers to include during agent invocation.
             kwargs: any additional keyword arguments.
                 Unused, can be used by subclasses of this Agent.
         """
@@ -258,6 +261,8 @@ class ChatAgent(BaseAgent):
 
         kwargs.update(additional_properties or {})
 
+        aggregate_context_providers = self._prepare_context_providers(context_providers)
+
         # We ignore the MCP Servers here and store them separately,
         # we add their functions to the tools list at runtime
         normalized_tools = [] if tools is None else tools if isinstance(tools, list) else [tools]
@@ -266,6 +271,7 @@ class ChatAgent(BaseAgent):
         args: dict[str, Any] = {
             "chat_client": chat_client,
             "chat_message_store_factory": chat_message_store_factory,
+            "context_providers": aggregate_context_providers,
             "chat_options": ChatOptions(
                 ai_model_id=model,
                 frequency_penalty=frequency_penalty,
@@ -301,12 +307,16 @@ class ChatAgent(BaseAgent):
     async def __aenter__(self) -> "Self":
         """Async context manager entry.
 
-        If either the chat_client or the local_mcp_tools are context managers,
+        If any of the chat_client, local_mcp_tools, or context_providers are context managers,
         they will be entered into the async exit stack to ensure proper cleanup.
 
         This list might be extended in the future.
         """
-        for context_manager in chain([self.chat_client], self._local_mcp_tools):
+        context_managers = [self.chat_client, *self._local_mcp_tools]
+        if self.context_providers:
+            context_managers.append(self.context_providers)
+
+        for context_manager in context_managers:
             if isinstance(context_manager, AbstractAsyncContextManager):
                 await self._async_exit_stack.enter_async_context(context_manager)
         return self
@@ -388,7 +398,10 @@ class ChatAgent(BaseAgent):
                 will only be passed to functions that are called.
         """
         input_messages = self._normalize_messages(messages)
-        thread, thread_messages = await self._prepare_thread_and_messages(thread=thread, input_messages=input_messages)
+        context = await self.context_providers.model_invoking(input_messages) if self.context_providers else None
+        thread, thread_messages = await self._prepare_thread_and_messages(
+            thread=thread, context=context, input_messages=input_messages
+        )
         agent_name = self._get_agent_name()
 
         # Resolve final tool list (runtime provided tools + local MCP server tools)
@@ -441,6 +454,10 @@ class ChatAgent(BaseAgent):
         await self._notify_thread_of_new_messages(thread, input_messages)
         await self._notify_thread_of_new_messages(thread, response.messages)
 
+        if self.context_providers:
+            await self.context_providers.thread_created(response.conversation_id)
+            await self.context_providers.messages_adding(thread.service_thread_id, input_messages + response.messages)
+
         return AgentRunResponse(
             messages=response.messages,
             response_id=response.response_id,
@@ -509,7 +526,10 @@ class ChatAgent(BaseAgent):
 
         """
         input_messages = self._normalize_messages(messages)
-        thread, thread_messages = await self._prepare_thread_and_messages(thread=thread, input_messages=input_messages)
+        context = await self.context_providers.model_invoking(input_messages) if self.context_providers else None
+        thread, thread_messages = await self._prepare_thread_and_messages(
+            thread=thread, context=context, input_messages=input_messages
+        )
         agent_name = self._get_agent_name()
         response_updates: list[ChatResponseUpdate] = []
 
@@ -575,6 +595,10 @@ class ChatAgent(BaseAgent):
         await self._notify_thread_of_new_messages(thread, input_messages)
         await self._notify_thread_of_new_messages(thread, response.messages)
 
+        if self.context_providers:
+            await self.context_providers.thread_created(response.conversation_id)
+            await self.context_providers.messages_adding(thread.service_thread_id, input_messages + response.messages)
+
     def get_new_thread(self) -> AgentThread:
         message_store: ChatMessageStore | None = None
 
@@ -617,12 +641,14 @@ class ChatAgent(BaseAgent):
         self,
         *,
         thread: AgentThread | None,
+        context: Context | None,
         input_messages: list[ChatMessage] | None = None,
     ) -> tuple[AgentThread, list[ChatMessage]]:
         """Prepare the messages for agent execution.
 
         Args:
             thread: The conversation thread.
+            context: Context to include in messages.
             input_messages: Messages to process.
 
         Returns:
@@ -636,6 +662,8 @@ class ChatAgent(BaseAgent):
         messages: list[ChatMessage] = []
         if self.instructions:
             messages.append(ChatMessage(role=Role.SYSTEM, text=self.instructions))
+        if context and context.contents:
+            messages.append(ChatMessage(role=Role.SYSTEM, contents=context.contents))
         if thread.message_store:
             messages.extend(await thread.message_store.list_messages() or [])
         messages.extend(input_messages or [])
@@ -658,3 +686,18 @@ class ChatAgent(BaseAgent):
 
     def _get_agent_name(self) -> str:
         return self.name or "UnnamedAgent"
+
+    def _prepare_context_providers(
+        self,
+        context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
+    ) -> AggregateContextProvider | None:
+        if not context_providers:
+            return None
+
+        if isinstance(context_providers, AggregateContextProvider):
+            return context_providers
+
+        if isinstance(context_providers, ContextProvider):
+            return AggregateContextProvider([context_providers])
+
+        return AggregateContextProvider(context_providers)

python/packages/main/agent_framework/_clients.py

@@ -9,6 +9,7 @@ from pydantic import BaseModel, Field
 
 from ._logging import get_logger
 from ._mcp import MCPTool
+from ._memory import AggregateContextProvider, ContextProvider
 from ._pydantic import AFBaseModel
 from ._threads import ChatMessageStore
 from ._tools import ToolProtocol
@@ -463,6 +464,7 @@ class BaseChatClient(AFBaseModel, ABC):
         | list[ToolProtocol | Callable[..., Any] | MutableMapping[str, Any]]
         | None = None,
         chat_message_store_factory: Callable[[], ChatMessageStore] | None = None,
+        context_providers: ContextProvider | list[ContextProvider] | AggregateContextProvider | None = None,
         **kwargs: Any,
     ) -> "ChatAgent":
         """Create an agent with the given name and instructions.
@@ -473,6 +475,7 @@ class BaseChatClient(AFBaseModel, ABC):
             tools: Optional list of tools to associate with the agent.
             chat_message_store_factory: Factory function to create an instance of ChatMessageStore. If not provided,
                 the default in-memory store will be used.
+            context_providers: Context providers to include during agent invocation.
             **kwargs: Additional keyword arguments to pass to the agent.
                 See ChatAgent for all the available options.
 
@@ -487,6 +490,7 @@ class BaseChatClient(AFBaseModel, ABC):
             instructions=instructions,
             tools=tools,
             chat_message_store_factory=chat_message_store_factory,
+            context_providers=context_providers,
             **kwargs,
         )
 

python/packages/main/agent_framework/_memory.py

@@ -0,0 +1,188 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import sys
+from abc import ABC, abstractmethod
+from collections.abc import MutableSequence, Sequence
+from contextlib import AsyncExitStack
+from types import TracebackType
+
+from ._pydantic import AFBaseModel
+from ._types import ChatMessage, Contents
+
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
+# region Context
+
+
+class Context(AFBaseModel):
+    """A class containing any context that should be provided to the AI model as supplied by an ContextProvider.
+
+    Each ContextProvider has the ability to provide its own context for each invocation.
+    The Context class contains the additional context supplied by the ContextProvider.
+    This context will be combined with context supplied by other providers before being passed to the AI model.
+    This context is per invocation, and will not be stored as part of the chat history.
+    """
+
+    contents: list[Contents] | None = None
+    """
+    Any content to pass to the AI model in addition to any other prompts
+    that it may already have (in the case of an agent), or chat history that may already exist.
+    """
+
+
+# region ContextProvider
+
+
+class ContextProvider(AFBaseModel, ABC):
+    """Base class for all context providers.
+
+    A context provider is a component that can be used to enhance the AI's context management.
+    It can listen to changes in the conversation and provide additional context to the AI model
+    just before invocation.
+    """
+
+    async def thread_created(self, thread_id: str | None) -> None:
+        """Called just after a new thread is created.
+
+        Implementers can use this method to do any operations required at the creation of a new thread.
+        For example, checking long term storage for any data that is relevant
+        to the current session based on the input text.
+
+        Args:
+            thread_id: The ID of the new thread.
+        """
+        pass
+
+    async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
+        """Called just before messages are added to the chat by any participant.
+
+        Inheritors can use this method to update their context based on new messages.
+
+        Args:
+            thread_id: The ID of the thread for the new message.
+            new_messages: New messages to add.
+        """
+        pass
+
+    @abstractmethod
+    async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
+        """Called just before the Model/Agent/etc. is invoked.
+
+        Implementers can load any additional context required at this time,
+        and they should return any context that should be passed to the agent.
+
+        Args:
+            messages: The most recent messages that the agent is being invoked with.
+        """
+        pass
+
+    async def __aenter__(self) -> "Self":
+        """Async context manager entry.
+
+        Override this method to perform any setup operations when the context provider is entered.
+
+        Returns:
+            Self for chaining.
+        """
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Async context manager exit.
+
+        Override this method to perform any cleanup operations when the context provider is exited.
+
+        Args:
+            exc_type: Exception type if an exception occurred, None otherwise.
+            exc_val: Exception value if an exception occurred, None otherwise.
+            exc_tb: Exception traceback if an exception occurred, None otherwise.
+        """
+        pass
+
+
+# region AggregateContextProvider
+
+
+class AggregateContextProvider(ContextProvider):
+    """A ContextProvider that contains multiple context providers.
+
+    It delegates events to multiple context providers and aggregates responses from those events before returning.
+    """
+
+    providers: list[ContextProvider]
+    """List of registered context providers."""
+
+    def __init__(self, context_providers: Sequence[ContextProvider] | None = None) -> None:
+        """Initialize AggregateContextProvider with context providers.
+
+        Args:
+            context_providers: Context providers to add.
+        """
+        super().__init__(providers=list(context_providers or []))  # type: ignore
+        self._exit_stack: AsyncExitStack | None = None
+
+    def add(self, context_provider: ContextProvider) -> None:
+        """Adds new context provider.
+
+        Args:
+            context_provider: Context provider to add.
+        """
+        self.providers.append(context_provider)
+
+    async def thread_created(self, thread_id: str | None = None) -> None:
+        await asyncio.gather(*[x.thread_created(thread_id) for x in self.providers])
+
+    async def messages_adding(self, thread_id: str | None, new_messages: ChatMessage | Sequence[ChatMessage]) -> None:
+        await asyncio.gather(*[x.messages_adding(thread_id, new_messages) for x in self.providers])
+
+    async def model_invoking(self, messages: ChatMessage | MutableSequence[ChatMessage]) -> Context:
+        sub_contexts = await asyncio.gather(*[x.model_invoking(messages) for x in self.providers])
+        combined_context = Context()
+        # Flatten the list of lists and filter out None values
+        all_contents = []
+        for ctx in sub_contexts:
+            if ctx.contents:
+                all_contents.extend(ctx.contents)
+
+        combined_context.contents = all_contents if all_contents else None
+        return combined_context
+
+    async def __aenter__(self) -> "Self":
+        """Enter async context manager and set up all providers.
+
+        Returns:
+            Self for chaining.
+        """
+        self._exit_stack = AsyncExitStack()
+        await self._exit_stack.__aenter__()
+
+        # Enter all context providers
+        for provider in self.providers:
+            await self._exit_stack.enter_async_context(provider)
+
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        """Exit async context manager and clean up all providers.
+
+        Args:
+            exc_type: Exception type if an exception occurred, None otherwise.
+            exc_val: Exception value if an exception occurred, None otherwise.
+            exc_tb: Exception traceback if an exception occurred, None otherwise.
+        """
+        if self._exit_stack is not None:
+            await self._exit_stack.__aexit__(exc_type, exc_val, exc_tb)
+            self._exit_stack = None

python/packages/main/agent_framework/mem0/__init__.py

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import importlib
+from typing import Any
+
+PACKAGE_NAME = "agent_framework_mem0"
+PACKAGE_EXTRA = "mem0"
+_IMPORTS = ["__version__", "Mem0Provider"]
+
+
+def __getattr__(name: str) -> Any:
+    if name in _IMPORTS:
+        try:
+            return getattr(importlib.import_module(PACKAGE_NAME), name)
+        except ModuleNotFoundError as exc:
+            raise ModuleNotFoundError(
+                f"The '{PACKAGE_EXTRA}' extra is not installed, "
+                f"please do `pip install agent-framework[{PACKAGE_EXTRA}]`"
+            ) from exc
+    raise AttributeError(f"Module {PACKAGE_NAME} has no attribute {name}.")
+
+
+def __dir__() -> list[str]:
+    return _IMPORTS

python/packages/main/agent_framework/mem0/__init__.pyi

@@ -0,0 +1,5 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from agent_framework_mem0 import Mem0Provider, __version__
+
+__all__ = ["Mem0Provider", "__version__"]

python/packages/main/agent_framework/openai/_responses_client.py

@@ -455,7 +455,7 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                         )
                         response_tools.append(
                             WebSearchToolParam(
-                                type="web_search",
+                                type="web_search_preview",
                                 user_location=WebSearchUserLocation(
                                     type="approximate",
                                     city=location.get("city", None),