Beta 2 - Azure AI Projects SDK (Azure#38578)

Author: dargilco

sdk/ai/azure-ai-projects/CHANGELOG.md

@@ -1,12 +1,14 @@
 # Release History
 
-## 1.0.0b2 (Unreleased)
+## 1.0.0b2 (2024-12-03)
 
 ### Bugs Fixed
 
-* Fixed bugs that were causing issues with tracing agent asynchronous functionality
-* Fix a bug causing warning about unclosed session, shown when using asynchronous credentials to create agent. 
-* Fix a bug that would cause agent function tool related function names and parameters to be included in traces even when content recording is not enabled
+* Fix a bug in the `.inference` operations when Entra ID authentication is used by the default connection.
+* Fixed bugs occurring during streaming in function tool calls by asynchronous agents.
+* Fixed bugs that were causing issues with tracing agent asynchronous functionality.
+* Fix a bug causing warning about unclosed session, shown when using asynchronous credentials to create agent.
+* Fix a bug that would cause agent function tool related function names and parameters to be included in traces even when content recording is not enabled.
 
 ## 1.0.0b1 (2024-11-15)
 

sdk/ai/azure-ai-projects/assets.json

@@ -2,5 +2,5 @@
   "AssetsRepo": "Azure/azure-sdk-assets",
   "AssetsRepoPrefixPath": "python",
   "TagPrefix": "python/ai/azure-ai-projects",
-  "Tag": "python/ai/azure-ai-projects_a5498ae251"
+  "Tag": "python/ai/azure-ai-projects_84b9cfdbbd"
 }

sdk/ai/azure-ai-projects/azure/ai/projects/aio/operations/_patch.py

@@ -97,6 +97,7 @@ async def get_chat_completions_client(self, **kwargs) -> "ChatCompletionsClient"
             connection = await self._outer_instance.connections.get_default(
                 connection_type=ConnectionType.AZURE_AI_SERVICES, include_credentials=True, **kwargs
             )
+        logger.debug("[InferenceOperations.get_chat_completions_client] connection = %s", str(connection))
 
         try:
             from azure.ai.inference.aio import ChatCompletionsClient
@@ -107,11 +108,10 @@ async def get_chat_completions_client(self, **kwargs) -> "ChatCompletionsClient"
 
         if use_serverless_connection:
             endpoint = connection.endpoint_url
+            credential_scopes = ["https://ml.azure.com/.default"]
         else:
