fix: robust metadata yielding (#52)

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@magicul/react-chat-stream",
   "description": "A React hook that lets you easily integrate your custom ChatGPT-like chat in React.",
-  "version": "0.5.1",
+  "version": "0.5.2",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "homepage": "https://github.com/XD2Sketch/react-chat-stream#readme",

src/hooks/useChatStream.ts

@@ -1,7 +1,7 @@
 import { ChangeEvent, FormEvent, useState } from 'react';
 import { decodeStreamToJson, getStream } from '../utils/streams';
 import { UseChatStreamChatMessage, UseChatStreamInput } from '../types';
-import { extractJsonFromEnd } from '../utils/json';
+import { getJsonObjectsFromChunks } from '../utils/json';
 
 const BOT_ERROR_MESSAGE = 'Something went wrong fetching AI response.';
 
@@ -10,7 +10,9 @@ const useChatStream = (input: UseChatStreamInput) => {
   const [formInput, setFormInput] = useState('');
   const [isStreaming, setIsStreaming] = useState(false);
 
-  const handleInputChange = (e: ChangeEvent<HTMLInputElement> | ChangeEvent<HTMLTextAreaElement>) => {
+  const handleInputChange = (
+    e: ChangeEvent<HTMLInputElement> | ChangeEvent<HTMLTextAreaElement>,
+  ) => {
     setFormInput(e.target.value);
   };
 
@@ -37,38 +39,47 @@ const useChatStream = (input: UseChatStreamInput) => {
     });
   };
 
-  const fetchAndUpdateAIResponse = async (message: string) => {
+  const fetchAndUpdateAIResponse = async (message: string, useMetadata: boolean) => {
     const charactersPerSecond = input.options.fakeCharactersPerSecond;
     const stream = await getStream(message, input.options, input.method);
     const initialMessage = addMessage({ content: '', role: 'bot' });
     let response = '';
+    let metadata = null;
 
-    for await (const chunk of decodeStreamToJson(stream)) {
+    const processContent = async (content: string) => {
       if (!charactersPerSecond) {
-        appendMessageToChat(chunk);
-        response += chunk;
-        continue;
+        appendMessageToChat(content);
+        response += content;
+        return;
       }
 
-      if (input.options.useMetadata) {
-        const metadata = extractJsonFromEnd(chunk);
-        if (metadata) {
-          return { ...initialMessage, content: response, metadata: metadata };
-        }
-      }
-
-      // Stream characters one by one based on the characters per second that is set.
-      for (const char of chunk) {
+      for (const char of content) {
         appendMessageToChat(char);
         response += char;
 
         if (charactersPerSecond > 0) {
           await new Promise(resolve => setTimeout(resolve, 1000 / charactersPerSecond));
         }
       }
+    };
+
+    for await (const chunk of decodeStreamToJson(stream)) {
+      if (useMetadata) {
+        const jsonObjects = getJsonObjectsFromChunks(chunk);
+
+        for (const parsedChunk of jsonObjects) {
+          if (parsedChunk.type === 'content') {
+            await processContent(parsedChunk.data);
+          } else if (parsedChunk.type === 'metadata') {
+            metadata = parsedChunk.data;
+          }
+        }
+      } else {
+        await processContent(chunk);
+      }
     }
 
-    return { ...initialMessage, content: response };
+    return { ...initialMessage, content: response, ...(useMetadata && { metadata }) };
   };
 
   const submitMessage = async (message: string) => resetInputAndGetResponse(message);
@@ -80,15 +91,18 @@ const useChatStream = (input: UseChatStreamInput) => {
     setFormInput('');
 
     try {
-      const addedMessage = await fetchAndUpdateAIResponse(message ?? formInput);
+      const addedMessage = await fetchAndUpdateAIResponse(
+        message ?? formInput,
+        input.options.useMetadata ?? false,
+      );
       await input.handlers.onMessageAdded?.(addedMessage);
     } catch {
       const addedMessage = addMessage({ content: BOT_ERROR_MESSAGE, role: 'bot' });
       await input.handlers.onMessageAdded?.(addedMessage);
     } finally {
       setIsStreaming(false);
     }
-  }
+  };
 
   return {
     messages,

src/utils/json.ts

@@ -1,20 +1,69 @@
-export const extractJsonFromEnd = (chunk: string) => {
-  const chunkTrimmed = chunk.trim();
+/**
+ * Extracts JSON objects from a string containing one or more dumps of JSON objects.
+ *
+ * This function is used to parse the stream only when the `useMetadata` option is enabled.
+ * Then, this hook expects to receive JSON dumps instead of plain text. These JSON dumps are of
+ * the following format:
+ *
+ * - For content to be used in the chat:
+ * ```json
+ * {
+ *   "type": "content",
+ *   "data": "Hello, world!"
+ * }
+ * ```
+ *
+ * - For metadata:
+ * ```json
+ * {
+ *   "type": "metadata",
+ *   "data": {
+ *     "key": "value",
+ *     "key2": "value2",
+ *     ...
+ *   }
+ * }
+ * ```
+ *
+ * @param chunk - The string containing one or more JSON object dumps.
+ * @returns An array of parsed JSON objects.
+ *
+ * @example
+ * ```typescript
+ * const chunk = '{"type": "content", "data": "Hello, world!"}{"type": "metadata", "data": {"key": "value"}}';
+ * const jsonObjects = getJsonObjectsFromChunks(chunk);
+ * console.log(jsonObjects);
+ * // Output: [
+ * //   { type: 'content', data: 'Hello, world!' },
+ * //   { type: 'metadata', data: { key: 'value' } }
+ * // ]
+ * ```
+ */
+export const getJsonObjectsFromChunks = (chunk: string) => {
+  const jsonObjects = [];
+  const braceStack = [];
+  let currentJsonStart = null;
 
-  const jsonObjectRegex = /({[^]*})\s*$/;
-  const match = chunkTrimmed.match(jsonObjectRegex);
+  for (let i = 0; i < chunk.length; i++) {
+    const char = chunk[i];
 
-  if (!match) {
-    return null;
-  }
-
-  const jsonStr = match[1];
-  try {
-    const parsedData = JSON.parse(jsonStr);
-    if (typeof parsedData === 'object' && parsedData !== null && !Array.isArray(parsedData)) {
-      return parsedData;
+    if (char === '{') {
+      if (braceStack.length === 0) {
+        currentJsonStart = i;
+      }
+      braceStack.push('{');
+    } else if (char === '}') {
+      braceStack.pop();
+      if (braceStack.length === 0 && currentJsonStart !== null) {
+        const potentialJson = chunk.substring(currentJsonStart, i + 1);
+        try {
+          const parsedJson = JSON.parse(potentialJson);
+          jsonObjects.push(parsedJson);
+        } catch {}
+        currentJsonStart = null;
+      }
     }
-  } catch {}
+  }
 
-  return null;
+  return jsonObjects;
 };