From 2f073e4e6c9cd0ff3ad434907da710704765a005 Mon Sep 17 00:00:00 2001
From: Stainless Bot <107565488+stainless-bot@users.noreply.github.com>
Date: Tue, 23 Jan 2024 18:09:39 +0100
Subject: [PATCH] chore(internal): minor streaming updates (#647)

---
 src/core.ts      | 7 +++++++
 src/streaming.ts | 6 +++---
 2 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/src/core.ts b/src/core.ts
index f431b7d58..d7c1fb831 100644
--- a/src/core.ts
+++ b/src/core.ts
@@ -44,6 +44,11 @@ async function defaultParseResponse<T>(props: APIResponseProps): Promise<T> {
 
     // Note: there is an invariant here that isn't represented in the type system
     // that if you set `stream: true` the response type must also be `Stream<T>`
+
+    if (props.options.__streamClass) {
+      return props.options.__streamClass.fromSSEResponse(response, props.controller) as any;
+    }
+
     return Stream.fromSSEResponse(response, props.controller) as any;
   }
 
@@ -743,6 +748,7 @@ export type RequestOptions<Req = unknown | Record<string, unknown> | Readable> =
   idempotencyKey?: string;
 
   __binaryResponse?: boolean | undefined;
+  __streamClass?: typeof Stream;
 };
 
 // This is required so that we can determine if a given object matches the RequestOptions
@@ -763,6 +769,7 @@ const requestOptionsKeys: KeysEnum<RequestOptions> = {
   idempotencyKey: true,
 
   __binaryResponse: true,
+  __streamClass: true,
 };
 
 export const isRequestOptions = (obj: unknown): obj is RequestOptions => {
diff --git a/src/streaming.ts b/src/streaming.ts
index d1ee70c2e..3ad11fb54 100644
--- a/src/streaming.ts
+++ b/src/streaming.ts
@@ -3,9 +3,9 @@ import { OpenAIError } from './error';
 
 import { APIError } from 'openai/error';
 
-type Bytes = string | ArrayBuffer | Uint8Array | Buffer | null | undefined;
+export type Bytes = string | ArrayBuffer | Uint8Array | Buffer | null | undefined;
 
-type ServerSentEvent = {
+export type ServerSentEvent = {
   event: string | null;
   data: string;
   raw: string[];
@@ -381,7 +381,7 @@ function partition(str: string, delimiter: string): [string, string, string] {
  *
  * This polyfill was pulled from https://github.com/MattiasBuelens/web-streams-polyfill/pull/122#issuecomment-1627354490
  */
-function readableStreamAsyncIterable<T>(stream: any): AsyncIterableIterator<T> {
+export function readableStreamAsyncIterable<T>(stream: any): AsyncIterableIterator<T> {
   if (stream[Symbol.asyncIterator]) return stream;
 
   const reader = stream.getReader();