Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 16 additions & 1 deletion python/packages/azure-ai/agent_framework_azure_ai/_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@ def __init__(
project_endpoint: str | None = None,
model_deployment_name: str | None = None,
async_credential: AsyncTokenCredential | None = None,
use_latest_version: bool | None = None,
env_file_path: str | None = None,
env_file_encoding: str | None = None,
**kwargs: Any,
Expand All @@ -76,6 +77,8 @@ def __init__(
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.
use_latest_version: Boolean flag that indicates whether to use latest agent version
if it exists in the service.
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.
Expand Down Expand Up @@ -139,6 +142,7 @@ def __init__(
# Initialize instance variables
self.agent_name = agent_name
self.agent_version = agent_version
self.use_latest_version = use_latest_version
self.project_client = project_client
self.credential = async_credential
self.model_id = azure_ai_settings.model_deployment_name
Expand Down Expand Up @@ -188,8 +192,19 @@ async def _get_agent_reference_or_create(
"""
agent_name = self.agent_name or "UnnamedAgent"

# If no agent_version is provided, create a new agent
# If no agent_version is provided, either use latest version or create a new agent:
if self.agent_version is None:
# Try to use latest version if requested and agent exists
if self.use_latest_version:
try:
existing_agent = await self.project_client.agents.retrieve(agent_name)
self.agent_name = existing_agent.name
self.agent_version = existing_agent.versions.latest.version
return {"name": self.agent_name, "version": self.agent_version, "type": "agent_reference"}
except ResourceNotFoundError:
# Agent doesn't exist, fall through to creation logic
pass

if "model" not in run_options or not run_options["model"]:
raise ServiceInitializationError(
"Model deployment name is required for agent creation, "
Expand Down
94 changes: 94 additions & 0 deletions python/packages/azure-ai/tests/test_azure_ai_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ def create_test_azure_ai_client(
conversation_id: str | None = None,
azure_ai_settings: AzureAISettings | None = None,
should_close_client: bool = False,
use_latest_version: bool | None = None,
) -> AzureAIClient:
"""Helper function to create AzureAIClient instances for testing, bypassing normal validation."""
if azure_ai_settings is None:
Expand All @@ -36,6 +37,7 @@ def create_test_azure_ai_client(
client.credential = None
client.agent_name = agent_name
client.agent_version = agent_version
client.use_latest_version = use_latest_version
client.model_id = azure_ai_settings.model_deployment_name
client.conversation_id = conversation_id
client._should_close_client = should_close_client # type: ignore
Expand Down Expand Up @@ -437,6 +439,98 @@ async def test_azure_ai_client_agent_creation_with_tools(
assert call_args[1]["definition"].tools == test_tools


async def test_azure_ai_client_use_latest_version_existing_agent(
mock_project_client: MagicMock,
) -> None:
"""Test _get_agent_reference_or_create when use_latest_version=True and agent exists."""
client = create_test_azure_ai_client(mock_project_client, agent_name="existing-agent", use_latest_version=True)

# Mock existing agent response
mock_existing_agent = MagicMock()
mock_existing_agent.name = "existing-agent"
mock_existing_agent.versions.latest.version = "2.5"
mock_project_client.agents.retrieve = AsyncMock(return_value=mock_existing_agent)

run_options = {"model": "test-model"}
agent_ref = await client._get_agent_reference_or_create(run_options, None) # type: ignore

# Verify existing agent was retrieved and used
mock_project_client.agents.retrieve.assert_called_once_with("existing-agent")
mock_project_client.agents.create_version.assert_not_called()

assert agent_ref == {"name": "existing-agent", "version": "2.5", "type": "agent_reference"}
assert client.agent_name == "existing-agent"
assert client.agent_version == "2.5"


async def test_azure_ai_client_use_latest_version_agent_not_found(
mock_project_client: MagicMock,
) -> None:
"""Test _get_agent_reference_or_create when use_latest_version=True but agent doesn't exist."""
from azure.core.exceptions import ResourceNotFoundError

client = create_test_azure_ai_client(mock_project_client, agent_name="non-existing-agent", use_latest_version=True)

# Mock ResourceNotFoundError when trying to retrieve agent
mock_project_client.agents.retrieve = AsyncMock(side_effect=ResourceNotFoundError("Agent not found"))

# Mock agent creation response for fallback
mock_created_agent = MagicMock()
mock_created_agent.name = "non-existing-agent"
mock_created_agent.version = "1.0"
mock_project_client.agents.create_version = AsyncMock(return_value=mock_created_agent)

run_options = {"model": "test-model"}
agent_ref = await client._get_agent_reference_or_create(run_options, None) # type: ignore

# Verify retrieval was attempted and creation was used as fallback
mock_project_client.agents.retrieve.assert_called_once_with("non-existing-agent")
mock_project_client.agents.create_version.assert_called_once()

assert agent_ref == {"name": "non-existing-agent", "version": "1.0", "type": "agent_reference"}
assert client.agent_name == "non-existing-agent"
assert client.agent_version == "1.0"


async def test_azure_ai_client_use_latest_version_false(
mock_project_client: MagicMock,
) -> None:
"""Test _get_agent_reference_or_create when use_latest_version=False (default behavior)."""
client = create_test_azure_ai_client(mock_project_client, agent_name="test-agent", use_latest_version=False)

# Mock agent creation response
mock_created_agent = MagicMock()
mock_created_agent.name = "test-agent"
mock_created_agent.version = "1.0"
mock_project_client.agents.create_version = AsyncMock(return_value=mock_created_agent)

run_options = {"model": "test-model"}
agent_ref = await client._get_agent_reference_or_create(run_options, None) # type: ignore

# Verify retrieval was not attempted and creation was used directly
mock_project_client.agents.retrieve.assert_not_called()
mock_project_client.agents.create_version.assert_called_once()

assert agent_ref == {"name": "test-agent", "version": "1.0", "type": "agent_reference"}


async def test_azure_ai_client_use_latest_version_with_existing_agent_version(
mock_project_client: MagicMock,
) -> None:
"""Test that use_latest_version is ignored when agent_version is already provided."""
client = create_test_azure_ai_client(
mock_project_client, agent_name="test-agent", agent_version="3.0", use_latest_version=True
)

agent_ref = await client._get_agent_reference_or_create({}, None) # type: ignore

# Verify neither retrieval nor creation was attempted since version is already set
mock_project_client.agents.retrieve.assert_not_called()
mock_project_client.agents.create_version.assert_not_called()

assert agent_ref == {"name": "test-agent", "version": "3.0", "type": "agent_reference"}


@pytest.fixture
def mock_project_client() -> MagicMock:
"""Fixture that provides a mock AIProjectClient."""
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
# Copyright (c) Microsoft. All rights reserved.

import asyncio
from random import randint
from typing import Annotated

from agent_framework.azure import AzureAIClient
from azure.identity.aio import AzureCliCredential
from pydantic import Field

"""
Azure AI Agent Basic Example

This sample demonstrates how to reuse the latest version of an existing agent instead of creating a new agent version on each instantiation.
The first call creates a new agent, while subsequent calls with `use_latest_version=True` reuse the latest agent version.
"""


def get_weather(
location: Annotated[str, Field(description="The location to get the weather for.")],
) -> str:
"""Get the weather for a given location."""
conditions = ["sunny", "cloudy", "rainy", "stormy"]
return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C."


async def main() -> None:
# For authentication, run `az login` command in terminal or replace AzureCliCredential with preferred
# authentication option.
async with AzureCliCredential() as credential:
async with (
AzureAIClient(
async_credential=credential,
).create_agent(
name="MyWeatherAgent",
instructions="You are a helpful weather agent.",
tools=get_weather,
) as agent,
):
# First query will create a new agent
query = "What's the weather like in Seattle?"
print(f"User: {query}")
result = await agent.run(query)
print(f"Agent: {result}\n")

# Create a new agent instance
async with (
AzureAIClient(
async_credential=credential,
# This parameter will allow to re-use latest agent version
# instead of creating a new one
use_latest_version=True,
).create_agent(
name="MyWeatherAgent",
instructions="You are a helpful weather agent.",
tools=get_weather,
) as agent,
):
query = "What's the weather like in Tokyo?"
print(f"User: {query}")
result = await agent.run(query)
print(f"Agent: {result}\n")


if __name__ == "__main__":
asyncio.run(main())