-            # Be sure to use the Azure resource name here, not the connection name. Connection name is something that
-            # admins can pick when they manually create a new connection (or use bicep). Get the Azure resource name
-            # from the end of the connection id.
-            endpoint = f"https://{connection.id.split('/')[-1]}.services.ai.azure.com/models"
+            endpoint = f"{connection.endpoint_url}/models"
+            credential_scopes = ["https://cognitiveservices.azure.com/.default"]
 
         if connection.authentication_type == AuthenticationType.API_KEY:
             logger.debug(
@@ -122,12 +122,13 @@ async def get_chat_completions_client(self, **kwargs) -> "ChatCompletionsClient"
 
             client = ChatCompletionsClient(endpoint=endpoint, credential=AzureKeyCredential(connection.key))
         elif connection.authentication_type == AuthenticationType.ENTRA_ID:
-            # MaaS models do not yet support EntraID auth
             logger.debug(
                 "[InferenceOperations.get_chat_completions_client]"
                 + " Creating ChatCompletionsClient using Entra ID authentication"
             )
-            client = ChatCompletionsClient(endpoint=endpoint, credential=connection.properties.token_credential)
+            client = ChatCompletionsClient(
+                endpoint=endpoint, credential=connection.token_credential, credential_scopes=credential_scopes
+            )
         elif connection.authentication_type == AuthenticationType.SAS:
             logger.debug(
                 "[InferenceOperations.get_chat_completions_client] "
@@ -172,6 +173,7 @@ async def get_embeddings_client(self, **kwargs) -> "EmbeddingsClient":
             connection = await self._outer_instance.connections.get_default(
                 connection_type=ConnectionType.AZURE_AI_SERVICES, include_credentials=True, **kwargs
             )
+        logger.debug("[InferenceOperations.get_embeddings_client] connection = %s", str(connection))
 
         try:
             from azure.ai.inference.aio import EmbeddingsClient
@@ -182,11 +184,10 @@ async def get_embeddings_client(self, **kwargs) -> "EmbeddingsClient":
 
         if use_serverless_connection:
             endpoint = connection.endpoint_url
+            credential_scopes = ["https://ml.azure.com/.default"]
         else:
-            # Be sure to use the Azure resource name here, not the connection name. Connection name is something that
-            # admins can pick when they manually create a new connection (or use bicep). Get the Azure resource name
-            # from the end of the connection id.
-            endpoint = f"https://{connection.id.split('/')[-1]}.services.ai.azure.com/models"
+            endpoint = f"{connection.endpoint_url}/models"
+            credential_scopes = ["https://cognitiveservices.azure.com/.default"]
 
         if connection.authentication_type == AuthenticationType.API_KEY:
             logger.debug(
@@ -196,11 +197,12 @@ async def get_embeddings_client(self, **kwargs) -> "EmbeddingsClient":
 
             client = EmbeddingsClient(endpoint=endpoint, credential=AzureKeyCredential(connection.key))
         elif connection.authentication_type == AuthenticationType.ENTRA_ID:
-            # MaaS models do not yet support EntraID auth
             logger.debug(
                 "[InferenceOperations.get_embeddings_client] Creating EmbeddingsClient using Entra ID authentication"
             )
-            client = EmbeddingsClient(endpoint=endpoint, credential=connection.properties.token_credential)
+            client = EmbeddingsClient(
+                endpoint=endpoint, credential=connection.token_credential, credential_scopes=credential_scopes
+            )
         elif connection.authentication_type == AuthenticationType.SAS:
             logger.debug(
                 "[InferenceOperations.get_embeddings_client] Creating EmbeddingsClient using SAS authentication"
@@ -235,6 +237,7 @@ async def get_azure_openai_client(self, *, api_version: Optional[str] = None, **
         connection = await self._outer_instance.connections.get_default(
             connection_type=ConnectionType.AZURE_OPEN_AI, include_credentials=True, **kwargs
         )
+        logger.debug("[InferenceOperations.get_azure_openai_client] connection = %s", str(connection))
 
         try:
             from openai import AsyncAzureOpenAI
@@ -250,25 +253,30 @@ async def get_azure_openai_client(self, *, api_version: Optional[str] = None, **
             client = AsyncAzureOpenAI(
                 api_key=connection.key, azure_endpoint=connection.endpoint_url, api_version=api_version
             )
-        elif connection.authentication_type in {AuthenticationType.ENTRA_ID, AuthenticationType.SAS}:
+        elif connection.authentication_type == AuthenticationType.ENTRA_ID:
+            logger.debug(
+                "[InferenceOperations.get_azure_openai_client] " + "Creating AzureOpenAI using Entra ID authentication"
+            )
             try:
                 from azure.identity.aio import get_bearer_token_provider
             except ModuleNotFoundError as e:
                 raise ModuleNotFoundError(
                     "azure.identity package not installed. Please install it using 'pip install azure-identity'"
                 ) from e
-            if connection.authentication_type == AuthenticationType.ENTRA_ID:
-                auth = "Creating AzureOpenAI using Entra ID authentication"
-            else:
-                auth = "Creating AzureOpenAI using SAS authentication"
-            logger.debug("[InferenceOperations.get_azure_openai_client] %s", auth)
             client = AsyncAzureOpenAI(
                 azure_ad_token_provider=get_bearer_token_provider(
                     connection.token_credential, "https://cognitiveservices.azure.com/.default"
                 ),
                 azure_endpoint=connection.endpoint_url,
                 api_version=api_version,
             )
+        elif connection.authentication_type == AuthenticationType.SAS:
+            logger.debug(
+                "[InferenceOperations.get_azure_openai_client] " + "Creating AzureOpenAI using SAS authentication"
+            )
+            raise ValueError(
+                "Getting an AzureOpenAI client from a connection with SAS authentication is not yet supported"
+            )
         else:
             raise ValueError("Unknown authentication type")
 
@@ -299,16 +307,17 @@ async def get_default(
         if not connection_type:
             raise ValueError("You must specify an connection type")
         # Since there is no notion of default connection at the moment, list all connections in the category
-        # and return the first one
+        # and return the first one (index 0), unless overridden by the environment variable DEFAULT_CONNECTION_INDEX.
         connection_properties_list = await self.list(connection_type=connection_type, **kwargs)
         if len(connection_properties_list) > 0:
+            default_connection_index = int(os.getenv("DEFAULT_CONNECTION_INDEX", "0"))
             if include_credentials:
                 return await self.get(
-                    connection_name=connection_properties_list[0].name,
+                    connection_name=connection_properties_list[default_connection_index].name,
                     include_credentials=include_credentials,
                     **kwargs,
                 )
-            return connection_properties_list[0]
+            return connection_properties_list[default_connection_index]
         raise ResourceNotFoundError(f"No connection of type {connection_type} found")
 
     @distributed_trace_async

sdk/ai/azure-ai-projects/azure/ai/projects/operations/_patch.py

@@ -87,6 +87,7 @@ def get_chat_completions_client(self, **kwargs) -> "ChatCompletionsClient":
             connection = self._outer_instance.connections.get_default(
                 connection_type=ConnectionType.AZURE_AI_SERVICES, include_credentials=True, **kwargs
             )
+        logger.debug("[InferenceOperations.get_chat_completions_client] connection = %s", str(connection))
 
         try:
             from azure.ai.inference import ChatCompletionsClient
@@ -97,11 +98,10 @@ def get_chat_completions_client(self, **kwargs) -> "ChatCompletionsClient":
 
         if use_serverless_connection:
             endpoint = connection.endpoint_url
+            credential_scopes = ["https://ml.azure.com/.default"]
         else:
-            # Be sure to use the Azure resource name here, not the connection name. Connection name is something that
-            # admins can pick when they manually create a new connection (or use bicep). Get the Azure resource name
-            # from the end of the connection id.
-            endpoint = f"https://{connection.id.split('/')[-1]}.services.ai.azure.com/models"
+            endpoint = f"{connection.endpoint_url}/models"
+            credential_scopes = ["https://cognitiveservices.azure.com/.default"]
 
         if connection.authentication_type == AuthenticationType.API_KEY:
             logger.debug(
@@ -112,12 +112,13 @@ def get_chat_completions_client(self, **kwargs) -> "ChatCompletionsClient":
 
             client = ChatCompletionsClient(endpoint=endpoint, credential=AzureKeyCredential(connection.key))
         elif connection.authentication_type == AuthenticationType.ENTRA_ID:
-            # MaaS models do not yet support EntraID auth
             logger.debug(
                 "[InferenceOperations.get_chat_completions_client] "
                 + "Creating ChatCompletionsClient using Entra ID authentication"
             )
-            client = ChatCompletionsClient(endpoint=endpoint, credential=connection.properties.token_credential)
+            client = ChatCompletionsClient(
+                endpoint=endpoint, credential=connection.token_credential, credential_scopes=credential_scopes
+            )
         elif connection.authentication_type == AuthenticationType.SAS:
             logger.debug(
                 "[InferenceOperations.get_chat_completions_client] "
@@ -161,6 +162,7 @@ def get_embeddings_client(self, **kwargs) -> "EmbeddingsClient":
             connection = self._outer_instance.connections.get_default(
                 connection_type=ConnectionType.AZURE_AI_SERVICES, include_credentials=True, **kwargs
             )
+        logger.debug("[InferenceOperations.get_embeddings_client] connection = %s", str(connection))
 
         try:
             from azure.ai.inference import EmbeddingsClient
@@ -171,11 +173,10 @@ def get_embeddings_client(self, **kwargs) -> "EmbeddingsClient":
 
         if use_serverless_connection:
             endpoint = connection.endpoint_url
+            credential_scopes = ["https://ml.azure.com/.default"]
         else:
-            # Be sure to use the Azure resource name here, not the connection name. Connection name is something that
-            # admins can pick when they manually create a new connection (or use bicep). Get the Azure resource name
-            # from the end of the connection id.
-            endpoint = f"https://{connection.id.split('/')[-1]}.services.ai.azure.com/models"
+            endpoint = f"{connection.endpoint_url}/models"
+            credential_scopes = ["https://cognitiveservices.azure.com/.default"]
 
         if connection.authentication_type == AuthenticationType.API_KEY:
             logger.debug(
@@ -185,11 +186,12 @@ def get_embeddings_client(self, **kwargs) -> "EmbeddingsClient":
 
             client = EmbeddingsClient(endpoint=endpoint, credential=AzureKeyCredential(connection.key))
         elif connection.authentication_type == AuthenticationType.ENTRA_ID:
-            # MaaS models do not yet support EntraID auth
             logger.debug(
                 "[InferenceOperations.get_embeddings_client] Creating EmbeddingsClient using Entra ID authentication"
             )
-            client = EmbeddingsClient(endpoint=endpoint, credential=connection.properties.token_credential)
+            client = EmbeddingsClient(
+                endpoint=endpoint, credential=connection.token_credential, credential_scopes=credential_scopes
+            )
         elif connection.authentication_type == AuthenticationType.SAS:
             logger.debug(
                 "[InferenceOperations.get_embeddings_client] Creating EmbeddingsClient using SAS authentication"
@@ -225,6 +227,7 @@ def get_azure_openai_client(self, *, api_version: Optional[str] = None, **kwargs
         connection = self._outer_instance.connections.get_default(
             connection_type=ConnectionType.AZURE_OPEN_AI, include_credentials=True, **kwargs
         )
+        logger.debug("[InferenceOperations.get_azure_openai_client] connection = %s", str(connection))
 
         try:
             from openai import AzureOpenAI
@@ -240,18 +243,16 @@ def get_azure_openai_client(self, *, api_version: Optional[str] = None, **kwargs
             client = AzureOpenAI(
                 api_key=connection.key, azure_endpoint=connection.endpoint_url, api_version=api_version
             )
-        elif connection.authentication_type == {AuthenticationType.ENTRA_ID, AuthenticationType.SAS}:
+        elif connection.authentication_type == AuthenticationType.ENTRA_ID:
+            logger.debug(
+                "[InferenceOperations.get_azure_openai_client] " + "Creating AzureOpenAI using Entra ID authentication"
+            )
             try:
                 from azure.identity import get_bearer_token_provider
             except ModuleNotFoundError as e:
                 raise ModuleNotFoundError(
                     "azure.identity package not installed. Please install it using 'pip install azure.identity'"
                 ) from e
-            if connection.authentication_type == AuthenticationType.ENTRA_ID:
-                auth = "Creating AzureOpenAI using Entra ID authentication"
-            else:
-                auth = "Creating AzureOpenAI using SAS authentication"
-            logger.debug("[InferenceOperations.get_azure_openai_client] %s", auth)
             client = AzureOpenAI(
                 # See https://learn.microsoft.com/python/api/azure-identity/azure.identity?view=azure-python#azure-identity-get-bearer-token-provider # pylint: disable=line-too-long
                 azure_ad_token_provider=get_bearer_token_provider(
@@ -260,6 +261,13 @@ def get_azure_openai_client(self, *, api_version: Optional[str] = None, **kwargs
                 azure_endpoint=connection.endpoint_url,
                 api_version=api_version,
             )
+        elif connection.authentication_type == AuthenticationType.SAS:
+            logger.debug(
+                "[InferenceOperations.get_azure_openai_client] " + "Creating AzureOpenAI using SAS authentication"
+            )
+            raise ValueError(
+                "Getting an AzureOpenAI client from a connection with SAS authentication is not yet supported"
+            )
         else:
             raise ValueError("Unknown authentication type")
 
@@ -290,16 +298,17 @@ def get_default(
         if not connection_type:
             raise ValueError("You must specify an connection type")
         # Since there is no notion of default connection at the moment, list all connections in the category
-        # and return the first one
+        # and return the first one (index 0), unless overridden by the environment variable DEFAULT_CONNECTION_INDEX.
         connection_properties_list = self.list(connection_type=connection_type, **kwargs)
         if len(connection_properties_list) > 0:
+            default_connection_index = int(os.getenv("DEFAULT_CONNECTION_INDEX", "0"))
             if include_credentials:
                 return self.get(
-                    connection_name=connection_properties_list[0].name,
+                    connection_name=connection_properties_list[default_connection_index].name,
                     include_credentials=include_credentials,
                     **kwargs,
                 )
-            return connection_properties_list[0]
+            return connection_properties_list[default_connection_index]
         raise ResourceNotFoundError(f"No connection of type {connection_type} found")
 
     @distributed_trace
@@ -2081,7 +2090,7 @@ def submit_tool_outputs_to_stream(
 
         elif tool_outputs is not _Unset:
             response = super().submit_tool_outputs_to_run(
-                thread_id, run_id, tool_outputs=tool_outputs, stream_parameter=False, stream=False, **kwargs
+                thread_id, run_id, tool_outputs=tool_outputs, stream_parameter=True, stream=True, **kwargs
             )
 
         elif isinstance(body, io.IOBase):

sdk/ai/azure-ai-projects/pyproject.toml

@@ -17,7 +17,7 @@ exclude = [
   "sample_chat_completions_with_azure_ai_inference_client_and_console_tracing\\.py",
   "sample_chat_completions_with_azure_ai_inference_client_and_azure_monitor_tracing\\.py",
   "sample_chat_completions_with_azure_ai_inference_client\\.py",
-  "sample_connections.py"
+  "sample_inference_client_from_connection.py"
 ]
 warn_unused_configs = true
 ignore_missing_imports = true