initial stream handling

examples/extract_user copy/index.ts

@@ -0,0 +1,31 @@
+import Instructor from "@/instructor"
+import OpenAI from "openai"
+import { z } from "zod"
+
+const UserSchema = z.object({
+  age: z.number(),
+  name: z.string().refine(name => name.includes(" "), {
+    message: "Name must contain a space"
+  })
+})
+
+type User = z.infer<typeof UserSchema>
+
+const oai = new OpenAI({
+  apiKey: process.env.OPENAI_API_KEY ?? undefined,
+  organization: process.env.OPENAI_ORG_ID ?? undefined
+})
+
+const client = Instructor({
+  client: oai,
+  mode: "FUNCTIONS"
+})
+
+const user: User = await client.chat.completions.create({
+  messages: [{ role: "user", content: "Jason Liu is 30 years old" }],
+  model: "gpt-3.5-turbo",
+  response_model: UserSchema,
+  max_retries: 3
+})
+
+console.log(user)

examples/extract_user_stream/index.ts

@@ -0,0 +1,53 @@
+import Instructor from "@/instructor"
+import OpenAI from "openai"
+import { z } from "zod"
+
+const UserSchema = z.object({
+  age: z.number(),
+  name: z.string()
+})
+
+type User = z.infer<typeof UserSchema>
+
+const oai = new OpenAI({
+  apiKey: process.env.OPENAI_API_KEY ?? undefined,
+  organization: process.env.OPENAI_ORG_ID ?? undefined
+})
+
+const client = Instructor({
+  client: oai,
+  mode: "FUNCTIONS"
+})
+
+const userStream = await client.chat.completions.create({
+  messages: [{ role: "user", content: "Jason Liu is 30 years old" }],
+  model: "gpt-3.5-turbo",
+  response_model: UserSchema,
+  max_retries: 3,
+  stream: true
+})
+
+const reader = userStream.readable.getReader()
+const decoder = new TextDecoder()
+
+let result = {}
+let done = false
+
+while (!done) {
+  try {
+    const { value, done: doneReading } = await reader.read()
+    done = doneReading
+
+    if (done) {
+      break
+    }
+
+    const chunkValue = decoder.decode(value)
+    result = JSON.parse(chunkValue)
+    console.log(result)
+  } catch (e) {
+    done = true
+    console.log(e)
+    break
+  }
+}

package.json

