feat: stream support (#14)

Author: lino-levan

deno.json

@@ -1,3 +1,4 @@
 {
-  "tasks": {}
+  "tasks": {},
+  "lock": false
 }

examples/chatCompletionStream.ts

@@ -0,0 +1,18 @@
+import { OpenAI } from "../mod.ts";
+
+const openAI = new OpenAI(Deno.env.get("YOUR_API_KEY")!);
+
+await openAI.createChatCompletionStream({
+  model: "gpt-3.5-turbo",
+  messages: [
+    { "role": "system", "content": "You are a helpful assistant." },
+    { "role": "user", "content": "Who won the world series in 2020?" },
+    {
+      "role": "assistant",
+      "content": "The Los Angeles Dodgers won the World Series in 2020.",
+    },
+    { "role": "user", "content": "Where was it played?" },
+  ],
+}, (chunk) => {
+  console.log(chunk);
+});

examples/completionStream.ts

@@ -0,0 +1,10 @@
+import { OpenAI } from "../mod.ts";
+
+const openAI = new OpenAI(Deno.env.get("YOUR_API_KEY")!);
+
+openAI.createCompletionStream({
+  model: "davinci",
+  prompt: "The meaning of life is",
+}, (chunk) => {
+  console.log(chunk);
+});

src/openai.ts

@@ -1,8 +1,12 @@
+import { basename } from "https://deno.land/std@0.189.0/path/mod.ts";
+import { decodeStream, throwError } from "./util.ts";
 import type {
   ChatCompletion,
   ChatCompletionOptions,
+  ChatCompletionStream,
   Completion,
   CompletionOptions,
+  CompletionStream,
   DeletedFile,
   DeletedFineTune,
   Edit,
@@ -29,7 +33,6 @@ import type {
   Translation,
   TranslationOptions,
 } from "./types.ts";
-import { basename } from "https://deno.land/std@0.187.0/path/mod.ts";
 
 const defaultBaseUrl = "https://api.openai.com/v1";
 
@@ -67,17 +70,7 @@ export class OpenAI {
     );
     const data = await response.json();
 
-    if (data.error) {
-      let errorMessage = `${data.error.type}`;
-      if (data.error.message) {
-        errorMessage += ": " + data.error.message;
-      }
-      if (data.error.code) {
-        errorMessage += ` (${data.error.code})`;
-      }
-      console.log(data.error);
-      throw new Error(errorMessage);
-    }
+    throwError(data);
 
     return data;
   }
@@ -108,7 +101,6 @@ export class OpenAI {
    * https://platform.openai.com/docs/api-reference/completions/create
    */
   async createCompletion(options: CompletionOptions): Promise<Completion> {
-    // TODO: make options.stream work
     return await this.#request(`/completions`, {
       model: options.model,
       prompt: options.prompt,
@@ -117,7 +109,6 @@ export class OpenAI {
       temperature: options.temperature,
       top_p: options.topP,
       n: options.n,
-      stream: options.stream,
       logprobs: options.logprobs,
       echo: options.echo,
       stop: options.stop,
@@ -129,6 +120,46 @@ export class OpenAI {
     });
   }
 
+  /**
+   * Creates a completion stream for the provided prompt and parameters
+   *
+   * https://platform.openai.com/docs/api-reference/completions/create
+   */
+  async createCompletionStream(
+    options: Omit<CompletionOptions, "bestOf">,
+    callback: (chunk: CompletionStream) => void,
+  ): Promise<void> {
+    const res = await fetch(
+      `${this.#baseUrl}/completions`,
+      {
+        method: "POST",
+        headers: {
+          Authorization: `Bearer ${this.#privateKey}`,
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+          model: options.model,
+          prompt: options.prompt,
+          suffix: options.suffix,
+          max_tokens: options.maxTokens,
+          temperature: options.temperature,
+          top_p: options.topP,
+          n: options.n,
+          stream: true,
+          logprobs: options.logprobs,
+          echo: options.echo,
+          stop: options.stop,
+          presence_penalty: options.presencePenalty,
+          frequency_penalty: options.frequencyPenalty,
+          logit_bias: options.logitBias,
+          user: options.user,
+        }),
+      },
+    );
+
+    await decodeStream(res, callback);
+  }
+
   /**
    * Creates a completion for the chat message
    *
@@ -143,7 +174,6 @@ export class OpenAI {
       temperature: options.temperature,
       top_p: options.topP,
       n: options.n,
-      stream: options.stream,
       stop: options.stop,
       max_tokens: options.maxTokens,
       presence_penalty: options.presencePenalty,
@@ -153,6 +183,43 @@ export class OpenAI {
     });
   }
 
+  /**
+   * Creates a completion stream for the chat message
+   *
+   * https://platform.openai.com/docs/api-reference/chat/create
+   */
+  async createChatCompletionStream(
+    options: ChatCompletionOptions,
+    callback: (chunk: ChatCompletionStream) => void,
+  ): Promise<void> {
+    const res = await fetch(
+      `${this.#baseUrl}/chat/completions`,
+      {
+        method: "POST",
+        headers: {
+          Authorization: `Bearer ${this.#privateKey}`,
+          "Content-Type": "application/json",
+        },
+        body: JSON.stringify({
+          model: options.model,
+          messages: options.messages,
+          temperature: options.temperature,
+          top_p: options.topP,
+          n: options.n,
+          stream: true,
+          stop: options.stop,
+          max_tokens: options.maxTokens,
+          presence_penalty: options.presencePenalty,
+          frequency_penalty: options.frequencyPenalty,
+          logit_bias: options.logitBias,
+          user: options.user,
+        }),
+      },
+    );
+
+    await decodeStream(res, callback);
+  }
+
   /**
    * Creates a new edit for the provided input, instruction, and parameters.
    *

src/types.ts

@@ -52,13 +52,6 @@ export interface CompletionOptions {
    */
   n?: number;
 
-  /**
-   * Whether to stream back partial progress.
-   * If set, tokens will be sent as data-only server-sent events as they become available, with the stream terminated by a data: [DONE] message.
-   * https://platform.openai.com/docs/api-reference/completions/create#completions/create-stream
-   */
-  stream?: boolean;
-
   /**
    * Include the log probabilities on the logprobs most likely tokens, as well the chosen tokens.
    * For example, if logprobs is 5, the API will return a list of the 5 most likely tokens.
@@ -161,13 +154,6 @@ export interface ChatCompletionOptions {
    */
   n?: number;
 
-  /**
-   * If set, partial message deltas will be sent, like in ChatGPT.
-   * Tokens will be sent as data-only server-sent events as they become available, with the stream terminated by a data: [DONE] message.
-   * https://platform.openai.com/docs/api-reference/chat/create#chat/create-stream
-   */
-  stream?: boolean;
-
   /**
    * Up to 4 sequences where the API will stop generating further tokens.
    * https://platform.openai.com/docs/api-reference/chat/create#chat/create-stop
@@ -607,6 +593,19 @@ export interface Completion {
   };
 }
 
+export interface CompletionStream {
+  id: string;
+  object: "text_completion";
+  created: number;
+  model: string;
+  choices: {
+    text: string;
+    index: number;
+    logprobs: number | null;
+    finish_reason: string;
+  }[];
+}
+
 export interface ChatCompletion {
   id: string;
   object: "chat.completion";
@@ -627,6 +626,21 @@ export interface ChatCompletion {
   };
 }
 
+export interface ChatCompletionStream {
+  id: string;
+  object: "chat.completion.chunk";
+  created: number;
+  choices: {
+    index: number;
+    delta: {
+      name?: string;
+      role?: "system" | "assistant" | "user";
+      content?: string;
+    };
+    finish_reason: string;
+  }[];
+}
+
 export interface Edit {
   object: "edit";
   created: number;

src/util.ts

@@ -0,0 +1,40 @@
+import { TextDelimiterStream } from "https://deno.land/std@0.189.0/streams/mod.ts";
+
+export function throwError(
+  data: { error?: { type: string; message: string; code: string } },
+) {
+  if (data.error) {
+    let errorMessage = `${data.error.type}`;
+    if (data.error.message) {
+      errorMessage += ": " + data.error.message;
+    }
+    if (data.error.code) {
+      errorMessage += ` (${data.error.code})`;
+    }
+    // console.log(data.error);
+    throw new Error(errorMessage);
+  }
+}
+
+// deno-lint-ignore no-explicit-any
+export async function decodeStream(
+  res: Response,
+  callback: (data: any) => void,
+) {
+  const chunks = res.body!
+    .pipeThrough(new TextDecoderStream())
+    .pipeThrough(new TextDelimiterStream("\n\n"));
+
+  for await (const chunk of chunks) {
+    let data;
+    try {
+      data = JSON.parse(chunk);
+    } catch {
+      // no-op (just checking if error message)
+    }
+    if (data) throwError(data);
+
+    if (chunk === "data: [DONE]") break;
+    callback(JSON.parse(chunk.slice(6)));
+  }
+}