Skip to content

Commit

Permalink
fix(azure): update build script (#825)
Browse files Browse the repository at this point in the history
  • Loading branch information
stainless-bot committed May 6, 2024
1 parent 625426c commit 258c191
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 19 deletions.
2 changes: 1 addition & 1 deletion scripts/utils/fix-index-exports.cjs
Original file line number Diff line number Diff line change
Expand Up @@ -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');
28 changes: 17 additions & 11 deletions src/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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}`,
Expand All @@ -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);
}
}

Expand Down
19 changes: 12 additions & 7 deletions tests/lib/azure.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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://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://localhost:5000/',
azureADTokenProvider: () => 'my token',
azureADTokenProvider: async () => 'my token',
apiKey: 'My API Key',
apiVersion,
}),
Expand Down

0 comments on commit 258c191

Please sign in to comment.