feat: managed identity integration

Author: sk-portkey

src/handlers/handlerUtils.ts

@@ -1011,11 +1011,15 @@ export function constructConfigFromRequestHeaders(
     resourceName: requestHeaders[`x-${POWERED_BY}-azure-resource-name`],
     deploymentId: requestHeaders[`x-${POWERED_BY}-azure-deployment-id`],
     apiVersion: requestHeaders[`x-${POWERED_BY}-azure-api-version`],
-    azureModelName: requestHeaders[`x-${POWERED_BY}-azure-model-name`],
+    azureAdToken: requestHeaders[`x-${POWERED_BY}-azure-ad-token`],
+    azureAuthMode: requestHeaders[`x-${POWERED_BY}-azure-auth-mode`],
+    azureManagedClientId:
+      requestHeaders[`x-${POWERED_BY}-azure-managed-client-id`],
     azureEntraClientId: requestHeaders[`x-${POWERED_BY}-azure-entra-client-id`],
     azureEntraClientSecret:
       requestHeaders[`x-${POWERED_BY}-azure-entra-client-secret`],
     azureEntraTenantId: requestHeaders[`x-${POWERED_BY}-azure-entra-tenant-id`],
+    azureModelName: requestHeaders[`x-${POWERED_BY}-azure-model-name`],
   };
 
   const stabilityAiConfig = {
@@ -1034,10 +1038,6 @@ export function constructConfigFromRequestHeaders(
       requestHeaders[`x-${POWERED_BY}-azure-deployment-type`],
     azureApiVersion: requestHeaders[`x-${POWERED_BY}-azure-api-version`],
     azureEndpointName: requestHeaders[`x-${POWERED_BY}-azure-endpoint-name`],
-    azureEntraClientId: requestHeaders[`x-${POWERED_BY}-azure-entra-client-id`],
-    azureEntraClientSecret:
-      requestHeaders[`x-${POWERED_BY}-azure-entra-client-secret`],
-    azureEntraTenantId: requestHeaders[`x-${POWERED_BY}-azure-entra-tenant-id`],
   };
 
   const bedrockConfig = {

src/providers/azure-ai-inference/api.ts

@@ -1,5 +1,4 @@
 import { GITHUB } from '../../globals';
-import { getAccessTokenFromEntraId } from '../azure-openai/utils';
 import { ProviderAPIConfig } from '../types';
 
 const AzureAIInferenceAPI: ProviderAPIConfig = {
@@ -21,30 +20,15 @@ const AzureAIInferenceAPI: ProviderAPIConfig = {
     return `https://${azureEndpointName}.${azureRegion}.inference.ml.azure.com/score`;
   },
   headers: async ({ providerOptions }) => {
-    const {
-      apiKey,
-      azureDeploymentType,
-      azureDeploymentName,
-      azureEntraClientId,
-      azureEntraClientSecret,
-      azureEntraTenantId,
-    } = providerOptions;
+    const { apiKey, azureDeploymentType, azureDeploymentName } =
+      providerOptions;
     const headers: Record<string, string> = {
       Authorization: `Bearer ${apiKey}`,
       'extra-parameters': 'ignore',
     };
     if (azureDeploymentType === 'managed' && azureDeploymentName) {
       headers['azureml-model-deployment'] = azureDeploymentName;
     }
-    if (azureEntraClientId && azureEntraClientSecret && azureEntraTenantId) {
-      const accessToken = await getAccessTokenFromEntraId(
-        azureEntraTenantId,
-        azureEntraClientId,
-        azureEntraClientSecret,
-        'https://cognitiveservices.azure.com/decision/.default'
-      );
-      headers['Authorization'] = `Bearer ${accessToken}`;
-    }
     return headers;
   },
   getEndpoint: ({ providerOptions, fn }) => {

src/providers/azure-openai/api.ts

@@ -1,29 +1,46 @@
 import { ProviderAPIConfig } from '../types';
-import { getAccessTokenFromEntraId } from './utils';
+import { getAccessTokenFromEntraId, getAzureManagedIdentityToken } from './utils';
 
 const AzureOpenAIAPIConfig: ProviderAPIConfig = {
   getBaseURL: ({ providerOptions }) => {
     const { resourceName, deploymentId } = providerOptions;
     return `https://${resourceName}.openai.azure.com/openai/deployments/${deploymentId}`;
   },
   headers: async ({ providerOptions, fn }) => {
+    const { apiKey, azureAuthMode } = providerOptions;
+
+    if (azureAuthMode === 'entra') {
+      const { azureEntraTenantId, azureEntraClientId, azureEntraClientSecret } =
+        providerOptions;
+      if (azureEntraTenantId && azureEntraClientId && azureEntraClientSecret) {
+        const scope = 'https://cognitiveservices.azure.com/.default';
+        const accessToken = await getAccessTokenFromEntraId(
+          azureEntraTenantId,
+          azureEntraClientId,
+          azureEntraClientSecret,
+          scope
+        );
+        return {
+          Authorization: `Bearer ${accessToken}`,
+        };
+      }
+    }
+    if (azureAuthMode === 'managed') {
+      const { azureManagedClientId } = providerOptions;
+      const resource = 'https://cognitiveservices.azure.com/';
+      const accessToken = await getAzureManagedIdentityToken(
+        resource,
+        azureManagedClientId
+      );
+      return {
+        Authorization: `Bearer ${accessToken}`,
+      };
+    }
     const headersObj: Record<string, string> = {
-      'api-key': `${providerOptions.apiKey}`,
+      'api-key': `${apiKey}`,
     };
     if (fn === 'createTranscription' || fn === 'createTranslation')
       headersObj['Content-Type'] = 'multipart/form-data';
-    const { azureEntraClientId, azureEntraClientSecret, azureEntraTenantId } =
-      providerOptions;
-    if (azureEntraClientId && azureEntraClientSecret && azureEntraTenantId) {
-      const accessToken = await getAccessTokenFromEntraId(
-        azureEntraTenantId,
-        azureEntraClientId,
-        azureEntraClientSecret,
-        'https://cognitiveservices.azure.com/decision/.default'
-      );
-      headersObj['Authorization'] = `Bearer ${accessToken}`;
-    }
-
     return headersObj;
   },
   getEndpoint: ({ providerOptions, fn }) => {

src/providers/azure-openai/utils.ts

@@ -2,28 +2,59 @@ export async function getAccessTokenFromEntraId(
   tenantId: string,
   clientId: string,
   clientSecret: string,
-  scope = 'https://openai.azure.com/.default'
+  scope = 'https://cognitiveservices.azure.com/.default'
 ) {
-  const url = `https://login.microsoftonline.com/${tenantId}/oauth2/v2.0/token`;
+  try {
+    const url = `https://login.microsoftonline.com/${tenantId}/oauth2/v2.0/token`;
+    const params = new URLSearchParams({
+      client_id: clientId,
+      client_secret: clientSecret,
+      scope: scope,
+      grant_type: 'client_credentials',
+    });
 
-  const params = new URLSearchParams({
-    client_id: clientId,
-    client_secret: clientSecret,
-    scope: scope,
-    grant_type: 'client_credentials',
-  });
+    const response = await fetch(url, {
+      method: 'POST',
+      headers: {
+        'Content-Type': 'application/x-www-form-urlencoded',
+      },
+      body: params,
+    });
 
-  const response = await fetch(url, {
-    method: 'POST',
-    headers: {
-      'Content-Type': 'application/x-www-form-urlencoded',
-    },
-    body: params,
-  });
+    if (!response.ok) {
+      const errorMessage = await response.text();
+      console.log({ message: `Error from Entra ${errorMessage}` });
+      return undefined;
+    }
+    const data: { access_token: string } = await response.json();
+    return data.access_token;
+  } catch (error) {
+    console.log(error);
+  }
+}
 
-  if (!response.ok) {
-    throw new Error(`Error fetching access token: ${response.statusText}`);
+export async function getAzureManagedIdentityToken(
+  resource: string,
+  clientId?: string
+) {
+  try {
+    const response = await fetch(
+      `http://169.254.169.254/metadata/identity/oauth2/token?api-version=2018-02-01&resource=${encodeURIComponent(resource)}${clientId ? `&client_id=${encodeURIComponent(clientId)}` : ''}`,
+      {
+        method: 'GET',
+        headers: {
+          Metadata: 'true',
+        },
+      }
+    );
+    if (!response.ok) {
+      const errorMessage = await response.text();
+      console.log({ message: `Error from Managed ${errorMessage}` });
+      return undefined;
+    }
+    const data: { access_token: string } = await response.json();
+    return data.access_token;
+  } catch (error) {
+    console.log({ error });
   }
-  const data: { access_token: string } = await response.json();
-  return data.access_token;
 }

src/types/requestBody.ts

@@ -60,6 +60,8 @@ export interface Options {
   apiVersion?: string;
   adAuth?: string;
   azureModelName?: string;
+  azureAuthMode?: string; // can be entra or managed
+  azureManagedClientId?: string;
   azureEntraClientId?: string;
   azureEntraClientSecret?: string;
   azureEntraTenantId?: string;
@@ -143,6 +145,12 @@ export interface Targets {
   deploymentId?: string;
   apiVersion?: string;
   adAuth?: string;
+  azureAuthMode?: string;
+  azureManagedClientId?: string;
+  azureEntraClientId?: string;
+  azureEntraClientSecret?: string;
+  azureEntraTenantId?: string;
+  azureModelName?: string;
   /** provider option index picked based on weight in loadbalance mode */
   index?: number;
   cache?: CacheSettings | string;
@@ -353,6 +361,11 @@ export interface ShortConfig {
   azureModelName?: string;
   workersAiAccountId?: string;
   apiVersion?: string;
+  azureAuthMode?: string;
+  azureManagedClientId?: string;
+  azureEntraClientId?: string;
+  azureEntraClientSecret?: string;
+  azureEntraTenantId?: string;
   customHost?: string;
   // Google Vertex AI specific
   vertexRegion?: string;