feat(client): add ._request_id property to object responses

RobertCraigie · RobertCraigie · commit 90f9eddee57e · 2024-09-16T20:16:04.000+01:00
diff --git a/src/core.ts b/src/core.ts
@@ -37,7 +37,7 @@ type APIResponseProps = {
   controller: AbortController;
 };
 
-async function defaultParseResponse<T>(props: APIResponseProps): Promise<T> {
+async function defaultParseResponse<T>(props: APIResponseProps): Promise<WithRequestID<T>> {
   const { response } = props;
   if (props.options.stream) {
     debug('response', response.status, response.url, response.headers, response.body);
@@ -54,11 +54,11 @@ async function defaultParseResponse<T>(props: APIResponseProps): Promise<T> {
 
   // fetch refuses to read the body when the status code is 204.
   if (response.status === 204) {
-    return null as T;
+    return null as WithRequestID<T>;
   }
 
   if (props.options.__binaryResponse) {
-    return response as unknown as T;
+    return response as unknown as WithRequestID<T>;
   }
 
   const contentType = response.headers.get('content-type');
@@ -69,26 +69,44 @@ async function defaultParseResponse<T>(props: APIResponseProps): Promise<T> {
 
     debug('response', response.status, response.url, response.headers, json);
 
-    return json as T;
+    return _addRequestID(json, response);
   }
 
   const text = await response.text();
   debug('response', response.status, response.url, response.headers, text);
 
   // TODO handle blob, arraybuffer, other content types, etc.
-  return text as unknown as T;
+  return text as unknown as WithRequestID<T>;
+}
+
+type WithRequestID<T> =
+  T extends Array<any> | Response | AbstractPage<any> ? T
+  : T extends Record<string, any> ? T & { _request_id?: string | null }
+  : T;
+
+function _addRequestID<T>(value: T, response: Response): WithRequestID<T> {
+  if (!value || typeof value !== 'object' || Array.isArray(value)) {
+    return value as WithRequestID<T>;
+  }
+
+  return Object.defineProperty(value, '_request_id', {
+    value: response.headers.get('x-request-id'),
+    enumerable: false,
+  }) as WithRequestID<T>;
 }
 
 /**
  * A subclass of `Promise` providing additional helper methods
  * for interacting with the SDK.
  */
-export class APIPromise<T> extends Promise<T> {
-  private parsedPromise: Promise<T> | undefined;
+export class APIPromise<T> extends Promise<WithRequestID<T>> {
+  private parsedPromise: Promise<WithRequestID<T>> | undefined;
 
   constructor(
     private responsePromise: Promise<APIResponseProps>,
-    private parseResponse: (props: APIResponseProps) => PromiseOrValue<T> = defaultParseResponse,
+    private parseResponse: (
+      props: APIResponseProps,
+    ) => PromiseOrValue<WithRequestID<T>> = defaultParseResponse,
   ) {
     super((resolve) => {
       // this is maybe a bit weird but this has to be a no-op to not implicitly
@@ -99,7 +117,9 @@ export class APIPromise<T> extends Promise<T> {
   }
 
   _thenUnwrap<U>(transform: (data: T) => U): APIPromise<U> {
-    return new APIPromise(this.responsePromise, async (props) => transform(await this.parseResponse(props)));
+    return new APIPromise(this.responsePromise, async (props) =>
+      _addRequestID(transform(await this.parseResponse(props)), props.response),
+    );
   }
 
   /**
@@ -136,27 +156,27 @@ export class APIPromise<T> extends Promise<T> {
     return { data, response };
   }
 
-  private parse(): Promise<T> {
+  private parse(): Promise<WithRequestID<T>> {
     if (!this.parsedPromise) {
-      this.parsedPromise = this.responsePromise.then(this.parseResponse);
+      this.parsedPromise = this.responsePromise.then(this.parseResponse) as any as Promise<WithRequestID<T>>;
     }
     return this.parsedPromise;
   }
 
-  override then<TResult1 = T, TResult2 = never>(
-    onfulfilled?: ((value: T) => TResult1 | PromiseLike<TResult1>) | undefined | null,
+  override then<TResult1 = WithRequestID<T>, TResult2 = never>(
+    onfulfilled?: ((value: WithRequestID<T>) => TResult1 | PromiseLike<TResult1>) | undefined | null,
     onrejected?: ((reason: any) => TResult2 | PromiseLike<TResult2>) | undefined | null,
   ): Promise<TResult1 | TResult2> {
     return this.parse().then(onfulfilled, onrejected);
   }
 
   override catch<TResult = never>(
     onrejected?: ((reason: any) => TResult | PromiseLike<TResult>) | undefined | null,
-  ): Promise<T | TResult> {
+  ): Promise<WithRequestID<T> | TResult> {
     return this.parse().catch(onrejected);
   }
 
-  override finally(onfinally?: (() => void) | undefined | null): Promise<T> {
+  override finally(onfinally?: (() => void) | undefined | null): Promise<WithRequestID<T>> {
     return this.parse().finally(onfinally);
   }
 }
@@ -706,7 +726,13 @@ export class PagePromise<
   ) {
     super(
       request,
-      async (props) => new Page(client, props.response, await defaultParseResponse(props), props.options),
+      async (props) =>
+        new Page(
+          client,
+          props.response,
+          await defaultParseResponse(props),
+          props.options,
+        ) as WithRequestID<PageClass>,
     );
   }
 
diff --git a/tests/responses.test.ts b/tests/responses.test.ts
@@ -1,5 +1,8 @@
-import { createResponseHeaders } from 'openai/core';
+import { APIPromise, createResponseHeaders } from 'openai/core';
+import OpenAI from 'openai/index';
 import { Headers } from 'openai/_shims/index';
+import { Response } from 'node-fetch';
+import { compareType } from './utils/typing';
 
 describe('response parsing', () => {
   // TODO: test unicode characters
@@ -23,3 +26,99 @@ describe('response parsing', () => {
     expect(headers['content-type']).toBe('text/xml, application/json');
   });
 });
+
+describe('request id', () => {
+  test('types', () => {
+    compareType<Awaited<APIPromise<string>>, string>(true);
+    compareType<Awaited<APIPromise<number>>, number>(true);
+    compareType<Awaited<APIPromise<null>>, null>(true);
+    compareType<Awaited<APIPromise<void>>, void>(true);
+    compareType<Awaited<APIPromise<Response>>, Response>(true);
+    compareType<Awaited<APIPromise<Response>>, Response>(true);
+    compareType<Awaited<APIPromise<{ foo: string }>>, { foo: string } & { _request_id?: string | null }>(
+      true,
+    );
+    compareType<Awaited<APIPromise<Array<{ foo: string }>>>, Array<{ foo: string }>>(true);
+  });
+
+  test('object response', async () => {
+    const client = new OpenAI({
+      fetch: async () =>
+        new Response(JSON.stringify({ id: 'bar' }), {
+          headers: { 'x-request-id': 'req_id_xxx', 'content-type': 'application/json' },
+        }),
+    });
+
+    const rsp = await client.chat.completions.create({ messages: [], model: 'gpt-4' });
+    expect(rsp.id).toBe('bar');
+    expect(rsp._request_id).toBe('req_id_xxx');
+  });
+
+  test('envelope response', async () => {
+    const promise = new APIPromise<{ data: { foo: string } }>(
+      (async () => {
+        return {
+          response: new Response(JSON.stringify({ data: { foo: 'bar' } }), {
+            headers: { 'x-request-id': 'req_id_xxx', 'content-type': 'application/json' },
+          }),
+          controller: {} as any,
+          options: {} as any,
+        };
+      })(),
+    )._thenUnwrap((d) => d.data);
+
+    const rsp = await promise;
+    expect(rsp.foo).toBe('bar');
+    expect(rsp._request_id).toBe('req_id_xxx');
+  });
+
+  test('page response', async () => {
+    const client = new OpenAI({
+      fetch: async () =>
+        new Response(JSON.stringify({ data: [{ foo: 'bar' }] }), {
+          headers: { 'x-request-id': 'req_id_xxx', 'content-type': 'application/json' },
+        }),
+    });
+
+    const page = await client.fineTuning.jobs.list();
+    expect(page.data).toMatchObject([{ foo: 'bar' }]);
+    expect((page as any)._request_id).toBeUndefined();
+  });
+
+  test('array response', async () => {
+    const promise = new APIPromise<Array<{ foo: string }>>(
+      (async () => {
+        return {
+          response: new Response(JSON.stringify([{ foo: 'bar' }]), {
+            headers: { 'x-request-id': 'req_id_xxx', 'content-type': 'application/json' },
+          }),
+          controller: {} as any,
+          options: {} as any,
+        };
+      })(),
+    );
+
+    const rsp = await promise;
+    expect(rsp.length).toBe(1);
+    expect(rsp[0]).toMatchObject({ foo: 'bar' });
+    expect((rsp as any)._request_id).toBeUndefined();
+  });
+
+  test('string response', async () => {
+    const promise = new APIPromise<string>(
+      (async () => {
+        return {
+          response: new Response('hello world', {
+            headers: { 'x-request-id': 'req_id_xxx', 'content-type': 'application/text' },
+          }),
+          controller: {} as any,
+          options: {} as any,
+        };
+      })(),
+    );
+
+    const result = await promise;
+    expect(result).toBe('hello world');
+    expect((result as any)._request_id).toBeUndefined();
+  });
+});
diff --git a/tests/utils/typing.ts b/tests/utils/typing.ts
@@ -0,0 +1,9 @@
+type Equal<X, Y> = (<T>() => T extends X ? 1 : 2) extends <T>() => T extends Y ? 1 : 2 ? true : false;
+
+export const expectType = <T>(_expression: T): void => {
+  return;
+};
+
+export const compareType = <T1, T2>(_expression: Equal<T1, T2>): void => {
+  return;
+};
