Removed automatic agent cleanup in AzureAIAgentClient

This commit is contained in:
Dmytro Struk
2025-11-03 17:43:31 -08:00
Unverified
parent 50d9b13bfc
commit 89846c7212
3 changed files with 7 additions and 95 deletions
@@ -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:
@@ -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()
@@ -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 (