fix(azure): update build script

scripts/utils/fix-index-exports.cjs

@@ -9,6 +9,6 @@ const indexJs =
 let before = fs.readFileSync(indexJs, 'utf8');
 let after = before.replace(
  /^\s*exports\.default\s*=\s*(\w+)/m,
- 'exports = module.exports = $1;\nexports.default = $1',
+ 'exports = module.exports = $1;\nmodule.exports.AzureOpenAI = AzureOpenAI;\nexports.default = $1',
 );
 fs.writeFileSync(indexJs, after, 'utf8');

src/index.ts

@@ -339,12 +339,12 @@ export interface AzureClientOptions extends ClientOptions {
  * A function that returns an access token for Microsoft Entra (formerly known as Azure Active Directory),
  * which will be invoked on every request.
  */
- azureADTokenProvider?: (() => string) | undefined;
+ azureADTokenProvider?: (() => Promise<string>) | undefined;
 }
 
 /** API Client for interfacing with the Azure OpenAI API. */
 export class AzureOpenAI extends OpenAI {
- private _azureADTokenProvider: (() => string) | undefined;
+ private _azureADTokenProvider: (() => Promise<string>) | undefined;
  apiVersion: string = '';
  /**
  * API Client for interfacing with the Azure OpenAI API.
@@ -451,9 +451,9 @@ export class AzureOpenAI extends OpenAI {
  return super.buildRequest(options);
  }
 
- private _getAzureADToken(): string | undefined {
+ private async _getAzureADToken(): Promise<string | undefined> {
  if (typeof this._azureADTokenProvider === 'function') {
- const token = this._azureADTokenProvider();
+ const token = await this._azureADTokenProvider();
  if (!token || typeof token !== 'string') {
  throw new Errors.OpenAIError(
  `Expected 'azureADTokenProvider' argument to return a string but it returned ${token}`,
@@ -465,17 +465,23 @@ export class AzureOpenAI extends OpenAI {
  }
 
  protected override authHeaders(opts: Core.FinalRequestOptions): Core.Headers {
+ return {};
+ }
+
+ protected override async prepareOptions(opts: Core.FinalRequestOptions<unknown>): Promise<void> {
  if (opts.headers?.['Authorization'] || opts.headers?.['api-key']) {
- return {};
+ return super.prepareOptions(opts);
  }
- const token = this._getAzureADToken();
+ const token = await this._getAzureADToken();
+ opts.headers ??= {};
  if (token) {
- return { Authorization: `Bearer ${token}` };
- }
- if (this.apiKey !== API_KEY_SENTINEL) {
- return { 'api-key': this.apiKey };
+ opts.headers['Authorization'] = `Bearer ${token}`;
+ } else if (this.apiKey !== API_KEY_SENTINEL) {
+ opts.headers['api-key'] = this.apiKey;
+ } else {
+ throw new Errors.OpenAIError('Unable to handle auth');
  }
- throw new Errors.OpenAIError('Unable to handle auth');
+ return super.prepareOptions(opts);
  }
 }
 

tests/lib/azure.test.ts

@@ -222,24 +222,29 @@ describe('instantiate azure client', () => {
  });
 
  describe('Azure Active Directory (AD)', () => {
- test('with azureADTokenProvider', () => {
+ test('with azureADTokenProvider', async () => {
+ const testFetch = async (url: RequestInfo, { headers }: RequestInit = {}): Promise<Response> => {
+ return new Response(JSON.stringify({ a: 1 }), { headers });
+ };
  const client = new AzureOpenAI({
  baseURL: 'http:https://localhost:5000/',
- azureADTokenProvider: () => 'my token',
+ azureADTokenProvider: async () => 'my token',
  apiVersion,
+ fetch: testFetch,
  });
- expect(client.buildRequest({ method: 'post', path: 'https://example.com' }).req.headers).toHaveProperty(
- 'authorization',
- 'Bearer my token',
- );
+ expect(
+ (await client.request({ method: 'post', path: 'https://example.com' }).asResponse()).headers.get(
+ 'authorization',
+ ),
+ ).toEqual('Bearer my token');
  });
 
  test('apiKey and azureADTokenProvider cant be combined', () => {
  expect(
  () =>
  new AzureOpenAI({
  baseURL: 'http:https://localhost:5000/',
- azureADTokenProvider: () => 'my token',
+ azureADTokenProvider: async () => 'my token',
  apiKey: 'My API Key',
  apiVersion,
  }),