feat: custom fetch param for node envs (#46)

- Adds `{ fetch }` param to `OpenAIOptions` for `node-fetch` support

Author: ctjlewis

package.json

@@ -52,6 +52,7 @@
   },
   "dependencies": {
     "eventsource-parser": "^1.0.0",
+    "node-fetch": "^3.3.1",
     "openai": "^3.2.1",
     "yield-stream": "^3.0.0"
   },

pnpm-lock.yaml

@@ -1,19 +1,28 @@
 /* eslint-disable no-console */
+import nodeFetch, { RequestInfo as NodeFetchRequestInfo, RequestInit as NodeFetchRequestInit } from "node-fetch";
+
 export interface BackoffOptions {
   maxRetries: number;
   delay: number;
 }
 
+const globalFetch = typeof fetch === "undefined" ? nodeFetch : fetch;
+
 export const fetchWithBackoff = async (
-  input: RequestInfo,
-  init?: RequestInit,
+  input: RequestInfo & NodeFetchRequestInfo,
+  init?: RequestInit & NodeFetchRequestInit,
+  fetch = globalFetch,
   { delay, maxRetries }: BackoffOptions = {
     delay: 500,
     maxRetries: 7
   }
 ) => {
   for (let i = 0; i <= maxRetries; i++) {
     try {
+      if (!fetch) {
+        throw new Error("No fetch implementation.");
+      }
+
       const response = await fetch(input, init);
 
       if (!response.ok) {

src/lib/backoff.ts

@@ -28,7 +28,8 @@ export const OpenAI: OpenAIEdgeClient = async (
     apiHeaders = {},
     controller,
     onDone,
-    onParse
+    onParse,
+    fetch
   } = {}
 ) => {
   if (!apiKey) {
@@ -50,7 +51,7 @@ export const OpenAI: OpenAIEdgeClient = async (
       ...apiHeaders,
     },
     signal: controller?.signal,
-  });
+  }, fetch);
 
   switch (response.status) {
     case 401:

src/lib/openai/edge.ts

@@ -3,6 +3,7 @@ import { ChatParser, TokenParser } from "./transforms";
 
 import { createParser } from "eventsource-parser";
 import { Transform, pipeline, yieldStream } from "yield-stream";
+import { yieldStream as yieldStreamNode } from "yield-stream/node";
 import { OpenAIError } from "../errors";
 
 export type StreamMode = "raw" | "tokens";
@@ -26,7 +27,7 @@ export interface OpenAIStreamOptions {
 }
 
 export type OpenAIStream = (
-  stream: ReadableStream<Uint8Array>,
+  stream: NodeJS.ReadableStream | ReadableStream<Uint8Array>,
   options: OpenAIStreamOptions
 ) => ReadableStream<Uint8Array>;
 
@@ -85,10 +86,17 @@ export const EventStream: OpenAIStream = (
           }
         }
       });
+
+      // Check if the stream is a NodeJS stream or a browser stream.
+      // @ts-ignore - TS doesn't know about `pipe` on streams.
+      const isNodeJsStream = typeof stream.pipe === "function";
+
       /**
        * Feed the parser with decoded chunks from the raw stream.
        */
-      for await (const chunk of yieldStream(stream)) {
+      for await (const chunk of isNodeJsStream
+        ? yieldStreamNode<Uint8Array>(stream as NodeJS.ReadableStream)
+        : yieldStream(stream as ReadableStream<Uint8Array>)) {
         const decoded = DECODER.decode(chunk);
 
         try {

src/lib/streaming/streams.ts

@@ -9,6 +9,8 @@ import type {
 
 import { OpenAIStreamOptions } from "./streaming";
 
+import nodeFetch from "node-fetch";
+
 export const OpenAIAPIEndpoints = {
   chat: "chat/completions",
   completions: "completions",
@@ -56,6 +58,11 @@ export interface OpenAIOptions extends OpenAIStreamOptions {
    * mid-flight.
    */
   controller?: AbortController;
+  /**
+   * An optional fetch implementation, which can be used to replace the default global fetch call used for making
+   * API requests. This is useful in environments like node where a global fetch is not provided by default.
+   */
+  fetch?: typeof fetch | typeof nodeFetch;
 }
 
 /**

src/lib/types.ts

@@ -5,6 +5,7 @@ import test from "ava";
 import { OpenAI } from "../src";
 import { yieldStream } from "yield-stream";
 import { DECODER } from "../src/globs/shared";
+import nodeFetch from "node-fetch";
 
 test.serial("'completions' endpoint", async (t) => {
   const stream = await OpenAI("completions", {
@@ -237,3 +238,28 @@ test.serial("cancelling streams", async (t) => {
 
   t.is(i, 5, "Stream should have been cancelled after 5 chunks.");
 });
+
+test.serial("should work with custom fetch", async (t) => {
+  let didUseMock = false;
+
+  const mockFetch: typeof nodeFetch = (...params) => {
+    didUseMock = true;
+    return nodeFetch(...params);
+  };
+
+  const stream = await OpenAI("chat", {
+    model: "gpt-3.5-turbo",
+    messages: [
+      { role: "system", content: "You are a helpful assistant." },
+      { role: "user", content: "This is a test message, say hello!" },
+    ],
+  }, { fetch: mockFetch });
+
+  const DECODER = new TextDecoder();
+  for await (const serialized of yieldStream(stream)) {
+    const string = DECODER.decode(serialized);
+    console.log(string);
+  }
+
+  t.true(didUseMock, "Did not use custom fetch.");
+});