Merge pull request #1148 from b4s36t4/feat/azure-ai-openai-endpoints

VisargD · web-flow · commit cccd1501b326 · 2025-06-25T13:30:43.000+05:30
feat: support extra endpoints for azure-ai provider
diff --git a/src/providers/azure-ai-inference/api.ts b/src/providers/azure-ai-inference/api.ts
@@ -5,9 +5,28 @@ import {
 } from '../azure-openai/utils';
 import { ProviderAPIConfig } from '../types';
 
+const NON_INFERENCE_ENDPOINTS = [
+  'createBatch',
+  'retrieveBatch',
+  'cancelBatch',
+  'getBatchOutput',
+  'listBatches',
+  'uploadFile',
+  'listFiles',
+  'retrieveFile',
+  'deleteFile',
+  'retrieveFileContent',
+];
+
 const AzureAIInferenceAPI: ProviderAPIConfig = {
-  getBaseURL: ({ providerOptions }) => {
+  getBaseURL: ({ providerOptions, fn }) => {
     const { provider, azureFoundryUrl } = providerOptions;
+
+    // Azure Foundry URL includes `/deployments/<deployment>`, strip out and append openai for batches/finetunes
+    if (fn && NON_INFERENCE_ENDPOINTS.includes(fn)) {
+      return new URL(azureFoundryUrl ?? '').origin + '/openai';
+    }
+
     if (provider === GITHUB) {
       return 'https://models.inference.ai.azure.com';
     }
@@ -17,7 +36,7 @@ const AzureAIInferenceAPI: ProviderAPIConfig = {
 
     return '';
   },
-  headers: async ({ providerOptions }) => {
+  headers: async ({ providerOptions, fn }) => {
     const {
       apiKey,
       azureExtraParams,
@@ -31,6 +50,13 @@ const AzureAIInferenceAPI: ProviderAPIConfig = {
       ...(azureDeploymentName && {
         'azureml-model-deployment': azureDeploymentName,
       }),
+      ...(['createTranscription', 'createTranslation', 'uploadFile'].includes(
+        fn
+      )
+        ? {
+            'Content-Type': 'multipart/form-data',
+          }
+        : {}),
     };
     if (azureAdToken) {
       headers['Authorization'] =
@@ -70,14 +96,37 @@ const AzureAIInferenceAPI: ProviderAPIConfig = {
     }
     return headers;
   },
-  getEndpoint: ({ providerOptions, fn }) => {
+  getEndpoint: ({ providerOptions, fn, gatewayRequestURL }) => {
     const { azureApiVersion, urlToFetch } = providerOptions;
     let mappedFn = fn;
 
+    const urlObj = new URL(gatewayRequestURL);
+    const path = urlObj.pathname.replace('/v1', '');
+    const searchParams = urlObj.searchParams;
+
+    if (azureApiVersion) {
+      searchParams.set('api-version', azureApiVersion);
+    }
+
     const ENDPOINT_MAPPING: Record<string, string> = {
       complete: '/completions',
       chatComplete: '/chat/completions',
       embed: '/embeddings',
+      realtime: '/realtime',
+      imageGenerate: '/images/generations',
+      createSpeech: '/audio/speech',
+      createTranscription: '/audio/transcriptions',
+      createTranslation: '/audio/translations',
+      uploadFile: path,
+      retrieveFile: path,
+      listFiles: path,
+      deleteFile: path,
+      retrieveFileContent: path,
+      listBatches: path,
+      retrieveBatch: path,
+      cancelBatch: path,
+      getBatchOutput: path,
+      createBatch: path,
     };
 
     const isGithub = providerOptions.provider === GITHUB;
@@ -92,23 +141,40 @@ const AzureAIInferenceAPI: ProviderAPIConfig = {
       }
     }
 
-    const apiVersion = azureApiVersion ? `?api-version=${azureApiVersion}` : '';
+    const searchParamsString = searchParams.toString();
     switch (mappedFn) {
       case 'complete': {
         return isGithub
           ? ENDPOINT_MAPPING[mappedFn]
-          : `${ENDPOINT_MAPPING[mappedFn]}${apiVersion}`;
+          : `${ENDPOINT_MAPPING[mappedFn]}?${searchParamsString}`;
       }
       case 'chatComplete': {
         return isGithub
           ? ENDPOINT_MAPPING[mappedFn]
-          : `${ENDPOINT_MAPPING[mappedFn]}${apiVersion}`;
+          : `${ENDPOINT_MAPPING[mappedFn]}?${searchParamsString}`;
       }
       case 'embed': {
         return isGithub
           ? ENDPOINT_MAPPING[mappedFn]
-          : `${ENDPOINT_MAPPING[mappedFn]}${apiVersion}`;
+          : `${ENDPOINT_MAPPING[mappedFn]}?${searchParamsString}`;
       }
+      case 'realtime':
+      case 'imageGenerate':
+      case 'createSpeech':
+      case 'createTranscription':
+      case 'createTranslation':
+      case 'cancelBatch':
+      case 'createBatch':
+      case 'getBatchOutput':
+      case 'retrieveBatch':
+      case 'listBatches':
+      case 'retrieveFile':
+      case 'listFiles':
+      case 'deleteFile':
+      case 'retrieveFileContent': {
+        return `${ENDPOINT_MAPPING[mappedFn]}?${searchParamsString}`;
+      }
+
       default:
         return '';
     }
diff --git a/src/providers/azure-ai-inference/getBatchOutput.ts b/src/providers/azure-ai-inference/getBatchOutput.ts
@@ -0,0 +1,124 @@
+import { Context } from 'hono';
+import AzureAIInferenceAPI from './api';
+import { Options } from '../../types/requestBody';
+import { RetrieveBatchResponse } from '../types';
+import { AZURE_OPEN_AI } from '../../globals';
+
+// Return a ReadableStream containing batches output data
+export const AzureAIInferenceGetBatchOutputRequestHandler = async ({
+  c,
+  providerOptions,
+  requestURL,
+}: {
+  c: Context;
+  providerOptions: Options;
+  requestURL: string;
+}) => {
+  // get batch details which has ouptut file id
+  // get file content as ReadableStream
+  // return file content
+  const baseUrl = AzureAIInferenceAPI.getBaseURL({
+    providerOptions,
+    fn: 'retrieveBatch',
+    c,
+    gatewayRequestURL: requestURL,
+  });
+  const retrieveBatchRequestURL = requestURL.replace('/output', '');
+  const retrieveBatchURL =
+    baseUrl +
+    AzureAIInferenceAPI.getEndpoint({
+      providerOptions,
+      fn: 'retrieveBatch',
+      gatewayRequestURL: retrieveBatchRequestURL,
+      c,
+      gatewayRequestBodyJSON: {},
+      gatewayRequestBody: {},
+    });
+  const retrieveBatchesHeaders = await AzureAIInferenceAPI.headers({
+    c,
+    providerOptions,
+    fn: 'retrieveBatch',
+    transformedRequestBody: {},
+    transformedRequestUrl: retrieveBatchURL,
+    gatewayRequestBody: {},
+  });
+  try {
+    const retrieveBatchesResponse = await fetch(retrieveBatchURL, {
+      method: 'GET',
+      headers: retrieveBatchesHeaders,
+    });
+
+    if (!retrieveBatchesResponse.ok) {
+      const error = await retrieveBatchesResponse.text();
+      return new Response(
+        JSON.stringify({
+          error: error || 'error fetching batch output',
+          provider: AZURE_OPEN_AI,
+          param: null,
+        }),
+        {
+          status: 500,
+        }
+      );
+    }
+
+    const batchDetails: RetrieveBatchResponse =
+      await retrieveBatchesResponse.json();
+
+    const outputFileId =
+      batchDetails.output_file_id || batchDetails.error_file_id;
+    if (!outputFileId) {
+      const errors = batchDetails.errors;
+      if (errors) {
+        return new Response(JSON.stringify(errors), {
+          status: 200,
+        });
+      }
+      return new Response(
+        JSON.stringify({
+          error: 'invalid response output format',
+          provider_response: batchDetails,
+          provider: AZURE_OPEN_AI,
+        }),
+        {
+          status: 400,
+        }
+      );
+    }
+    const retrieveFileContentRequestURL = `https://api.portkey.ai/v1/files/${outputFileId}/content`; // construct the entire url instead of the path of sanity sake
+    const retrieveFileContentURL =
+      baseUrl +
+      AzureAIInferenceAPI.getEndpoint({
+        providerOptions,
+        fn: 'retrieveFileContent',
+        gatewayRequestURL: retrieveFileContentRequestURL,
+        c,
+        gatewayRequestBodyJSON: {},
+        gatewayRequestBody: {},
+      });
+    const retrieveFileContentHeaders = await AzureAIInferenceAPI.headers({
+      c,
+      providerOptions,
+      fn: 'retrieveFileContent',
+      transformedRequestBody: {},
+      transformedRequestUrl: retrieveFileContentURL,
+      gatewayRequestBody: {},
+    });
+    const response = fetch(retrieveFileContentURL, {
+      method: 'GET',
+      headers: retrieveFileContentHeaders,
+    });
+    return response;
+  } catch (e) {
+    return new Response(
+      JSON.stringify({
+        error: 'error fetching batch output',
+        provider: AZURE_OPEN_AI,
+        param: null,
+      }),
+      {
+        status: 500,
+      }
+    );
+  }
+};
diff --git a/src/providers/azure-ai-inference/index.ts b/src/providers/azure-ai-inference/index.ts
@@ -13,17 +13,58 @@ import {
   AzureAIInferenceChatCompleteResponseTransform,
 } from './chatComplete';
 import { AZURE_AI_INFERENCE, GITHUB } from '../../globals';
+import { AzureOpenAIImageGenerateConfig } from '../azure-openai/imageGenerate';
+import { AzureOpenAICreateSpeechConfig } from '../azure-openai/createSpeech';
+import { OpenAICreateFinetuneConfig } from '../openai/createFinetune';
+import { AzureOpenAICreateBatchConfig } from '../azure-openai/createBatch';
+import { AzureAIInferenceGetBatchOutputRequestHandler } from './getBatchOutput';
+import { OpenAIFileUploadRequestTransform } from '../openai/uploadFile';
+import {
+  AzureAIInferenceCreateSpeechResponseTransform,
+  AzureAIInferenceCreateTranscriptionResponseTransform,
+  AzureAIInferenceCreateTranslationResponseTransform,
+  AzureAIInferenceResponseTransform,
+} from './utils';
 
 const AzureAIInferenceAPIConfig: ProviderConfigs = {
   complete: AzureAIInferenceCompleteConfig,
   embed: AzureAIInferenceEmbedConfig,
   api: AzureAIInferenceAPI,
   chatComplete: AzureAIInferenceChatCompleteConfig,
+  imageGenerate: AzureOpenAIImageGenerateConfig,
+  createSpeech: AzureOpenAICreateSpeechConfig,
+  createFinetune: OpenAICreateFinetuneConfig,
+  createTranscription: {},
+  createTranslation: {},
+  realtime: {},
+  cancelBatch: {},
+  createBatch: AzureOpenAICreateBatchConfig,
+  cancelFinetune: {},
+  requestHandlers: {
+    getBatchOutput: AzureAIInferenceGetBatchOutputRequestHandler,
+  },
+  requestTransforms: {
+    uploadFile: OpenAIFileUploadRequestTransform,
+  },
   responseTransforms: {
     complete: AzureAIInferenceCompleteResponseTransform(AZURE_AI_INFERENCE),
     chatComplete:
       AzureAIInferenceChatCompleteResponseTransform(AZURE_AI_INFERENCE),
     embed: AzureAIInferenceEmbedResponseTransform(AZURE_AI_INFERENCE),
+    imageGenerate: AzureAIInferenceResponseTransform,
+    createSpeech: AzureAIInferenceCreateSpeechResponseTransform,
+    createTranscription: AzureAIInferenceCreateTranscriptionResponseTransform,
+    createTranslation: AzureAIInferenceCreateTranslationResponseTransform,
+    realtime: {},
+    createBatch: AzureAIInferenceResponseTransform,
+    retrieveBatch: AzureAIInferenceResponseTransform,
+    cancelBatch: AzureAIInferenceResponseTransform,
+    listBatches: AzureAIInferenceResponseTransform,
+    uploadFile: AzureAIInferenceResponseTransform,
+    listFiles: AzureAIInferenceResponseTransform,
+    retrieveFile: AzureAIInferenceResponseTransform,
+    deleteFile: AzureAIInferenceResponseTransform,
+    retrieveFileContent: AzureAIInferenceResponseTransform,
   },
 };
 
diff --git a/src/providers/azure-ai-inference/utils.ts b/src/providers/azure-ai-inference/utils.ts
@@ -0,0 +1,47 @@
+import { AZURE_AI_INFERENCE } from '../../globals';
+import { OpenAIErrorResponseTransform } from '../openai/utils';
+import { ErrorResponse } from '../types';
+
+export const AzureAIInferenceResponseTransform = (
+  response: any,
+  responseStatus: number
+) => {
+  if (responseStatus !== 200 && 'error' in response) {
+    return OpenAIErrorResponseTransform(response, AZURE_AI_INFERENCE);
+  }
+
+  return { ...response, provider: AZURE_AI_INFERENCE };
+};
+
+export const AzureAIInferenceCreateSpeechResponseTransform = (
+  response: any,
+  responseStatus: number
+) => {
+  if (responseStatus !== 200 && 'error' in response) {
+    return OpenAIErrorResponseTransform(response, AZURE_AI_INFERENCE);
+  }
+
+  return { ...response, provider: AZURE_AI_INFERENCE };
+};
+
+export const AzureAIInferenceCreateTranscriptionResponseTransform = (
+  response: any,
+  responseStatus: number
+) => {
+  if (responseStatus !== 200 && 'error' in response) {
+    return OpenAIErrorResponseTransform(response, AZURE_AI_INFERENCE);
+  }
+
+  return { ...response, provider: AZURE_AI_INFERENCE };
+};
+
+export const AzureAIInferenceCreateTranslationResponseTransform = (
+  response: any,
+  responseStatus: number
+) => {
+  if (responseStatus !== 200 && 'error' in response) {
+    return OpenAIErrorResponseTransform(response, AZURE_AI_INFERENCE);
+  }
+
+  return { ...response, provider: AZURE_AI_INFERENCE };
+};
