All files / api/src/ai-providers/azure-ai-foundry api.ts

100% Statements 76/76
100% Branches 21/21
100% Functions 3/3
100% Lines 76/76

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89    1x 1x 1x         1x 1x 1x 1x 1x   1x 1x 5x 5x 2x 2x   5x 1x 1x   2x 5x 1x 12x 12x   12x 12x 12x 12x 2x 2x 2x   12x 3x 3x 3x 3x 3x 3x 3x 2x 1x 3x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 3x 12x 2x 2x 2x 2x 2x 2x 2x 2x 2x   12x 4x 4x 4x 3x 12x 1x 6x 6x 6x 1x 1x 5x 6x 1x  
import type { InternalProviderAPIConfig } from '@shared/types/ai-providers/config';
 
import { FunctionName } from '@shared/types/api/request';
import { AIProvider } from '@shared/types/constants';
import {
  getAccessTokenFromEntraId,
  getAzureManagedIdentityToken,
} from '../azure-openai/utils';
 
const EndpointMap: Partial<Record<FunctionName, string>> = {
  [FunctionName.COMPLETE]: '/models/completions',
  [FunctionName.CHAT_COMPLETE]: '/models/chat/completions',
  [FunctionName.EMBED]: '/models/embeddings',
};
 
export const azureAIInferenceAPI: InternalProviderAPIConfig = {
  getBaseURL: ({ saTarget }) => {
    const { configuration, azure_ai_foundry_config } = saTarget;
    if (configuration.ai_provider === AIProvider.GITHUB) {
      return 'https://models.inference.ai.azure.com';
    }
 
    if (!azure_ai_foundry_config) {
      throw new Error('`azure_ai_foundry_config` is required in target');
    }
 
    return azure_ai_foundry_config.url;
  },
  headers: async ({ saTarget: providerOptions }) => {
    const { api_key, azure_extra_params, azure_ad_token, azure_auth_mode } =
      providerOptions;
 
    const headers: Record<string, string> = {
      'extra-parameters': azure_extra_params ?? 'drop',
    };
    if (azure_ad_token) {
      headers.Authorization = `Bearer ${azure_ad_token?.replace('Bearer ', '')}`;
      return headers;
    }
 
    if (azure_auth_mode === 'entra') {
      const {
        azure_entra_tenant_id,
        azure_entra_client_id,
        azure_entra_client_secret,
      } = providerOptions;
      if (
        azure_entra_tenant_id &&
        azure_entra_client_id &&
        azure_entra_client_secret
      ) {
        const scope = 'https://cognitiveservices.azure.com/.default';
        const accessToken = await getAccessTokenFromEntraId(
          azure_entra_tenant_id,
          azure_entra_client_id,
          azure_entra_client_secret,
          scope,
        );
        headers.Authorization = `Bearer ${accessToken}`;
        return headers;
      }
    }
    if (azure_auth_mode === 'managed') {
      const { azure_managed_client_id } = providerOptions;
      const resource = 'https://cognitiveservices.azure.com/';
      const accessToken = await getAzureManagedIdentityToken(
        resource,
        azure_managed_client_id,
      );
      headers.Authorization = `Bearer ${accessToken}`;
      return headers;
    }
 
    if (api_key) {
      headers.Authorization = `Bearer ${api_key}`;
      return headers;
    }
    return headers;
  },
  getEndpoint: ({ saRequestData }) => {
    const fn = saRequestData.functionName;
    const endpoint = EndpointMap[fn];
    if (!endpoint) {
      throw new Error(`Endpoint not found for function ${fn}`);
    }
    return endpoint;
  },
};