Update default headers azure & Add Token Provider Check (#5379)

* Update Default Headers for Azure OpenAI Requests

* nit

* Update user agent string

* Add Token Provider check

* Update libs/langchain-openai/src/azure/llms.ts

Co-authored-by: Yohan Lasorsa <[email protected]>

* Update libs/langchain-openai/src/azure/embeddings.ts

Co-authored-by: Yohan Lasorsa <[email protected]>

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

Co-authored-by: Yohan Lasorsa <[email protected]>

* Format the code

---------

Co-authored-by: Yohan Lasorsa <[email protected]>

Author: 2 people

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

@@ -50,6 +50,7 @@ export class AzureChatOpenAI extends ChatOpenAI {
         azureOpenAIApiInstanceName: this.azureOpenAIApiInstanceName,
         azureOpenAIApiKey: this.azureOpenAIApiKey,
         azureOpenAIBasePath: this.azureOpenAIBasePath,
+        azureADTokenProvider: this.azureADTokenProvider,
         baseURL: this.clientConfig.baseURL,
       };
 
@@ -70,6 +71,13 @@ export class AzureChatOpenAI extends ChatOpenAI {
         delete params.baseURL;
       }
 
+      params.defaultHeaders = {
+        ...params.defaultHeaders,
+        "User-Agent": params.defaultHeaders?.["User-Agent"]
+          ? `${params.defaultHeaders["User-Agent"]}: langchainjs-azure-openai-v2`
+          : `langchainjs-azure-openai-v2`,
+      };
+
       this.client = new AzureOpenAIClient({
         apiVersion: this.azureOpenAIApiVersion,
         azureADTokenProvider: this.azureADTokenProvider,

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

@@ -44,6 +44,7 @@ export class AzureOpenAIEmbeddings extends OpenAIEmbeddings {
         azureOpenAIApiInstanceName: this.azureOpenAIApiInstanceName,
         azureOpenAIApiKey: this.azureOpenAIApiKey,
         azureOpenAIBasePath: this.azureOpenAIBasePath,
+        azureADTokenProvider: this.azureADTokenProvider,
         baseURL: this.clientConfig.baseURL,
       };
 
@@ -64,6 +65,13 @@ export class AzureOpenAIEmbeddings extends OpenAIEmbeddings {
         delete params.baseURL;
       }
 
+      params.defaultHeaders = {
+        ...params.defaultHeaders,
+        "User-Agent": params.defaultHeaders?.["User-Agent"]
+          ? `${params.defaultHeaders["User-Agent"]}: langchainjs-azure-openai-v2`
+          : `langchainjs-azure-openai-v2`,
+      };
+
       this.client = new AzureOpenAIClient({
         apiVersion: this.azureOpenAIApiVersion,
         azureADTokenProvider: this.azureADTokenProvider,

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

@@ -46,6 +46,7 @@ export class AzureOpenAI extends OpenAI {
         azureOpenAIApiInstanceName: this.azureOpenAIApiInstanceName,
         azureOpenAIApiKey: this.azureOpenAIApiKey,
         azureOpenAIBasePath: this.azureOpenAIBasePath,
+        azureADTokenProvider: this.azureADTokenProvider,
         baseURL: this.clientConfig.baseURL,
       };
 
@@ -66,6 +67,13 @@ export class AzureOpenAI extends OpenAI {
         delete params.baseURL;
       }
 
+      params.defaultHeaders = {
+        ...params.defaultHeaders,
+        "User-Agent": params.defaultHeaders?.["User-Agent"]
+          ? `${params.defaultHeaders["User-Agent"]}: langchainjs-azure-openai-v2`
+          : `langchainjs-azure-openai-v2`,
+      };
+
       this.client = new AzureOpenAIClient({
         apiVersion: this.azureOpenAIApiVersion,
         azureADTokenProvider: this.azureADTokenProvider,

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

@@ -2,6 +2,7 @@ export interface OpenAIEndpointConfig {
   azureOpenAIApiDeploymentName?: string;
   azureOpenAIApiInstanceName?: string;
   azureOpenAIApiKey?: string;
+  azureADTokenProvider?: () => Promise<string>;
   azureOpenAIBasePath?: string;
   baseURL?: string | null;
 }
@@ -35,6 +36,7 @@ export function getEndpoint(config: OpenAIEndpointConfig) {
     azureOpenAIApiKey,
     azureOpenAIBasePath,
     baseURL,
+    azureADTokenProvider,
   } = config;
 
   if (
@@ -45,7 +47,7 @@ export function getEndpoint(config: OpenAIEndpointConfig) {
     return `${azureOpenAIBasePath}/${azureOpenAIApiDeploymentName}`;
   }
 
-  if (azureOpenAIApiKey) {
+  if (azureOpenAIApiKey || azureADTokenProvider) {
     if (!azureOpenAIApiInstanceName) {
       throw new Error(
         "azureOpenAIApiInstanceName is required when using azureOpenAIApiKey"