openai[patch]: add AZURE_OPENAI_ENDPOINT supoprt (langchain-ai#6868)

Co-authored-by: jacoblee93 <[email protected]>

Author: baskaryan

libs/langchain-openai/src/azure/chat_models.ts

@@ -434,6 +434,11 @@ export class AzureChatOpenAI extends ChatOpenAI {
       openAIApiKey: "openai_api_key",
       openAIApiVersion: "openai_api_version",
       openAIBasePath: "openai_api_base",
+      deploymentName: "deployment_name",
+      azureOpenAIEndpoint: "azure_endpoint",
+      azureOpenAIApiVersion: "openai_api_version",
+      azureOpenAIBasePath: "openai_api_base",
+      azureOpenAIApiDeploymentName: "deployment_name",
     };
   }
 
@@ -477,6 +482,7 @@ export class AzureChatOpenAI extends ChatOpenAI {
         azureOpenAIBasePath: this.azureOpenAIBasePath,
         azureADTokenProvider: this.azureADTokenProvider,
         baseURL: this.clientConfig.baseURL,
+        azureOpenAIEndpoint: this.azureOpenAIEndpoint,
       };
 
       const endpoint = getEndpoint(openAIEndpointConfig);
@@ -541,6 +547,44 @@ export class AzureChatOpenAI extends ChatOpenAI {
       delete json.kwargs.azure_openai_api_key;
       delete json.kwargs.azure_openai_api_version;
       delete json.kwargs.azure_open_ai_base_path;
+
+      if (!json.kwargs.azure_endpoint && this.azureOpenAIEndpoint) {
+        json.kwargs.azure_endpoint = this.azureOpenAIEndpoint;
+      }
+      if (!json.kwargs.azure_endpoint && this.azureOpenAIBasePath) {
+        const parts = this.azureOpenAIBasePath.split("/openai/deployments/");
+        if (parts.length === 2 && parts[0].startsWith("http")) {
+          const [endpoint] = parts;
+          json.kwargs.azure_endpoint = endpoint;
+        }
+      }
+      if (!json.kwargs.azure_endpoint && this.azureOpenAIApiInstanceName) {
+        json.kwargs.azure_endpoint = `https://${this.azureOpenAIApiInstanceName}.openai.azure.com/`;
+      }
+      if (!json.kwargs.deployment_name && this.azureOpenAIApiDeploymentName) {
+        json.kwargs.deployment_name = this.azureOpenAIApiDeploymentName;
+      }
+      if (!json.kwargs.deployment_name && this.azureOpenAIBasePath) {
+        const parts = this.azureOpenAIBasePath.split("/openai/deployments/");
+        if (parts.length === 2) {
+          const [, deployment] = parts;
+          json.kwargs.deployment_name = deployment;
+        }
+      }
+
+      if (
+        json.kwargs.azure_endpoint &&
+        json.kwargs.deployment_name &&
+        json.kwargs.openai_api_base
+      ) {
+        delete json.kwargs.openai_api_base;
+      }
+      if (
+        json.kwargs.azure_openai_api_instance_name &&
+        json.kwargs.azure_endpoint
+      ) {
+        delete json.kwargs.azure_openai_api_instance_name;
+      }
     }
 
     return json;

libs/langchain-openai/src/chat_models.ts

@@ -933,6 +933,8 @@ export class ChatOpenAI<
 
   azureOpenAIBasePath?: string;
 
+  azureOpenAIEndpoint?: string;
+
   organization?: string;
 
   __includeRawResponse?: boolean;
@@ -993,6 +995,10 @@ export class ChatOpenAI<
       fields?.configuration?.organization ??
       getEnvironmentVariable("OPENAI_ORGANIZATION");
 
+    this.azureOpenAIEndpoint =
+      fields?.azureOpenAIEndpoint ??
+      getEnvironmentVariable("AZURE_OPENAI_ENDPOINT");
+
     this.modelName = fields?.model ?? fields?.modelName ?? this.model;
     this.model = this.modelName;
     this.modelKwargs = fields?.modelKwargs ?? {};
