mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Merge branch 'main' into feature-python-foundry-agents
This commit is contained in:
@@ -120,6 +120,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
project_endpoint: str | None = None,
|
||||
model_deployment_name: str | None = None,
|
||||
async_credential: AsyncTokenCredential | None = None,
|
||||
should_cleanup_agent: bool = True,
|
||||
env_file_path: str | None = None,
|
||||
env_file_encoding: str | None = None,
|
||||
**kwargs: Any,
|
||||
@@ -140,6 +141,9 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
model_deployment_name: The model deployment name to use for agent creation.
|
||||
Can also be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME.
|
||||
async_credential: Azure async credential to use for authentication.
|
||||
should_cleanup_agent: Whether to cleanup (delete) agents created by this client when
|
||||
the client is closed or context is exited. Defaults to True. Only affects agents
|
||||
created by this client instance; existing agents passed via agent_id are never deleted.
|
||||
env_file_path: Path to environment file for loading settings.
|
||||
env_file_encoding: Encoding of the environment file.
|
||||
kwargs: Additional keyword arguments passed to the parent class.
|
||||
@@ -212,7 +216,8 @@ 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_cleanup_agent = should_cleanup_agent # Track whether we should delete the agent
|
||||
self._agent_created = False # Track whether agent was created inside this class
|
||||
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
|
||||
|
||||
@@ -245,6 +250,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
agent_name=settings.get("agent_name"),
|
||||
credential=settings.get("credential"),
|
||||
env_file_path=settings.get("env_file_path"),
|
||||
should_cleanup_agent=settings.get("should_cleanup_agent", True),
|
||||
)
|
||||
|
||||
async def _inner_get_response(
|
||||
@@ -323,7 +329,7 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
|
||||
self.agent_id = str(created_agent.id)
|
||||
self._agent_definition = created_agent
|
||||
self._should_delete_agent = True
|
||||
self._agent_created = True
|
||||
|
||||
return self.agent_id
|
||||
|
||||
@@ -655,10 +661,10 @@ class AzureAIAgentClient(BaseChatClient):
|
||||
|
||||
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:
|
||||
if self._agent_created and self.should_cleanup_agent and self.agent_id is not None:
|
||||
await self.agents_client.delete_agent(self.agent_id)
|
||||
self.agent_id = None
|
||||
self._should_delete_agent = False
|
||||
self._agent_created = False
|
||||
|
||||
async def _load_agent_definition_if_needed(self) -> Agent | None:
|
||||
"""Load and cache agent details if not already loaded."""
|
||||
|
||||
@@ -71,7 +71,7 @@ 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,
|
||||
should_cleanup_agent: bool = True,
|
||||
agent_name: str | None = None,
|
||||
) -> AzureAIAgentClient:
|
||||
"""Helper function to create AzureAIAgentClient instances for testing, bypassing normal validation."""
|
||||
@@ -88,9 +88,10 @@ 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._agent_definition = None # type: ignore
|
||||
client.should_cleanup_agent = should_cleanup_agent
|
||||
client._agent_created = False
|
||||
client._should_close_client = False
|
||||
client._agent_definition = None
|
||||
client.additional_properties = {}
|
||||
client.middleware = None
|
||||
|
||||
@@ -125,7 +126,6 @@ def test_azure_ai_chat_client_init_with_client(mock_agents_client: MagicMock) ->
|
||||
assert chat_client.agents_client is mock_agents_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)
|
||||
|
||||
|
||||
@@ -141,7 +141,6 @@ def test_azure_ai_chat_client_init_auto_create_client(
|
||||
chat_client.agents_client = mock_agents_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
|
||||
@@ -151,7 +150,6 @@ def test_azure_ai_chat_client_init_auto_create_client(
|
||||
|
||||
assert chat_client.agents_client is mock_agents_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:
|
||||
@@ -302,7 +300,7 @@ 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
|
||||
assert not chat_client._agent_created
|
||||
|
||||
|
||||
async def test_azure_ai_chat_client_get_agent_id_or_create_create_new(
|
||||
@@ -316,7 +314,7 @@ 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
|
||||
assert chat_client._agent_created
|
||||
|
||||
|
||||
async def test_azure_ai_chat_client_thread_management_through_public_api(mock_agents_client: MagicMock) -> None:
|
||||
@@ -364,74 +362,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_agents_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test _cleanup_agent_if_needed when agent should be deleted."""
|
||||
chat_client = create_test_azure_ai_chat_client(
|
||||
mock_agents_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_agents_client.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_agents_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test _cleanup_agent_if_needed when agent should not be deleted."""
|
||||
chat_client = create_test_azure_ai_chat_client(
|
||||
mock_agents_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_agents_client.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_agents_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test _cleanup_agent_if_needed propagates exceptions (it doesn't handle them)."""
|
||||
chat_client = create_test_azure_ai_chat_client(
|
||||
mock_agents_client, agent_id="agent-to-delete", should_delete_agent=True
|
||||
)
|
||||
mock_agents_client.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_agents_client: MagicMock) -> None:
|
||||
"""Test aclose method calls cleanup."""
|
||||
chat_client = create_test_azure_ai_chat_client(
|
||||
mock_agents_client, agent_id="agent-to-delete", should_delete_agent=True
|
||||
)
|
||||
|
||||
await chat_client.close()
|
||||
|
||||
# Verify agent deletion was called
|
||||
mock_agents_client.delete_agent.assert_called_once_with("agent-to-delete")
|
||||
|
||||
|
||||
async def test_azure_ai_chat_client_async_context_manager(mock_agents_client: MagicMock) -> None:
|
||||
"""Test async context manager functionality."""
|
||||
chat_client = create_test_azure_ai_chat_client(
|
||||
mock_agents_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_agents_client.delete_agent.assert_called_once_with("agent-to-delete")
|
||||
|
||||
|
||||
async def test_azure_ai_chat_client_create_run_options_basic(mock_agents_client: MagicMock) -> None:
|
||||
"""Test _create_run_options with basic ChatOptions."""
|
||||
chat_client = create_test_azure_ai_chat_client(mock_agents_client)
|
||||
@@ -1347,23 +1277,6 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_with_tool_resources(
|
||||
assert call_kwargs["tool_resources"] == {"vector_store_ids": ["vs-123"]}
|
||||
|
||||
|
||||
async def test_azure_ai_chat_client_close_method(mock_agents_client: MagicMock) -> None:
|
||||
"""Test close method."""
|
||||
chat_client = create_test_azure_ai_chat_client(mock_agents_client, should_delete_agent=True)
|
||||
chat_client._should_close_client = True # type: ignore
|
||||
chat_client.agent_id = "test-agent"
|
||||
|
||||
# Mock cleanup methods
|
||||
mock_agents_client.delete_agent = AsyncMock()
|
||||
mock_agents_client.close = AsyncMock()
|
||||
|
||||
await chat_client.close()
|
||||
|
||||
# Verify cleanup was called
|
||||
mock_agents_client.delete_agent.assert_called_once_with("test-agent")
|
||||
mock_agents_client.close.assert_called_once()
|
||||
|
||||
|
||||
async def test_azure_ai_chat_client_create_agent_stream_submit_tool_outputs(
|
||||
mock_agents_client: MagicMock,
|
||||
) -> None:
|
||||
@@ -1837,3 +1750,57 @@ async def test_azure_ai_chat_client_agent_chat_options_agent_level() -> None:
|
||||
assert isinstance(response, AgentRunResponse)
|
||||
assert response.text is not None
|
||||
assert len(response.text) > 0
|
||||
|
||||
|
||||
async def test_azure_ai_chat_client_cleanup_agent_when_enabled_and_created(
|
||||
mock_agents_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test that agent is cleaned up when should_cleanup_agent=True and agent was created by client."""
|
||||
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id=None, should_cleanup_agent=True)
|
||||
|
||||
# Simulate agent creation
|
||||
chat_client.agent_id = "created-agent-id"
|
||||
chat_client._agent_created = True # type: ignore
|
||||
|
||||
await chat_client._cleanup_agent_if_needed() # type: ignore
|
||||
|
||||
# Verify agent was deleted
|
||||
mock_agents_client.delete_agent.assert_called_once_with("created-agent-id")
|
||||
assert chat_client.agent_id is None
|
||||
assert chat_client._agent_created is False # type: ignore
|
||||
|
||||
|
||||
async def test_azure_ai_chat_client_no_cleanup_when_disabled(
|
||||
mock_agents_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test that agent is not cleaned up when should_cleanup_agent=False."""
|
||||
chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id=None, should_cleanup_agent=False)
|
||||
|
||||
# Simulate agent creation
|
||||
chat_client.agent_id = "created-agent-id"
|
||||
chat_client._agent_created = True
|
||||
|
||||
await chat_client._cleanup_agent_if_needed() # type: ignore
|
||||
|
||||
# Verify agent was NOT deleted
|
||||
mock_agents_client.delete_agent.assert_not_called()
|
||||
assert chat_client.agent_id == "created-agent-id"
|
||||
assert chat_client._agent_created is True
|
||||
|
||||
|
||||
async def test_azure_ai_chat_client_no_cleanup_when_agent_not_created_by_client(
|
||||
mock_agents_client: MagicMock,
|
||||
) -> None:
|
||||
"""Test that agent is not cleaned up when it was not created by this client instance."""
|
||||
chat_client = create_test_azure_ai_chat_client(
|
||||
mock_agents_client, agent_id="existing-agent-id", should_cleanup_agent=True
|
||||
)
|
||||
|
||||
# Agent exists but was not created by this client (_agent_created = False)
|
||||
assert chat_client._agent_created is False # type: ignore
|
||||
|
||||
await chat_client._cleanup_agent_if_needed() # type: ignore
|
||||
|
||||
# Verify agent was NOT deleted
|
||||
mock_agents_client.delete_agent.assert_not_called()
|
||||
assert chat_client.agent_id == "existing-agent-id"
|
||||
|
||||
+1
@@ -41,6 +41,7 @@ async def main() -> None:
|
||||
model_deployment_name=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"],
|
||||
async_credential=credential,
|
||||
agent_name="WeatherAgent",
|
||||
should_cleanup_agent=True, # Set to False if you want to disable automatic agent cleanup
|
||||
),
|
||||
instructions="You are a helpful weather agent.",
|
||||
tools=get_weather,
|
||||
|
||||
Reference in New Issue
Block a user