Merge branch 'main' into feature-python-foundry-agents

This commit is contained in:
Dmytro Struk
2025-11-05 09:21:26 -08:00
Unverified
3 changed files with 72 additions and 98 deletions
@@ -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"
@@ -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,