Merge pull request #1182 from utsavtulsyan/feature/azureopenai-managed-identity

feat(dspy): added support for managed identity in AzureOpenAI

Author: arnavsinghvi11

docs/api/language_model_clients/AzureOpenAI.md

@@ -29,30 +29,32 @@ class AzureOpenAI(LM):
     ):
 ```
 
+**Parameters:**
 
-
-**Parameters:** 
 - `api_base` (str): Azure Base URL.
 - `api_version` (str): Version identifier for Azure OpenAI API.
 - `api_key` (_Optional[str]_, _optional_): API provider authentication token. Retrieves from `AZURE_OPENAI_KEY` environment variable if None.
 - `model_type` (_Literal["chat", "text"]_): Specified model type to use, defaults to 'chat'.
+- `azure_ad_token_provider` (_Optional[AzureADTokenProvider]_, _optional_): Pass the bearer token provider created by _get_bearer_token_provider()_ when using DefaultAzureCredential, alternative to api token.
 - `**kwargs`: Additional language model arguments to pass to the API provider.
 
 ### Methods
 
 #### `__call__(self, prompt: str, only_completed: bool = True, return_sorted: bool = False, **kwargs) -> List[Dict[str, Any]]`
 
-Retrieves completions from Azure OpenAI Endpoints by calling `request`. 
+Retrieves completions from Azure OpenAI Endpoints by calling `request`.
 
 Internally, the method handles the specifics of preparing the request prompt and corresponding payload to obtain the response.
 
 After generation, the completions are post-processed based on the `model_type` parameter. If the parameter is set to 'chat', the generated content look like `choice["message"]["content"]`. Otherwise, the generated text will be `choice["text"]`.
 
 **Parameters:**
+
 - `prompt` (_str_): Prompt to send to Azure OpenAI.
 - `only_completed` (_bool_, _optional_): Flag to return only completed responses and ignore completion due to length. Defaults to True.
 - `return_sorted` (_bool_, _optional_): Flag to sort the completion choices using the returned averaged log-probabilities. Defaults to False.
 - `**kwargs`: Additional keyword arguments for completion request.
 
 **Returns:**
+
 - `List[Dict[str, Any]]`: List of completion choices.

dsp/modules/azure_openai.py

@@ -1,7 +1,7 @@
 import functools
 import json
 import logging
-from typing import Any, Literal, Optional, cast
+from typing import Any, Callable, Literal, Optional, cast
 
 import backoff
 import openai
@@ -37,6 +37,9 @@ def backoff_hdlr(details):
     )
 
 
+AzureADTokenProvider = Callable[[], str]
+
+
 class AzureOpenAI(LM):
     """Wrapper around Azure's API for OpenAI.
 
@@ -57,6 +60,7 @@ def __init__(
         api_key: Optional[str] = None,
         model_type: Literal["chat", "text"] = "chat",
         system_prompt: Optional[str] = None,
+        azure_ad_token_provider: Optional[AzureADTokenProvider] = None,
         **kwargs,
     ):
         super().__init__(model)
@@ -75,6 +79,7 @@ def __init__(
             openai.api_key = api_key
             openai.api_type = "azure"
             openai.api_version = api_version
+            openai.azure_ad_token_provider = azure_ad_token_provider
 
             self.client = None
 
@@ -83,6 +88,7 @@ def __init__(
                 azure_endpoint=api_base,
                 api_key=api_key,
                 api_version=api_version,
+                azure_ad_token_provider=azure_ad_token_provider,
             )
 
             self.client = client