@@ -41,6 +41,7 @@
   },
   "homepage": "https://github.com/jxnl/instructor-js#readme",
   "dependencies": {
+    "schema-stream": "^1.1.0",
     "zod-to-json-schema": "^3.22.3"
   },
   "peerDependencies": {

src/instructor.ts

@@ -8,9 +8,11 @@ import {
   OAIResponseJSONStringParser,
   OAIResponseToolArgsParser
 } from "@/oai/parser"
+import { OAIStream } from "@/oai/stream"
 import OpenAI from "openai"
-import { ChatCompletion, ChatCompletionCreateParamsNonStreaming } from "openai/resources/index.mjs"
-import { ZodObject } from "zod"
+import { ChatCompletion, ChatCompletionCreateParams } from "openai/resources/index.mjs"
+import { SchemaStream } from "schema-stream"
+import { z, ZodObject } from "zod"
 import zodToJsonSchema from "zod-to-json-schema"
 
 import { MODE } from "@/constants/modes"
@@ -31,7 +33,7 @@ const MODE_TO_PARAMS = {
   [MODE.JSON_SCHEMA]: OAIBuildMessageBasedParams
 }
 
-type PatchedChatCompletionCreateParams = ChatCompletionCreateParamsNonStreaming & {
+type PatchedChatCompletionCreateParams = ChatCompletionCreateParams & {
   //eslint-disable-next-line @typescript-eslint/no-explicit-any
   response_model?: ZodObject<any>
   max_retries?: number
@@ -84,9 +86,14 @@ class Instructor {
         }
 
         const completion = await this.client.chat.completions.create(resolvedParams)
-        const response = this.parseOAIResponse(completion)
+        const parser = MODE_TO_PARSER[this.mode]
 
-        return response
+        if ("choices" in completion) {
+          const parsedCompletion = parser(completion)
+          return JSON.parse(parsedCompletion)
+        } else {
+          return OAIStream({ res: completion, parser })
+        }
       } catch (error) {
         throw error
       }
@@ -95,6 +102,15 @@ class Instructor {
     const makeCompletionCallWithRetries = async () => {
       try {
         const data = await makeCompletionCall()
+
+        //short circuit if this is a stream for now
+        if (params.stream) {
+          return this.partialStreamResponse({
+            stream: data,
+            schema: params.response_model
+          })
+        }
+
         const validation = params.response_model.safeParse(data)
 
         if (!validation.success) {
@@ -125,15 +141,30 @@ class Instructor {
     return await makeCompletionCallWithRetries()
   }
 
+  private async partialStreamResponse({ stream, schema }) {
+    const streamParser = new SchemaStream(schema, {
+      onKeyComplete: console.log
+    })
+
+    const parser = streamParser.parse({
+      stringStreaming: true,
+      handleUnescapedNewLines: true
+    })
+
+    stream.pipeThrough(parser)
+
+    return parser
+  }
+
   /**
    * Builds the chat completion parameters.
    * @param {PatchedChatCompletionCreateParams} params - The parameters for chat completion.
-   * @returns {ChatCompletionCreateParamsNonStreaming} The chat completion parameters.
+   * @returns {ChatCompletionCreateParams} The chat completion parameters.
    */
   private buildChatCompletionParams = ({
     response_model,
     ...params
-  }: PatchedChatCompletionCreateParams): ChatCompletionCreateParamsNonStreaming => {
+  }: PatchedChatCompletionCreateParams): ChatCompletionCreateParams => {
     const jsonSchema = zodToJsonSchema(response_model, "response_model")
 
     const definition = {
@@ -149,17 +180,6 @@ class Instructor {
     }
   }
 
-  /**
-   * Parses the OAI response.
-   * @param {ChatCompletion} response - The response from the chat completion.
-   * @returns {any} The parsed response.
-   */
-  private parseOAIResponse = (response: ChatCompletion) => {
-    const parser = MODE_TO_PARSER[this.mode]
-
-    return parser(response)
-  }
-
   /**
    * Public chat interface.
    */

src/oai/parser.ts

@@ -16,9 +16,9 @@ export function OAIResponseTextParser(
 ) {
   const parsedData = typeof data === "string" ? JSON.parse(data) : data
 
-  const text = parsedData?.choices[0]?.message?.content ?? "{}"
+  const text = parsedData?.choices[0]?.message?.content ?? null
 
-  return JSON.parse(text)
+  return text
 }
 
 /**
@@ -36,9 +36,12 @@ export function OAIResponseFnArgsParser(
 ) {
   const parsedData = typeof data === "string" ? JSON.parse(data) : data
 
-  const text = parsedData.choices?.[0]?.message?.function_call?.arguments ?? "{}"
+  const text =
+    parsedData.choices?.[0].delta?.function_call?.arguments ??
+    parsedData.choices?.[0]?.message?.function_call?.arguments ??
+    null
 
-  return JSON.parse(text)
+  return text
 }
 
 /**
@@ -56,9 +59,12 @@ export function OAIResponseToolArgsParser(
 ) {
   const parsedData = typeof data === "string" ? JSON.parse(data) : data
 
-  const text = parsedData.choices?.[0]?.message?.tool_call?.function?.arguments ?? "{}"
+  const text =
+    parsedData.choices?.[0].delta?.tool_call?.function?.arguments ??
+    parsedData.choices?.[0]?.message?.tool_call?.function?.arguments ??
+    null
 
-  return JSON.parse(text)
+  return text
 }
 
 /**
@@ -75,7 +81,8 @@ export function OAIResponseJSONStringParser(
     | OpenAI.Chat.Completions.ChatCompletion
 ) {
   const parsedData = typeof data === "string" ? JSON.parse(data) : data
-  const text = parsedData?.choices[0]?.message?.content ?? "{}"
+  const text =
+    parsedData.choices?.[0].delta?.content ?? parsedData?.choices[0]?.message?.content ?? null
 
-  return JSON.parse(text)
+  return text
 }

src/oai/stream.ts

@@ -0,0 +1,51 @@
+import OpenAI from "openai"
+import { Stream } from "openai/streaming"
+
+interface OaiStreamArgs {
+  res: Stream<OpenAI.Chat.Completions.ChatCompletionChunk>
+  parser
+}
+
+/**
+ * `OaiStream` creates a ReadableStream that parses the SSE response from OAI
+ * and returns a parsed string from the response.
+ *
+ * @param {OaiStreamArgs} args - The arguments for the function.
+ * @returns {ReadableStream<string>} - The created ReadableStream.
+ */
+export function OAIStream({ res, parser }: OaiStreamArgs): ReadableStream<Uint8Array> {
+  const encoder = new TextEncoder()
+  let cancelGenerator: () => void
+
+  async function* generateStream(res): AsyncGenerator<string> {
+    cancelGenerator = () => {
+      return
+    }
+
+    for await (const part of res) {
+      if (part?.choices?.[0]?.finish_reason === "stop") {
+        cancelGenerator()
+        break
+      }
+
+      yield parser(part)
+    }
+  }
+
+  const generator = generateStream(res)
+
+  return new ReadableStream({
+    async start(controller) {
+      for await (const parsedData of generator) {
+        controller.enqueue(encoder.encode(parsedData))
+      }
+
+      controller.close()
+    },
+    cancel() {
+      if (cancelGenerator) {
+        cancelGenerator()
+      }
+    }
+  })
+}

tests/functions.test.ts

@@ -26,6 +26,7 @@ async function extractUser() {
   const user: User = await client.chat.completions.create({
     messages: [{ role: "user", content: "Jason Liu is 30 years old" }],
     model: "gpt-3.5-turbo",
+
     response_model: UserSchema,
     max_retries: 3
   })