mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Removed automatic agent cleanup in AzureAIAgentClient
This commit is contained in:
@@ -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 (
|
||||
|
||||
Reference in New Issue
Block a user