chore: import upstream snapshot with attribution

python/packages/mem0/agent_framework_mem0/__init__.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import importlib.metadata
+import os
+
+# Disable Mem0 telemetry by default to prevent usage data from being sent to telemetry provider.
+# Users can opt-in by setting MEM0_TELEMETRY=true before importing this package.
+if os.environ.get("MEM0_TELEMETRY") is None:
+    os.environ["MEM0_TELEMETRY"] = "false"
+
+from ._context_provider import Mem0ContextProvider
+
+try:
+    __version__ = importlib.metadata.version(__name__)
+except importlib.metadata.PackageNotFoundError:
+    __version__ = "0.0.0"  # Fallback for development mode
+
+__all__ = [
+    "Mem0ContextProvider",
+    "__version__",
+]

python/packages/mem0/agent_framework_mem0/_context_provider.py

@@ -0,0 +1,271 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+"""New-pattern Mem0 context provider using ContextProvider.
+
+This module provides ``Mem0ContextProvider``, built on the new
+:class:`ContextProvider` hooks pattern.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import sys
+from collections.abc import Awaitable
+from contextlib import AbstractAsyncContextManager
+from typing import TYPE_CHECKING, Any, ClassVar, TypeAlias, TypedDict
+
+from agent_framework import Message
+from agent_framework._sessions import AgentSession, ContextProvider, SessionContext
+from mem0 import AsyncMemory, AsyncMemoryClient
+
+if sys.version_info >= (3, 11):
+    from typing import Self  # pragma: no cover
+else:
+    from typing_extensions import Self  # pragma: no cover
+
+if TYPE_CHECKING:
+    from agent_framework._agents import SupportsAgentRun
+
+logger = logging.getLogger(__name__)
+MemoryRecord: TypeAlias = dict[str, object]
+
+
+class SearchResults(TypedDict):
+    results: list[MemoryRecord]
+
+
+SearchResponse: TypeAlias = list[MemoryRecord] | SearchResults
+
+
+class Mem0ContextProvider(ContextProvider):
+    """Mem0 context provider using the new ContextProvider hooks pattern.
+
+    Integrates Mem0 for persistent semantic memory, searching and storing
+    memories via the Mem0 API.
+    """
+
+    DEFAULT_CONTEXT_PROMPT = "## Memories\nConsider the following memories when answering user questions:"
+    DEFAULT_SOURCE_ID: ClassVar[str] = "mem0"
+
+    def __init__(
+        self,
+        source_id: str = DEFAULT_SOURCE_ID,
+        mem0_client: AsyncMemory | AsyncMemoryClient | None = None,
+        api_key: str | None = None,
+        application_id: str | None = None,
+        agent_id: str | None = None,
+        user_id: str | None = None,
+        *,
+        context_prompt: str | None = None,
+    ) -> None:
+        """Initialize the Mem0 context provider.
+
+        Args:
+            source_id: Unique identifier for this provider instance.
+            mem0_client: A pre-created Mem0 MemoryClient or None to create a default client.
+            api_key: The API key for authenticating with the Mem0 API.
+            application_id: The application ID for scoping memories. Platform-only:
+                the OSS ``AsyncMemory`` client does not recognize an application
+                scope (it scopes only by user_id/agent_id in this provider), so
+                application_id cannot be used with an OSS client.
+            agent_id: The agent ID for scoping memories.
+            user_id: The user ID for scoping memories.
+            context_prompt: The prompt to prepend to retrieved memories.
+        """
+        super().__init__(source_id)
+        should_close_client = False
+        if mem0_client is None:
+            mem0_client = AsyncMemoryClient(api_key=api_key)
+            should_close_client = True
+
+        self.api_key = api_key
+        self.application_id = application_id
+        self.agent_id = agent_id
+        self.user_id = user_id
+        self.context_prompt = context_prompt or self.DEFAULT_CONTEXT_PROMPT
+        self.mem0_client = mem0_client
+        self._should_close_client = should_close_client
+
+    async def __aenter__(self) -> Self:
+        """Async context manager entry."""
+        if self.mem0_client and isinstance(self.mem0_client, AbstractAsyncContextManager):
+            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 and isinstance(self.mem0_client, AbstractAsyncContextManager):
+            await self.mem0_client.__aexit__(exc_type, exc_val, exc_tb)  # pyright: ignore[reportUnknownMemberType]
+
+    # -- Hooks pattern ---------------------------------------------------------
+
+    async def before_run(
+        self,
+        *,
+        agent: SupportsAgentRun,
+        session: AgentSession,
+        context: SessionContext,
+        state: dict[str, Any],
+    ) -> None:
+        """Search Mem0 for relevant memories and add to the session context."""
+        self._validate_filters()
+        input_text = "\n".join(msg.text for msg in context.input_messages if msg and msg.text and msg.text.strip())
+        if not input_text.strip():
+            return
+
+        # Query entity partitions independently to bypass strict logical AND limitations
+        # Mem0 OSS and Platform SDKs expose inconsistent search typings.
+        search_tasks: list[Awaitable[Any]] = []
+
+        # 1. Query User partition independently
+        if self.user_id:
+            user_kwargs = self._build_search_kwargs(input_text, "user_id", self.user_id)
+            search_tasks.append(self.mem0_client.search(**user_kwargs))  # type: ignore[reportUnknownMemberType, reportUnknownArgumentType]
+
+        # 2. Query Agent partition independently
+        if self.agent_id:
+            agent_kwargs = self._build_search_kwargs(input_text, "agent_id", self.agent_id)
+            search_tasks.append(self.mem0_client.search(**agent_kwargs))  # type: ignore[reportUnknownMemberType, reportUnknownArgumentType]
+
+        # Fall back to an app-scoped search when only application_id is configured.
+        if not search_tasks and self.application_id:
+            app_kwargs: dict[str, Any] = {"query": input_text, "filters": self._build_filters()}
+            search_tasks.append(self.mem0_client.search(**app_kwargs))  # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType]
+        if not search_tasks:
+            return
+
+        results: list[SearchResponse | BaseException] = await asyncio.gather(*search_tasks, return_exceptions=True)
+
+        # Merge and deduplicate results
+        memories: list[MemoryRecord] = []
+        seen_memory_ids: set[str] = set()
+        failed_tasks_count: int = 0
+
+        for search_response in results:
+            if isinstance(search_response, asyncio.CancelledError):
+                raise search_response
+
+            if isinstance(search_response, BaseException):
+                failed_tasks_count += 1
+                logger.error(
+                    "Mem0 partition search task failed: %s",
+                    search_response,
+                    exc_info=(type(search_response), search_response, search_response.__traceback__),
+                )
+                continue
+
+            current_memories: list[MemoryRecord] = []
+            if isinstance(search_response, list):
+                current_memories = [mem for mem in search_response if isinstance(mem, dict)]
+            elif isinstance(search_response, dict):
+                results_field = search_response.get("results")
+                if isinstance(results_field, list):
+                    current_memories = [item for item in results_field if isinstance(item, dict)]
+                else:
+                    logger.warning(
+                        "Unexpected Mem0 search response format: %s",
+                        type(results_field).__name__,
+                    )
+
+            for mem in current_memories:
+                mem_id = mem.get("id")
+                if mem_id is not None and not isinstance(mem_id, str):
+                    mem_id = str(mem_id)
+
+                if mem_id is not None and mem_id in seen_memory_ids:
+                    continue
+
+                if mem_id is not None:
+                    seen_memory_ids.add(mem_id)
+
+                memories.append(mem)
+
+        if failed_tasks_count == len(search_tasks):
+            logger.error("All Mem0 retrieval tasks failed. Context provider is unable to verify memory state.")
+
+        line_separated_memories = "\n".join(str(memory.get("memory", "")) for memory in memories)
+        if line_separated_memories:
+            context.extend_messages(
+                self.source_id,
+                [Message(role="user", contents=[f"{self.context_prompt}\n{line_separated_memories}"])],
+            )
+
+    async def after_run(
+        self,
+        *,
+        agent: SupportsAgentRun,
+        session: AgentSession,
+        context: SessionContext,
+        state: dict[str, Any],
+    ) -> None:
+        """Store request/response messages to Mem0 for future retrieval."""
+        self._validate_filters()
+
+        messages_to_store: list[Message] = list(context.input_messages)
+        if context.response and context.response.messages:
+            messages_to_store.extend(context.response.messages)
+
+        def get_role_value(role: Any) -> str:
+            return role.value if hasattr(role, "value") else str(role)
+
+        messages: list[dict[str, str]] = [
+            {"role": get_role_value(message.role), "content": message.text}
+            for message in messages_to_store
+            if get_role_value(message.role) in {"user", "assistant", "system"} and message.text and message.text.strip()
+        ]
+
+        if messages:
+            add_kwargs: dict[str, Any] = {
+                "messages": messages,
+            }
+
+            if isinstance(self.mem0_client, AsyncMemory):
+                add_kwargs["user_id"] = self.user_id
+                add_kwargs["agent_id"] = self.agent_id
+            else:
+                add_kwargs["filters"] = self._build_filters()
+
+            await self.mem0_client.add(**add_kwargs)  # type: ignore[misc, call-arg]
+
+    # -- Internal methods ------------------------------------------------------
+
+    def _validate_filters(self) -> None:
+        """Validates that at least one usable filter is provided for the configured client."""
+        if not self.agent_id and not self.user_id and not self.application_id:
+            raise ValueError("At least one of the filters: agent_id, user_id, or application_id is required.")
+        if isinstance(self.mem0_client, AsyncMemory) and self.application_id:
+            raise ValueError(
+                "application_id is not supported by the OSS AsyncMemory client, which scopes "
+                "memories only by user_id/agent_id. Remove application_id or use AsyncMemoryClient."
+            )
+
+    def _build_search_kwargs(self, input_text: str, entity_key: str, entity_value: str) -> dict[str, Any]:
+        """Build search keyword arguments formatted for OSS vs Platform clients."""
+        filters: dict[str, Any] = {"query": input_text}
+
+        if self.application_id and isinstance(self.mem0_client, AsyncMemory):
+            raise ValueError(
+                "application_id is not supported by the OSS AsyncMemory client, which scopes "
+                "memories only by user_id/agent_id. Remove application_id or use AsyncMemoryClient."
+            )
+
+        filters["filters"] = {entity_key: entity_value}
+        if self.application_id and not isinstance(self.mem0_client, AsyncMemory):
+            filters["filters"]["app_id"] = self.application_id
+
+        return filters
+
+    def _build_filters(self) -> dict[str, Any]:
+        """Build identity filters from initialization parameters."""
+        filters: dict[str, Any] = {}
+        if self.user_id:
+            filters["user_id"] = self.user_id
+        if self.agent_id:
+            filters["agent_id"] = self.agent_id
+        if self.application_id:
+            filters["app_id"] = self.application_id
+        return filters
+
+
+__all__ = ["Mem0ContextProvider"]