feat: streaming improvements (#411)

openai · Oct 30, 2023 · 37b622c · 37b622c
1 parent 510c1f3
commit 37b622c
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 37 deletions.
diff --git a/src/lib/ChatCompletionRunner.ts b/src/lib/ChatCompletionRunner.ts
@@ -3,7 +3,6 @@ import {
   type Completions,
   type ChatCompletionMessage,
   type ChatCompletionMessageParam,
-  type ChatCompletionCreateParams,
   type ChatCompletionCreateParamsNonStreaming,
 } from 'openai/resources/chat/completions';
 import { type RunnableFunctions, type BaseFunctionsArgs } from './RunnableFunction';
@@ -34,16 +33,6 @@ export class ChatCompletionRunner extends AbstractChatCompletionRunner<ChatCompl
     return runner;
   }
 
-  static createChatCompletion(
-    completions: Completions,
-    params: ChatCompletionCreateParams,
-    options?: Core.RequestOptions,
-  ): ChatCompletionRunner {
-    const runner = new ChatCompletionRunner();
-    runner._run(() => runner._runChatCompletion(completions, params, options));
-    return runner;
-  }
-
   override _addMessage(message: ChatCompletionMessage | ChatCompletionMessageParam) {
     super._addMessage(message);
     if (message.role === 'assistant' && message.content) {

diff --git a/src/lib/ChatCompletionStream.ts b/src/lib/ChatCompletionStream.ts
@@ -5,7 +5,7 @@ import {
   type ChatCompletion,
   type ChatCompletionChunk,
   type ChatCompletionCreateParams,
-  type ChatCompletionCreateParamsStreaming,
+  ChatCompletionCreateParamsBase,
 } from 'openai/resources/chat/completions';
 import {
   AbstractChatCompletionRunner,
@@ -19,7 +19,9 @@ export interface ChatCompletionStreamEvents extends AbstractChatCompletionRunner
   chunk: (chunk: ChatCompletionChunk, snapshot: ChatCompletionSnapshot) => void;
 }
 
-export type ChatCompletionStreamParams = ChatCompletionCreateParamsStreaming;
+export type ChatCompletionStreamParams = Omit<ChatCompletionCreateParamsBase, 'stream'> & {
+  stream?: true;
+};
 
 export class ChatCompletionStream
   extends AbstractChatCompletionRunner<ChatCompletionStreamEvents>
@@ -31,6 +33,13 @@ export class ChatCompletionStream
     return this.#currentChatCompletionSnapshot;
   }
 
+  /**
+   * Intended for use on the frontend, consuming a stream produced with
+   * `.toReadableStream()` on the backend.
+   *
+   * Note that messages sent to the model do not appear in `.on('message')`
+   * in this context.
+   */
   static fromReadableStream(stream: ReadableStream): ChatCompletionStream {
     const runner = new ChatCompletionStream();
     runner._run(() => runner._fromReadableStream(stream));
@@ -39,11 +48,11 @@ export class ChatCompletionStream
 
   static createChatCompletion(
     completions: Completions,
-    params: ChatCompletionCreateParams,
+    params: ChatCompletionStreamParams,
     options?: Core.RequestOptions,
   ): ChatCompletionStream {
     const runner = new ChatCompletionStream();
-    runner._run(() => runner._runChatCompletion(completions, params, options));
+    runner._run(() => runner._runChatCompletion(completions, { ...params, stream: true }, options));
     return runner;
   }
 
@@ -110,8 +119,15 @@ export class ChatCompletionStream
     this.#beginRequest();
     this._connected();
     const stream = Stream.fromReadableStream<ChatCompletionChunk>(readableStream, this.controller);
+    let chatId;
     for await (const chunk of stream) {
+      if (chatId && chatId !== chunk.id) {
+        // A new request has been made.
+        this._addChatCompletion(this.#endRequest());
+      }
+
       this.#addChunk(chunk);
+      chatId = chunk.id;
     }
     if (stream.controller.signal?.aborted) {
       throw new APIUserAbortError();

diff --git a/src/lib/ChatCompletionStreamingRunner.ts b/src/lib/ChatCompletionStreamingRunner.ts
@@ -2,7 +2,6 @@ import * as Core from 'openai/core';
 import {
   Completions,
   type ChatCompletionChunk,
-  type ChatCompletionCreateParams,
   type ChatCompletionCreateParamsStreaming,
 } from 'openai/resources/chat/completions';
 import { type AbstractChatCompletionRunnerEvents } from './AbstractChatCompletionRunner';
@@ -41,14 +40,4 @@ export class ChatCompletionStreamingRunner
     runner._run(() => runner._runFunctions(completions, params, options));
     return runner;
   }
-
-  static override createChatCompletion(
-    completions: Completions,
-    params: ChatCompletionCreateParams,
-    options?: Core.RequestOptions,
-  ): ChatCompletionStreamingRunner {
-    const runner = new ChatCompletionStreamingRunner();
-    runner._run(() => runner._runChatCompletion(completions, params, options));
-    return runner;
-  }
 }
diff --git a/src/resources/beta/chat/completions.ts b/src/resources/beta/chat/completions.ts
@@ -20,8 +20,8 @@ export {
   RunnableFunctionWithoutParse,
   ParsingFunction,
 } from 'openai/lib/RunnableFunction';
-import { ChatCompletionStream } from 'openai/lib/ChatCompletionStream';
-import { ChatCompletionCreateParamsStreaming } from 'openai/resources/chat/completions';
+import { ChatCompletionStream, type ChatCompletionStreamParams } from 'openai/lib/ChatCompletionStream';
+export { ChatCompletionStream, type ChatCompletionStreamParams } from 'openai/lib/ChatCompletionStream';
 
 export class Completions extends APIResource {
   /**
@@ -64,14 +64,7 @@ export class Completions extends APIResource {
   /**
    * Creates a chat completion stream
    */
-  stream(
-    body: Omit<ChatCompletionCreateParamsStreaming, 'stream'> & { stream?: true },
-    options?: Core.RequestOptions,
-  ): ChatCompletionStream {
-    return ChatCompletionStream.createChatCompletion(
-      this.client.chat.completions,
-      { ...body, stream: true },
-      options,
-    );
+  stream(body: ChatCompletionStreamParams, options?: Core.RequestOptions): ChatCompletionStream {
+    return ChatCompletionStream.createChatCompletion(this.client.chat.completions, body, options);
   }
 }