fix(azure): include retry count header

openai · Jan 22, 2025 · 3e0ba40 · 3e0ba40
1 parent b7ab6bb
commit 3e0ba40
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/src/index.ts b/src/index.ts
@@ -577,7 +577,10 @@ export class AzureOpenAI extends OpenAI {
     this._deployment = deployment;
   }
 
-  override buildRequest(options: Core.FinalRequestOptions<unknown>): {
+  override buildRequest(
+    options: Core.FinalRequestOptions<unknown>,
+    props: { retryCount?: number } = {},
+  ): {
     req: RequestInit;
     url: string;
     timeout: number;
@@ -591,7 +594,7 @@ export class AzureOpenAI extends OpenAI {
         options.path = `/deployments/${model}${options.path}`;
       }
     }
-    return super.buildRequest(options);
+    return super.buildRequest(options, props);
   }
 
   private async _getAzureADToken(): Promise<string | undefined> {

diff --git a/tests/lib/azure.test.ts b/tests/lib/azure.test.ts
@@ -51,6 +51,18 @@ describe('instantiate azure client', () => {
       });
       expect(req.headers as Headers).not.toHaveProperty('x-my-default-header');
     });
+
+    test('includes retry count', () => {
+      const { req } = client.buildRequest(
+        {
+          path: '/foo',
+          method: 'post',
+          headers: { 'X-My-Default-Header': null },
+        },
+        { retryCount: 1 },
+      );
+      expect((req.headers as Headers)['x-stainless-retry-count']).toEqual('1');
+    });
   });
 
   describe('defaultQuery', () => {