@@ -1013,9 +1019,21 @@ export class ChatOpenAI<
     this.__includeRawResponse = fields?.__includeRawResponse;
 
     if (this.azureOpenAIApiKey || this.azureADTokenProvider) {
-      if (!this.azureOpenAIApiInstanceName && !this.azureOpenAIBasePath) {
+      if (
+        !this.azureOpenAIApiInstanceName &&
+        !this.azureOpenAIBasePath &&
+        !this.azureOpenAIEndpoint
+      ) {
         throw new Error("Azure OpenAI API instance name not found");
       }
+
+      if (!this.azureOpenAIApiDeploymentName && this.azureOpenAIBasePath) {
+        const parts = this.azureOpenAIBasePath.split("/openai/deployments/");
+        if (parts.length === 2) {
+          const [, deployment] = parts;
+          this.azureOpenAIApiDeploymentName = deployment;
+        }
+      }
       if (!this.azureOpenAIApiDeploymentName) {
         throw new Error("Azure OpenAI API deployment name not found");
       }
@@ -1654,6 +1672,7 @@ export class ChatOpenAI<
         azureOpenAIApiKey: this.azureOpenAIApiKey,
         azureOpenAIBasePath: this.azureOpenAIBasePath,
         baseURL: this.clientConfig.baseURL,
+        azureOpenAIEndpoint: this.azureOpenAIEndpoint,
       };
 
       const endpoint = getEndpoint(openAIEndpointConfig);

libs/langchain-openai/src/tests/azure/chat_models.test.ts

@@ -0,0 +1,37 @@
+import { AzureChatOpenAI } from "../../azure/chat_models.js";
+
+test("Test Azure OpenAI serialization from azure endpoint", async () => {
+  const chat = new AzureChatOpenAI({
+    azureOpenAIEndpoint: "https://foobar.openai.azure.com/",
+    azureOpenAIApiDeploymentName: "gpt-4o",
+    azureOpenAIApiVersion: "2024-08-01-preview",
+    azureOpenAIApiKey: "foo",
+  });
+  expect(JSON.stringify(chat)).toEqual(
+    `{"lc":1,"type":"constructor","id":["langchain","chat_models","azure_openai","AzureChatOpenAI"],"kwargs":{"azure_endpoint":"https://foobar.openai.azure.com/","openai_api_key":{"lc":1,"type":"secret","id":["OPENAI_API_KEY"]},"deployment_name":"gpt-4o"}}`
+  );
+});
+
+test("Test Azure OpenAI serialization from base path", async () => {
+  const chat = new AzureChatOpenAI({
+    azureOpenAIBasePath:
+      "https://foobar.openai.azure.com/openai/deployments/gpt-4o",
+    azureOpenAIApiVersion: "2024-08-01-preview",
+    azureOpenAIApiKey: "foo",
+  });
+  expect(JSON.stringify(chat)).toEqual(
+    `{"lc":1,"type":"constructor","id":["langchain","chat_models","azure_openai","AzureChatOpenAI"],"kwargs":{"openai_api_key":{"lc":1,"type":"secret","id":["OPENAI_API_KEY"]},"azure_endpoint":"https://foobar.openai.azure.com","deployment_name":"gpt-4o"}}`
+  );
+});
+
+test("Test Azure OpenAI serialization from instance name", async () => {
+  const chat = new AzureChatOpenAI({
+    azureOpenAIApiInstanceName: "foobar",
+    azureOpenAIApiDeploymentName: "gpt-4o",
+    azureOpenAIApiVersion: "2024-08-01-preview",
+    azureOpenAIApiKey: "foo",
+  });
+  expect(JSON.stringify(chat)).toEqual(
+    `{"lc":1,"type":"constructor","id":["langchain","chat_models","azure_openai","AzureChatOpenAI"],"kwargs":{"openai_api_key":{"lc":1,"type":"secret","id":["OPENAI_API_KEY"]},"azure_endpoint":"https://foobar.openai.azure.com/","deployment_name":"gpt-4o"}}`
+  );
+});

libs/langchain-openai/src/types.ts

@@ -216,12 +216,19 @@ export declare interface AzureOpenAIInput {
   azureOpenAIApiCompletionsDeploymentName?: string;
 
   /**
-   * Custom endpoint for Azure OpenAI API. This is useful in case you have a deployment in another region.
+   * Custom base url for Azure OpenAI API. This is useful in case you have a deployment in another region.
    * e.g. setting this value to "https://westeurope.api.cognitive.microsoft.com/openai/deployments"
    * will be result in the endpoint URL: https://westeurope.api.cognitive.microsoft.com/openai/deployments/{DeploymentName}/
    */
   azureOpenAIBasePath?: string;
 
+  /**
+   * Custom endpoint for Azure OpenAI API. This is useful in case you have a deployment in another region.
+   * e.g. setting this value to "https://westeurope.api.cognitive.microsoft.com/"
+   * will be result in the endpoint URL: https://westeurope.api.cognitive.microsoft.com/openai/deployments/{DeploymentName}/
+   */
+  azureOpenAIEndpoint?: string;
+
   /**
    * A function that returns an access token for Microsoft Entra (formerly known as Azure Active Directory),
    * which will be invoked on every request.

libs/langchain-openai/src/utils/azure.ts

@@ -5,6 +5,7 @@ export interface OpenAIEndpointConfig {
   azureADTokenProvider?: () => Promise<string>;
   azureOpenAIBasePath?: string;
   baseURL?: string | null;
+  azureOpenAIEndpoint?: string;
 }
 
 /**
@@ -14,13 +15,15 @@ export interface OpenAIEndpointConfig {
  * @param {OpenAIEndpointConfig} config - The configuration object for the (Azure) endpoint.
  *
  * @property {string} config.azureOpenAIApiDeploymentName - The deployment name of Azure OpenAI.
- * @property {string} config.azureOpenAIApiInstanceName - The instance name of Azure OpenAI.
+ * @property {string} config.azureOpenAIApiInstanceName - The instance name of Azure OpenAI, e.g. `example-resource`.
  * @property {string} config.azureOpenAIApiKey - The API Key for Azure OpenAI.
- * @property {string} config.azureOpenAIBasePath - The base path for Azure OpenAI.
+ * @property {string} config.azureOpenAIBasePath - The base path for Azure OpenAI, e.g. `https://example-resource.azure.openai.com/openai/deployments/`.
  * @property {string} config.baseURL - Some other custom base path URL.
+ * @property {string} config.azureOpenAIEndpoint - The endpoint for the Azure OpenAI instance, e.g. `https://example-resource.azure.openai.com/`.
  *
  * The function operates as follows:
  * - If both `azureOpenAIBasePath` and `azureOpenAIApiDeploymentName` (plus `azureOpenAIApiKey`) are provided, it returns an URL combining these two parameters (`${azureOpenAIBasePath}/${azureOpenAIApiDeploymentName}`).
+ * - If both `azureOpenAIEndpoint` and `azureOpenAIApiDeploymentName` (plus `azureOpenAIApiKey`) are provided, it returns an URL combining these two parameters (`${azureOpenAIEndpoint}/openai/deployments/${azureOpenAIApiDeploymentName}`).
  * - If `azureOpenAIApiKey` is provided, it checks for `azureOpenAIApiInstanceName` and `azureOpenAIApiDeploymentName` and throws an error if any of these is missing. If both are provided, it generates an URL incorporating these parameters.
  * - If none of the above conditions are met, return any custom `baseURL`.
  * - The function returns the generated URL as a string, or undefined if no custom paths are specified.
@@ -37,6 +40,7 @@ export function getEndpoint(config: OpenAIEndpointConfig) {
     azureOpenAIBasePath,
     baseURL,
     azureADTokenProvider,
+    azureOpenAIEndpoint,
   } = config;
 
   if (
@@ -46,6 +50,13 @@ export function getEndpoint(config: OpenAIEndpointConfig) {
   ) {
     return `${azureOpenAIBasePath}/${azureOpenAIApiDeploymentName}`;
   }
+  if (
+    (azureOpenAIApiKey || azureADTokenProvider) &&
+    azureOpenAIEndpoint &&
+    azureOpenAIApiDeploymentName
+  ) {
+    return `${azureOpenAIEndpoint}/openai/deployments/${azureOpenAIApiDeploymentName}`;
+  }
 
   if (azureOpenAIApiKey || azureADTokenProvider) {
     if (!azureOpenAIApiInstanceName) {