diff --git a/python/packages/azure-ai/agent_framework_azure_ai/__init__.py b/python/packages/azure-ai/agent_framework_azure_ai/__init__.py index cbd836099e..6e6ac7a5e5 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/__init__.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/__init__.py @@ -2,7 +2,8 @@ import importlib.metadata -from ._chat_client import AzureAIAgentClient, AzureAISettings +from ._chat_client import AzureAIAgentClient +from ._shared import AzureAISettings try: __version__ = importlib.metadata.version(__name__) diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py index 834f8debcc..f16560e517 100644 --- a/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py +++ b/python/packages/azure-ai/agent_framework_azure_ai/_chat_client.py @@ -40,9 +40,9 @@ from agent_framework import ( use_chat_middleware, use_function_invocation, ) -from agent_framework._pydantic import AFBaseSettings from agent_framework.exceptions import ServiceInitializationError, ServiceResponseException from agent_framework.observability import use_observability +from azure.ai.agents.aio import AgentsClient from azure.ai.agents.models import ( Agent, AgentsNamedToolChoice, @@ -85,11 +85,11 @@ from azure.ai.agents.models import ( ToolDefinition, ToolOutput, ) -from azure.ai.projects.aio import AIProjectClient from azure.core.credentials_async import AsyncTokenCredential -from azure.core.exceptions import HttpResponseError, ResourceNotFoundError from pydantic import ValidationError +from ._shared import AzureAISettings + if sys.version_info >= (3, 11): from typing import Self # pragma: no cover else: @@ -99,47 +99,6 @@ else: logger = get_logger("agent_framework.azure") -class AzureAISettings(AFBaseSettings): - """Azure AI Project settings. - - The settings are first loaded from environment variables with the prefix 'AZURE_AI_'. - If the environment variables are not found, the settings can be loaded from a .env file - with the encoding 'utf-8'. If the settings are not found in the .env file, the settings - are ignored; however, validation will fail alerting that the settings are missing. - - Keyword Args: - project_endpoint: The Azure AI Project endpoint URL. - Can be set via environment variable AZURE_AI_PROJECT_ENDPOINT. - model_deployment_name: The name of the model deployment to use. - Can be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME. - env_file_path: If provided, the .env settings are read from this file path location. - env_file_encoding: The encoding of the .env file, defaults to 'utf-8'. - - Examples: - .. code-block:: python - - from agent_framework_azure_ai import AzureAISettings - - # Using environment variables - # Set AZURE_AI_PROJECT_ENDPOINT=https://your-project.cognitiveservices.azure.com - # Set AZURE_AI_MODEL_DEPLOYMENT_NAME=gpt-4 - settings = AzureAISettings() - - # Or passing parameters directly - settings = AzureAISettings( - project_endpoint="https://your-project.cognitiveservices.azure.com", model_deployment_name="gpt-4" - ) - - # Or loading from a .env file - settings = AzureAISettings(env_file_path="path/to/.env") - """ - - env_prefix: ClassVar[str] = "AZURE_AI_" - - project_endpoint: str | None = None - model_deployment_name: str | None = None - - TAzureAIAgentClient = TypeVar("TAzureAIAgentClient", bound="AzureAIAgentClient") @@ -154,7 +113,7 @@ class AzureAIAgentClient(BaseChatClient): def __init__( self, *, - project_client: AIProjectClient | None = None, + agents_client: AgentsClient | None = None, agent_id: str | None = None, agent_name: str | None = None, thread_id: str | None = None, @@ -169,16 +128,16 @@ class AzureAIAgentClient(BaseChatClient): """Initialize an Azure AI Agent client. 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 + agents_client: An existing AgentsClient to use. If not provided, one will be created. + agent_id: The ID of an existing agent to use. If not provided and agents_client is provided, + a new agent will be created (and deleted after the request). If neither agents_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 conversation_id property when making a request. project_endpoint: The Azure AI Project endpoint URL. Can also be set via environment variable AZURE_AI_PROJECT_ENDPOINT. - Ignored when a project_client is passed. + Ignored when a agents_client is passed. 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. @@ -221,9 +180,9 @@ class AzureAIAgentClient(BaseChatClient): except ValidationError as ex: raise ServiceInitializationError("Failed to create Azure AI settings.", ex) from ex - # If no project_client is provided, create one + # If no agents_client is provided, create one should_close_client = False - if project_client is None: + if agents_client is None: if not azure_ai_settings.project_endpoint: raise ServiceInitializationError( "Azure AI project endpoint is required. Set via 'project_endpoint' parameter " @@ -238,8 +197,8 @@ class AzureAIAgentClient(BaseChatClient): # Use provided credential if not async_credential: - raise ServiceInitializationError("Azure credential is required when project_client is not provided.") - project_client = AIProjectClient( + raise ServiceInitializationError("Azure credential is required when agents_client is not provided.") + agents_client = AgentsClient( endpoint=azure_ai_settings.project_endpoint, credential=async_credential, user_agent=AGENT_FRAMEWORK_USER_AGENT, @@ -250,7 +209,7 @@ class AzureAIAgentClient(BaseChatClient): super().__init__(**kwargs) # Initialize instance variables - self.project_client = project_client + self.agents_client = agents_client self.credential = async_credential self.agent_id = agent_id self.agent_name = agent_name @@ -261,27 +220,6 @@ class AzureAIAgentClient(BaseChatClient): 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 - async def setup_azure_ai_observability(self, enable_sensitive_data: bool | None = None) -> None: - """Use this method to setup tracing in your Azure AI Project. - - This will take the connection string from the project project_client. - It will override any connection string that is set in the environment variables. - It will disable any OTLP endpoint that might have been set. - """ - try: - conn_string = await self.project_client.telemetry.get_application_insights_connection_string() - except ResourceNotFoundError: - logger.warning( - "No Application Insights connection string found for the Azure AI Project, " - "please call setup_observability() manually." - ) - return - from agent_framework.observability import setup_observability - - setup_observability( - applicationinsights_connection_string=conn_string, enable_sensitive_data=enable_sensitive_data - ) - async def __aenter__(self) -> "Self": """Async context manager entry.""" return self @@ -291,7 +229,7 @@ class AzureAIAgentClient(BaseChatClient): await self.close() async def close(self) -> None: - """Close the project_client and clean up any agents we created.""" + """Close the agents_client and clean up any agents we created.""" await self._cleanup_agent_if_needed() await self._close_client_if_needed() @@ -303,7 +241,7 @@ class AzureAIAgentClient(BaseChatClient): settings: A dictionary of settings for the service. """ return cls( - project_client=settings.get("project_client"), + agents_client=settings.get("agents_client"), agent_id=settings.get("agent_id"), thread_id=settings.get("thread_id"), project_endpoint=settings.get("project_endpoint"), @@ -380,11 +318,14 @@ class AzureAIAgentClient(BaseChatClient): args["instructions"] = run_options["instructions"] if "response_format" in run_options: args["response_format"] = run_options["response_format"] + if "temperature" in run_options: args["temperature"] = run_options["temperature"] if "top_p" in run_options: args["top_p"] = run_options["top_p"] - created_agent = await self.project_client.agents.create_agent(**args) + + created_agent = await self.agents_client.create_agent(**args) + self.agent_id = str(created_agent.id) self._agent_definition = created_agent self._agent_created = True @@ -428,7 +369,7 @@ class AzureAIAgentClient(BaseChatClient): args["tool_outputs"] = tool_outputs if tool_approvals: args["tool_approvals"] = tool_approvals - await self.project_client.agents.runs.submit_tool_outputs_stream(**args) # type: ignore[reportUnknownMemberType] + await self.agents_client.runs.submit_tool_outputs_stream(**args) # type: ignore[reportUnknownMemberType] # Pass the handler to the stream to continue processing stream = handler # type: ignore final_thread_id = thread_run.thread_id @@ -438,7 +379,7 @@ class AzureAIAgentClient(BaseChatClient): # Now create a new run and stream the results. run_options.pop("conversation_id", None) - stream = await self.project_client.agents.runs.stream( # type: ignore[reportUnknownMemberType] + stream = await self.agents_client.runs.stream( # type: ignore[reportUnknownMemberType] final_thread_id, agent_id=agent_id, **run_options ) @@ -449,9 +390,7 @@ class AzureAIAgentClient(BaseChatClient): if thread_id is None: return None - async for run in self.project_client.agents.runs.list( - thread_id=thread_id, limit=1, order=ListSortOrder.DESCENDING - ): # type: ignore[reportUnknownMemberType] + async for run in self.agents_client.runs.list(thread_id=thread_id, limit=1, order=ListSortOrder.DESCENDING): # type: ignore[reportUnknownMemberType] if run.status not in [ RunStatus.COMPLETED, RunStatus.CANCELLED, @@ -468,12 +407,12 @@ class AzureAIAgentClient(BaseChatClient): if thread_id is not None: if thread_run is not None: # There was an active run; we need to cancel it before starting a new run. - await self.project_client.agents.runs.cancel(thread_id, thread_run.id) + await self.agents_client.runs.cancel(thread_id, thread_run.id) return thread_id # No thread ID was provided, so create a new thread. - thread = await self.project_client.agents.threads.create( + thread = await self.agents_client.threads.create( tool_resources=run_options.get("tool_resources"), metadata=run_options.get("metadata") ) thread_id = thread.id @@ -482,7 +421,7 @@ class AzureAIAgentClient(BaseChatClient): # once fixed, in the function above, readd: # `messages=run_options.pop("additional_messages")` for msg in run_options.pop("additional_messages", []): - await self.project_client.agents.messages.create( + await self.agents_client.messages.create( thread_id=thread_id, role=msg.role, content=msg.content, metadata=msg.metadata ) # and remove until here. @@ -715,21 +654,21 @@ class AzureAIAgentClient(BaseChatClient): return [] async def _close_client_if_needed(self) -> None: - """Close project_client session if we created it.""" + """Close agents_client session if we created it.""" if self._should_close_client: - await self.project_client.close() + await self.agents_client.close() async def _cleanup_agent_if_needed(self) -> None: """Clean up the agent if we created it.""" if self._agent_created and self.should_cleanup_agent and self.agent_id is not None: - await self.project_client.agents.delete_agent(self.agent_id) + await self.agents_client.delete_agent(self.agent_id) self.agent_id = None self._agent_created = 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: - self._agent_definition = await self.project_client.agents.get_agent(self.agent_id) + self._agent_definition = await self.agents_client.get_agent(self.agent_id) return self._agent_definition def _prepare_tool_choice(self, chat_options: ChatOptions) -> None: @@ -919,59 +858,34 @@ class AzureAIAgentClient(BaseChatClient): config_args["market"] = market if set_lang := additional_props.get("set_lang"): config_args["set_lang"] = set_lang - # Bing Grounding (support both connection_id and connection_name) + # Bing Grounding connection_id = additional_props.get("connection_id") or os.getenv("BING_CONNECTION_ID") - connection_name = additional_props.get("connection_name") or os.getenv("BING_CONNECTION_NAME") # Custom Bing Search - custom_connection_name = additional_props.get("custom_connection_name") or os.getenv( - "BING_CUSTOM_CONNECTION_NAME" + custom_connection_id = additional_props.get("custom_connection_id") or os.getenv( + "BING_CUSTOM_CONNECTION_ID" ) - custom_configuration_name = additional_props.get("custom_instance_name") or os.getenv( + custom_instance_name = additional_props.get("custom_instance_name") or os.getenv( "BING_CUSTOM_INSTANCE_NAME" ) bing_search: BingGroundingTool | BingCustomSearchTool | None = None - if ( - (connection_id or connection_name) - and not custom_connection_name - and not custom_configuration_name - ): + if (connection_id) and not custom_connection_id and not custom_instance_name: if connection_id: conn_id = connection_id - elif connection_name: - try: - bing_connection = await self.project_client.connections.get(name=connection_name) - except HttpResponseError as err: - raise ServiceInitializationError( - f"Bing connection '{connection_name}' not found in the Azure AI Project.", - err, - ) from err - else: - conn_id = bing_connection.id else: - raise ServiceInitializationError("Neither connection_id nor connection_name provided.") + raise ServiceInitializationError("Parameter connection_id is not provided.") bing_search = BingGroundingTool(connection_id=conn_id, **config_args) - if custom_connection_name and custom_configuration_name: - try: - bing_custom_connection = await self.project_client.connections.get( - name=custom_connection_name - ) - except HttpResponseError as err: - raise ServiceInitializationError( - f"Bing custom connection '{custom_connection_name}' not found in the Azure AI Project.", - err, - ) from err - else: - bing_search = BingCustomSearchTool( - connection_id=bing_custom_connection.id, - instance_name=custom_configuration_name, - **config_args, - ) + if custom_connection_id and custom_instance_name: + bing_search = BingCustomSearchTool( + connection_id=custom_connection_id, + instance_name=custom_instance_name, + **config_args, + ) if not bing_search: raise ServiceInitializationError( - "Bing search tool requires either 'connection_id' or 'connection_name' for Bing Grounding " - "or both 'custom_connection_name' and 'custom_instance_name' for Custom Bing Search. " + "Bing search tool requires either 'connection_id' for Bing Grounding " + "or both 'custom_connection_id' and 'custom_instance_name' for Custom Bing Search. " "These can be provided via additional_properties or environment variables: " - "'BING_CONNECTION_ID', 'BING_CONNECTION_NAME', 'BING_CUSTOM_CONNECTION_NAME', " + "'BING_CONNECTION_ID', 'BING_CUSTOM_CONNECTION_ID', " "'BING_CUSTOM_INSTANCE_NAME'" ) tool_definitions.extend(bing_search.definitions) @@ -1062,4 +976,4 @@ class AzureAIAgentClient(BaseChatClient): Returns: The service URL for the chat client, or None if not set. """ - return self.project_client._config.endpoint + return self.agents_client._config.endpoint # type: ignore diff --git a/python/packages/azure-ai/agent_framework_azure_ai/_shared.py b/python/packages/azure-ai/agent_framework_azure_ai/_shared.py new file mode 100644 index 0000000000..a120e9f92e --- /dev/null +++ b/python/packages/azure-ai/agent_framework_azure_ai/_shared.py @@ -0,0 +1,46 @@ +# Copyright (c) Microsoft. All rights reserved. + +from typing import ClassVar + +from agent_framework._pydantic import AFBaseSettings + + +class AzureAISettings(AFBaseSettings): + """Azure AI Project settings. + + The settings are first loaded from environment variables with the prefix 'AZURE_AI_'. + If the environment variables are not found, the settings can be loaded from a .env file + with the encoding 'utf-8'. If the settings are not found in the .env file, the settings + are ignored; however, validation will fail alerting that the settings are missing. + + Keyword Args: + project_endpoint: The Azure AI Project endpoint URL. + Can be set via environment variable AZURE_AI_PROJECT_ENDPOINT. + model_deployment_name: The name of the model deployment to use. + Can be set via environment variable AZURE_AI_MODEL_DEPLOYMENT_NAME. + env_file_path: If provided, the .env settings are read from this file path location. + env_file_encoding: The encoding of the .env file, defaults to 'utf-8'. + + Examples: + .. code-block:: python + + from agent_framework.azure import AzureAISettings + + # Using environment variables + # Set AZURE_AI_PROJECT_ENDPOINT=https://your-project.cognitiveservices.azure.com + # Set AZURE_AI_MODEL_DEPLOYMENT_NAME=gpt-4 + settings = AzureAISettings() + + # Or passing parameters directly + settings = AzureAISettings( + project_endpoint="https://your-project.cognitiveservices.azure.com", model_deployment_name="gpt-4" + ) + + # Or loading from a .env file + settings = AzureAISettings(env_file_path="path/to/.env") + """ + + env_prefix: ClassVar[str] = "AZURE_AI_" + + project_endpoint: str | None = None + model_deployment_name: str | None = None diff --git a/python/packages/azure-ai/tests/conftest.py b/python/packages/azure-ai/tests/conftest.py index 13764c24f0..3f635e8213 100644 --- a/python/packages/azure-ai/tests/conftest.py +++ b/python/packages/azure-ai/tests/conftest.py @@ -44,31 +44,30 @@ def azure_ai_unit_test_env(monkeypatch, exclude_list, override_env_param_dict): @fixture -def mock_ai_project_client() -> MagicMock: - """Fixture that provides a mock AIProjectClient.""" +def mock_agents_client() -> MagicMock: + """Fixture that provides a mock AgentsClient.""" mock_client = MagicMock() # Mock agents property - mock_client.agents = MagicMock() - mock_client.agents.create_agent = AsyncMock() - mock_client.agents.delete_agent = AsyncMock() + mock_client.create_agent = AsyncMock() + mock_client.delete_agent = AsyncMock() # Mock agent creation response mock_agent = MagicMock() mock_agent.id = "test-agent-id" - mock_client.agents.create_agent.return_value = mock_agent + mock_client.create_agent.return_value = mock_agent # Mock threads property - mock_client.agents.threads = MagicMock() - mock_client.agents.threads.create = AsyncMock() - mock_client.agents.messages.create = AsyncMock() + mock_client.threads = MagicMock() + mock_client.threads.create = AsyncMock() + mock_client.messages.create = AsyncMock() # Mock runs property - mock_client.agents.runs = MagicMock() - mock_client.agents.runs.list = AsyncMock() - mock_client.agents.runs.cancel = AsyncMock() - mock_client.agents.runs.stream = AsyncMock() - mock_client.agents.runs.submit_tool_outputs_stream = AsyncMock() + mock_client.runs = MagicMock() + mock_client.runs.list = AsyncMock() + mock_client.runs.cancel = AsyncMock() + mock_client.runs.stream = AsyncMock() + mock_client.runs.submit_tool_outputs_stream = AsyncMock() return mock_client diff --git a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py index 9a243382e2..555d27d560 100644 --- a/python/packages/azure-ai/tests/test_azure_ai_agent_client.py +++ b/python/packages/azure-ai/tests/test_azure_ai_agent_client.py @@ -52,7 +52,6 @@ from azure.ai.agents.models import ( VectorStore, ) from azure.core.credentials_async import AsyncTokenCredential -from azure.core.exceptions import HttpResponseError, ResourceNotFoundError from azure.identity.aio import AzureCliCredential from pydantic import BaseModel, Field, ValidationError @@ -68,7 +67,7 @@ skip_if_azure_ai_integration_tests_disabled = pytest.mark.skipif( def create_test_azure_ai_chat_client( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, agent_id: str | None = None, thread_id: str | None = None, azure_ai_settings: AzureAISettings | None = None, @@ -83,7 +82,7 @@ def create_test_azure_ai_chat_client( client = object.__new__(AzureAIAgentClient) # Set attributes directly - client.project_client = mock_ai_project_client + client.agents_client = mock_agents_client client.credential = None client.agent_id = agent_id client.agent_name = agent_name @@ -118,13 +117,13 @@ def test_azure_ai_settings_init_with_explicit_values() -> None: assert settings.model_deployment_name == "custom-model" -def test_azure_ai_chat_client_init_with_client(mock_ai_project_client: MagicMock) -> None: - """Test AzureAIAgentClient initialization with existing project_client.""" +def test_azure_ai_chat_client_init_with_client(mock_agents_client: MagicMock) -> None: + """Test AzureAIAgentClient initialization with existing agents_client.""" chat_client = create_test_azure_ai_chat_client( - mock_ai_project_client, agent_id="existing-agent-id", thread_id="test-thread-id" + mock_agents_client, agent_id="existing-agent-id", thread_id="test-thread-id" ) - assert chat_client.project_client is mock_ai_project_client + 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 isinstance(chat_client, ChatClientProtocol) @@ -132,14 +131,14 @@ def test_azure_ai_chat_client_init_with_client(mock_ai_project_client: MagicMock def test_azure_ai_chat_client_init_auto_create_client( azure_ai_unit_test_env: dict[str, str], - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: - """Test AzureAIAgentClient initialization with auto-created project_client.""" + """Test AzureAIAgentClient initialization with auto-created agents_client.""" azure_ai_settings = AzureAISettings(**azure_ai_unit_test_env) # type: ignore # Create client instance directly chat_client = object.__new__(AzureAIAgentClient) - chat_client.project_client = mock_ai_project_client + chat_client.agents_client = mock_agents_client chat_client.agent_id = None chat_client.thread_id = None chat_client._should_close_client = False # type: ignore @@ -149,12 +148,12 @@ def test_azure_ai_chat_client_init_auto_create_client( chat_client.additional_properties = {} chat_client.middleware = None - assert chat_client.project_client is mock_ai_project_client + assert chat_client.agents_client is mock_agents_client assert chat_client.agent_id is None def test_azure_ai_chat_client_init_missing_project_endpoint() -> None: - """Test AzureAIAgentClient initialization when project_endpoint is missing and no project_client provided.""" + """Test AzureAIAgentClient initialization when project_endpoint is missing and no agents_client provided.""" # Mock AzureAISettings to return settings with None project_endpoint with patch("agent_framework_azure_ai._chat_client.AzureAISettings") as mock_settings: mock_settings_instance = MagicMock() @@ -165,7 +164,7 @@ def test_azure_ai_chat_client_init_missing_project_endpoint() -> None: with pytest.raises(ServiceInitializationError, match="project endpoint is required"): AzureAIAgentClient( - project_client=None, + agents_client=None, agent_id=None, project_endpoint=None, # Missing endpoint model_deployment_name="test-model", @@ -185,7 +184,7 @@ def test_azure_ai_chat_client_init_missing_model_deployment_for_agent_creation() with pytest.raises(ServiceInitializationError, match="model deployment name is required"): AzureAIAgentClient( - project_client=None, + agents_client=None, agent_id=None, # No existing agent project_endpoint="https://test.com", model_deployment_name=None, # Missing for agent creation @@ -193,10 +192,10 @@ def test_azure_ai_chat_client_init_missing_model_deployment_for_agent_creation() ) -def test_azure_ai_chat_client_from_dict(mock_ai_project_client: MagicMock) -> None: +def test_azure_ai_chat_client_from_dict(mock_agents_client: MagicMock) -> None: """Test AzureAIAgentClient.from_dict method.""" settings = { - "project_client": mock_ai_project_client, + "agents_client": mock_agents_client, "agent_id": "test-agent-id", "thread_id": "test-thread-id", "project_endpoint": "https://test-endpoint.com/", @@ -210,24 +209,24 @@ def test_azure_ai_chat_client_from_dict(mock_ai_project_client: MagicMock) -> No ) chat_client: AzureAIAgentClient = create_test_azure_ai_chat_client( - mock_ai_project_client, + mock_agents_client, agent_id=settings["agent_id"], # type: ignore thread_id=settings["thread_id"], # type: ignore azure_ai_settings=azure_ai_settings, ) - assert chat_client.project_client is mock_ai_project_client + assert chat_client.agents_client is mock_agents_client assert chat_client.agent_id == "test-agent-id" assert chat_client.thread_id == "test-thread-id" def test_azure_ai_chat_client_init_missing_credential(azure_ai_unit_test_env: dict[str, str]) -> None: - """Test AzureAIAgentClient.__init__ when async_credential is missing and no project_client provided.""" + """Test AzureAIAgentClient.__init__ when async_credential is missing and no agents_client provided.""" with pytest.raises( - ServiceInitializationError, match="Azure credential is required when project_client is not provided" + ServiceInitializationError, match="Azure credential is required when agents_client is not provided" ): AzureAIAgentClient( - project_client=None, + agents_client=None, agent_id="existing-agent", project_endpoint=azure_ai_unit_test_env["AZURE_AI_PROJECT_ENDPOINT"], model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"], @@ -251,9 +250,9 @@ def test_azure_ai_chat_client_init_validation_error(mock_azure_credential: Magic def test_azure_ai_chat_client_from_settings() -> None: """Test from_settings class method.""" - mock_project_client = MagicMock() + mock_agents_client = MagicMock() settings = { - "project_client": mock_project_client, + "agents_client": mock_agents_client, "agent_id": "test-agent", "thread_id": "test-thread", "project_endpoint": "https://test.com", @@ -263,18 +262,18 @@ def test_azure_ai_chat_client_from_settings() -> None: client = AzureAIAgentClient.from_settings(settings) - assert client.project_client is mock_project_client + assert client.agents_client is mock_agents_client assert client.agent_id == "test-agent" assert client.thread_id == "test-thread" assert client.agent_name == "TestAgent" async def test_azure_ai_chat_client_get_agent_id_or_create_with_temperature_and_top_p( - mock_ai_project_client: MagicMock, azure_ai_unit_test_env: dict[str, str] + mock_agents_client: MagicMock, azure_ai_unit_test_env: dict[str, str] ) -> None: """Test _get_agent_id_or_create with temperature and top_p in run_options.""" azure_ai_settings = AzureAISettings(model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"]) - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, azure_ai_settings=azure_ai_settings) + chat_client = create_test_azure_ai_chat_client(mock_agents_client, azure_ai_settings=azure_ai_settings) run_options = { "model": azure_ai_settings.model_deployment_name, @@ -286,17 +285,17 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_with_temperature_and_ assert agent_id == "test-agent-id" # Verify create_agent was called with temperature and top_p parameters - mock_ai_project_client.agents.create_agent.assert_called_once() - call_kwargs = mock_ai_project_client.agents.create_agent.call_args[1] + mock_agents_client.create_agent.assert_called_once() + call_kwargs = mock_agents_client.create_agent.call_args[1] assert call_kwargs["temperature"] == 0.7 assert call_kwargs["top_p"] == 0.9 async def test_azure_ai_chat_client_get_agent_id_or_create_existing_agent( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _get_agent_id_or_create when agent_id is already provided.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="existing-agent-id") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="existing-agent-id") agent_id = await chat_client._get_agent_id_or_create() # type: ignore @@ -305,12 +304,12 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_existing_agent( async def test_azure_ai_chat_client_get_agent_id_or_create_create_new( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, azure_ai_unit_test_env: dict[str, str], ) -> None: """Test _get_agent_id_or_create when creating a new agent.""" azure_ai_settings = AzureAISettings(model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"]) - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, azure_ai_settings=azure_ai_settings) + chat_client = create_test_azure_ai_chat_client(mock_agents_client, azure_ai_settings=azure_ai_settings) agent_id = await chat_client._get_agent_id_or_create(run_options={"model": azure_ai_settings.model_deployment_name}) # type: ignore @@ -318,19 +317,19 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_create_new( assert chat_client._agent_created -async def test_azure_ai_chat_client_thread_management_through_public_api(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_thread_management_through_public_api(mock_agents_client: MagicMock) -> None: """Test thread creation and management through public API.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Mock get_agent to avoid the async error - mock_ai_project_client.agents.get_agent = AsyncMock(return_value=None) + mock_agents_client.get_agent = AsyncMock(return_value=None) mock_thread = MagicMock() mock_thread.id = "new-thread-456" - mock_ai_project_client.agents.threads.create = AsyncMock(return_value=mock_thread) + mock_agents_client.threads.create = AsyncMock(return_value=mock_thread) mock_stream = AsyncMock() - mock_ai_project_client.agents.runs.stream = AsyncMock(return_value=mock_stream) + mock_agents_client.runs.stream = AsyncMock(return_value=mock_stream) # Create an async iterator that yields nothing (empty stream) async def empty_async_iter(): @@ -349,23 +348,23 @@ async def test_azure_ai_chat_client_thread_management_through_public_api(mock_ai pass # Verify thread creation was called - mock_ai_project_client.agents.threads.create.assert_called_once() + mock_agents_client.threads.create.assert_called_once() @pytest.mark.parametrize("exclude_list", [["AZURE_AI_MODEL_DEPLOYMENT_NAME"]], indirect=True) async def test_azure_ai_chat_client_get_agent_id_or_create_missing_model( - mock_ai_project_client: MagicMock, azure_ai_unit_test_env: dict[str, str] + mock_agents_client: MagicMock, azure_ai_unit_test_env: dict[str, str] ) -> None: """Test _get_agent_id_or_create when model_deployment_name is missing.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) with pytest.raises(ServiceInitializationError, match="Model deployment name is required"): await chat_client._get_agent_id_or_create() # type: ignore -async def test_azure_ai_chat_client_create_run_options_basic(mock_ai_project_client: MagicMock) -> None: +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_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) messages = [ChatMessage(role=Role.USER, text="Hello")] chat_options = ChatOptions(max_tokens=100, temperature=0.7) @@ -376,9 +375,9 @@ async def test_azure_ai_chat_client_create_run_options_basic(mock_ai_project_cli assert tool_results is None -async def test_azure_ai_chat_client_create_run_options_no_chat_options(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_create_run_options_no_chat_options(mock_agents_client: MagicMock) -> None: """Test _create_run_options with no ChatOptions.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) messages = [ChatMessage(role=Role.USER, text="Hello")] @@ -388,13 +387,13 @@ async def test_azure_ai_chat_client_create_run_options_no_chat_options(mock_ai_p assert tool_results is None -async def test_azure_ai_chat_client_create_run_options_with_image_content(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_create_run_options_with_image_content(mock_agents_client: MagicMock) -> None: """Test _create_run_options with image content.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Mock get_agent - mock_ai_project_client.agents.get_agent = AsyncMock(return_value=None) + mock_agents_client.get_agent = AsyncMock(return_value=None) image_content = UriContent(uri="https://example.com/image.jpg", media_type="image/jpeg") messages = [ChatMessage(role=Role.USER, contents=[image_content])] @@ -408,9 +407,9 @@ async def test_azure_ai_chat_client_create_run_options_with_image_content(mock_a assert len(message.content) == 1 -def test_azure_ai_chat_client_convert_function_results_to_tool_output_none(mock_ai_project_client: MagicMock) -> None: +def test_azure_ai_chat_client_convert_function_results_to_tool_output_none(mock_agents_client: MagicMock) -> None: """Test _convert_required_action_to_tool_output with None input.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) run_id, tool_outputs, tool_approvals = chat_client._convert_required_action_to_tool_output(None) # type: ignore @@ -419,31 +418,31 @@ def test_azure_ai_chat_client_convert_function_results_to_tool_output_none(mock_ assert tool_approvals is None -async def test_azure_ai_chat_client_close_client_when_should_close_true(mock_ai_project_client: MagicMock) -> None: - """Test _close_client_if_needed closes project_client when should_close_client is True.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) +async def test_azure_ai_chat_client_close_client_when_should_close_true(mock_agents_client: MagicMock) -> None: + """Test _close_client_if_needed closes agents_client when should_close_client is True.""" + chat_client = create_test_azure_ai_chat_client(mock_agents_client) chat_client._should_close_client = True # type: ignore - mock_ai_project_client.close = AsyncMock() + mock_agents_client.close = AsyncMock() await chat_client._close_client_if_needed() # type: ignore - mock_ai_project_client.close.assert_called_once() + mock_agents_client.close.assert_called_once() -async def test_azure_ai_chat_client_close_client_when_should_close_false(mock_ai_project_client: MagicMock) -> None: - """Test _close_client_if_needed does not close project_client when should_close_client is False.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) +async def test_azure_ai_chat_client_close_client_when_should_close_false(mock_agents_client: MagicMock) -> None: + """Test _close_client_if_needed does not close agents_client when should_close_client is False.""" + chat_client = create_test_azure_ai_chat_client(mock_agents_client) chat_client._should_close_client = False # type: ignore await chat_client._close_client_if_needed() # type: ignore - mock_ai_project_client.close.assert_not_called() + mock_agents_client.close.assert_not_called() -def test_azure_ai_chat_client_update_agent_name_when_current_is_none(mock_ai_project_client: MagicMock) -> None: +def test_azure_ai_chat_client_update_agent_name_when_current_is_none(mock_agents_client: MagicMock) -> None: """Test _update_agent_name updates name when current agent_name is None.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) chat_client.agent_name = None # type: ignore chat_client._update_agent_name("NewAgentName") # type: ignore @@ -451,9 +450,9 @@ def test_azure_ai_chat_client_update_agent_name_when_current_is_none(mock_ai_pro assert chat_client.agent_name == "NewAgentName" -def test_azure_ai_chat_client_update_agent_name_when_current_exists(mock_ai_project_client: MagicMock) -> None: +def test_azure_ai_chat_client_update_agent_name_when_current_exists(mock_agents_client: MagicMock) -> None: """Test _update_agent_name does not update when current agent_name exists.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) chat_client.agent_name = "ExistingName" # type: ignore chat_client._update_agent_name("NewAgentName") # type: ignore @@ -461,9 +460,9 @@ def test_azure_ai_chat_client_update_agent_name_when_current_exists(mock_ai_proj assert chat_client.agent_name == "ExistingName" -def test_azure_ai_chat_client_update_agent_name_with_none_input(mock_ai_project_client: MagicMock) -> None: +def test_azure_ai_chat_client_update_agent_name_with_none_input(mock_agents_client: MagicMock) -> None: """Test _update_agent_name with None input.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) chat_client.agent_name = None # type: ignore chat_client._update_agent_name(None) # type: ignore @@ -471,9 +470,9 @@ def test_azure_ai_chat_client_update_agent_name_with_none_input(mock_ai_project_ assert chat_client.agent_name is None -async def test_azure_ai_chat_client_create_run_options_with_messages(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_create_run_options_with_messages(mock_agents_client: MagicMock) -> None: """Test _create_run_options with different message types.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) # Test with system message (becomes instruction) messages = [ @@ -489,9 +488,9 @@ async def test_azure_ai_chat_client_create_run_options_with_messages(mock_ai_pro assert len(run_options["additional_messages"]) == 1 # Only user message -async def test_azure_ai_chat_client_inner_get_response(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_inner_get_response(mock_agents_client: MagicMock) -> None: """Test _inner_get_response method.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") messages = [ChatMessage(role=Role.USER, text="Hello")] chat_options = ChatOptions() @@ -512,11 +511,11 @@ async def test_azure_ai_chat_client_inner_get_response(mock_ai_project_client: M async def test_azure_ai_chat_client_get_agent_id_or_create_with_run_options( - mock_ai_project_client: MagicMock, azure_ai_unit_test_env: dict[str, str] + mock_agents_client: MagicMock, azure_ai_unit_test_env: dict[str, str] ) -> None: """Test _get_agent_id_or_create with run_options containing tools and instructions.""" azure_ai_settings = AzureAISettings(model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"]) - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, azure_ai_settings=azure_ai_settings) + chat_client = create_test_azure_ai_chat_client(mock_agents_client, azure_ai_settings=azure_ai_settings) run_options = { "tools": [{"type": "function", "function": {"name": "test_tool"}}], @@ -529,16 +528,16 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_with_run_options( assert agent_id == "test-agent-id" # Verify create_agent was called with run_options parameters - mock_ai_project_client.agents.create_agent.assert_called_once() - call_args = mock_ai_project_client.agents.create_agent.call_args[1] + mock_agents_client.create_agent.assert_called_once() + call_args = mock_agents_client.create_agent.call_args[1] assert "tools" in call_args assert "instructions" in call_args assert "response_format" in call_args -async def test_azure_ai_chat_client_prepare_thread_cancels_active_run(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_prepare_thread_cancels_active_run(mock_agents_client: MagicMock) -> None: """Test _prepare_thread cancels active thread run when provided.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") mock_thread_run = MagicMock() mock_thread_run.id = "run_123" @@ -549,12 +548,12 @@ async def test_azure_ai_chat_client_prepare_thread_cancels_active_run(mock_ai_pr result = await chat_client._prepare_thread("test-thread", mock_thread_run, run_options) # type: ignore assert result == "test-thread" - mock_ai_project_client.agents.runs.cancel.assert_called_once_with("test-thread", "run_123") + mock_agents_client.runs.cancel.assert_called_once_with("test-thread", "run_123") -def test_azure_ai_chat_client_create_function_call_contents_basic(mock_ai_project_client: MagicMock) -> None: +def test_azure_ai_chat_client_create_function_call_contents_basic(mock_agents_client: MagicMock) -> None: """Test _create_function_call_contents with basic function call.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) mock_tool_call = MagicMock(spec=RequiredFunctionToolCall) mock_tool_call.id = "call_123" @@ -575,9 +574,9 @@ def test_azure_ai_chat_client_create_function_call_contents_basic(mock_ai_projec assert result[0].call_id == '["response_123", "call_123"]' -def test_azure_ai_chat_client_create_function_call_contents_no_submit_action(mock_ai_project_client: MagicMock) -> None: +def test_azure_ai_chat_client_create_function_call_contents_no_submit_action(mock_agents_client: MagicMock) -> None: """Test _create_function_call_contents when required_action is not SubmitToolOutputsAction.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) mock_event_data = MagicMock(spec=ThreadRun) mock_event_data.required_action = MagicMock() @@ -588,10 +587,10 @@ def test_azure_ai_chat_client_create_function_call_contents_no_submit_action(moc def test_azure_ai_chat_client_create_function_call_contents_non_function_tool_call( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _create_function_call_contents with non-function tool call.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) mock_tool_call = MagicMock() @@ -607,10 +606,10 @@ def test_azure_ai_chat_client_create_function_call_contents_non_function_tool_ca async def test_azure_ai_chat_client_create_run_options_with_none_tool_choice( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _create_run_options with tool_choice set to 'none'.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) chat_options = ChatOptions() chat_options.tool_choice = "none" @@ -623,10 +622,10 @@ async def test_azure_ai_chat_client_create_run_options_with_none_tool_choice( async def test_azure_ai_chat_client_create_run_options_with_auto_tool_choice( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _create_run_options with tool_choice set to 'auto'.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) chat_options = ChatOptions() chat_options.tool_choice = "auto" @@ -639,10 +638,10 @@ async def test_azure_ai_chat_client_create_run_options_with_auto_tool_choice( async def test_azure_ai_chat_client_prepare_tool_choice_none_string( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _prepare_tool_choice when tool_choice is string 'none'.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) # Create a mock tool for testing mock_tool = MagicMock() @@ -657,10 +656,10 @@ async def test_azure_ai_chat_client_prepare_tool_choice_none_string( async def test_azure_ai_chat_client_create_run_options_tool_choice_required_specific_function( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _create_run_options with ToolMode.REQUIRED specifying a specific function name.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) required_tool_mode = ToolMode.REQUIRED("specific_function_name") @@ -680,10 +679,10 @@ async def test_azure_ai_chat_client_create_run_options_tool_choice_required_spec async def test_azure_ai_chat_client_create_run_options_with_response_format( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _create_run_options with response_format configured.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) class TestResponseModel(BaseModel): name: str = Field(description="Test name") @@ -698,19 +697,19 @@ async def test_azure_ai_chat_client_create_run_options_with_response_format( assert response_format.json_schema.name == "TestResponseModel" -def test_azure_ai_chat_client_service_url_method(mock_ai_project_client: MagicMock) -> None: +def test_azure_ai_chat_client_service_url_method(mock_agents_client: MagicMock) -> None: """Test service_url method returns endpoint.""" - mock_ai_project_client._config.endpoint = "https://test-endpoint.com/" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + mock_agents_client._config.endpoint = "https://test-endpoint.com/" + chat_client = create_test_azure_ai_chat_client(mock_agents_client) url = chat_client.service_url() assert url == "https://test-endpoint.com/" -async def test_azure_ai_chat_client_prep_tools_ai_function(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_prep_tools_ai_function(mock_agents_client: MagicMock) -> None: """Test _prep_tools with AIFunction tool.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Create a mock AIFunction mock_ai_function = MagicMock(spec=AIFunction) @@ -723,10 +722,10 @@ async def test_azure_ai_chat_client_prep_tools_ai_function(mock_ai_project_clien mock_ai_function.to_json_schema_spec.assert_called_once() -async def test_azure_ai_chat_client_prep_tools_code_interpreter(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_prep_tools_code_interpreter(mock_agents_client: MagicMock) -> None: """Test _prep_tools with HostedCodeInterpreterTool.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") code_interpreter_tool = HostedCodeInterpreterTool() @@ -736,10 +735,10 @@ async def test_azure_ai_chat_client_prep_tools_code_interpreter(mock_ai_project_ assert isinstance(result[0], CodeInterpreterToolDefinition) -async def test_azure_ai_chat_client_prep_tools_mcp_tool(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_prep_tools_mcp_tool(mock_agents_client: MagicMock) -> None: """Test _prep_tools with HostedMCPTool.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") mcp_tool = HostedMCPTool(name="Test MCP Tool", url="https://example.com/mcp", allowed_tools=["tool1", "tool2"]) @@ -761,9 +760,9 @@ async def test_azure_ai_chat_client_prep_tools_mcp_tool(mock_ai_project_client: assert set(call_args["allowed_tools"]) == {"tool1", "tool2"} -async def test_azure_ai_chat_client_create_run_options_mcp_never_require(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_create_run_options_mcp_never_require(mock_agents_client: MagicMock) -> None: """Test _create_run_options with HostedMCPTool having never_require approval mode.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) mcp_tool = HostedMCPTool(name="Test MCP Tool", url="https://example.com/mcp", approval_mode="never_require") @@ -790,9 +789,9 @@ async def test_azure_ai_chat_client_create_run_options_mcp_never_require(mock_ai assert mcp_resource["require_approval"] == "never" -async def test_azure_ai_chat_client_create_run_options_mcp_with_headers(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_create_run_options_mcp_with_headers(mock_agents_client: MagicMock) -> None: """Test _create_run_options with HostedMCPTool having headers.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client) + chat_client = create_test_azure_ai_chat_client(mock_agents_client) # Test with headers headers = {"Authorization": "Bearer DUMMY_TOKEN", "X-API-Key": "DUMMY_KEY"} @@ -822,14 +821,14 @@ async def test_azure_ai_chat_client_create_run_options_mcp_with_headers(mock_ai_ assert mcp_resource["headers"] == headers -async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding(mock_agents_client: MagicMock) -> None: """Test _prep_tools with HostedWebSearchTool using Bing Grounding.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") web_search_tool = HostedWebSearchTool( additional_properties={ - "connection_name": "test-connection-name", + "connection_id": "test-connection-id", "count": 5, "freshness": "Day", "market": "en-US", @@ -837,11 +836,6 @@ async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding(mock_ai } ) - # Mock connection get - mock_connection = MagicMock() - mock_connection.id = "test-connection-id" - mock_ai_project_client.connections.get = AsyncMock(return_value=mock_connection) - # Mock BingGroundingTool with patch("agent_framework_azure_ai._chat_client.BingGroundingTool") as mock_bing_grounding: mock_bing_tool = MagicMock() @@ -861,11 +855,11 @@ async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding(mock_ai async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding_with_connection_id( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _prep_tools with HostedWebSearchTool using Bing Grounding with connection_id (no HTTP call).""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") web_search_tool = HostedWebSearchTool( additional_properties={ @@ -884,29 +878,22 @@ async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding_with_co assert len(result) == 1 assert result[0] == {"type": "bing_grounding"} - # Verify that connection_id was used directly (no HTTP call to connections.get) - mock_ai_project_client.connections.get.assert_not_called() mock_bing_grounding.assert_called_once_with(connection_id="direct-connection-id", count=3) -async def test_azure_ai_chat_client_prep_tools_web_search_custom_bing(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_prep_tools_web_search_custom_bing(mock_agents_client: MagicMock) -> None: """Test _prep_tools with HostedWebSearchTool using Custom Bing Search.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") web_search_tool = HostedWebSearchTool( additional_properties={ - "custom_connection_name": "custom-bing-connection", + "custom_connection_id": "custom-connection-id", "custom_instance_name": "custom-instance", "count": 10, } ) - # Mock connection get - mock_connection = MagicMock() - mock_connection.id = "custom-connection-id" - mock_ai_project_client.connections.get = AsyncMock(return_value=mock_connection) - # Mock BingCustomSearchTool with patch("agent_framework_azure_ai._chat_client.BingCustomSearchTool") as mock_custom_bing: mock_custom_tool = MagicMock() @@ -917,65 +904,14 @@ async def test_azure_ai_chat_client_prep_tools_web_search_custom_bing(mock_ai_pr assert len(result) == 1 assert result[0] == {"type": "bing_custom_search"} - mock_ai_project_client.connections.get.assert_called_once_with(name="custom-bing-connection") - mock_custom_bing.assert_called_once_with( - connection_id="custom-connection-id", instance_name="custom-instance", count=10 - ) - - -async def test_azure_ai_chat_client_prep_tools_web_search_custom_bing_connection_error( - mock_ai_project_client: MagicMock, -) -> None: - """Test _prep_tools with HostedWebSearchTool when custom connection is not found.""" - - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") - - web_search_tool = HostedWebSearchTool( - additional_properties={ - "custom_connection_name": "nonexistent-connection", - "custom_instance_name": "custom-instance", - } - ) - - # Mock connection get to raise HttpResponseError - mock_ai_project_client.connections.get = AsyncMock(side_effect=HttpResponseError("Connection not found")) - - with pytest.raises( - ServiceInitializationError, - match="Bing custom connection 'nonexistent-connection' not found in the Azure AI Project", - ): - await chat_client._prep_tools([web_search_tool]) # type: ignore - - -async def test_azure_ai_chat_client_prep_tools_web_search_bing_grounding_connection_error( - mock_ai_project_client: MagicMock, -) -> None: - """Test _prep_tools with HostedWebSearchTool when Bing Grounding connection is not found.""" - - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") - - web_search_tool = HostedWebSearchTool( - additional_properties={ - "connection_name": "nonexistent-bing-connection", - } - ) - - # Mock connection get to raise HttpResponseError - mock_ai_project_client.connections.get = AsyncMock(side_effect=HttpResponseError("Connection not found")) - - with pytest.raises( - ServiceInitializationError, - match="Bing connection 'nonexistent-bing-connection' not found in the Azure AI Project", - ): - await chat_client._prep_tools([web_search_tool]) # type: ignore async def test_azure_ai_chat_client_prep_tools_file_search_with_vector_stores( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _prep_tools with HostedFileSearchTool using vector stores.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") vector_store_input = HostedVectorStoreContent(vector_store_id="vs-123") file_search_tool = HostedFileSearchTool(inputs=[vector_store_input]) @@ -996,33 +932,11 @@ async def test_azure_ai_chat_client_prep_tools_file_search_with_vector_stores( mock_file_search.assert_called_once_with(vector_store_ids=["vs-123"]) -async def test_azure_ai_chat_client_setup_azure_ai_observability_success( - mock_ai_project_client: MagicMock, -) -> None: - """Test setup_azure_ai_observability success case with connection string.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") - - # Mock successful connection string retrieval - mock_ai_project_client.telemetry.get_application_insights_connection_string = AsyncMock( - return_value="InstrumentationKey=test-key;IngestionEndpoint=https://test.com" - ) - - # Mock setup_observability function (it's imported inside the method) - with patch("agent_framework.observability.setup_observability") as mock_setup: - await chat_client.setup_azure_ai_observability(enable_sensitive_data=True) - - # Verify setup_observability was called with the correct parameters - mock_setup.assert_called_once_with( - applicationinsights_connection_string="InstrumentationKey=test-key;IngestionEndpoint=https://test.com", - enable_sensitive_data=True, - ) - - async def test_azure_ai_chat_client_create_agent_stream_submit_tool_approvals( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _create_agent_stream with tool approvals submission path.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Mock active thread run that matches the tool run ID mock_thread_run = MagicMock() @@ -1041,7 +955,7 @@ async def test_azure_ai_chat_client_create_agent_stream_submit_tool_approvals( # Mock submit_tool_outputs_stream mock_handler = MagicMock() - mock_ai_project_client.agents.runs.submit_tool_outputs_stream = AsyncMock() + mock_agents_client.runs.submit_tool_outputs_stream = AsyncMock() with patch("azure.ai.agents.models.AsyncAgentEventHandler", return_value=mock_handler): stream, final_thread_id = await chat_client._create_agent_stream( # type: ignore @@ -1052,16 +966,16 @@ async def test_azure_ai_chat_client_create_agent_stream_submit_tool_approvals( assert final_thread_id == "test-thread" # Verify submit_tool_outputs_stream was called with approvals - mock_ai_project_client.agents.runs.submit_tool_outputs_stream.assert_called_once() - call_args = mock_ai_project_client.agents.runs.submit_tool_outputs_stream.call_args[1] + mock_agents_client.runs.submit_tool_outputs_stream.assert_called_once() + call_args = mock_agents_client.runs.submit_tool_outputs_stream.call_args[1] assert "tool_approvals" in call_args assert call_args["tool_approvals"][0].tool_call_id == "test-call-id" assert call_args["tool_approvals"][0].approve is True -async def test_azure_ai_chat_client_prep_tools_dict_tool(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_prep_tools_dict_tool(mock_agents_client: MagicMock) -> None: """Test _prep_tools with dictionary tool definition.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") dict_tool = {"type": "custom_tool", "config": {"param": "value"}} @@ -1071,9 +985,9 @@ async def test_azure_ai_chat_client_prep_tools_dict_tool(mock_ai_project_client: assert result[0] == dict_tool -async def test_azure_ai_chat_client_prep_tools_unsupported_tool(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_prep_tools_unsupported_tool(mock_agents_client: MagicMock) -> None: """Test _prep_tools with unsupported tool type.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") unsupported_tool = "not_a_tool" @@ -1081,10 +995,10 @@ async def test_azure_ai_chat_client_prep_tools_unsupported_tool(mock_ai_project_ await chat_client._prep_tools([unsupported_tool]) # type: ignore -async def test_azure_ai_chat_client_get_active_thread_run_with_active_run(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_get_active_thread_run_with_active_run(mock_agents_client: MagicMock) -> None: """Test _get_active_thread_run when there's an active run.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Mock an active run mock_run = MagicMock() @@ -1093,17 +1007,17 @@ async def test_azure_ai_chat_client_get_active_thread_run_with_active_run(mock_a async def mock_list_runs(*args, **kwargs): # type: ignore yield mock_run - mock_ai_project_client.agents.runs.list = mock_list_runs + mock_agents_client.runs.list = mock_list_runs result = await chat_client._get_active_thread_run("thread-123") # type: ignore assert result == mock_run -async def test_azure_ai_chat_client_get_active_thread_run_no_active_run(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_get_active_thread_run_no_active_run(mock_agents_client: MagicMock) -> None: """Test _get_active_thread_run when there's no active run.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Mock a completed run (not active) mock_run = MagicMock() @@ -1112,32 +1026,32 @@ async def test_azure_ai_chat_client_get_active_thread_run_no_active_run(mock_ai_ async def mock_list_runs(*args, **kwargs): # type: ignore yield mock_run - mock_ai_project_client.agents.runs.list = mock_list_runs + mock_agents_client.runs.list = mock_list_runs result = await chat_client._get_active_thread_run("thread-123") # type: ignore assert result is None -async def test_azure_ai_chat_client_get_active_thread_run_no_thread(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_get_active_thread_run_no_thread(mock_agents_client: MagicMock) -> None: """Test _get_active_thread_run with None thread_id.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") result = await chat_client._get_active_thread_run(None) # type: ignore assert result is None # Should not call list since thread_id is None - mock_ai_project_client.agents.runs.list.assert_not_called() + mock_agents_client.runs.list.assert_not_called() -async def test_azure_ai_chat_client_service_url(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_service_url(mock_agents_client: MagicMock) -> None: """Test service_url method.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Mock the config endpoint mock_config = MagicMock() mock_config.endpoint = "https://test-endpoint.com/" - mock_ai_project_client._config = mock_config + mock_agents_client._config = mock_config result = chat_client.service_url() @@ -1145,10 +1059,10 @@ async def test_azure_ai_chat_client_service_url(mock_ai_project_client: MagicMoc async def test_azure_ai_chat_client_convert_required_action_to_tool_output_function_result( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _convert_required_action_to_tool_output with FunctionResultContent.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Test with simple result function_result = FunctionResultContent(call_id='["run_123", "call_456"]', result="Simple result") @@ -1163,10 +1077,10 @@ async def test_azure_ai_chat_client_convert_required_action_to_tool_output_funct assert tool_outputs[0].output == "Simple result" -async def test_azure_ai_chat_client_convert_required_action_invalid_call_id(mock_ai_project_client: MagicMock) -> None: +async def test_azure_ai_chat_client_convert_required_action_invalid_call_id(mock_agents_client: MagicMock) -> None: """Test _convert_required_action_to_tool_output with invalid call_id format.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Invalid call_id format - should raise JSONDecodeError function_result = FunctionResultContent(call_id="invalid_json", result="result") @@ -1176,10 +1090,10 @@ async def test_azure_ai_chat_client_convert_required_action_invalid_call_id(mock async def test_azure_ai_chat_client_convert_required_action_invalid_structure( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _convert_required_action_to_tool_output with invalid call_id structure.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Valid JSON but invalid structure (missing second element) function_result = FunctionResultContent(call_id='["run_123"]', result="result") @@ -1193,7 +1107,7 @@ async def test_azure_ai_chat_client_convert_required_action_invalid_structure( async def test_azure_ai_chat_client_convert_required_action_serde_model_results( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _convert_required_action_to_tool_output with BaseModel results.""" @@ -1202,7 +1116,7 @@ async def test_azure_ai_chat_client_convert_required_action_serde_model_results( self.name = name self.value = value - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Test with BaseModel result mock_result = MockResult(name="test", value=42) @@ -1221,7 +1135,7 @@ async def test_azure_ai_chat_client_convert_required_action_serde_model_results( async def test_azure_ai_chat_client_convert_required_action_multiple_results( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _convert_required_action_to_tool_output with multiple results.""" @@ -1229,7 +1143,7 @@ async def test_azure_ai_chat_client_convert_required_action_multiple_results( def __init__(self, data: str): self.data = data - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Test with multiple results - mix of BaseModel and regular objects mock_basemodel = MockResult(data="model_data") @@ -1254,10 +1168,10 @@ async def test_azure_ai_chat_client_convert_required_action_multiple_results( async def test_azure_ai_chat_client_convert_required_action_approval_response( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _convert_required_action_to_tool_output with FunctionApprovalResponseContent.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Test with approval response - need to provide required fields approval_response = FunctionApprovalResponseContent( @@ -1277,10 +1191,10 @@ async def test_azure_ai_chat_client_convert_required_action_approval_response( async def test_azure_ai_chat_client_create_function_call_contents_approval_request( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _create_function_call_contents with approval action.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Mock SubmitToolApprovalAction with RequiredMcpToolCall mock_tool_call = MagicMock(spec=RequiredMcpToolCall) @@ -1304,11 +1218,11 @@ async def test_azure_ai_chat_client_create_function_call_contents_approval_reque async def test_azure_ai_chat_client_get_agent_id_or_create_with_agent_name( - mock_ai_project_client: MagicMock, azure_ai_unit_test_env: dict[str, str] + mock_agents_client: MagicMock, azure_ai_unit_test_env: dict[str, str] ) -> None: """Test _get_agent_id_or_create uses default name when no agent_name set.""" azure_ai_settings = AzureAISettings(model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"]) - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, azure_ai_settings=azure_ai_settings) + chat_client = create_test_azure_ai_chat_client(mock_agents_client, azure_ai_settings=azure_ai_settings) # Ensure agent_name is None to test the default chat_client.agent_name = None # type: ignore @@ -1317,17 +1231,17 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_with_agent_name( assert agent_id == "test-agent-id" # Verify create_agent was called with default "UnnamedAgent" - mock_ai_project_client.agents.create_agent.assert_called_once() - call_kwargs = mock_ai_project_client.agents.create_agent.call_args[1] + mock_agents_client.create_agent.assert_called_once() + call_kwargs = mock_agents_client.create_agent.call_args[1] assert call_kwargs["name"] == "UnnamedAgent" async def test_azure_ai_chat_client_get_agent_id_or_create_with_response_format( - mock_ai_project_client: MagicMock, azure_ai_unit_test_env: dict[str, str] + mock_agents_client: MagicMock, azure_ai_unit_test_env: dict[str, str] ) -> None: """Test _get_agent_id_or_create with response_format in run_options.""" azure_ai_settings = AzureAISettings(model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"]) - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, azure_ai_settings=azure_ai_settings) + chat_client = create_test_azure_ai_chat_client(mock_agents_client, azure_ai_settings=azure_ai_settings) # Test with response_format in run_options run_options = {"response_format": {"type": "json_object"}, "model": azure_ai_settings.model_deployment_name} @@ -1336,17 +1250,17 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_with_response_format( assert agent_id == "test-agent-id" # Verify create_agent was called with response_format - mock_ai_project_client.agents.create_agent.assert_called_once() - call_kwargs = mock_ai_project_client.agents.create_agent.call_args[1] + mock_agents_client.create_agent.assert_called_once() + call_kwargs = mock_agents_client.create_agent.call_args[1] assert call_kwargs["response_format"] == {"type": "json_object"} async def test_azure_ai_chat_client_get_agent_id_or_create_with_tool_resources( - mock_ai_project_client: MagicMock, azure_ai_unit_test_env: dict[str, str] + mock_agents_client: MagicMock, azure_ai_unit_test_env: dict[str, str] ) -> None: """Test _get_agent_id_or_create with tool_resources in run_options.""" azure_ai_settings = AzureAISettings(model_deployment_name=azure_ai_unit_test_env["AZURE_AI_MODEL_DEPLOYMENT_NAME"]) - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, azure_ai_settings=azure_ai_settings) + chat_client = create_test_azure_ai_chat_client(mock_agents_client, azure_ai_settings=azure_ai_settings) # Test with tool_resources in run_options run_options = { @@ -1358,16 +1272,16 @@ async def test_azure_ai_chat_client_get_agent_id_or_create_with_tool_resources( assert agent_id == "test-agent-id" # Verify create_agent was called with tool_resources - mock_ai_project_client.agents.create_agent.assert_called_once() - call_kwargs = mock_ai_project_client.agents.create_agent.call_args[1] + mock_agents_client.create_agent.assert_called_once() + call_kwargs = mock_agents_client.create_agent.call_args[1] assert call_kwargs["tool_resources"] == {"vector_store_ids": ["vs-123"]} async def test_azure_ai_chat_client_create_agent_stream_submit_tool_outputs( - mock_ai_project_client: MagicMock, + mock_agents_client: MagicMock, ) -> None: """Test _create_agent_stream with tool outputs submission path.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Mock active thread run that matches the tool run ID mock_thread_run = MagicMock() @@ -1380,7 +1294,7 @@ async def test_azure_ai_chat_client_create_agent_stream_submit_tool_outputs( # Mock submit_tool_outputs_stream mock_handler = MagicMock() - mock_ai_project_client.agents.runs.submit_tool_outputs_stream = AsyncMock() + mock_agents_client.runs.submit_tool_outputs_stream = AsyncMock() with patch("azure.ai.agents.models.AsyncAgentEventHandler", return_value=mock_handler): stream, final_thread_id = await chat_client._create_agent_stream( # type: ignore @@ -1388,13 +1302,13 @@ async def test_azure_ai_chat_client_create_agent_stream_submit_tool_outputs( ) # Should call submit_tool_outputs_stream since we have matching run ID - mock_ai_project_client.agents.runs.submit_tool_outputs_stream.assert_called_once() + mock_agents_client.runs.submit_tool_outputs_stream.assert_called_once() assert final_thread_id == "test-thread" -def test_azure_ai_chat_client_extract_url_citations_with_citations(mock_ai_project_client: MagicMock) -> None: +def test_azure_ai_chat_client_extract_url_citations_with_citations(mock_agents_client: MagicMock) -> None: """Test _extract_url_citations with MessageDeltaChunk containing URL citations.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") + chat_client = create_test_azure_ai_chat_client(mock_agents_client, agent_id="test-agent") # Create mock URL citation annotation mock_url_citation = MagicMock() @@ -1437,28 +1351,6 @@ def test_azure_ai_chat_client_extract_url_citations_with_citations(mock_ai_proje assert citation.annotated_regions[0].end_index == 20 -async def test_azure_ai_chat_client_setup_azure_ai_observability_resource_not_found( - mock_ai_project_client: MagicMock, -) -> None: - """Test setup_azure_ai_observability when Application Insights connection string is not found.""" - chat_client = create_test_azure_ai_chat_client(mock_ai_project_client, agent_id="test-agent") - - # Mock telemetry.get_application_insights_connection_string to raise ResourceNotFoundError - mock_ai_project_client.telemetry.get_application_insights_connection_string = AsyncMock( - side_effect=ResourceNotFoundError("No Application Insights found") - ) - - # Mock logger.warning to capture the warning message - with patch("agent_framework_azure_ai._chat_client.logger") as mock_logger: - await chat_client.setup_azure_ai_observability() - - # Verify warning was logged - mock_logger.warning.assert_called_once_with( - "No Application Insights connection string found for the Azure AI Project, " - "please call setup_observability() manually." - ) - - def get_weather( location: Annotated[str, Field(description="The location to get the weather for.")], ) -> str: @@ -1483,7 +1375,7 @@ async def test_azure_ai_chat_client_get_response() -> None: ) messages.append(ChatMessage(role="user", text="What's the weather like today?")) - # Test that the project_client can be used to get a response + # Test that the agents_client can be used to get a response response = await azure_ai_chat_client.get_response(messages=messages) assert response is not None @@ -1501,7 +1393,7 @@ async def test_azure_ai_chat_client_get_response_tools() -> None: messages: list[ChatMessage] = [] messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?")) - # Test that the project_client can be used to get a response + # Test that the agents_client can be used to get a response response = await azure_ai_chat_client.get_response( messages=messages, tools=[get_weather], @@ -1530,7 +1422,7 @@ async def test_azure_ai_chat_client_streaming() -> None: ) messages.append(ChatMessage(role="user", text="What's the weather like today?")) - # Test that the project_client can be used to get a response + # Test that the agents_client can be used to get a response response = azure_ai_chat_client.get_streaming_response(messages=messages) full_message: str = "" @@ -1554,7 +1446,7 @@ async def test_azure_ai_chat_client_streaming_tools() -> None: messages: list[ChatMessage] = [] messages.append(ChatMessage(role="user", text="What's the weather like in Seattle?")) - # Test that the project_client can be used to get a response + # Test that the agents_client can be used to get a response response = azure_ai_chat_client.get_streaming_response( messages=messages, tools=[get_weather], @@ -1704,10 +1596,8 @@ async def test_azure_ai_chat_client_agent_file_search(): try: # 1. Read and upload the test file to the Azure AI agent service test_file_path = Path(__file__).parent / "resources" / "employees.pdf" - file = await client.project_client.agents.files.upload_and_poll( - file_path=str(test_file_path), purpose="assistants" - ) - vector_store = await client.project_client.agents.vector_stores.create_and_poll( + file = await client.agents_client.files.upload_and_poll(file_path=str(test_file_path), purpose="assistants") + vector_store = await client.agents_client.vector_stores.create_and_poll( file_ids=[file.id], name="test_employees_vectorstore" ) @@ -1732,9 +1622,9 @@ async def test_azure_ai_chat_client_agent_file_search(): # 4. Cleanup: Delete the vector store and file try: if vector_store: - await client.project_client.agents.vector_stores.delete(vector_store.id) + await client.agents_client.vector_stores.delete(vector_store.id) if file: - await client.project_client.agents.files.delete(file.id) + await client.agents_client.files.delete(file.id) except Exception: # Ignore cleanup errors to avoid masking the actual test failure pass @@ -1863,10 +1753,10 @@ async def test_azure_ai_chat_client_agent_chat_options_agent_level() -> None: async def test_azure_ai_chat_client_cleanup_agent_when_enabled_and_created( - mock_ai_project_client: MagicMock, + 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_ai_project_client, agent_id=None, should_cleanup_agent=True) + 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" @@ -1875,16 +1765,16 @@ async def test_azure_ai_chat_client_cleanup_agent_when_enabled_and_created( await chat_client._cleanup_agent_if_needed() # type: ignore # Verify agent was deleted - mock_ai_project_client.agents.delete_agent.assert_called_once_with("created-agent-id") + 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_ai_project_client: MagicMock, + 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_ai_project_client, agent_id=None, 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" @@ -1893,17 +1783,17 @@ async def test_azure_ai_chat_client_no_cleanup_when_disabled( await chat_client._cleanup_agent_if_needed() # type: ignore # Verify agent was NOT deleted - mock_ai_project_client.agents.delete_agent.assert_not_called() + 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_ai_project_client: MagicMock, + 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_ai_project_client, agent_id="existing-agent-id", should_cleanup_agent=True + mock_agents_client, agent_id="existing-agent-id", should_cleanup_agent=True ) # Agent exists but was not created by this client (_agent_created = False) @@ -1912,5 +1802,5 @@ async def test_azure_ai_chat_client_no_cleanup_when_agent_not_created_by_client( await chat_client._cleanup_agent_if_needed() # type: ignore # Verify agent was NOT deleted - mock_ai_project_client.agents.delete_agent.assert_not_called() + mock_agents_client.delete_agent.assert_not_called() assert chat_client.agent_id == "existing-agent-id" diff --git a/python/packages/lab/gaia/samples/azure_ai_agent.py b/python/packages/lab/gaia/samples/azure_ai_agent.py index 2fa4ce7e92..78b9c49327 100644 --- a/python/packages/lab/gaia/samples/azure_ai_agent.py +++ b/python/packages/lab/gaia/samples/azure_ai_agent.py @@ -10,8 +10,6 @@ Required Environment Variables: AZURE_AI_MODEL_DEPLOYMENT_NAME: Name of the model deployment to use Optional Environment Variables: - BING_CONNECTION_NAME: Name of the Bing connection for web search - OR BING_CONNECTION_ID: ID of the Bing connection for web search Authentication: @@ -21,7 +19,7 @@ Authentication: Example: export AZURE_AI_PROJECT_ENDPOINT="https://your-project.azure.com" export AZURE_AI_MODEL_DEPLOYMENT_NAME="gpt-4o" - export BING_CONNECTION_NAME="bing-grounding-connection" + export BING_CONNECTION_ID="connection-id" az login """ diff --git a/python/samples/getting_started/agents/azure_ai/README.md b/python/samples/getting_started/agents/azure_ai/README.md index 77faa26192..375f682474 100644 --- a/python/samples/getting_started/agents/azure_ai/README.md +++ b/python/samples/getting_started/agents/azure_ai/README.md @@ -38,10 +38,8 @@ Before running the examples, you need to set up your environment variables. You AZURE_AI_MODEL_DEPLOYMENT_NAME="your-model-deployment-name" ``` -3. For samples using Bing Grounding search (like `azure_ai_with_bing_grounding.py` and `azure_ai_with_multiple_tools.py`), you'll also need either: +3. For samples using Bing Grounding search (like `azure_ai_with_bing_grounding.py` and `azure_ai_with_multiple_tools.py`), you'll also need: ``` - BING_CONNECTION_NAME="bing-grounding-connection" - # OR BING_CONNECTION_ID="your-bing-connection-id" ``` @@ -49,7 +47,7 @@ Before running the examples, you need to set up your environment variables. You - Go to [Azure AI Foundry portal](https://ai.azure.com) - Navigate to your project's "Connected resources" section - Add a new connection for "Grounding with Bing Search" - - Copy either the connection name or ID + - Copy the ID ### Option 2: Using environment variables directly @@ -58,9 +56,7 @@ Set the environment variables in your shell: ```bash export AZURE_AI_PROJECT_ENDPOINT="your-project-endpoint" export AZURE_AI_MODEL_DEPLOYMENT_NAME="your-model-deployment-name" -export BING_CONNECTION_NAME="your-bing-connection-name" # Optional, only needed for web search samples -# OR -export BING_CONNECTION_ID="your-bing-connection-id" # Alternative to BING_CONNECTION_NAME +export BING_CONNECTION_ID="your-bing-connection-id" ``` ### Required Variables @@ -70,4 +66,4 @@ export BING_CONNECTION_ID="your-bing-connection-id" # Alternative to BING_CONNE ### Optional Variables -- `BING_CONNECTION_NAME` or `BING_CONNECTION_ID`: Your Bing connection name or ID (required for `azure_ai_with_bing_grounding.py` and `azure_ai_with_multiple_tools.py`) +- `BING_CONNECTION_ID`: Your Bing connection ID (required for `azure_ai_with_bing_grounding.py` and `azure_ai_with_multiple_tools.py`) diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_azure_ai_search.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_azure_ai_search.py index c925514c89..7d094089bc 100644 --- a/python/samples/getting_started/agents/azure_ai/azure_ai_with_azure_ai_search.py +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_azure_ai_search.py @@ -5,6 +5,7 @@ import os from agent_framework import ChatAgent, CitationAnnotation from agent_framework.azure import AzureAIAgentClient +from azure.ai.agents.aio import AgentsClient from azure.ai.projects.aio import AIProjectClient from azure.ai.projects.models import ConnectionType from azure.identity.aio import AzureCliCredential @@ -38,16 +39,17 @@ async def main() -> None: # Create the client and manually create an agent with Azure AI Search tool async with ( AzureCliCredential() as credential, - AIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as client, + AIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as project_client, + AgentsClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as agents_client, ): ai_search_conn_id = "" - async for connection in client.connections.list(): + async for connection in project_client.connections.list(): if connection.type == ConnectionType.AZURE_AI_SEARCH: ai_search_conn_id = connection.id break # 1. Create Azure AI agent with the search tool - azure_ai_agent = await client.agents.create_agent( + azure_ai_agent = await project_client.agents.create_agent( model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], name="HotelSearchAgent", instructions=( @@ -69,7 +71,7 @@ async def main() -> None: ) # 2. Create chat client with the existing agent - chat_client = AzureAIAgentClient(project_client=client, agent_id=azure_ai_agent.id) + chat_client = AzureAIAgentClient(agents_client=agents_client, agent_id=azure_ai_agent.id) try: async with ChatAgent( @@ -112,7 +114,7 @@ async def main() -> None: finally: # Clean up the agent manually - await client.agents.delete_agent(azure_ai_agent.id) + await project_client.agents.delete_agent(azure_ai_agent.id) if __name__ == "__main__": diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_bing_grounding.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_bing_grounding.py index 6972eab820..462bf26f7b 100644 --- a/python/samples/getting_started/agents/azure_ai/azure_ai_with_bing_grounding.py +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_bing_grounding.py @@ -12,8 +12,7 @@ uses Bing Grounding search to find real-time information from the web. Prerequisites: 1. A connected Grounding with Bing Search resource in your Azure AI project -2. Set either BING_CONNECTION_NAME or BING_CONNECTION_ID environment variable - Example: BING_CONNECTION_NAME="bing-grounding-connection" +2. Set BING_CONNECTION_ID environment variable Example: BING_CONNECTION_ID="your-bing-connection-id" To set up Bing Grounding: @@ -27,7 +26,7 @@ To set up Bing Grounding: async def main() -> None: """Main function demonstrating Azure AI agent with Bing Grounding search.""" # 1. Create Bing Grounding search tool using HostedWebSearchTool - # The connection_name or ID will be automatically picked up from environment variable + # The connection ID will be automatically picked up from environment variable bing_search_tool = HostedWebSearchTool( name="Bing Grounding Search", description="Search the web for current information using Bing", diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_existing_agent.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_existing_agent.py index cbbfca486e..f0fc2c79fc 100644 --- a/python/samples/getting_started/agents/azure_ai/azure_ai_with_existing_agent.py +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_existing_agent.py @@ -5,6 +5,7 @@ import os from agent_framework import ChatAgent from agent_framework.azure import AzureAIAgentClient +from azure.ai.agents.aio import AgentsClient from azure.ai.projects.aio import AIProjectClient from azure.identity.aio import AzureCliCredential @@ -22,16 +23,17 @@ async def main() -> None: # Create the client async with ( AzureCliCredential() as credential, - AIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as client, + AIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as project_client, + AgentsClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as agents_client, ): - azure_ai_agent = await client.agents.create_agent( + azure_ai_agent = await project_client.agents.create_agent( model=os.environ["AZURE_AI_MODEL_DEPLOYMENT_NAME"], # Create remote agent with default instructions # These instructions will persist on created agent for every run. instructions="End each response with [END].", ) - chat_client = AzureAIAgentClient(project_client=client, agent_id=azure_ai_agent.id) + chat_client = AzureAIAgentClient(agents_client=agents_client, agent_id=azure_ai_agent.id) try: async with ChatAgent( @@ -50,7 +52,7 @@ async def main() -> None: print(f"Agent: {result}\n") finally: # Clean up the agent manually - await client.agents.delete_agent(azure_ai_agent.id) + await project_client.agents.delete_agent(azure_ai_agent.id) if __name__ == "__main__": diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_existing_thread.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_existing_thread.py index 4ebd462656..b96b6e5686 100644 --- a/python/samples/getting_started/agents/azure_ai/azure_ai_with_existing_thread.py +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_existing_thread.py @@ -7,7 +7,7 @@ from typing import Annotated from agent_framework import ChatAgent from agent_framework.azure import AzureAIAgentClient -from azure.ai.projects.aio import AIProjectClient +from azure.ai.agents.aio import AgentsClient from azure.identity.aio import AzureCliCredential from pydantic import Field @@ -33,16 +33,16 @@ async def main() -> None: # Create the client async with ( AzureCliCredential() as credential, - AIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as client, + AgentsClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as agents_client, ): # Create an thread that will persist - created_thread = await client.agents.threads.create() + created_thread = await agents_client.threads.create() try: async with ChatAgent( # passing in the client is optional here, so if you take the agent_id from the portal # you can use it directly without the two lines above. - chat_client=AzureAIAgentClient(project_client=client), + chat_client=AzureAIAgentClient(agents_client=agents_client), instructions="You are a helpful weather agent.", tools=get_weather, ) as agent: @@ -52,7 +52,7 @@ async def main() -> None: print(f"Result: {result}\n") finally: # Clean up the thread manually - await client.agents.threads.delete(created_thread.id) + await agents_client.threads.delete(created_thread.id) if __name__ == "__main__": diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_hosted_mcp.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_hosted_mcp.py index dde7836445..fb9d13323e 100644 --- a/python/samples/getting_started/agents/azure_ai/azure_ai_with_hosted_mcp.py +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_hosted_mcp.py @@ -44,8 +44,6 @@ async def main() -> None: AzureCliCredential() as credential, AzureAIAgentClient(async_credential=credential) as chat_client, ): - # enable azure-ai observability - await chat_client.setup_azure_ai_observability() agent = chat_client.create_agent( name="DocsAgent", instructions="You are a helpful assistant that can help with microsoft documentation questions.", diff --git a/python/samples/getting_started/agents/azure_ai/azure_ai_with_multiple_tools.py b/python/samples/getting_started/agents/azure_ai/azure_ai_with_multiple_tools.py index 9e3d088343..7d8a226f80 100644 --- a/python/samples/getting_started/agents/azure_ai/azure_ai_with_multiple_tools.py +++ b/python/samples/getting_started/agents/azure_ai/azure_ai_with_multiple_tools.py @@ -69,8 +69,6 @@ async def main() -> None: AzureCliCredential() as credential, AzureAIAgentClient(async_credential=credential) as chat_client, ): - # enable azure-ai observability - await chat_client.setup_azure_ai_observability() agent = chat_client.create_agent( name="DocsAgent", instructions="You are a helpful assistant that can help with microsoft documentation questions.", diff --git a/python/samples/getting_started/observability/azure_ai_agent_observability.py b/python/samples/getting_started/observability/azure_ai_agent_observability.py index d30b7569c7..e236e43ef2 100644 --- a/python/samples/getting_started/observability/azure_ai_agent_observability.py +++ b/python/samples/getting_started/observability/azure_ai_agent_observability.py @@ -9,7 +9,9 @@ import dotenv from agent_framework import ChatAgent from agent_framework.azure import AzureAIAgentClient from agent_framework.observability import get_tracer +from azure.ai.agents.aio import AgentsClient from azure.ai.projects.aio import AIProjectClient +from azure.core.exceptions import ResourceNotFoundError from azure.identity.aio import AzureCliCredential from opentelemetry.trace import SpanKind from opentelemetry.trace.span import format_trace_id @@ -38,16 +40,36 @@ async def get_weather( return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." +async def setup_azure_ai_observability( + project_client: AIProjectClient, enable_sensitive_data: bool | None = None +) -> None: + """Use this method to setup tracing in your Azure AI Project. + + This will take the connection string from the AIProjectClient. + It will override any connection string that is set in the environment variables. + It will disable any OTLP endpoint that might have been set. + """ + try: + conn_string = await project_client.telemetry.get_application_insights_connection_string() + except ResourceNotFoundError: + print("No Application Insights connection string found for the Azure AI Project.") + return + from agent_framework.observability import setup_observability + + setup_observability(applicationinsights_connection_string=conn_string, enable_sensitive_data=enable_sensitive_data) + + async def main(): async with ( AzureCliCredential() as credential, - AIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as project, - AzureAIAgentClient(project_client=project) as client, + AIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as project_client, + AgentsClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as agents_client, + AzureAIAgentClient(agents_client=agents_client) as client, ): # This will enable tracing and configure the application to send telemetry data to the # Application Insights instance attached to the Azure AI project. # This will override any existing configuration. - await client.setup_azure_ai_observability() + await setup_azure_ai_observability(project_client) questions = ["What's the weather in Amsterdam?", "and in Paris, and which is better?", "Why is the sky blue?"] diff --git a/python/samples/getting_started/observability/azure_ai_chat_client_with_observability.py b/python/samples/getting_started/observability/azure_ai_chat_client_with_observability.py index 2dafa753da..492e10fe51 100644 --- a/python/samples/getting_started/observability/azure_ai_chat_client_with_observability.py +++ b/python/samples/getting_started/observability/azure_ai_chat_client_with_observability.py @@ -9,7 +9,9 @@ import dotenv from agent_framework import HostedCodeInterpreterTool from agent_framework.azure import AzureAIAgentClient from agent_framework.observability import get_tracer +from azure.ai.agents.aio import AgentsClient from azure.ai.projects.aio import AIProjectClient +from azure.core.exceptions import ResourceNotFoundError from azure.identity.aio import AzureCliCredential from opentelemetry.trace import SpanKind from opentelemetry.trace.span import format_trace_id @@ -42,6 +44,25 @@ async def get_weather( return f"The weather in {location} is {conditions[randint(0, 3)]} with a high of {randint(10, 30)}°C." +async def setup_azure_ai_observability( + project_client: AIProjectClient, enable_sensitive_data: bool | None = None +) -> None: + """Use this method to setup tracing in your Azure AI Project. + + This will take the connection string from the AIProjectClient instance. + It will override any connection string that is set in the environment variables. + It will disable any OTLP endpoint that might have been set. + """ + try: + conn_string = await project_client.telemetry.get_application_insights_connection_string() + except ResourceNotFoundError: + print("No Application Insights connection string found for the Azure AI Project.") + return + from agent_framework.observability import setup_observability + + setup_observability(applicationinsights_connection_string=conn_string, enable_sensitive_data=enable_sensitive_data) + + async def main() -> None: """Run an AI service. @@ -62,13 +83,14 @@ async def main() -> None: ] async with ( AzureCliCredential() as credential, - AIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as project, - AzureAIAgentClient(project_client=project) as client, + AIProjectClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as project_client, + AgentsClient(endpoint=os.environ["AZURE_AI_PROJECT_ENDPOINT"], credential=credential) as agents_client, + AzureAIAgentClient(agents_client=agents_client) as client, ): # This will enable tracing and configure the application to send telemetry data to the # Application Insights instance attached to the Azure AI project. # This will override any existing configuration. - await client.setup_azure_ai_observability() + await setup_azure_ai_observability(project_client) with get_tracer().start_as_current_span( name="Foundry Telemetry from Agent Framework", kind=SpanKind.CLIENT