From a883dd37637d074dfe3a6c823ece3702e1a0ba56 Mon Sep 17 00:00:00 2001
From: Umut Hope YILDIRIM <umut475@gmail.com>
Date: Wed, 5 Feb 2025 17:12:28 -0500
Subject: [PATCH] feat: allow azure openai to pass custom uri

---
 src/handlers/handlerUtils.ts      |  1 +
 src/providers/azure-openai/api.ts | 31 ++++++++++++++++++++++---------
 src/types/requestBody.ts          |  1 +
 3 files changed, 24 insertions(+), 9 deletions(-)

diff --git a/src/handlers/handlerUtils.ts b/src/handlers/handlerUtils.ts
index b8c8419d..f0e4b2e3 100644
--- a/src/handlers/handlerUtils.ts
+++ b/src/handlers/handlerUtils.ts
@@ -896,6 +896,7 @@ export function constructConfigFromRequestHeaders(
     openaiBeta:
       requestHeaders[`x-${POWERED_BY}-openai-beta`] ||
       requestHeaders[`openai-beta`],
+    azureCustomUri: requestHeaders[`x-${POWERED_BY}-azure-custom-uri`],
   };
 
   const stabilityAiConfig = {
diff --git a/src/providers/azure-openai/api.ts b/src/providers/azure-openai/api.ts
index 0786e2df..2ee29783 100644
--- a/src/providers/azure-openai/api.ts
+++ b/src/providers/azure-openai/api.ts
@@ -6,7 +6,15 @@ import {
 
 const AzureOpenAIAPIConfig: ProviderAPIConfig = {
   getBaseURL: ({ providerOptions }) => {
-    const { resourceName } = providerOptions;
+    const { resourceName, azureCustomUri } = providerOptions;
+    /**
+     * Uses the standard OpenAI endpoint (openai.azure.com) by default.
+     * If a custom subdomain is provided, that endpoint is used instead.
+     * Example: DeepSeek-R1 uses `https://{resourceName}.services.ai.azure.com/models`.
+     */
+    if (azureCustomUri) {
+      return `https://${resourceName}.${azureCustomUri}`;
+    }
     return `https://${resourceName}.openai.azure.com/openai`;
   },
   headers: async ({ providerOptions, fn }) => {
@@ -55,7 +63,8 @@ const AzureOpenAIAPIConfig: ProviderAPIConfig = {
     return headersObj;
   },
   getEndpoint: ({ providerOptions, fn, gatewayRequestURL }) => {
-    const { apiVersion, urlToFetch, deploymentId } = providerOptions;
+    const { apiVersion, urlToFetch, deploymentId, azureCustomUri } =
+      providerOptions;
     let mappedFn = fn;
 
     if (fn === 'proxy' && urlToFetch) {
@@ -79,27 +88,31 @@ const AzureOpenAIAPIConfig: ProviderAPIConfig = {
     const segments = gatewayRequestURL?.split('/');
     const id = segments?.at(-1) ?? '';
 
+    const deploymentPrefix = azureCustomUri
+      ? ''
+      : `/deployments/${deploymentId}`;
+
     switch (mappedFn) {
       case 'complete': {
-        return `/deployments/${deploymentId}/completions?api-version=${apiVersion}`;
+        return `${deploymentPrefix}/completions?api-version=${apiVersion}`;
       }
       case 'chatComplete': {
-        return `/deployments/${deploymentId}/chat/completions?api-version=${apiVersion}`;
+        return `${deploymentPrefix}/chat/completions?api-version=${apiVersion}`;
       }
       case 'embed': {
-        return `/deployments/${deploymentId}/embeddings?api-version=${apiVersion}`;
+        return `${deploymentPrefix}/embeddings?api-version=${apiVersion}`;
       }
       case 'imageGenerate': {
-        return `/deployments/${deploymentId}/images/generations?api-version=${apiVersion}`;
+        return `${deploymentPrefix}/images/generations?api-version=${apiVersion}`;
       }
       case 'createSpeech': {
-        return `/deployments/${deploymentId}/audio/speech?api-version=${apiVersion}`;
+        return `${deploymentPrefix}/audio/speech?api-version=${apiVersion}`;
       }
       case 'createTranscription': {
-        return `/deployments/${deploymentId}/audio/transcriptions?api-version=${apiVersion}`;
+        return `${deploymentPrefix}/audio/transcriptions?api-version=${apiVersion}`;
       }
       case 'createTranslation': {
-        return `/deployments/${deploymentId}/audio/translations?api-version=${apiVersion}`;
+        return `${deploymentPrefix}/audio/translations?api-version=${apiVersion}`;
       }
       case 'realtime': {
         return `/realtime?api-version=${apiVersion}&deployment=${providerOptions.deploymentId}`;
diff --git a/src/types/requestBody.ts b/src/types/requestBody.ts
index f4eeac8a..0d58d925 100644
--- a/src/types/requestBody.ts
+++ b/src/types/requestBody.ts
@@ -65,6 +65,7 @@ export interface Options {
   azureEntraClientId?: string;
   azureEntraClientSecret?: string;
   azureEntraTenantId?: string;
+  azureCustomUri?: string;
   /** Workers AI specific */
   workersAiAccountId?: string;
   /** The parameter to set custom base url */