feat: update azure foundry provider

src/handlers/handlerUtils.ts

@@ -973,13 +973,9 @@ export function constructConfigFromRequestHeaders(
   };
 
   const azureAiInferenceConfig = {
-    azureDeploymentName:
-      requestHeaders[`x-${POWERED_BY}-azure-deployment-name`],
-    azureRegion: requestHeaders[`x-${POWERED_BY}-azure-region`],
-    azureDeploymentType:
-      requestHeaders[`x-${POWERED_BY}-azure-deployment-type`],
     azureApiVersion: requestHeaders[`x-${POWERED_BY}-azure-api-version`],
     azureEndpointName: requestHeaders[`x-${POWERED_BY}-azure-endpoint-name`],
+    azureFoundryUrl: requestHeaders[`x-${POWERED_BY}-azure-foundry-url`],
     azureExtraParams: requestHeaders[`x-${POWERED_BY}-azure-extra-params`],
   };
 

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

@@ -1,37 +1,72 @@
 import { GITHUB } from '../../globals';
+import {
+  getAccessTokenFromEntraId,
+  getAzureManagedIdentityToken,
+} from '../azure-openai/utils';
 import { ProviderAPIConfig } from '../types';
 
 const AzureAIInferenceAPI: ProviderAPIConfig = {
   getBaseURL: ({ providerOptions }) => {
-    const {
-      azureDeploymentName,
-      azureRegion,
-      azureDeploymentType,
-      provider,
-      azureEndpointName,
-    } = providerOptions;
+    const { provider, azureFoundryUrl } = providerOptions;
     if (provider === GITHUB) {
       return 'https://models.inference.ai.azure.com';
     }
-    if (azureDeploymentType === 'serverless') {
-      return `https://${azureDeploymentName?.toLowerCase()}.${azureRegion}.models.ai.azure.com`;
+    if (azureFoundryUrl) {
+      return azureFoundryUrl;
     }
 
-    return `https://${azureEndpointName}.${azureRegion}.inference.ml.azure.com/score`;
+    return '';
   },
-  headers: ({ providerOptions }) => {
+  headers: async ({ providerOptions }) => {
     const {
       apiKey,
-      azureDeploymentType,
-      azureDeploymentName,
       azureExtraParams,
+      azureDeploymentName,
+      azureAdToken,
+      azureAuthMode,
     } = providerOptions;
+
     const headers: Record<string, string> = {
-      Authorization: `Bearer ${apiKey}`,
-      'extra-parameters': azureExtraParams || 'pass-through',
+      'extra-parameters': azureExtraParams ?? 'drop',
+      ...(azureDeploymentName && {
+        'azureml-model-deployment': azureDeploymentName,
+      }),
     };
-    if (azureDeploymentType === 'managed' && azureDeploymentName) {
-      headers['azureml-model-deployment'] = azureDeploymentName;
+    if (azureAdToken) {
+      headers['Authorization'] =
+        `Bearer ${azureAdToken?.replace('Bearer ', '')}`;
+      return headers;
+    }
+
+    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
+        );
+        headers['Authorization'] = `Bearer ${accessToken}`;
+        return headers;
+      }
+    }
+    if (azureAuthMode === 'managed') {
+      const { azureManagedClientId } = providerOptions;
+      const resource = 'https://cognitiveservices.azure.com/';
+      const accessToken = await getAzureManagedIdentityToken(
+        resource,
+        azureManagedClientId
+      );
+      headers['Authorization'] = `Bearer ${accessToken}`;
+      return headers;
+    }
+
+    if (apiKey) {
+      headers['Authorization'] = `Bearer ${apiKey}`;
+      return headers;
     }
     return headers;
   },
@@ -57,21 +92,22 @@ const AzureAIInferenceAPI: ProviderAPIConfig = {
       }
     }
 
+    const apiVersion = azureApiVersion ? `?api-version=${azureApiVersion}` : '';
     switch (mappedFn) {
       case 'complete': {
         return isGithub
           ? ENDPOINT_MAPPING[mappedFn]
-          : `${ENDPOINT_MAPPING[mappedFn]}?api-version=${azureApiVersion}`;
+          : `${ENDPOINT_MAPPING[mappedFn]}${apiVersion}`;
       }
       case 'chatComplete': {
         return isGithub
           ? ENDPOINT_MAPPING[mappedFn]
-          : `${ENDPOINT_MAPPING[mappedFn]}?api-version=${azureApiVersion}`;
+          : `${ENDPOINT_MAPPING[mappedFn]}${apiVersion}`;
       }
       case 'embed': {
         return isGithub
           ? ENDPOINT_MAPPING[mappedFn]
-          : `${ENDPOINT_MAPPING[mappedFn]}?api-version=${azureApiVersion}`;
+          : `${ENDPOINT_MAPPING[mappedFn]}${apiVersion}`;
       }
       default:
         return '';

src/types/requestBody.ts

@@ -135,12 +135,10 @@ export interface Options {
   openaiBeta?: string;
 
   /** Azure Inference Specific */
-  azureRegion?: string;
   azureDeploymentName?: string;
-  azureDeploymentType?: 'managed' | 'serverless';
-  azureEndpointName?: string;
   azureApiVersion?: string;
   azureExtraParams?: string;
+  azureFoundryUrl?: string;
 
   /** The parameter to determine if extra non-openai compliant fields should be returned in response */
   strictOpenAiCompliance?: boolean;