From 89846c72128c22167c3fb6dbcaf0687aa75175a9 Mon Sep 17 00:00:00 2001 From: Dmytro Struk <13853051+dmytrostruk@users.noreply.github.com> Date: Mon, 3 Nov 2025 17:43:31 -0800 Subject: [PATCH] Removed automatic agent cleanup in AzureAIAgentClient --- .../agent_framework_azure_ai/_chat_client.py | 12 +-- .../tests/test_azure_ai_agent_client.py | 83 +------------------ .../agents/azure_ai/azure_ai_basic.py | 7 +- 3 files changed, 7 insertions(+), 95 deletions(-) diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py index 0f35158da7..7664db8efe 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py @@ -170,7 +170,7 @@ class AzureAIAgentClient(BaseChatClient): Keyword Args: project_client: An existing AIProjectClient to use. If not provided, one will be created. agent_id: The ID of an existing agent to use. If not provided and project_client is provided, - a new agent will be created (and deleted after the request). If neither project_client + a new agent will be created. If neither project_client nor agent_id is provided, both will be created and managed automatically. agent_name: The name to use when creating new agents. thread_id: Default thread ID to use for conversations. Can be overridden by @@ -252,7 +252,6 @@ class AzureAIAgentClient(BaseChatClient): self.agent_name = agent_name self.model_id = azure_ai_settings.model_deployment_name self.thread_id = thread_id - self._should_delete_agent = False # Track whether we should delete the agent self._should_close_client = should_close_client # Track whether we should close client connection self._agent_definition: Agent | None = None # Cached definition for existing agent @@ -287,7 +286,6 @@ class AzureAIAgentClient(BaseChatClient): async def close(self) -> None: """Close the project_client and clean up any agents we created.""" - await self._cleanup_agent_if_needed() await self._close_client_if_needed() @classmethod @@ -381,7 +379,6 @@ class AzureAIAgentClient(BaseChatClient): created_agent = await self.project_client.agents.create_agent(**args) self.agent_id = str(created_agent.id) self._agent_definition = created_agent - self._should_delete_agent = True return self.agent_id @@ -713,13 +710,6 @@ class AzureAIAgentClient(BaseChatClient): if self._should_close_client: await self.project_client.close() - async def _cleanup_agent_if_needed(self) -> None: - """Clean up the agent if we created it.""" - if self._should_delete_agent and self.agent_id is not None: - await self.project_client.agents.delete_agent(self.agent_id) - self.agent_id = None - self._should_delete_agent = False - async def _load_agent_definition_if_needed(self) -> Agent | None: """Load and cache agent details if not already loaded.""" if self._agent_definition is None and self.agent_id is not None: diff --git a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py index 9f3a06ebee..92ed8bcf79 100644 --- a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py +++ b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py @@ -72,8 +72,8 @@ def create_test_azure_ai_chat_client( agent_id: str | None = None, thread_id: str | None = None, azure_ai_settings: AzureAISettings | None = None, - should_delete_agent: bool = False, agent_name: str | None = None, + should_close_client: bool = False, ) -> AzureAIAgentClient: """Helper function to create AzureAIAgentClient instances for testing, bypassing normal validation.""" if azure_ai_settings is None: @@ -89,8 +89,7 @@ def create_test_azure_ai_chat_client( client.agent_name = agent_name client.model_id = azure_ai_settings.model_deployment_name client.thread_id = thread_id - client._should_delete_agent = should_delete_agent # type: ignore - client._should_close_client = False # type: ignore + client._should_close_client = should_close_client # type: ignore client._agent_definition = None # type: ignore client.additional_properties = {} client.middleware = None @@ -126,7 +125,6 @@ def test_azure_ai_chat_client_init_with_client(mock_ai_project_client: MagicMock assert chat_client.project_client is mock_ai_project_client assert chat_client.agent_id == "existing-agent-id" assert chat_client.thread_id == "test-thread-id" - assert not chat_client._should_delete_agent # type: ignore assert isinstance(chat_client, ChatClientProtocol) @@ -142,7 +140,6 @@ def test_azure_ai_chat_client_init_auto_create_client( chat_client.project_client = mock_ai_project_client chat_client.agent_id = None chat_client.thread_id = None - chat_client._should_delete_agent = False # type: ignore chat_client._should_close_client = False # type: ignore chat_client.credential = None chat_client.model_id = azure_ai_settings.model_deployment_name @@ -152,7 +149,6 @@ def test_azure_ai_chat_client_init_auto_create_client( assert chat_client.project_client is mock_ai_project_client assert chat_client.agent_id is None - assert not chat_client._should_delete_agent # type: ignore def test_azure_ai_chat_client_init_missing_project_endpoint() -> None: @@ -303,7 +299,6 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_existing_agent( agent_id = await chat_client._get_agent_id_or_create() # type: ignore assert agent_id == "existing-agent-id" - assert not chat_client._should_delete_agent # type: ignore async def test_azure_ai_chat_client_get_agent_id_or_create_create_new( @@ -317,7 +312,6 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_create_new( agent_id = await chat_client._get_agent_id_or_create(run_options={"model": azure_ai_settings.model_deployment_name}) # type: ignore assert agent_id == "test-agent-id" - assert chat_client._should_delete_agent # type: ignore async def test_azure_ai_chat_client_thread_management_through_public_api(mock_ai_project_client: MagicMock) -> None: @@ -365,74 +359,6 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_missing_model( await chat_client._get_agent_id_or_create() # type: ignore -async def test_azure_ai_chat_client_cleanup_agent_if_needed_should_delete( - mock_ai_project_client: MagicMock, -) -> None: - """Test _cleanup_agent_if_needed when agent should be deleted.""" - chat_client = create_test_azure_ai_chat_client( - mock_ai_project_client, agent_id="agent-to-delete", should_delete_agent=True - ) - - await chat_client._cleanup_agent_if_needed() # type: ignore - # Verify agent deletion was called - mock_ai_project_client.agents.delete_agent.assert_called_once_with("agent-to-delete") - assert not chat_client._should_delete_agent # type: ignore - - -async def test_azure_ai_chat_client_cleanup_agent_if_needed_should_not_delete( - mock_ai_project_client: MagicMock, -) -> None: - """Test _cleanup_agent_if_needed when agent should not be deleted.""" - chat_client = create_test_azure_ai_chat_client( - mock_ai_project_client, agent_id="agent-to-keep", should_delete_agent=False - ) - - await chat_client._cleanup_agent_if_needed() # type: ignore - - # Verify agent deletion was not called - mock_ai_project_client.agents.delete_agent.assert_not_called() - assert not chat_client._should_delete_agent # type: ignore - - -async def test_azure_ai_chat_client_cleanup_agent_if_needed_exception_handling( - mock_ai_project_client: MagicMock, -) -> None: - """Test _cleanup_agent_if_needed propagates exceptions (it doesn't handle them).""" - chat_client = create_test_azure_ai_chat_client( - mock_ai_project_client, agent_id="agent-to-delete", should_delete_agent=True - ) - mock_ai_project_client.agents.delete_agent.side_effect = Exception("Deletion failed") - - with pytest.raises(Exception, match="Deletion failed"): - await chat_client._cleanup_agent_if_needed() # type: ignore - - -async def test_azure_ai_chat_client_aclose(mock_ai_project_client: MagicMock) -> None: - """Test aclose method calls cleanup.""" - chat_client = create_test_azure_ai_chat_client( - mock_ai_project_client, agent_id="agent-to-delete", should_delete_agent=True - ) - - await chat_client.close() - - # Verify agent deletion was called - mock_ai_project_client.agents.delete_agent.assert_called_once_with("agent-to-delete") - - -async def test_azure_ai_chat_client_async_context_manager(mock_ai_project_client: MagicMock) -> None: - """Test async context manager functionality.""" - chat_client = create_test_azure_ai_chat_client( - mock_ai_project_client, agent_id="agent-to-delete", should_delete_agent=True - ) - - # Test context manager - async with chat_client: - pass # Just test that we can enter and exit - - # Verify cleanup was called on exit - mock_ai_project_client.agents.delete_agent.assert_called_once_with("agent-to-delete") - - async def test_azure_ai_chat_client_create_run_options_basic(mock_ai_project_client: MagicMock) -> None: """Test _create_run_options with basic ChatOptions.""" chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) @@ -1435,18 +1361,15 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_with_tool_resources( async def test_azure_ai_chat_client_close_method(mock_ai_project_client: MagicMock) -> None: """Test close method.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, should_delete_agent=True) - chat_client._should_close_client = True # type: ignore + chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, should_close_client=True) chat_client.agent_id = "test-agent" # Mock cleanup methods - mock_ai_project_client.agents.delete_agent = AsyncMock() mock_ai_project_client.close = AsyncMock() await chat_client.close() # Verify cleanup was called - mock_ai_project_client.agents.delete_agent.assert_called_once_with("test-agent") mock_ai_project_client.close.assert_called_once() diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_basic.py b/python/samples/getting_started/agents/azure_ai/azure_ai_basic.py index 633b5b9daa..d448079b29 100644 --- a/python/samples/getting_started/agents/azure_ai/azure_ai_basic.py +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_basic.py @@ -11,8 +11,8 @@ from pydantic import Field """ Azure AI Agent Basic Example -This sample demonstrates basic usage of AzureAIAgentClient to create agents with automatic -lifecycle management. Shows both streaming and non-streaming responses with function tools. +This sample demonstrates basic usage of AzureAIAgentClient. +Shows both streaming and non-streaming responses with function tools. """ @@ -28,8 +28,7 @@ async def non_streaming_example() -> None: """Example of non-streaming response (get the complete result at once).""" print("=== Non-streaming Response Example ===") - # Since no Agent ID is provided, the agent will be automatically created - # and deleted after getting a response + # Since no Agent ID is provided, the agent will be automatically created. # For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred # authentication option. async with (