From a719a884fd24024e3e2dc9a415717060022a815a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Germ=C3=A1n=20Calcedo?= <49457798+gcalcedo@users.noreply.github.com> Date: Thu, 28 Nov 2024 14:39:39 +0100 Subject: [PATCH] fix: robust metadata yielding (#52) --- package.json | 2 +- src/hooks/useChatStream.ts | 52 ++++++++++++++++--------- src/utils/json.ts | 79 ++++++++++++++++++++++++++++++-------- 3 files changed, 98 insertions(+), 35 deletions(-) diff --git a/package.json b/package.json index 331ad4d..5f205dc 100644 --- a/package.json +++ b/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", diff --git a/src/hooks/useChatStream.ts b/src/hooks/useChatStream.ts index 423c492..65c83ed 100644 --- a/src/hooks/useChatStream.ts +++ b/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 | ChangeEvent) => { + const handleInputChange = ( + e: ChangeEvent | ChangeEvent, + ) => { setFormInput(e.target.value); }; @@ -37,28 +39,21 @@ 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; @@ -66,9 +61,25 @@ const useChatStream = (input: UseChatStreamInput) => { 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,7 +91,10 @@ 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' }); @@ -88,7 +102,7 @@ const useChatStream = (input: UseChatStreamInput) => { } finally { setIsStreaming(false); } - } + }; return { messages, diff --git a/src/utils/json.ts b/src/utils/json.ts index 215d0fd..81ec96b 100644 --- a/src/utils/json.ts +++ b/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; };