[PY] feat: Add managed identity auth support to AssistantsPlanner (#2153)

singhk97 · web-flow · commit 3509d0848e07 · 2024-10-31T13:16:35.000-07:00
## Linked issues closes: #minor tracking: #1918 ## Details Add ability to authenticate `AssistantsPlanner` using Azure Managed Identity by providing a `azure_ad_token_bearer` field in the initializer options. #### Change details Updated samples - math bot - order bot ## Attestation Checklist - [x] My code follows the style guidelines of this project - I have checked for/fixed spelling, linting, and other errors - I have commented my code for clarity - I have made corresponding changes to the documentation (updating the doc strings in the code is sufficient) - My changes generate no new warnings - I have added tests that validates my changes, and provides sufficient test coverage. I have tested with: - Local testing - E2E testing in Teams - New and existing unit tests pass locally with my changes
diff --git a/python/packages/ai/teams/ai/planners/assistants_planner.py b/python/packages/ai/teams/ai/planners/assistants_planner.py
@@ -9,7 +9,7 @@
 import json
 from dataclasses import dataclass
 from importlib.metadata import version
-from typing import Dict, Generic, List, Optional, TypeVar, Union
+from typing import Callable, Dict, Generic, List, Optional, TypeVar, Union
 
 import openai
 from botbuilder.core import TurnContext
@@ -49,10 +49,6 @@ class AzureOpenAIAssistantsOptions:
     """
     Options for configuring the AssistantsPlanner for AzureOpenAI.
     """
-
-    api_key: str
-    "The AzureOpenAI API key."
-
     default_model: str
     "Default name of the Azure OpenAI deployment (model) to use."
 
@@ -62,6 +58,14 @@ class AzureOpenAIAssistantsOptions:
     endpoint: str
     "Deployment endpoint to use."
 
+    api_key: Optional[str] = None
+    "The AzureOpenAI API key."
+
+    azure_ad_token_provider: Optional[Callable[..., str]] = None
+    """Optional. A function that returns an access token for Microsoft Entra 
+    (formerly known as Azure Active Directory), which will be invoked in every request.
+    """
+
     polling_interval: float = DEFAULT_POLLING_INTERVAL
     "Optional. Polling interval in seconds. Defaults to 1 second"
 
@@ -140,6 +144,7 @@ def __init__(
             self._client = openai.AsyncAzureOpenAI(
                 api_key=options.api_key,
                 api_version=options.api_version,
+                azure_ad_token_provider=options.azure_ad_token_provider,
                 azure_endpoint=options.endpoint,
                 organization=options.organization if options.organization else None,
                 default_headers={"User-Agent": self.user_agent},
@@ -196,7 +201,8 @@ async def continue_task(self, context: TurnContext, state: TurnState) -> Plan:
 
     @staticmethod
     async def create_assistant(
-        api_key: str,
+        api_key: Optional[str],
+        azure_ad_token_provider: Optional[Callable[..., str]],
         api_version: Optional[str],
         organization: Optional[str],
         endpoint: Optional[str],
@@ -221,6 +227,7 @@ async def create_assistant(
             user_agent = f"teamsai-py/{version('teams-ai')}"
             client = openai.AsyncAzureOpenAI(
                 api_key=api_key,
+                azure_ad_token_provider=azure_ad_token_provider,
                 api_version=api_version if api_version else "2024-02-15-preview",
                 azure_endpoint=endpoint,
                 organization=organization if organization else None,
diff --git a/python/packages/ai/tests/ai/planners/test_assistants_planner.py b/python/packages/ai/tests/ai/planners/test_assistants_planner.py
@@ -560,7 +560,12 @@ async def test_create_openai_assistant(self, mock_async_openai):
         params = beta.AssistantCreateParams(model="123")
 
         assistant = await AssistantsPlanner.create_assistant(
-            api_key="", api_version="", organization="", endpoint="", request=params
+            api_key="",
+            azure_ad_token_provider=None,
+            api_version="",
+            organization="",
+            endpoint="",
+            request=params
         )
 
         self.assertTrue(mock_async_openai.called)
@@ -573,6 +578,24 @@ async def test_create_azure_openai_assistant(self, mock_async_azure_openai):
 
         assistant = await AssistantsPlanner.create_assistant(
             api_key="",
+            azure_ad_token_provider=None,
+            api_version="",
+            organization="",
+            endpoint="this is my endpoint",
+            request=params,
+        )
+
+        self.assertTrue(mock_async_azure_openai.called)
+        self.assertEqual(assistant.id, ASSISTANT_ID)
+        self.assertEqual(assistant.model, ASSISTANT_MODEL)
+    
+    @mock.patch("openai.AsyncAzureOpenAI", return_value=MockAsyncOpenAI())
+    async def test_create_azure_openai_assistant_with_az_token_provider(self, mock_async_azure_openai):
+        params = beta.AssistantCreateParams(model="123")
+
+        assistant = await AssistantsPlanner.create_assistant(
+            api_key=None,
+            azure_ad_token_provider=lambda: "test-token",
             api_version="",
             organization="",
             endpoint="this is my endpoint",
diff --git a/python/samples/06.assistants.a.mathBot/pyproject.toml b/python/samples/06.assistants.a.mathBot/pyproject.toml
@@ -12,6 +12,7 @@ packages = [
 python = ">=3.8,<4.0"
 teams-ai = "^1.2.2"
 python-dotenv = "^1.0.1"
+azure-identity = "^1.19.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.0"
diff --git a/python/samples/06.assistants.a.mathBot/src/bot.py b/python/samples/06.assistants.a.mathBot/src/bot.py
@@ -22,27 +22,40 @@
 
 from config import Config
 from state import AppTurnState
+from azure.identity import get_bearer_token_provider, DefaultAzureCredential
 
 config = Config()
 
-if config.OPENAI_KEY is None and config.AZURE_OPENAI_KEY is None:
+if config.OPENAI_KEY is None and config.AZURE_OPENAI_ENDPOINT is None:
     raise RuntimeError(
-        "Missing environment variables - please check that OPENAI_KEY or AZURE_OPENAI_KEY is set."
+        "Missing environment variables - please check that OPENAI_KEY or AZURE_OPENAI_ENDPOINT is set."
     )
 
 planner: AssistantsPlanner
 
 # Create Assistant Planner
-if config.AZURE_OPENAI_KEY and config.AZURE_OPENAI_ENDPOINT:
-    planner = AssistantsPlanner[AppTurnState](
-        AzureOpenAIAssistantsOptions(
-            api_key=config.AZURE_OPENAI_KEY,
-            api_version="2024-02-15-preview",
-            endpoint=config.AZURE_OPENAI_ENDPOINT,
-            default_model="gpt-4",
-            assistant_id=config.ASSISTANT_ID,
+if config.AZURE_OPENAI_ENDPOINT:
+    if config.AZURE_OPENAI_KEY:
+        planner = AssistantsPlanner[AppTurnState](
+            AzureOpenAIAssistantsOptions(
+                api_key=config.AZURE_OPENAI_KEY,
+                api_version="2024-05-01-preview",
+                endpoint=config.AZURE_OPENAI_ENDPOINT,
+                default_model="gpt-4o",
+                assistant_id=config.ASSISTANT_ID,
+            )
+        )
+    else:
+        # Managed Identity Auth
+        planner = AssistantsPlanner[AppTurnState](
+            AzureOpenAIAssistantsOptions(
+                azure_ad_token_provider=get_bearer_token_provider(DefaultAzureCredential(), 'https://cognitiveservices.azure.com/.default'),
+                api_version="2024-05-01-preview",
+                endpoint=config.AZURE_OPENAI_ENDPOINT,
+                default_model="gpt-4o",
+                assistant_id=config.ASSISTANT_ID,
+            )
         )
-    )
 else:
     planner = AssistantsPlanner[AppTurnState](
         OpenAIAssistantsOptions(api_key=config.OPENAI_KEY, assistant_id=config.ASSISTANT_ID)
@@ -69,22 +82,34 @@ async def setup_assistant(context: TurnContext, state: AppTurnState):
             name="Math Tutor",
             instructions="You are a personal math tutor. Write and run code to answer math questions.",
             tools=[CodeInterpreterToolParam(type="code_interpreter")],
-            model="gpt-4",
+            model="gpt-4o",
         )
 
         assistant: Assistant
 
-        if config.AZURE_OPENAI_KEY and config.AZURE_OPENAI_ENDPOINT:
-            assistant = await AssistantsPlanner.create_assistant(
-                api_key=config.AZURE_OPENAI_KEY,
-                api_version="",
-                organization="",
-                endpoint=config.AZURE_OPENAI_ENDPOINT,
-                request=params,
-            )
+        if config.AZURE_OPENAI_ENDPOINT:
+            if config.AZURE_OPENAI_KEY:
+                assistant = await AssistantsPlanner.create_assistant(
+                    api_key=config.AZURE_OPENAI_KEY,
+                    azure_ad_token_provider=None,
+                    api_version="2024-05-01-preview",
+                    organization="",
+                    endpoint=config.AZURE_OPENAI_ENDPOINT,
+                    request=params,
+                )
+            else:
+                assistant = await AssistantsPlanner.create_assistant(
+                    api_key=None,
+                    azure_ad_token_provider=get_bearer_token_provider(DefaultAzureCredential(), 'https://cognitiveservices.azure.com/.default'),
+                    api_version="2024-05-01-preview",
+                    organization="",
+                    endpoint=config.AZURE_OPENAI_ENDPOINT,
+                    request=params,
+                )
         else:
             assistant = await AssistantsPlanner.create_assistant(
                 api_key=config.OPENAI_KEY,
+                azure_ad_token_provider=None,
                 api_version="",
                 organization="",
                 endpoint="",
diff --git a/python/samples/06.assistants.a.mathBot/teamsapp.yml b/python/samples/06.assistants.a.mathBot/teamsapp.yml
@@ -97,3 +97,4 @@ deploy:
       # You can replace it with your existing Azure Resource id
       # or add it to your environment variable file.
       resourceId: ${{BOT_AZURE_APP_SERVICE_RESOURCE_ID}}
+projectId: c5366cf6-846e-4dfb-bf8c-4fe0afa40c1d
diff --git a/python/samples/06.assistants.b.orderBot/pyproject.toml b/python/samples/06.assistants.b.orderBot/pyproject.toml
@@ -12,6 +12,7 @@ packages = [
 python = ">=3.8,<4.0"
 teams-ai = "^1.2.2"
 python-dotenv = "^1.0.1"
+azure-identity = "^1.19.0"
 
 [tool.poetry.group.dev.dependencies]
 pytest = "^7.4.0"
diff --git a/python/samples/06.assistants.b.orderBot/src/bot.py b/python/samples/06.assistants.b.orderBot/src/bot.py
@@ -23,6 +23,7 @@
     AzureOpenAIAssistantsOptions,
     OpenAIAssistantsOptions,
 )
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 
 from config import Config
 from food_order_card import generate_card_for_order
@@ -31,24 +32,36 @@
 
 config = Config()
 
-if config.OPENAI_KEY is None and config.AZURE_OPENAI_KEY is None:
+if config.OPENAI_KEY is None and config.AZURE_OPENAI_ENDPOINT is None:
     raise RuntimeError(
-        "Missing environment variables - please check that OPENAI_KEY or AZURE_OPENAI_KEY is set."
+        "Missing environment variables - please check that OPENAI_KEY or AZURE_OPENAI_ENDPOINT is set."
     )
 
 planner: AssistantsPlanner
 
 # Create Assistant Planner
-if config.AZURE_OPENAI_KEY and config.AZURE_OPENAI_ENDPOINT:
-    planner = AssistantsPlanner[AppTurnState](
-        AzureOpenAIAssistantsOptions(
-            api_key=config.AZURE_OPENAI_KEY,
-            api_version="2024-02-15-preview",
-            endpoint=config.AZURE_OPENAI_ENDPOINT,
-            default_model="gpt-4",
-            assistant_id=config.ASSISTANT_ID,
+if config.AZURE_OPENAI_ENDPOINT:
+    if config.AZURE_OPENAI_KEY:
+        planner = AssistantsPlanner[AppTurnState](
+            AzureOpenAIAssistantsOptions(
+                api_key=config.AZURE_OPENAI_KEY,
+                api_version="2024-05-01-preview",
+                endpoint=config.AZURE_OPENAI_ENDPOINT,
+                default_model="gpt-4o",
+                assistant_id=config.ASSISTANT_ID,
+            )
+        )
+    else:
+        # Managed Identity Auth
+        planner = AssistantsPlanner[AppTurnState](
+            AzureOpenAIAssistantsOptions(
+                azure_ad_token_provider=get_bearer_token_provider(DefaultAzureCredential(), 'https://cognitiveservices.azure.com/.default'),
+                api_version="2024-05-01-preview",
+                endpoint=config.AZURE_OPENAI_ENDPOINT,
+                default_model="gpt-4o",
+                assistant_id=config.ASSISTANT_ID,
+            )
         )
-    )
 else:
     planner = AssistantsPlanner[AppTurnState](
         OpenAIAssistantsOptions(api_key=config.OPENAI_KEY, assistant_id=config.ASSISTANT_ID)
@@ -96,22 +109,34 @@ async def setup_assistant(context: TurnContext, state: AppTurnState):
                         ),
                     )
                 ],
-                model="gpt-4",
+                model="gpt-4o",
             )
 
             assistant: Assistant
 
-            if config.AZURE_OPENAI_KEY and config.AZURE_OPENAI_ENDPOINT:
-                assistant = await AssistantsPlanner.create_assistant(
-                    api_key=config.AZURE_OPENAI_KEY,
-                    api_version="",
-                    organization="",
-                    endpoint=config.AZURE_OPENAI_ENDPOINT,
-                    request=params,
-                )
+            if config.AZURE_OPENAI_ENDPOINT:
+                if config.AZURE_OPENAI_KEY:
+                    assistant = await AssistantsPlanner.create_assistant(
+                        api_key=config.AZURE_OPENAI_KEY,
+                        azure_ad_token_provider=None,
+                        api_version="2024-05-01-preview",
+                        organization="",
+                        endpoint=config.AZURE_OPENAI_ENDPOINT,
+                        request=params,
+                    )
+                else:
+                    assistant = await AssistantsPlanner.create_assistant(
+                        api_key=None,
+                        azure_ad_token_provider=get_bearer_token_provider(DefaultAzureCredential(), 'https://cognitiveservices.azure.com/.default'),
+                        api_version="2024-05-01-preview",
+                        organization="",
+                        endpoint=config.AZURE_OPENAI_ENDPOINT,
+                        request=params,
+                    )
             else:
                 assistant = await AssistantsPlanner.create_assistant(
                     api_key=config.OPENAI_KEY,
+                    azure_ad_token_provider=None,
                     api_version="",
                     organization="",
                     endpoint="",
diff --git a/python/samples/06.assistants.b.orderBot/teamsapp.yml b/python/samples/06.assistants.b.orderBot/teamsapp.yml
@@ -97,3 +97,4 @@ deploy:
       # You can replace it with your existing Azure Resource id
       # or add it to your environment variable file.
       resourceId: ${{BOT_AZURE_APP_SERVICE_RESOURCE_ID}}
+projectId: 821a55d5-585b-4805-8558-793d811b1215
