⚡ (openai) Stream chat completion to avoid serverless timeout (#526)

Closes #520
2023-05-25 10:32:35 +02:00
parent 6bb6a2b0e3
commit 56364fd863
39 changed files with 556 additions and 121 deletions
--- a/apps/viewer/src/features/blocks/integrations/openai/getChatCompletionStream.ts
+++ b/apps/viewer/src/features/blocks/integrations/openai/getChatCompletionStream.ts
@@ -0,0 +1,103 @@
+import { parseVariableNumber } from '@/features/variables/parseVariableNumber'
+import { Connection } from '@planetscale/database'
+import { decrypt } from '@typebot.io/lib/api/encryption'
+import {
+  ChatCompletionOpenAIOptions,
+  OpenAICredentials,
+} from '@typebot.io/schemas/features/blocks/integrations/openai'
+import { SessionState } from '@typebot.io/schemas/features/chat'
+import {
+  ParsedEvent,
+  ReconnectInterval,
+  createParser,
+} from 'eventsource-parser'
+import type {
+  ChatCompletionRequestMessage,
+  CreateChatCompletionRequest,
+} from 'openai'
+
+export const getChatCompletionStream =
+  (conn: Connection) =>
+  async (
+    state: SessionState,
+    options: ChatCompletionOpenAIOptions,
+    messages: ChatCompletionRequestMessage[]
+  ) => {
+    if (!options.credentialsId) return
+    const credentials = (
+      await conn.execute('select data, iv from Credentials where id=?', [
+        options.credentialsId,
+      ])
+    ).rows.at(0) as { data: string; iv: string } | undefined
+    if (!credentials) {
+      console.error('Could not find credentials in database')
+      return
+    }
+    const { apiKey } = (await decrypt(
+      credentials.data,
+      credentials.iv
+    )) as OpenAICredentials['data']
+
+    const temperature = parseVariableNumber(state.typebot.variables)(
+      options.advancedSettings?.temperature
+    )
+
+    const encoder = new TextEncoder()
+    const decoder = new TextDecoder()
+
+    let counter = 0
+
+    const res = await fetch('https://api.openai.com/v1/chat/completions', {
+      headers: {
+        'Content-Type': 'application/json',
+        Authorization: `Bearer ${apiKey}`,
+      },
+      method: 'POST',
+      body: JSON.stringify({
+        messages,
+        model: options.model,
+        temperature,
+        stream: true,
+      } satisfies CreateChatCompletionRequest),
+    })
+
+    const stream = new ReadableStream({
+      async start(controller) {
+        function onParse(event: ParsedEvent | ReconnectInterval) {
+          if (event.type === 'event') {
+            const data = event.data
+            if (data === '[DONE]') {
+              controller.close()
+              return
+            }
+            try {
+              const json = JSON.parse(data) as {
+                choices: { delta: { content: string } }[]
+              }
+              const text = json.choices.at(0)?.delta.content
+              if (counter < 2 && (text?.match(/\n/) || []).length) {
+                return
+              }
+              const queue = encoder.encode(text)
+              controller.enqueue(queue)
+              counter++
+            } catch (e) {
+              controller.error(e)
+            }
+          }
+        }
+
+        // stream response (SSE) from OpenAI may be fragmented into multiple chunks
+        // this ensures we properly read chunks & invoke an event for each SSE event stream
+        const parser = createParser(onParse)
+
+        // https://web.dev/streams/#asynchronous-iteration
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        for await (const chunk of res.body as any) {
+          parser.feed(decoder.decode(chunk))
+        }
+      },
+    })
+
+    return stream
+  }