Repository: telekom/mltb2
Pull request #168: add option for `azure_ad_token_provider` to `OpenAiAzureChat`
Author: login: john-tlkm
Labels: name: enhancement
description: New feature or request
Created: 2025-03-14T15:48:42Z | Merged: 2025-03-21T15:02:01Z

## Description

we can now initialise the OpenAiAzureChat using an `azure_ad_token_provider`. Using the token provider we do not need to pass an AD token manually. This is especially useful when an azure ML job is queuing for a long time and the provided AD token has expired.

Here is an Example:
```
from mltb2.openai import OpenAiAzureChat
from azure.identity import DefaultAzureCredential, get_bearer_token_provider

# define resource
scopes = "https://cognitiveservices.azure.com/.default"
api_version = "2023-07-01-preview"
endpoint = "https://my-resource.openai.azure.com/"

# init token provider
token_provider = get_bearer_token_provider(DefaultAzureCredential(), scopes)

# init OpenAiAzureChat class
chat = OpenAiAzureChat(azure_ad_token_provider=token_provider, 
                azure_endpoint=endpoint, 
                model="GPT-4o",
                api_version=api_version)

# run a completion
chat.create_completions(prompt="hello")
```

## Pull request checklist

- tests
  - [x] Do we have tests? Should we add more?
  - [x] Should we parametrize tests?
  - [x] Should we add hypothesis tests?
- documentation
  - [x] add docstrings
  - [x] if a module is new: add API reference file for Sphinx in `docs/source/api-reference`
  - [x] check API doc of Sphinx - build and show with `make sphinx && make open-sphinx`
- [x] fix checks
  - start locale checks with `make check`
  - fix formatting with `make format`
- Python type info
  - [x] check and improve Python type info
  - [x] double check if return types are specified - including `-> None`
- [x] check TODO comments in code
- [x] check FIXME comments in code
- [x] check if we want to increase the version or rc number
- [x] check if the copyright needs an update
  - file header
  - README.md
  - LICENSE file



Diff:
diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml
index ebe8c51..545f095 100644
--- a/.github/workflows/pytest.yml
+++ b/.github/workflows/pytest.yml
@@ -18,7 +18,9 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11", "3.12"]
+        # only test high and low version numbers
+        python-version: ["3.9", "3.12"]
+        # python-version: ["3.9", "3.10", "3.11", "3.12"]
 
     steps:
     # https://github.com/actions/checkout
diff --git a/LICENSE b/LICENSE
index 5229464..41eaa93 100644
--- a/LICENSE
+++ b/LICENSE
@@ -3,6 +3,7 @@ MIT License
 Copyright (c) 2023-2024 Philip May
 Copyright (c) 2023-2024 Philip May, Deutsche Telekom AG
 Copyright (c) 2023-2024 Alaeddine Abdessalem, Deutsche Telekom AG
