set_current_environment to support openai > 1.0 (Azure#33636)

* set_current_environment to support openai 1.0

* Changelog

* Update azure-ai-ml version

* Update azure-ai-ml version

* Stop using pkg_resources

Author: luigiw

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

@@ -4,6 +4,8 @@
 
 ### Features Added
 
+- AzureOpenAIConnection.set_current_environment supports openai 1.0 and above.
+
 ### Breaking Changes
 
 ### Bugs Fixed

sdk/ai/azure-ai-resources/azure/ai/resources/entities/connection_subtypes.py

@@ -3,6 +3,7 @@
 # ---------------------------------------------------------
 
 
+import os
 from typing import Optional
 
 from azure.ai.ml.entities._credentials import ApiKeyConfiguration
@@ -99,16 +100,33 @@ def set_current_environment(self, credential: Optional[TokenCredential] = None):
         :type credential: :class:`~azure.core.credentials.TokenCredential`
         """
 
-        import os
-        def get_api_version_case_insensitive(connection):
-            if connection.api_version == None:
-                raise ValueError(f"Connection {connection.name} is being used to set environment variables, but lacks required api_version")
-            return connection.api_version.lower()
+        from importlib.metadata import version as get_version
+        from packaging.version import Version
 
-        try:
-            import openai
-        except ImportError:
-            raise Exception("OpenAI SDK not installed. Please install it using `pip install openai`")
+        openai_version_str = get_version("openai")
+        openai_version = Version(openai_version_str)
+        if openai_version >= Version("1.0.0"):
+            self._set_current_environment_new(credential)
+        else:
+            self._set_current_environment_old(credential)
+
+    def _get_api_version_case_insensitive(self, connection):
+        if connection.api_version == None:
+            raise ValueError(f"Connection {connection.name} is being used to set environment variables, but lacks required api_version")
+        return connection.api_version.lower()
+
+    def _set_current_environment_new(self, credential: Optional[TokenCredential] = None):
+        if not credential:
+            os.environ["AZURE_OPENAI_API_KEY"] = self._workspace_connection.credentials.key
+        else:
+            token = credential.get_token("https://cognitiveservices.azure.com/.default")
+            os.environ["AZURE_OPENAI_AD_TOKEN"] = token.token
+
+        os.environ["OPENAI_API_VERSION"] = self._get_api_version_case_insensitive(self._workspace_connection)
+        os.environ["AZURE_OPENAI_ENDPOINT"] = self._workspace_connection.target
+
+    def _set_current_environment_old(self, credential: Optional[TokenCredential] = None):
+        import openai
 
         if not credential:
             openai.api_type = "azure"
@@ -123,13 +141,12 @@ def get_api_version_case_insensitive(connection):
             openai.api_key = token.token
             os.environ["OPENAI_API_KEY"] = token.token
 
-        openai.api_version = get_api_version_case_insensitive(self._workspace_connection)
+        openai.api_version = self._get_api_version_case_insensitive(self._workspace_connection)
 
         openai.api_base = self._workspace_connection.target
 
         os.environ["OPENAI_API_BASE"] = self._workspace_connection.target
-        os.environ["OPENAI_API_VERSION"] = get_api_version_case_insensitive(self._workspace_connection)
-
+        os.environ["OPENAI_API_VERSION"] = self._get_api_version_case_insensitive(self._workspace_connection)
 
 class AzureAISearchConnection(BaseConnection):
     """A Connection for Azure AI Search

sdk/ai/azure-ai-resources/setup.py

@@ -66,7 +66,7 @@
     python_requires="<4.0,>=3.8",
     install_requires=[
         # NOTE: To avoid breaking changes in a major version bump, all dependencies should pin an upper bound if possible.
-        "azure-ai-ml>=1.12.0",
+        "azure-ai-ml>=1.13.0",
         "mlflow-skinny<3",
         "opencensus-ext-logging<=0.1.1",
         "azure-mgmt-resource<23.0.0,>=22.0.0",

sdk/ai/azure-ai-resources/tests/connections/unittests/test_connection.py

@@ -0,0 +1,59 @@
+import pytest
+import unittest
+from unittest.mock import Mock, patch
+
+from azure.ai.resources.entities import AzureOpenAIConnection
+from azure.ai.ml.entities._credentials import ApiKeyConfiguration
+
+
+@pytest.mark.unittest
+class TestConnection:
+
+    @pytest.fixture()
+    def mock_credentials(self):
+        mock_credentials = Mock()
+        mock_token = Mock()
+        mock_credentials.get_token.return_value = mock_token
+        mock_token.token = "my-aad-token"
+        return mock_credentials
+
+    @patch("importlib.metadata.version")
+    @patch("os.environ")
+    def test_set_environment(self, mock_os_environ, mock_get_version, mock_credentials):
+        mock_get_version.return_value = "0.28.1"
+        connection = AzureOpenAIConnection(name="my-connection", target="https://test.com", credentials=ApiKeyConfiguration(key="abc"), api_version="2021-01-01")
+
+        connection.set_current_environment()
+        assert mock_os_environ.__setitem__.call_count == 4
+        assert mock_os_environ.__setitem__.call_args_list[0][0] == ("OPENAI_API_KEY", "abc")
+        assert mock_os_environ.__setitem__.call_args_list[1][0] == ("OPENAI_API_TYPE", "azure")
+        assert mock_os_environ.__setitem__.call_args_list[2][0] == ("OPENAI_API_BASE", "https://test.com")
+        assert mock_os_environ.__setitem__.call_args_list[3][0] == ("OPENAI_API_VERSION", "2021-01-01")
+
+        connection.set_current_environment(mock_credentials)
+        assert mock_os_environ.__setitem__.call_count == 8 
+        assert mock_os_environ.__setitem__.call_args_list[4][0] == ("OPENAI_API_TYPE", "azure_ad")
+        assert mock_os_environ.__setitem__.call_args_list[5][0] == ("OPENAI_API_KEY", "my-aad-token")
+        assert mock_os_environ.__setitem__.call_args_list[6][0] == ("OPENAI_API_BASE", "https://test.com")
+        assert mock_os_environ.__setitem__.call_args_list[7][0] == ("OPENAI_API_VERSION", "2021-01-01")
+
+
+    @patch("importlib.metadata.version")
+    @patch("os.environ")
+    def test_set_environment_100(self, mock_os_environ, mock_get_version, mock_credentials):
+        mock_get_version.return_value = "1.0.0"
+        connection = AzureOpenAIConnection(name="my-connection", target="https://test.com", credentials=ApiKeyConfiguration(key="abc"), api_version="2021-01-01")
+
+        connection.set_current_environment()
+        # This is a hacky way to verity os.environ, any better way?
+        assert mock_os_environ.__setitem__.call_count == 3
+        # call_args[0] is args, call_args[1] is kwargs
+        assert mock_os_environ.__setitem__.call_args_list[0][0] == ("AZURE_OPENAI_API_KEY", "abc")
+        assert mock_os_environ.__setitem__.call_args_list[1][0] == ("OPENAI_API_VERSION", "2021-01-01")
+        assert mock_os_environ.__setitem__.call_args_list[2][0] == ("AZURE_OPENAI_ENDPOINT", "https://test.com")
+
+        connection.set_current_environment(mock_credentials)
+        assert mock_os_environ.__setitem__.call_count == 6
+        assert mock_os_environ.__setitem__.call_args_list[3][0] == ("AZURE_OPENAI_AD_TOKEN", "my-aad-token")
+        assert mock_os_environ.__setitem__.call_args_list[4][0] == ("OPENAI_API_VERSION", "2021-01-01")
+        assert mock_os_environ.__setitem__.call_args_list[5][0] == ("AZURE_OPENAI_ENDPOINT", "https://test.com")

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

@@ -1,6 +1,16 @@
 # Release History
 
-## 1.12.0 (unreleased)
+## 1.13.0 (unreleased)
+
+### Features Added
+
+### Bugs Fixed
+
+### Breaking Changes
+
+### Other Changes
+
+## 1.12.0 (2023-11-13)
 
 ### Features Added
 - Workspace Connections had 3 child classes added for open AI, cog search, and cog service connections.

sdk/ml/azure-ai-ml/azure/ai/ml/_version.py

@@ -2,4 +2,4 @@
 # Copyright (c) Microsoft Corporation. All rights reserved.
 # ---------------------------------------------------------
 
-VERSION = "1.12.0"
+VERSION = "1.13.0"

sdk/ml/azure-ai-ml/azure/ai/ml/entities/_resource.py

@@ -54,9 +54,7 @@ def __init__(
         self.tags = dict(tags) if tags else {}
         self.properties = dict(properties) if properties else {}
         # Conditional assignment to prevent entity bloat when unused.
-        print_as_yaml = kwargs.pop("print_as_yaml", in_jupyter_notebook())
-        if print_as_yaml:
-            self.print_as_yaml = True
+        self._print_as_yaml = kwargs.pop("print_as_yaml", False)
 
         # Hide read only properties in kwargs
         self._id = kwargs.pop("id", None)
@@ -194,7 +192,7 @@ def __repr__(self) -> str:
         return f"{self.__class__.__name__}({var_dict})"
 
     def __str__(self) -> str:
-        if hasattr(self, "print_as_yaml") and self.print_as_yaml:
+        if self._print_as_yaml or in_jupyter_notebook():
             # pylint: disable=no-member
             yaml_serialized = self._to_dict()
             return dump_yaml(yaml_serialized, default_flow_style=False)