fix: use new AI protocol format for friendly errors (#122)

oaknational · Sep 12, 2024 · 1ecce63 · 1ecce63
1 parent babe29d
commit 1ecce63
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 6 deletions.
diff --git a/apps/nextjs/src/app/api/chat/errorHandling.test.ts b/apps/nextjs/src/app/api/chat/errorHandling.test.ts
@@ -6,7 +6,10 @@ import { RateLimitExceededError } from "@oakai/core/src/utils/rateLimiting/userB
 import { PrismaClientWithAccelerate } from "@oakai/db";
 import invariant from "tiny-invariant";
 
-import { consumeStream } from "@/utils/testHelpers/consumeStream";
+import {
+  consumeStream,
+  extractStreamMessage,
+} from "@/utils/testHelpers/consumeStream";
 
 import { handleChatException } from "./errorHandling";
 
@@ -35,7 +38,8 @@ describe("handleChatException", () => {
       expect(response.status).toBe(200);
 
       invariant(response.body instanceof ReadableStream);
-      const message = JSON.parse(await consumeStream(response.body));
+      const message = extractStreamMessage(await consumeStream(response.body));
+
       expect(message).toEqual({
         type: "error",
         value: "Threat detected",
@@ -84,8 +88,8 @@ describe("handleChatException", () => {
       expect(response.status).toBe(200);
 
       const consumed = await consumeStream(response.body as ReadableStream);
-      expect(consumed).toMatch(/^\{/);
-      const message = JSON.parse(consumed);
+      const message = extractStreamMessage(consumed);
+
       expect(message).toEqual({
         type: "error",
         value: "Rate limit exceeded",
@@ -110,7 +114,7 @@ describe("handleChatException", () => {
 
       expect(response.status).toBe(200);
 
-      const message = JSON.parse(
+      const message = extractStreamMessage(
         await consumeStream(response.body as ReadableStream),
       );
       expect(message).toEqual({

diff --git a/apps/nextjs/src/app/api/chat/protocol.ts b/apps/nextjs/src/app/api/chat/protocol.ts
@@ -5,7 +5,9 @@ import {
 import { StreamingTextResponse } from "ai";
 
 export function streamingJSON(message: ErrorDocument | ActionDocument) {
-  const errorMessage = JSON.stringify(message);
+  const jsonContent = JSON.stringify(message);
+  const errorMessage = `0:"${jsonContent.replace(/"/g, '\\"')}"`;
+
   const errorEncoder = new TextEncoder();
 
   return new StreamingTextResponse(

diff --git a/apps/nextjs/src/utils/testHelpers/consumeStream.ts b/apps/nextjs/src/utils/testHelpers/consumeStream.ts
@@ -14,3 +14,12 @@ export async function consumeStream(
 
   return result;
 }
+
+export function extractStreamMessage(streamedText: string) {
+  const content = streamedText.match(/0:"(.*)"/);
+  if (!content?.[1]) {
+    throw new Error("No message found in streamed text");
+  }
+  const strippedContent = content[1].replace(/\\"/g, '"');
+  return JSON.parse(strippedContent);
+}