+Copyright (c) 2025 Sijun John Tu, Deutsche Telekom AG
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal
diff --git a/README.md b/README.md
index 54a0692..3c4733b 100644
--- a/README.md
+++ b/README.md
@@ -33,7 +33,8 @@ To install those module specific dependencies see
 
 Copyright (c) 2023-2024 [Philip May](https://philipmay.org)\
 Copyright (c) 2023-2024 [Philip May](https://philipmay.org), [Deutsche Telekom AG](https://www.telekom.de/)\
-Copyright (c) 2023-2024 Alaeddine Abdessalem, [Deutsche Telekom AG](https://www.telekom.de/)
+Copyright (c) 2023-2024 Alaeddine Abdessalem, [Deutsche Telekom AG](https://www.telekom.de/)\
+Copyright (c) 2025 Sijun John Tu, [Deutsche Telekom AG](https://www.telekom.de/)
 
 Licensed under the **MIT License** (the "License"); you may not use this file except in compliance with the License.
 You may obtain a copy of the License by reviewing the file
diff --git a/mltb2/openai.py b/mltb2/openai.py
index 77487e8..6311145 100644
--- a/mltb2/openai.py
+++ b/mltb2/openai.py
@@ -1,6 +1,7 @@
-# Copyright (c) 2023-2024 Philip May
-# Copyright (c) 2024 Philip May, Deutsche Telekom AG
+# Copyright (c) 2023-2025 Philip May
+# Copyright (c) 2024-2025 Philip May, Deutsche Telekom AG
 # Copyright (c) 2024 Alaeddine Abdessalem, Deutsche Telekom AG
+# Copyright (c) 2025 Sijun John Tu, Deutsche Telekom AG
 # This software is distributed under the terms of the MIT license
 # which is available at https://opensource.org/licenses/MIT
 
@@ -15,11 +16,13 @@
 import os
 from collections.abc import Iterable
 from dataclasses import dataclass, field
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, cast
 
 import tiktoken
 import yaml
+from azure.identity import DefaultAzureCredential, get_bearer_token_provider
 from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
+from openai.lib.azure import AzureADTokenProvider
 from openai.types.chat import ChatCompletion
 from tiktoken.core import Encoding
 from tqdm import tqdm
@@ -347,36 +350,59 @@ class OpenAiAzureChat(OpenAiChat, _OpenAiAzureChatBase):
     See Also:
         * OpenAI API reference: `Create chat completion <https://platform.openai.com/docs/api-reference/chat/create>`_
         * `Quickstart: Get started generating text using Azure OpenAI Service <https://learn.microsoft.com/en-us/azure/ai-services/openai/quickstart?tabs=command-line&pivots=programming-language-python>`_
+        * ```AzureADTokenProvider`` example <https://github.com/openai/openai-python/blob/main/examples/azure_ad.py>`_
 
     Args:
         api_key: The OpenAI API key.
         model: The OpenAI model name.
         api_version: The OpenAI API version.
             A common value for this is ``2023-05-15``.
+        azure_ad_token: The Azure Active Directory token.
+        azure_ad_token_provider: A function that returns an Azure Active Directory token,
+            which will be invoked on every request.
+            Or set to "auto" to use default credentials.
         azure_endpoint: The Azure endpoint.
     """
 
     api_version: Optional[str] = None
     api_key: Optional[str] = None
     azure_ad_token: Optional[str] = None
+    azure_ad_token_provider: Union[AzureADTokenProvider, str, None] = None
 
     def __post_init__(self) -> None:
         """Do post init."""
+        # init default token provider if azure_ad_token_provider=="auto"
+        if self.azure_ad_token_provider == "auto":  # NOQA: S105
+            self.azure_ad_token_provider = get_bearer_token_provider(
+                DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+            )
+
+        self.azure_ad_token_provider = cast("Optional[AzureADTokenProvider]", self.azure_ad_token_provider)
+
         self.client = AzureOpenAI(
             api_key=self.api_key,
             api_version=self.api_version,
             azure_endpoint=self.azure_endpoint,
             azure_ad_token=self.azure_ad_token,
+            azure_ad_token_provider=self.azure_ad_token_provider,
         )
         self.async_client = AsyncAzureOpenAI(
             api_key=self.api_key,
             api_version=self.api_version,
             azure_endpoint=self.azure_endpoint,
             azure_ad_token=self.azure_ad_token,
+            azure_ad_token_provider=self.azure_ad_token_provider,
         )
 
     @classmethod
-    def from_yaml(cls, yaml_file, api_key: Optional[str] = None, azure_ad_token: Optional[str] = None, **kwargs):
+    def from_yaml(
+        cls,
+        yaml_file,
+        api_key: Optional[str] = None,
+        azure_ad_token: Optional[str] = None,
+        azure_ad_token_provider: Union[AzureADTokenProvider, str, None] = None,
+        **kwargs,
+    ):
         """Construct this class from a yaml file.
 
         If the ``api_key`` is not set in the yaml file,
@@ -385,7 +411,10 @@ def from_yaml(cls, yaml_file, api_key: Optional[str] = None, azure_ad_token: Opt
         Args:
             yaml_file: The yaml file.
             api_key: The OpenAI API key.
-            azure_ad_token: Azure AD token
+            azure_ad_token: The Azure Active Directory token.
+            azure_ad_token_provider: A function that returns an Azure Active Directory token,
+                which will be invoked on every request.
+                Or set to "auto" to use default credentials.
             kwargs: extra kwargs to override parameters
         Returns:
             The constructed class.
@@ -396,4 +425,20 @@ def from_yaml(cls, yaml_file, api_key: Optional[str] = None, azure_ad_token: Opt
         # set azure_ad_token according to this priority:
         # method parameter > yaml > environment variable
         azure_ad_token = azure_ad_token or completion_kwargs.get("AZURE_AD_TOKEN") or os.getenv("AZURE_AD_TOKEN")
-        return super().from_yaml(yaml_file, api_key=api_key, azure_ad_token=azure_ad_token, **kwargs)
+
+        # init the token provider
+        ## method parameter > yaml
+        azure_ad_token_provider = azure_ad_token_provider or completion_kwargs.get("azure_ad_token_provider")
+        ## if token_provider==auto use default settings
+        if azure_ad_token_provider == "auto":  # NOQA: S105
+            azure_ad_token_provider = get_bearer_token_provider(
+                DefaultAzureCredential(), "https://cognitiveservices.azure.com/.default"
+            )
+
+        return super().from_yaml(
+            yaml_file,
+            api_key=api_key,
+            azure_ad_token=azure_ad_token,
+            azure_ad_token_provider=azure_ad_token_provider,
+            **kwargs,
+        )
diff --git a/pyproject.toml b/pyproject.toml
index 3951838..b356f03 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -84,6 +84,7 @@ python-arango = {version = "*", optional = true}
 jsonlines = {version = "*", optional = true}
 markdownify = {version = "*", optional = true}
 mdformat = {version = "*", optional = true}
+azure-identity = {version = "*", optional = true}
 
 [tool.poetry.extras]
 files = ["platformdirs", "scikit-learn", "joblib"]
@@ -95,7 +96,7 @@ somajo = ["SoMaJo"]
 transformers = ["scikit-learn", "torch", "transformers", "safetensors"]
 md = ["scikit-learn", "torch", "transformers", "safetensors"]
 somajo_transformers = ["SoMaJo", "scikit-learn", "torch", "transformers", "safetensors"]
-openai = ["tiktoken", "openai", "pyyaml"]
+openai = ["tiktoken", "openai", "pyyaml", "azure-identity"]
 arangodb = ["python-dotenv", "python-arango", "jsonlines"]
 bs = ["beautifulsoup4", "markdownify", "mdformat"]
 


Reviews:
user: login: PhilipMay
state: COMMENTED
author_association: MEMBER
submitted_at: 2025-03-15T06:40:45Z
user: login: alaeddine-13
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-18T09:19:07Z
user: login: john-tlkm
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-21T14:18:32Z

Review comments:
diff_hunk: @@ -19,8 +19,9 @@
 
 import tiktoken
 import yaml
-from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
+from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI 
path: mltb2/openai.py
user: login: PhilipMay
body: The trailing blankline should be removed. You can do this with `make format`. Please also see https://github.com/telekom/mltb2/blob/main/CONTRIBUTING.md#testing-linting-and-formatting
created_at: 2025-03-15T06:40:45Z
author_association: MEMBER
diff_hunk: @@ -21,6 +22,8 @@
 import yaml
 from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
 from openai.types.chat import ChatCompletion
+from openai.lib.azure import AzureADTokenProvider, AsyncAzureADTokenProvider
path: mltb2/openai.py
user: login: alaeddine-13
body: is AsyncAzureADTokenProvider being used somewhere ?
created_at: 2025-03-18T09:19:07Z
author_association: COLLABORATOR
diff_hunk: @@ -21,6 +22,8 @@
 import yaml
 from openai import AsyncAzureOpenAI, AsyncOpenAI, AzureOpenAI, OpenAI
 from openai.types.chat import ChatCompletion
+from openai.lib.azure import AzureADTokenProvider, AsyncAzureADTokenProvider
path: mltb2/openai.py
user: login: john-tlkm
body: actually not, I will check if its needed
created_at: 2025-03-21T14:18:32Z
author_association: CONTRIBUTOR

PR comments:
user: login: PhilipMay
created_at: 2025-03-15T06:37:19Z
body: @john-tlkm many thanks for this PR

Please so these additional changes:
- add yourself to the copyright header in `openai.py`
- also please add yourself to the copyright footer in `README.md`
- and the `LICENSE` file

author_association: MEMBER
user: login: PhilipMay
created_at: 2025-03-15T06:38:24Z
body: Some CI jobs fail. I will repair them in main.
author_association: MEMBER
user: login: PhilipMay
created_at: 2025-03-15T06:46:07Z
body: The `azure_ad_token_provider` should also be passable to the `from_yaml` function. Can you please extend that?
author_association: MEMBER
user: login: PhilipMay
created_at: 2025-03-15T08:07:39Z
body: @john-tlkm I fixed some build issues in main. Can you please rebase main into your branch and then push again?
See here: https://github.com/telekom/mltb2/blob/main/CONTRIBUTING.md#sync-your-branch-with-the-original-repository
author_association: MEMBER
user: login: PhilipMay
created_at: 2025-03-17T13:55:23Z
body: Oops - I saw there is an `AzureADTokenProvider` and an `AsyncAzureADTokenProvider`. This will be a bit tricky...
author_association: MEMBER
user: login: john-tlkm
created_at: 2025-03-17T16:27:14Z
body: > Oops - I saw there is an `AzureADTokenProvider` and an `AsyncAzureADTokenProvider`. This will be a bit tricky...

should not be a problem. the token provider should work the same in async
author_association: CONTRIBUTOR
user: login: PhilipMay
created_at: 2025-03-21T14:14:41Z
body: Not 100% sure what this exactly means...
https://github.com/openai/openai-python/blob/main/src/openai/lib/azure.py#L32-L33
author_association: MEMBER
user: login: john-tlkm
created_at: 2025-03-21T14:17:53Z
body: > Not 100% sure what this exactly means... https://github.com/openai/openai-python/blob/main/src/openai/lib/azure.py#L32-L33

also not sure about this, but the example here shows no difference between regular and async token provider: 
https://github.com/openai/openai-python/blob/main/examples/azure_ad.py
author_association: CONTRIBUTOR
user: login: PhilipMay
created_at: 2025-03-21T14:23:58Z
body: > Not 100% sure what this exactly means... https://github.com/openai/openai-python/blob/main/src/openai/lib/azure.py#L32-L33

Well. If we translate it it means:

`AzureADTokenProvider = Callable[[], str]` means that it is a callable with no input values which returns a string.
While `AsyncAzureADTokenProvider` is also a callable with no input values which returns a string or an Awaitable string.

This means an `AzureADTokenProvider` should also be an `AsyncAzureADTokenProvider` but not the other way around.
author_association: MEMBER