From a287cae1a12920e08102c0b3b2c2fa9a7b59096d Mon Sep 17 00:00:00 2001
From: Sarangan Rajamanickam <sarajama@microsoft.com>
Date: Wed, 15 May 2024 15:07:59 -0700
Subject: [PATCH] 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>
---
 libs/langchain-openai/src/azure/chat_models.ts | 8 ++++++++
 libs/langchain-openai/src/azure/embeddings.ts  | 8 ++++++++
 libs/langchain-openai/src/azure/llms.ts        | 8 ++++++++
 libs/langchain-openai/src/utils/azure.ts       | 4 +++-
 4 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/libs/langchain-openai/src/azure/chat_models.ts b/libs/langchain-openai/src/azure/chat_models.ts
index fc1a36f46581..49db684d53b0 100644
--- 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
index d5023bfbada9..81cdb181df21 100644
--- 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
index 546893113a06..08c87789692c 100644
--- 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
index d1af1b718f9f..d3db6386c62a 100644
--- 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"