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 <noda@free.fr> * Update libs/langchain-openai/src/azure/embeddings.ts Co-authored-by: Yohan Lasorsa <noda@free.fr> * Update libs/langchain-openai/src/azure/chat_models.ts Co-authored-by: Yohan Lasorsa <noda@free.fr> * Format the code --------- Co-authored-by: Yohan Lasorsa <noda@free.fr>
langchain-ai · May 15, 2024 · a287cae · a287cae
1 parent 69aca46
commit a287cae
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 1 deletion.
diff --git a/libs/langchain-openai/src/azure/chat_models.ts b/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,

diff --git a/libs/langchain-openai/src/azure/embeddings.ts b/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,

diff --git a/libs/langchain-openai/src/azure/llms.ts b/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,

diff --git a/libs/langchain-openai/src/utils/azure.ts b/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"