fix: prompt response stream (#190)

lunary-ai · Apr 8, 2024 · 031ef4a · 031ef4a
1 parent b7e0902
commit 031ef4a
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 27 deletions.
diff --git a/packages/backend/src/api/v1/orgs.ts b/packages/backend/src/api/v1/orgs.ts
@@ -226,28 +226,26 @@ orgs.post("/playground", async (ctx: Context) => {
   `
   const model = extra?.model || "gpt-3.5-turbo"
 
-  const res = await runAImodel(content, extra, variables, model, false)
-
-  ctx.body = res
-  // const stream = new PassThrough()
-  // stream.pipe(ctx.res)
-  // ctx.status = 200
-  // ctx.body = stream
-
-  // handleStream(
-  //   res,
-  //   (data) => {
-  //     stream.write(JSON.stringify(data) + "\n")
-  //   },
-  //   () => {
-  //     console.log("Stream finished")
-  //     stream.end()
-  //   },
-  //   (error) => {
-  //     throw new Error(error)
-  //     stream.end()
-  //   },
-  // )
+  const res = await runAImodel(content, extra, variables, model, true)
+
+  const stream = new PassThrough()
+  stream.pipe(ctx.res)
+  ctx.status = 200
+  ctx.body = stream
+
+  handleStream(
+    res,
+    (data) => {
+      stream.write(JSON.stringify(data) + "\n")
+    },
+    () => {
+      stream.end()
+    },
+    (error) => {
+      throw new Error(error)
+      stream.end()
+    },
+  )
 })
 
 export default orgs
diff --git a/packages/frontend/pages/prompts/[[...id]].tsx b/packages/frontend/pages/prompts/[[...id]].tsx
@@ -294,11 +294,15 @@ function Playground() {
           variables: templateVersion.testValues,
         },
         (chunk) => {
-          const parsedLine = JSON.parse(chunk)
-
-          setOutput(parsedLine.choices[0]?.message)
-          setOutputTokens(parsedLine.usage?.completion_tokens || 0)
-          setError(null)
+          try {
+            const parsedLine = JSON.parse(chunk)
+
+            setOutput(parsedLine.choices[0]?.message)
+            setOutputTokens(parsedLine.usage?.completion_tokens || 0)
+            setError(null)
+          } catch (error) {
+            console.error(error)
+          }
         },
       )
     } catch (e) {