From 21d964628ba6c031578b64ef2987ed4e9d14580a Mon Sep 17 00:00:00 2001
From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
Date: Tue, 3 Mar 2026 13:42:07 -0800
Subject: [PATCH 1/2] Upgraded azure-ai-projects to 2.0.0b4

---
 .../agent_framework_azure_ai/_client.py       | 31 ++++++++++------
 .../_foundry_memory_provider.py               | 21 ++++++-----
 .../_project_provider.py                      | 35 +++++++++++--------
 .../agent_framework_azure_ai/_shared.py       | 20 +++++------
 .../azure-ai/tests/test_azure_ai_client.py    | 20 +++++------
 .../packages/azure-ai/tests/test_provider.py  |  3 +-
 python/packages/core/pyproject.toml           |  3 +-
 .../azure_ai_foundry_memory.py                |  6 ++--
 .../azure_ai/azure_ai_provider_methods.py     |  8 ++---
 .../azure_ai/azure_ai_with_memory_search.py   |  6 ++--
 python/uv.lock                                | 19 +++++-----
 11 files changed, 93 insertions(+), 79 deletions(-)

diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_client.py
index 7c698847cc..d18b082a9e 100644
--- a/python/packages/azure-ai/agent_framework_azure_ai/_client.py
+++ b/python/packages/azure-ai/agent_framework_azure_ai/_client.py
@@ -37,12 +37,13 @@
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.projects.models import (
     ApproximateLocation,
+    CodeInterpreterContainerAuto,
     CodeInterpreterTool,
-    CodeInterpreterToolAuto,
+    FoundryFeaturesOptInKeys,
     ImageGenTool,
     MCPTool,
     PromptAgentDefinition,
-    PromptAgentDefinitionText,
+    PromptAgentDefinitionTextOptions,
     RaiConfig,
     Reasoning,
     WebSearchPreviewTool,
@@ -78,6 +79,9 @@ class AzureAIProjectAgentOptions(OpenAIResponsesOptions, total=False):
     reasoning: Reasoning  # type: ignore[misc]
     """Configuration for enabling reasoning capabilities (requires azure.ai.projects.models.Reasoning)."""
 
+    foundry_features: FoundryFeaturesOptInKeys | str
+    """Optional Foundry preview feature opt-in for agent version creation."""
+
 
 AzureAIClientOptionsT = TypeVar(
     "AzureAIClientOptionsT",
@@ -392,7 +396,7 @@ async def _get_agent_reference_or_create(
             # response_format is accessed from chat_options or additional_properties
             # since the base class excludes it from run_options
             if chat_options and (response_format := chat_options.get("response_format")):
-                args["text"] = PromptAgentDefinitionText(format=create_text_format_config(response_format))
+                args["text"] = PromptAgentDefinitionTextOptions(format=create_text_format_config(response_format))
 
             # Combine instructions from messages and options
             # instructions is accessed from chat_options since the base class excludes it from run_options
@@ -404,11 +408,15 @@ async def _get_agent_reference_or_create(
             if combined_instructions:
                 args["instructions"] = "".join(combined_instructions)
 
-            created_agent = await self.project_client.agents.create_version(
-                agent_name=self.agent_name,
-                definition=PromptAgentDefinition(**args),
-                description=self.agent_description,
-            )
+            create_version_kwargs: dict[str, Any] = {
+                "agent_name": self.agent_name,
+                "definition": PromptAgentDefinition(**args),
+                "description": self.agent_description,
+            }
+            if foundry_features := run_options.get("foundry_features"):
+                create_version_kwargs["foundry_features"] = foundry_features
+
+            created_agent = await self.project_client.agents.create_version(**create_version_kwargs)
 
             self.agent_version = created_agent.version
             self.warn_runtime_tools_and_structure_changed = True
@@ -500,6 +508,7 @@ def _remove_agent_level_run_options(
             "temperature": ("temperature",),
             "top_p": ("top_p",),
             "reasoning": ("reasoning",),
+            "foundry_features": ("foundry_features",),
         }
 
         for run_keys in agent_level_option_to_run_keys.values():
@@ -526,9 +535,9 @@ async def _prepare_options(
             run_options["input"] = self._transform_input_for_azure_ai(cast(list[dict[str, Any]], run_options["input"]))
 
         if not self._is_application_endpoint:
-            # Application-scoped response APIs do not support "agent" property.
+            # Application-scoped response APIs do not support "agent_reference" property.
             agent_reference = await self._get_agent_reference_or_create(run_options, instructions, options)
-            run_options["extra_body"] = {"agent": agent_reference}
+            run_options["extra_body"] = {"agent_reference": agent_reference}
 
         # Remove only keys that map to this client's declared options TypedDict.
         self._remove_agent_level_run_options(run_options, options)
@@ -857,7 +866,7 @@ def get_code_interpreter_tool(  # type: ignore[override]
         # Extract file_ids from container if provided as dict and file_ids not explicitly set
         if file_ids is None and isinstance(container, dict):
             file_ids = container.get("file_ids")
-        tool_container = CodeInterpreterToolAuto(file_ids=file_ids if file_ids else None)
+        tool_container = CodeInterpreterContainerAuto(file_ids=file_ids if file_ids else None)
         return CodeInterpreterTool(container=tool_container, **kwargs)
 
     @staticmethod
diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_foundry_memory_provider.py b/python/packages/azure-ai/agent_framework_azure_ai/_foundry_memory_provider.py
index eba210ff10..d02eb31bb6 100644
--- a/python/packages/azure-ai/agent_framework_azure_ai/_foundry_memory_provider.py
+++ b/python/packages/azure-ai/agent_framework_azure_ai/_foundry_memory_provider.py
@@ -18,7 +18,6 @@
 from agent_framework._settings import load_settings
 from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
 from azure.ai.projects.aio import AIProjectClient
-from azure.ai.projects.models import ItemParam, ResponsesAssistantMessageItemParam, ResponsesUserMessageItemParam
 
 from ._shared import AzureAISettings
 
@@ -149,7 +148,7 @@ async def before_run(
         # On first run, retrieve static memories (user profile memories)
         if not state.get("initialized"):
             try:
-                static_search_result = await self.project_client.memory_stores.search_memories(
+                static_search_result = await self.project_client.beta.memory_stores.search_memories(
                     name=self.memory_store_name,
                     scope=self.scope or context.session_id,  # type: ignore[arg-type]
                 )
@@ -169,15 +168,15 @@ async def before_run(
         if not has_input:
             return
 
-        # Convert input messages to ItemParam format for search
+        # Convert input messages to memory search item format
         items = [
-            ItemParam({"type": "text", "text": msg.text})
+            {"type": "text", "text": msg.text}
             for msg in context.input_messages
             if msg and msg.text and msg.text.strip()
         ]
 
         try:
-            search_result = await self.project_client.memory_stores.search_memories(
+            search_result = await self.project_client.beta.memory_stores.search_memories(
                 name=self.memory_store_name,
                 scope=self.scope or context.session_id,  # type: ignore[arg-type]
                 items=items,
@@ -224,24 +223,24 @@ async def after_run(
         if context.response and context.response.messages:
             messages_to_store.extend(context.response.messages)
 
-        # Filter and convert messages to ItemParam format
-        items: list[ResponsesUserMessageItemParam | ResponsesAssistantMessageItemParam] = []
+        # Filter and convert messages to memory update item format
+        items: list[dict[str, str]] = []
         for message in messages_to_store:
             if message.role in {"user", "assistant", "system"} and message.text and message.text.strip():
                 if message.role == "user":
-                    items.append(ResponsesUserMessageItemParam(content=message.text))
+                    items.append({"role": "user", "type": "message", "content": message.text})
                 elif message.role == "assistant":
-                    items.append(ResponsesAssistantMessageItemParam(content=message.text))
+                    items.append({"role": "assistant", "type": "message", "content": message.text})
 
         if not items:
             return
 
         try:
             # Fire and forget - don't wait for the update to complete
-            update_poller = await self.project_client.memory_stores.begin_update_memories(
+            update_poller = await self.project_client.beta.memory_stores.begin_update_memories(
                 name=self.memory_store_name,
                 scope=self.scope or context.session_id,  # type: ignore[arg-type]
-                items=items,  # type: ignore[arg-type]
+                items=items,
                 previous_update_id=state.get("previous_update_id"),
                 update_delay=self.update_delay,
             )
diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py b/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
index 81276d446b..d6b922db91 100644
--- a/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
+++ b/python/packages/azure-ai/agent_framework_azure_ai/_project_provider.py
@@ -4,7 +4,7 @@
 
 import logging
 import sys
-from collections.abc import Callable, MutableMapping, Sequence
+from collections.abc import Callable, Mapping, MutableMapping, Sequence
 from typing import Any, Generic
 
 from agent_framework import (
@@ -21,10 +21,9 @@
 from agent_framework.azure._entra_id_authentication import AzureCredentialTypes
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.projects.models import (
-    AgentReference,
     AgentVersionDetails,
     PromptAgentDefinition,
-    PromptAgentDefinitionText,
+    PromptAgentDefinitionTextOptions,
 )
 from azure.ai.projects.models import (
     FunctionTool as AzureFunctionTool,
@@ -200,13 +199,14 @@ async def create_agent(
         response_format = opts.get("response_format")
         rai_config = opts.get("rai_config")
         reasoning = opts.get("reasoning")
+        foundry_features = opts.get("foundry_features")
 
         args: dict[str, Any] = {"model": resolved_model}
 
         if instructions:
             args["instructions"] = instructions
         if response_format and isinstance(response_format, (type, dict)):
-            args["text"] = PromptAgentDefinitionText(
+            args["text"] = PromptAgentDefinitionTextOptions(
                 format=create_text_format_config(response_format)  # type: ignore[arg-type]
             )
         if rai_config:
@@ -241,11 +241,15 @@ async def create_agent(
         if all_tools_for_azure:
             args["tools"] = to_azure_ai_tools(all_tools_for_azure)
 
-        created_agent = await self._project_client.agents.create_version(
-            agent_name=name,
-            definition=PromptAgentDefinition(**args),
-            description=description,
-        )
+        create_version_kwargs: dict[str, Any] = {
+            "agent_name": name,
+            "definition": PromptAgentDefinition(**args),
+            "description": description,
+        }
+        if foundry_features:
+            create_version_kwargs["foundry_features"] = foundry_features
+
+        created_agent = await self._project_client.agents.create_version(**create_version_kwargs)
 
         return self._to_chat_agent_from_details(
             created_agent,
@@ -259,7 +263,7 @@ async def get_agent(
         self,
         *,
         name: str | None = None,
-        reference: AgentReference | None = None,
+        reference: Mapping[str, str | None] | None = None,
         tools: ToolTypes | Callable[..., Any] | Sequence[ToolTypes | Callable[..., Any]] | None = None,
         default_options: OptionsCoT | None = None,
         middleware: Sequence[MiddlewareTypes] | None = None,
@@ -272,7 +276,7 @@ async def get_agent(
 
         Args:
             name: The name of the agent to retrieve (fetches latest version).
-            reference: Reference containing the agent's name and optionally a specific version.
+            reference: Mapping containing the agent's ``name`` and optionally a specific ``version``.
             tools: Tools to make available to the agent. Required if the agent has function tools.
             default_options: A TypedDict containing default chat options for the agent.
                 These options are applied to every run unless overridden.
@@ -287,12 +291,15 @@ async def get_agent(
         """
         existing_agent: AgentVersionDetails
 
-        if reference and reference.version:
+        reference_name = str(reference.get("name")) if reference and reference.get("name") else None
+        reference_version = str(reference.get("version")) if reference and reference.get("version") else None
+
+        if reference_name and reference_version:
             # Fetch specific version
             existing_agent = await self._project_client.agents.get_version(
-                agent_name=reference.name, agent_version=reference.version
+                agent_name=reference_name, agent_version=reference_version
             )
-        elif agent_name := (reference.name if reference else name):
+        elif agent_name := (reference_name if reference_name else name):
             # Fetch latest version
             details = await self._project_client.agents.get(agent_name=agent_name)
             existing_agent = details.versions.latest
diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_shared.py b/python/packages/azure-ai/agent_framework_azure_ai/_shared.py
index 7dd1064bda..9612899d79 100644
--- a/python/packages/azure-ai/agent_framework_azure_ai/_shared.py
+++ b/python/packages/azure-ai/agent_framework_azure_ai/_shared.py
@@ -18,9 +18,9 @@
 from azure.ai.projects.models import (
     CodeInterpreterTool,
     MCPTool,
-    ResponseTextFormatConfigurationJsonObject,
-    ResponseTextFormatConfigurationJsonSchema,
-    ResponseTextFormatConfigurationText,
+    TextResponseFormatConfigurationResponseFormatJsonObject,
+    TextResponseFormatConfigurationResponseFormatText,
+    TextResponseFormatJsonSchema,
     Tool,
     WebSearchPreviewTool,
 )
@@ -421,9 +421,9 @@ def _prepare_mcp_tool_dict_for_azure_ai(tool_dict: dict[str, Any]) -> MCPTool:
 def create_text_format_config(
     response_format: type[BaseModel] | Mapping[str, Any],
 ) -> (
-    ResponseTextFormatConfigurationJsonSchema
-    | ResponseTextFormatConfigurationJsonObject
-    | ResponseTextFormatConfigurationText
+    TextResponseFormatJsonSchema
+    | TextResponseFormatConfigurationResponseFormatJsonObject
+    | TextResponseFormatConfigurationResponseFormatText
 ):
     """Convert response_format into Azure text format configuration."""
     if isinstance(response_format, type) and issubclass(response_format, BaseModel):
@@ -431,7 +431,7 @@ def create_text_format_config(
         # Ensure additionalProperties is explicitly false to satisfy Azure validation
         if isinstance(schema, dict):
             schema.setdefault("additionalProperties", False)
-        return ResponseTextFormatConfigurationJsonSchema(
+        return TextResponseFormatJsonSchema(
             name=response_format.__name__,
             schema=schema,
             strict=True,
@@ -452,11 +452,11 @@ def create_text_format_config(
                 config_kwargs["strict"] = format_config["strict"]
             if "description" in format_config:
                 config_kwargs["description"] = format_config["description"]
-            return ResponseTextFormatConfigurationJsonSchema(**config_kwargs)
+            return TextResponseFormatJsonSchema(**config_kwargs)
         if format_type == "json_object":
-            return ResponseTextFormatConfigurationJsonObject()
+            return TextResponseFormatConfigurationResponseFormatJsonObject()
         if format_type == "text":
-            return ResponseTextFormatConfigurationText()
+            return TextResponseFormatConfigurationResponseFormatText()
 
     raise IntegrationInvalidRequestException("response_format must be a Pydantic model or mapping.")
 
diff --git a/python/packages/azure-ai/tests/test_azure_ai_client.py b/python/packages/azure-ai/tests/test_azure_ai_client.py
index 4ec1b90971..519080894d 100644
--- a/python/packages/azure-ai/tests/test_azure_ai_client.py
+++ b/python/packages/azure-ai/tests/test_azure_ai_client.py
@@ -28,12 +28,12 @@
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.projects.models import (
     ApproximateLocation,
+    CodeInterpreterContainerAuto,
     CodeInterpreterTool,
-    CodeInterpreterToolAuto,
     FileSearchTool,
     ImageGenTool,
     MCPTool,
-    ResponseTextFormatConfigurationJsonSchema,
+    TextResponseFormatJsonSchema,
     WebSearchPreviewTool,
 )
 from azure.core.exceptions import ResourceNotFoundError
@@ -427,7 +427,7 @@ async def test_prepare_options_basic(mock_project_client: MagicMock) -> None:
         run_options = await client._prepare_options(messages, {})
 
         assert "extra_body" in run_options
-        assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+        assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
 
 
 @pytest.mark.parametrize(
@@ -465,7 +465,7 @@ async def test_prepare_options_with_application_endpoint(
 
     if expects_agent:
         assert "extra_body" in run_options
-        assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+        assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
     else:
         assert "extra_body" not in run_options
 
@@ -507,7 +507,7 @@ async def test_prepare_options_with_application_project_client(
 
     if expects_agent:
         assert "extra_body" in run_options
-        assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+        assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
     else:
         assert "extra_body" not in run_options
 
@@ -979,10 +979,10 @@ async def test_agent_creation_with_response_format(
     assert hasattr(created_definition, "text")
     assert created_definition.text is not None
 
-    # Check that the format is a ResponseTextFormatConfigurationJsonSchema
+    # Check that the format is a TextResponseFormatJsonSchema
     assert hasattr(created_definition.text, "format")
     format_config = created_definition.text.format
-    assert isinstance(format_config, ResponseTextFormatConfigurationJsonSchema)
+    assert isinstance(format_config, TextResponseFormatJsonSchema)
 
     # Check the schema name matches the model class name
     assert format_config.name == "ResponseFormatModel"
@@ -1040,7 +1040,7 @@ async def test_agent_creation_with_mapping_response_format(
     assert hasattr(created_definition, "text")
     assert created_definition.text is not None
     format_config = created_definition.text.format
-    assert isinstance(format_config, ResponseTextFormatConfigurationJsonSchema)
+    assert isinstance(format_config, TextResponseFormatJsonSchema)
     assert format_config.name == runtime_schema["title"]
     assert format_config.schema == runtime_schema
     assert format_config.strict is True
@@ -1110,7 +1110,7 @@ async def test_prepare_options_excludes_response_format(
         assert "text_format" not in run_options
         # But extra_body should contain agent reference
         assert "extra_body" in run_options
-        assert run_options["extra_body"]["agent"]["name"] == "test-agent"
+        assert run_options["extra_body"]["agent_reference"]["name"] == "test-agent"
 
 
 async def test_prepare_options_keeps_values_for_unsupported_option_keys(
@@ -1254,7 +1254,7 @@ def test_from_azure_ai_tools_mcp() -> None:
 
 def test_from_azure_ai_tools_code_interpreter() -> None:
     """Test from_azure_ai_tools with Code Interpreter tool."""
-    ci_tool = CodeInterpreterTool(container=CodeInterpreterToolAuto(file_ids=["file-1"]))
+    ci_tool = CodeInterpreterTool(container=CodeInterpreterContainerAuto(file_ids=["file-1"]))
     parsed_tools = from_azure_ai_tools([ci_tool])
     assert len(parsed_tools) == 1
     assert parsed_tools[0]["type"] == "code_interpreter"
diff --git a/python/packages/azure-ai/tests/test_provider.py b/python/packages/azure-ai/tests/test_provider.py
index 3765f17f1c..cb312983d4 100644
--- a/python/packages/azure-ai/tests/test_provider.py
+++ b/python/packages/azure-ai/tests/test_provider.py
@@ -8,7 +8,6 @@
 from agent_framework._mcp import MCPTool
 from azure.ai.projects.aio import AIProjectClient
 from azure.ai.projects.models import (
-    AgentReference,
     AgentVersionDetails,
     PromptAgentDefinition,
 )
@@ -345,7 +344,7 @@ async def test_provider_get_agent_with_reference(mock_project_client: MagicMock)
     mock_project_client.agents = AsyncMock()
     mock_project_client.agents.get_version.return_value = mock_agent_version
 
-    agent_reference = AgentReference(name="test-agent", version="1.0")
+    agent_reference = {"name": "test-agent", "version": "1.0"}
     agent = await provider.get_agent(reference=agent_reference)
 
     assert isinstance(agent, Agent)
diff --git a/python/packages/core/pyproject.toml b/python/packages/core/pyproject.toml
index a3c3f53ed6..0346257bb0 100644
--- a/python/packages/core/pyproject.toml
+++ b/python/packages/core/pyproject.toml
@@ -34,8 +34,7 @@ dependencies = [
     # connectors and functions
     "openai>=1.99.0",
     "azure-identity>=1,<2",
-    # Pinned to 2.0.0b3 - breaking changes in 2.0.0b4, unpin once upgrades complete
-    "azure-ai-projects == 2.0.0b3",
+    "azure-ai-projects == 2.0.0b4",
     "mcp[ws]>=1.24.0,<2",
     "packaging>=24.1",
 ]
diff --git a/python/samples/02-agents/context_providers/azure_ai_foundry_memory.py b/python/samples/02-agents/context_providers/azure_ai_foundry_memory.py
index f31e01ea1c..f7662d1e2f 100644
--- a/python/samples/02-agents/context_providers/azure_ai_foundry_memory.py
+++ b/python/samples/02-agents/context_providers/azure_ai_foundry_memory.py
@@ -61,7 +61,7 @@ async def main() -> None:
         print(f"Creating memory store '{memory_store_name}'...")
         try:
             # Create a memory store
-            memory_store = await project_client.memory_stores.create(
+            memory_store = await project_client.beta.memory_stores.create(
                 name=memory_store_name,
                 description="Memory store for Agent Framework with FoundryMemoryProvider",
                 definition=memory_store_definition,
@@ -126,7 +126,7 @@ async def main() -> None:
                 print(f"Agent: {result3}\n")
 
                 print(f"Stored memories from: {memory_store.name} ({memory_store.id})")
-                res = await project_client.memory_stores.search_memories(name=memory_store.name, scope="user_123")
+                res = await project_client.beta.memory_stores.search_memories(name=memory_store.name, scope="user_123")
                 for memory in res.memories:
                     print(f"Memory: {memory.memory_item.content}")
 
@@ -134,7 +134,7 @@ async def main() -> None:
                 print(f"An error occurred: {e}")
 
             finally:
-                await project_client.memory_stores.delete(memory_store_name)
+                await project_client.beta.memory_stores.delete(memory_store_name)
                 print("==========================================")
                 print("Memory store deleted")
 
diff --git a/python/samples/02-agents/providers/azure_ai/azure_ai_provider_methods.py b/python/samples/02-agents/providers/azure_ai/azure_ai_provider_methods.py
index e08dfcc1bc..9efa5592c7 100644
--- a/python/samples/02-agents/providers/azure_ai/azure_ai_provider_methods.py
+++ b/python/samples/02-agents/providers/azure_ai/azure_ai_provider_methods.py
@@ -8,7 +8,7 @@
 from agent_framework import tool
 from agent_framework.azure import AzureAIProjectAgentProvider
 from azure.ai.projects.aio import AIProjectClient
-from azure.ai.projects.models import AgentReference, PromptAgentDefinition
+from azure.ai.projects.models import PromptAgentDefinition
 from azure.identity.aio import AzureCliCredential
 from dotenv import load_dotenv
 from pydantic import Field
@@ -116,7 +116,7 @@ async def get_agent_by_name_example() -> None:
 async def get_agent_by_reference_example() -> None:
     """Example of using provider.get_agent(reference=...) to retrieve a specific agent version.
 
-    This method fetches a specific version of an agent using an AgentReference.
+    This method fetches a specific version of an agent using a reference mapping.
     Use this when you need to use a particular version of an agent.
     """
     print("=== provider.get_agent(reference=...) Example ===")
@@ -136,9 +136,9 @@ async def get_agent_by_reference_example() -> None:
         )
 
         try:
-            # Get the agent using an AgentReference with specific version
+            # Get the agent using a reference mapping with specific version
             provider = AzureAIProjectAgentProvider(project_client=project_client)
-            reference = AgentReference(name=created_agent.name, version=created_agent.version)
+            reference = {"name": created_agent.name, "version": created_agent.version}
             agent = await provider.get_agent(reference=reference)
 
             print(f"Retrieved agent: {agent.name} (version via reference)")
diff --git a/python/samples/02-agents/providers/azure_ai/azure_ai_with_memory_search.py b/python/samples/02-agents/providers/azure_ai/azure_ai_with_memory_search.py
index 2d1cb43c30..9377a78214 100644
--- a/python/samples/02-agents/providers/azure_ai/azure_ai_with_memory_search.py
+++ b/python/samples/02-agents/providers/azure_ai/azure_ai_with_memory_search.py
@@ -43,7 +43,7 @@ async def main() -> None:
                 options=MemoryStoreDefaultOptions(user_profile_enabled=True, chat_summary_enabled=True),
             )
 
-            memory_store = await project_client.memory_stores.create(
+            memory_store = await project_client.beta.memory_stores.create(
                 name=memory_store_name,
                 description="Memory store for Agent Framework conversations",
                 definition=memory_store_definition,
@@ -57,7 +57,7 @@ async def main() -> None:
                 instructions="""You are a helpful assistant that remembers past conversations.
                 Use the memory search tool to recall relevant information from previous interactions.""",
                 tools={
-                    "type": "memory_search",
+                    "type": "memory_search_preview",
                     "memory_store_name": memory_store.name,
                     "scope": "user_123",
                     "update_delay": 1,  # Wait 1 second before updating memories (use higher value in production)
@@ -84,7 +84,7 @@ async def main() -> None:
 
         # Clean up - delete the memory store
         async with AIProjectClient(endpoint=endpoint, credential=credential) as project_client:
-            await project_client.memory_stores.delete(memory_store_name)
+            await project_client.beta.memory_stores.delete(memory_store_name)
             print("Memory store deleted")
 
 
diff --git a/python/uv.lock b/python/uv.lock
index 15b3f18c46..415aa04f2b 100644
--- a/python/uv.lock
+++ b/python/uv.lock
@@ -401,7 +401,7 @@ requires-dist = [
     { name = "agent-framework-orchestrations", marker = "extra == 'all'", editable = "packages/orchestrations" },
     { name = "agent-framework-purview", marker = "extra == 'all'", editable = "packages/purview" },
     { name = "agent-framework-redis", marker = "extra == 'all'", editable = "packages/redis" },
-    { name = "azure-ai-projects", specifier = "==2.0.0b3" },
+    { name = "azure-ai-projects", specifier = "==2.0.0b4" },
     { name = "azure-identity", specifier = ">=1,<2" },
     { name = "mcp", extras = ["ws"], specifier = ">=1.24.0,<2" },
     { name = "openai", specifier = ">=1.99.0" },
@@ -1014,7 +1014,7 @@ wheels = [
 
 [[package]]
 name = "azure-ai-projects"
-version = "2.0.0b3"
+version = "2.0.0b4"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "azure-core", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
@@ -1022,10 +1022,11 @@ dependencies = [
     { name = "azure-storage-blob", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
     { name = "isodate", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
     { name = "openai", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/24/e0/3512d3f07e9dd2eb4af684387c31598c435bd87833b6a81850972963cb9c/azure_ai_projects-2.0.0b3.tar.gz", hash = "sha256:6d09ad110086e450a47b991ee8a3644f1be97fa3085d5981d543f900d78f4505", size = 431749, upload-time = "2026-01-06T05:31:25.849Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/24/e9/1cb8e95a19fbf174cfd7b30368a011b3e17503928b7801b8d9129b7cc59b/azure_ai_projects-2.0.0b4.tar.gz", hash = "sha256:b6082eacf0a11db59ad4c48cb7962f5204b9a0391000bc22421236f229ff783a", size = 477764, upload-time = "2026-02-24T17:57:52.489Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/4e/b6/8fbd4786bb5c0dd19eaff86ddce0fbfb53a6f90d712038272161067a076a/azure_ai_projects-2.0.0b3-py3-none-any.whl", hash = "sha256:3b3048a3ba3904d556ba392b7bd20b6e84c93bb39df6d43a6470cdb0ad08af8c", size = 240717, upload-time = "2026-01-06T05:31:27.716Z" },
+    { url = "https://files.pythonhosted.org/packages/27/6e/6445d510a8cb6a54f57e4344c14d825c37c5146fa69ccf9d9d15a29d23e2/azure_ai_projects-2.0.0b4-py3-none-any.whl", hash = "sha256:f4cf1615bd815744ddce304b97eea9456b7f6f0bd8725547c4e54e3a67534635", size = 231920, upload-time = "2026-02-24T17:57:53.917Z" },
 ]
 
 [[package]]
@@ -1408,7 +1409,7 @@ name = "clr-loader"
 version = "0.2.10"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "cffi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "cffi", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/18/24/c12faf3f61614b3131b5c98d3bf0d376b49c7feaa73edca559aeb2aee080/clr_loader-0.2.10.tar.gz", hash = "sha256:81f114afbc5005bafc5efe5af1341d400e22137e275b042a8979f3feb9fc9446", size = 83605, upload-time = "2026-01-03T23:13:06.984Z" }
 wheels = [
@@ -1887,7 +1888,7 @@ name = "exceptiongroup"
 version = "1.3.1"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "typing-extensions", marker = "(python_full_version < '3.13' and sys_platform == 'darwin') or (python_full_version < '3.13' and sys_platform == 'linux') or (python_full_version < '3.13' and sys_platform == 'win32')" },
+    { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'darwin') or (python_full_version < '3.11' and sys_platform == 'linux') or (python_full_version < '3.11' and sys_platform == 'win32')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" }
 wheels = [
@@ -4654,8 +4655,8 @@ name = "powerfx"
 version = "0.0.34"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "cffi", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
-    { name = "pythonnet", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "cffi", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
+    { name = "pythonnet", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/9f/fb/6c4bf87e0c74ca1c563921ce89ca1c5785b7576bca932f7255cdf81082a7/powerfx-0.0.34.tar.gz", hash = "sha256:956992e7afd272657ed16d80f4cad24ec95d9e4a79fb9dfa4a068a09e136af32", size = 3237555, upload-time = "2025-12-22T15:50:59.682Z" }
 wheels = [
@@ -5318,7 +5319,7 @@ name = "pythonnet"
 version = "3.0.5"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "clr-loader", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" },
+    { name = "clr-loader", marker = "(python_full_version < '3.14' and sys_platform == 'darwin') or (python_full_version < '3.14' and sys_platform == 'linux') or (python_full_version < '3.14' and sys_platform == 'win32')" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/9a/d6/1afd75edd932306ae9bd2c2d961d603dc2b52fcec51b04afea464f1f6646/pythonnet-3.0.5.tar.gz", hash = "sha256:48e43ca463941b3608b32b4e236db92d8d40db4c58a75ace902985f76dac21cf", size = 239212, upload-time = "2024-12-13T08:30:44.393Z" }
 wheels = [

From b0f93125c91c5324675bc770e1d714818cf4093e Mon Sep 17 00:00:00 2001
From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com>
Date: Tue, 3 Mar 2026 14:02:50 -0800
Subject: [PATCH 2/2] Fixed tests

---
 .../tests/test_foundry_memory_provider.py     | 53 ++++++++++---------
 1 file changed, 27 insertions(+), 26 deletions(-)

diff --git a/python/packages/azure-ai/tests/test_foundry_memory_provider.py b/python/packages/azure-ai/tests/test_foundry_memory_provider.py
index 9c2968a65e..943a528968 100644
--- a/python/packages/azure-ai/tests/test_foundry_memory_provider.py
+++ b/python/packages/azure-ai/tests/test_foundry_memory_provider.py
@@ -17,9 +17,10 @@
 def mock_project_client() -> AsyncMock:
     """Create a mock AIProjectClient."""
     mock_client = AsyncMock()
-    mock_client.memory_stores = AsyncMock()
-    mock_client.memory_stores.search_memories = AsyncMock()
-    mock_client.memory_stores.begin_update_memories = AsyncMock()
+    mock_client.beta = AsyncMock()
+    mock_client.beta.memory_stores = AsyncMock()
+    mock_client.beta.memory_stores.search_memories = AsyncMock()
+    mock_client.beta.memory_stores.begin_update_memories = AsyncMock()
     mock_client.__aenter__ = AsyncMock(return_value=mock_client)
     mock_client.__aexit__ = AsyncMock()
     return mock_client
@@ -146,7 +147,7 @@ async def test_retrieves_static_memories_on_first_run(self, mock_project_client:
         mem2.memory_item.content = "User is based in Seattle"
         mock_search_result = Mock()
         mock_search_result.memories = [mem1, mem2]
-        mock_project_client.memory_stores.search_memories.return_value = mock_search_result
+        mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -161,7 +162,7 @@ async def test_retrieves_static_memories_on_first_run(self, mock_project_client:
         )
 
         # Should call search_memories twice: once for static, once for contextual
-        assert mock_project_client.memory_stores.search_memories.call_count == 2
+        assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
         # Static memories should be cached
         assert len(session.state[provider.source_id]["static_memories"]) == 2
         assert session.state[provider.source_id]["initialized"] is True
@@ -181,7 +182,7 @@ async def test_contextual_memories_added_to_context(self, mock_project_client: A
         contextual_result.memories = [contextual_mem]
         contextual_result.search_id = "search-123"
 
-        mock_project_client.memory_stores.search_memories.side_effect = [static_result, contextual_result]
+        mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -208,7 +209,7 @@ async def test_empty_input_skips_contextual_search(self, mock_project_client: As
         """Empty input messages → only static search performed, no contextual search."""
         static_result = Mock()
         static_result.memories = []
-        mock_project_client.memory_stores.search_memories.return_value = static_result
+        mock_project_client.beta.memory_stores.search_memories.return_value = static_result
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -223,14 +224,14 @@ async def test_empty_input_skips_contextual_search(self, mock_project_client: As
         )
 
         # Should only call search_memories once for static memories
-        assert mock_project_client.memory_stores.search_memories.call_count == 1
+        assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
         assert provider.source_id not in ctx.context_messages
 
     async def test_empty_search_results_no_messages(self, mock_project_client: AsyncMock) -> None:
         """Empty search results → no messages added."""
         mock_search_result = Mock()
         mock_search_result.memories = []
-        mock_project_client.memory_stores.search_memories.return_value = mock_search_result
+        mock_project_client.beta.memory_stores.search_memories.return_value = mock_search_result
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -255,7 +256,7 @@ async def test_static_memories_only_retrieved_once(self, mock_project_client: As
         contextual_result = Mock()
         contextual_result.memories = []
 
-        mock_project_client.memory_stores.search_memories.side_effect = [static_result, contextual_result]
+        mock_project_client.beta.memory_stores.search_memories.side_effect = [static_result, contextual_result]
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -269,24 +270,24 @@ async def test_static_memories_only_retrieved_once(self, mock_project_client: As
         await provider.before_run(  # type: ignore[arg-type]
             agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
-        assert mock_project_client.memory_stores.search_memories.call_count == 2
+        assert mock_project_client.beta.memory_stores.search_memories.call_count == 2
 
         # Reset mock for second call
-        mock_project_client.memory_stores.search_memories.reset_mock()
+        mock_project_client.beta.memory_stores.search_memories.reset_mock()
         contextual_result2 = Mock()
         contextual_result2.memories = []
-        mock_project_client.memory_stores.search_memories.return_value = contextual_result2
+        mock_project_client.beta.memory_stores.search_memories.return_value = contextual_result2
 
         # Second call - should only search contextual, not static
         ctx2 = SessionContext(input_messages=[Message(role="user", text="World")], session_id="s1")
         await provider.before_run(  # type: ignore[arg-type]
             agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
         )
-        assert mock_project_client.memory_stores.search_memories.call_count == 1
+        assert mock_project_client.beta.memory_stores.search_memories.call_count == 1
 
     async def test_handles_search_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
         """Search exception is logged but doesn't fail the operation."""
-        mock_project_client.memory_stores.search_memories.side_effect = Exception("API error")
+        mock_project_client.beta.memory_stores.search_memories.side_effect = Exception("API error")
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -315,7 +316,7 @@ async def test_stores_input_and_response(self, mock_project_client: AsyncMock) -
         """Stores input+response messages via begin_update_memories."""
         mock_poller = Mock()
         mock_poller.update_id = "update-456"
-        mock_project_client.memory_stores.begin_update_memories.return_value = mock_poller
+        mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -330,8 +331,8 @@ async def test_stores_input_and_response(self, mock_project_client: AsyncMock) -
             agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        mock_project_client.memory_stores.begin_update_memories.assert_awaited_once()
-        call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+        mock_project_client.beta.memory_stores.begin_update_memories.assert_awaited_once()
+        call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
         assert call_kwargs["name"] == "test_store"
         assert call_kwargs["scope"] == "user_123"
         assert len(call_kwargs["items"]) == 2
@@ -342,7 +343,7 @@ async def test_stores_input_and_response(self, mock_project_client: AsyncMock) -
     async def test_only_stores_user_assistant_system(self, mock_project_client: AsyncMock) -> None:
         """Only stores user/assistant/system messages with text."""
         mock_poller = Mock()
-        mock_project_client.memory_stores.begin_update_memories.return_value = mock_poller
+        mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -363,7 +364,7 @@ async def test_only_stores_user_assistant_system(self, mock_project_client: Asyn
             agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+        call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
         items = call_kwargs["items"]
         assert len(items) == 2
         assert items[0]["content"] == "hello"
@@ -390,12 +391,12 @@ async def test_skips_empty_messages(self, mock_project_client: AsyncMock) -> Non
             agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        mock_project_client.memory_stores.begin_update_memories.assert_not_awaited()
+        mock_project_client.beta.memory_stores.begin_update_memories.assert_not_awaited()
 
     async def test_uses_configured_update_delay(self, mock_project_client: AsyncMock) -> None:
         """Uses the configured update_delay parameter."""
         mock_poller = Mock()
-        mock_project_client.memory_stores.begin_update_memories.return_value = mock_poller
+        mock_project_client.beta.memory_stores.begin_update_memories.return_value = mock_poller
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -411,7 +412,7 @@ async def test_uses_configured_update_delay(self, mock_project_client: AsyncMock
             agent=None, session=session, context=ctx, state=session.state.setdefault(provider.source_id, {})
         )
 
-        call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+        call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
         assert call_kwargs["update_delay"] == 60
 
     async def test_uses_previous_update_id_for_incremental_updates(self, mock_project_client: AsyncMock) -> None:
@@ -421,7 +422,7 @@ async def test_uses_previous_update_id_for_incremental_updates(self, mock_projec
         mock_poller2 = Mock()
         mock_poller2.update_id = "update-2"
 
-        mock_project_client.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
+        mock_project_client.beta.memory_stores.begin_update_memories.side_effect = [mock_poller1, mock_poller2]
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,
@@ -446,13 +447,13 @@ async def test_uses_previous_update_id_for_incremental_updates(self, mock_projec
             agent=None, session=session, context=ctx2, state=session.state.setdefault(provider.source_id, {})
         )
 
-        call_kwargs = mock_project_client.memory_stores.begin_update_memories.call_args.kwargs
+        call_kwargs = mock_project_client.beta.memory_stores.begin_update_memories.call_args.kwargs
         assert call_kwargs["previous_update_id"] == "update-1"
         assert session.state[provider.source_id]["previous_update_id"] == "update-2"
 
     async def test_handles_update_exception_gracefully(self, mock_project_client: AsyncMock) -> None:
         """Update exception is logged but doesn't fail the operation."""
-        mock_project_client.memory_stores.begin_update_memories.side_effect = Exception("API error")
+        mock_project_client.beta.memory_stores.begin_update_memories.side_effect = Exception("API error")
 
         provider = FoundryMemoryProvider(
             project_client=mock_project_client,