From 02b46a1784ec74c447d5237e8b759f30442f58b3 Mon Sep 17 00:00:00 2001 From: Debanjum Singh Solanky Date: Fri, 2 Aug 2024 02:33:04 +0530 Subject: [PATCH] Render references after chat response is streamed for smoother render Otherwise the Khoj's chat response is filling up in between the streamed message and already rendered references section at the bottom of the message Define OnlineContext type to simplify typing online context param across other interfaces and functions --- src/interface/web/app/chat/page.tsx | 19 +++-- src/interface/web/app/common/chatFunctions.ts | 76 ++++++++++--------- .../components/chatMessage/chatMessage.tsx | 14 ++-- .../referencePanel/referencePanel.tsx | 4 +- src/interface/web/app/factchecker/page.tsx | 6 +- 5 files changed, 60 insertions(+), 59 deletions(-) diff --git a/src/interface/web/app/chat/page.tsx b/src/interface/web/app/chat/page.tsx index 9a8ee838..1768a228 100644 --- a/src/interface/web/app/chat/page.tsx +++ b/src/interface/web/app/chat/page.tsx @@ -9,11 +9,11 @@ import NavMenu from '../components/navMenu/navMenu'; import { useSearchParams } from 'next/navigation' import Loading from '../components/loading/loading'; -import { convertMessageChunkToJson, handleImageResponse, processMessageChunk, RawReferenceData } from '../common/chatFunctions'; +import { processMessageChunk } from '../common/chatFunctions'; import 'katex/dist/katex.min.css'; -import { StreamMessage } from '../components/chatMessage/chatMessage'; +import { Context, OnlineContext, StreamMessage } from '../components/chatMessage/chatMessage'; import { useIPLocationData, welcomeConsole } from '../common/utils'; import ChatInputArea, { ChatOptions } from '../components/chatInputArea/chatInputArea'; import { useAuthenticatedData } from '../common/auth'; @@ -110,7 +110,6 @@ export default function Chat() { const [uploadedFiles, setUploadedFiles] = useState([]); const [isMobileWidth, setIsMobileWidth] = useState(false); const locationData = useIPLocationData(); - const authenticatedData = useAuthenticatedData(); useEffect(() => { @@ -169,8 +168,11 @@ export default function Chat() { const eventDelimiter = '␃🔚␗'; let buffer = ""; - while (true) { + // Track context used for chat response + let context: Context[] = []; + let onlineContext: OnlineContext = {}; + while (true) { const { done, value } = await reader.read(); if (done) { setQueryToProcess(''); @@ -179,7 +181,6 @@ export default function Chat() { } const chunk = decoder.decode(value, { stream: true }); - buffer += chunk; let newEventIndex; @@ -194,7 +195,8 @@ export default function Chat() { return; } - processMessageChunk(event, currentMessage); + // Track context used for chat response. References are rendered at the end of the chat + ({context, onlineContext} = processMessageChunk(event, currentMessage, context, onlineContext)); setMessages([...messages]); } @@ -222,10 +224,7 @@ export default function Chat() { setConversationID(newConversationId); }; - if (isLoading) { - return ; - } - + if (isLoading) return ; return (
diff --git a/src/interface/web/app/common/chatFunctions.ts b/src/interface/web/app/common/chatFunctions.ts index 17d6aec9..25742eec 100644 --- a/src/interface/web/app/common/chatFunctions.ts +++ b/src/interface/web/app/common/chatFunctions.ts @@ -1,17 +1,13 @@ -import { Context, OnlineContextData, StreamMessage } from "../components/chatMessage/chatMessage"; +import { Context, OnlineContext, StreamMessage } from "../components/chatMessage/chatMessage"; export interface RawReferenceData { context?: Context[]; - onlineContext?: { - [key: string]: OnlineContextData - } + onlineContext?: OnlineContext; } export interface ResponseWithReferences { context?: Context[]; - online?: { - [key: string]: OnlineContextData - } + online?: OnlineContext; response?: string; } @@ -50,57 +46,65 @@ export function convertMessageChunkToJson(chunk: string): MessageChunk { } } +function handleJsonResponse(chunkData: any) { + const jsonData = chunkData as any; + if (jsonData.image || jsonData.detail) { + let responseWithReference = handleImageResponse(chunkData, true); + if (responseWithReference.response) return responseWithReference.response; + } else if (jsonData.response) { + return jsonData.response; + } else { + throw new Error("Invalid JSON response"); + } +} + +export function processMessageChunk( + rawChunk: string, + currentMessage: StreamMessage, + context: Context[] = [], + onlineContext: OnlineContext = {}): { context: Context[], onlineContext: OnlineContext } { -export function processMessageChunk(rawChunk: string, currentMessage: StreamMessage) { const chunk = convertMessageChunkToJson(rawChunk); - if (!currentMessage) { - return; - } - - if (!chunk || !chunk.type) { - return; - } + if (!currentMessage || !chunk || !chunk.type) return {context, onlineContext}; if (chunk.type === "status") { + console.log(`status: ${chunk.data}`); const statusMessage = chunk.data as string; currentMessage.trainOfThought.push(statusMessage); } else if (chunk.type === "references") { const references = chunk.data as RawReferenceData; - if (references.context) { - currentMessage.context = references.context; - } - - if (references.onlineContext) { - currentMessage.onlineContext = references.onlineContext; - } + if (references.context) context = references.context; + if (references.onlineContext) onlineContext = references.onlineContext; + return {context, onlineContext} } else if (chunk.type === "message") { const chunkData = chunk.data; - if (chunkData !== null && typeof chunkData === 'object') { + currentMessage.rawResponse += handleJsonResponse(chunkData); + } else if (typeof chunkData === 'string' && chunkData.trim()?.startsWith("{") && chunkData.trim()?.endsWith("}")) { try { - const jsonData = chunkData as any; - if (jsonData.image || jsonData.detail) { - let responseWithReference = handleImageResponse(chunk.data, true); - if (responseWithReference.response) currentMessage.rawResponse = responseWithReference.response; - if (responseWithReference.online) currentMessage.onlineContext = responseWithReference.online; - if (responseWithReference.context) currentMessage.context = responseWithReference.context; - } else if (jsonData.response) { - currentMessage.rawResponse = jsonData.response; - } - else { - console.debug("any message", chunk); - } + const jsonData = JSON.parse(chunkData.trim()); + currentMessage.rawResponse += handleJsonResponse(jsonData); } catch (e) { - currentMessage.rawResponse += chunkData; + currentMessage.rawResponse += JSON.stringify(chunkData); } } else { currentMessage.rawResponse += chunkData; } + } else if (chunk.type === "start_llm_response") { + console.log(`Started streaming: ${new Date()}`); } else if (chunk.type === "end_llm_response") { + console.log(`Completed streaming: ${new Date()}`); + + // Append any references after all the data has been streamed + if (onlineContext) currentMessage.onlineContext = onlineContext; + if (context) currentMessage.context = context; + + // Mark current message streaming as completed currentMessage.completed = true; } + return {context, onlineContext}; } export function handleImageResponse(imageJson: any, liveStream: boolean): ResponseWithReferences { diff --git a/src/interface/web/app/components/chatMessage/chatMessage.tsx b/src/interface/web/app/components/chatMessage/chatMessage.tsx index b8187f4c..71a5b29e 100644 --- a/src/interface/web/app/components/chatMessage/chatMessage.tsx +++ b/src/interface/web/app/components/chatMessage/chatMessage.tsx @@ -33,6 +33,10 @@ export interface Context { file: string; } +export interface OnlineContext { + [key: string]: OnlineContextData; +} + export interface WebPage { link: string; query: string; @@ -85,11 +89,9 @@ export interface SingleChatMessage { automationId: string; by: string; message: string; - context: Context[]; created: string; - onlineContext: { - [key: string]: OnlineContextData - } + context: Context[]; + onlineContext: OnlineContext; rawQuery?: string; intent?: Intent; agent?: AgentData; @@ -99,9 +101,7 @@ export interface StreamMessage { rawResponse: string; trainOfThought: string[]; context: Context[]; - onlineContext: { - [key: string]: OnlineContextData - } + onlineContext: OnlineContext; completed: boolean; rawQuery: string; timestamp: string; diff --git a/src/interface/web/app/components/referencePanel/referencePanel.tsx b/src/interface/web/app/components/referencePanel/referencePanel.tsx index 75345904..5fc70406 100644 --- a/src/interface/web/app/components/referencePanel/referencePanel.tsx +++ b/src/interface/web/app/components/referencePanel/referencePanel.tsx @@ -11,7 +11,7 @@ const md = new markdownIt({ typographer: true }); -import { Context, WebPage, OnlineContextData } from "../chatMessage/chatMessage"; +import { Context, WebPage, OnlineContext } from "../chatMessage/chatMessage"; import { Card } from "@/components/ui/card"; import { @@ -161,7 +161,7 @@ function GenericOnlineReferenceCard(props: OnlineReferenceCardProps) { ) } -export function constructAllReferences(contextData: Context[], onlineData: { [key: string]: OnlineContextData }) { +export function constructAllReferences(contextData: Context[], onlineData: OnlineContext) { const onlineReferences: OnlineReferenceData[] = []; const contextReferences: NotesContextReferenceData[] = []; diff --git a/src/interface/web/app/factchecker/page.tsx b/src/interface/web/app/factchecker/page.tsx index 5f2a9845..bb6e5af2 100644 --- a/src/interface/web/app/factchecker/page.tsx +++ b/src/interface/web/app/factchecker/page.tsx @@ -4,7 +4,7 @@ import styles from './factChecker.module.css'; import { useAuthenticatedData } from '@/app/common/auth'; import { useState, useEffect } from 'react'; -import ChatMessage, { Context, OnlineContextData, WebPage } from '../components/chatMessage/chatMessage'; +import ChatMessage, { Context, OnlineContext, OnlineContextData, WebPage } from '../components/chatMessage/chatMessage'; import { ModelPicker, Model } from '../components/modelPicker/modelPicker'; import ShareLink from '../components/shareLink/shareLink'; @@ -47,9 +47,7 @@ interface SupplementReferences { interface ResponseWithReferences { context?: Context[]; - online?: { - [key: string]: OnlineContextData - } + online?: OnlineContext; response?: string; }