From 32c33ce9484ea07712ad9ca95a0f736968e339bf Mon Sep 17 00:00:00 2001 From: NolanTrem <34580718+NolanTrem@users.noreply.github.com> Date: Mon, 6 Jan 2025 17:21:44 -0800 Subject: [PATCH 1/2] Fix JSON parsing for large payloads --- package.json | 2 +- pnpm-lock.yaml | 32 +++--- src/components/ChatDemo/answer.tsx | 1 + src/components/ChatDemo/result.tsx | 169 +++++++++++++++++++---------- src/lib/utils.ts | 30 +++++ 5 files changed, 160 insertions(+), 74 deletions(-) diff --git a/package.json b/package.json index ef122f6..f094086 100644 --- a/package.json +++ b/package.json @@ -98,7 +98,7 @@ "postcss": "^8.4.49", "posthog-js": "^1.194.3", "r2r": "0.0.0-beta.0", - "r2r-js": "^0.4.7", + "r2r-js": "^0.4.8", "radix-ui": "^1.0.1", "react": "18.3.1", "react-chartjs-2": "^5.2.0", diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 137b11a..4b397f4 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -222,8 +222,8 @@ importers: specifier: 0.0.0-beta.0 version: 0.0.0-beta.0 r2r-js: - specifier: ^0.4.7 - version: 0.4.7 + specifier: ^0.4.8 + version: 0.4.8 radix-ui: specifier: ^1.0.1 version: 1.0.1(@types/react-dom@18.3.1)(@types/react@18.3.3)(react-dom@18.3.1(react@18.3.1))(react@18.3.1) @@ -5496,8 +5496,8 @@ packages: resolution: {integrity: sha512-ZBbGYOpact8QAH9RprFWL4RAESYwbDodxiuDjOnzwzzk9pBzKycoifGuUrHHcDixE/eLMKPHRaXenTgu1qXBqA==} hasBin: true - jsesc@3.0.2: - resolution: {integrity: sha512-xKqzzWXDttJuOcawBt4KnKHHIf5oQ/Cxax+0PWFG+DFDgHNAdi+TXECADI+RYiFUMmx8792xsMbbgXj4CwnP4g==} + jsesc@3.1.0: + resolution: {integrity: sha512-/sM3dO2FOzXjKQhJuo0Q173wf2KOo8t4I8vHy6lF9poUp7bKT0/NHE8fPX23PwfhnykfqnC2xRxOnVw5XuGIaA==} engines: {node: '>=6'} hasBin: true @@ -6406,8 +6406,8 @@ packages: posthog-js@1.194.3: resolution: {integrity: sha512-/YFpBMqZzRpywa07QeoaIojdrUDijFajT4gZBSCFUBuZA5BN5xr5S1spsvtpT7E4RjkQSVgRvUngI4W19csgQw==} - posthog-node@4.3.1: - resolution: {integrity: sha512-By9SEGZxBLC7GVyVb+HlJlpxM/xI4iLUgwtsBS8f4bZ0wqYKiNHoYcFEwMJnTM9xQcQZztr6dqLmsJ7jCv0vxA==} + posthog-node@4.3.2: + resolution: {integrity: sha512-vy8Mt9IEfniUgqQ1rOCQ31CBO1VNqDGd3ZtHlWR9/YfU6RiuK+9pUXPb4h6HTGzQmjL8NFnjd8K8NMXSX8S6MQ==} engines: {node: '>=15.0.0'} preact@10.25.1: @@ -6464,8 +6464,8 @@ packages: queue-microtask@1.2.3: resolution: {integrity: sha512-NuaNSa6flKT5JaSYQzJok04JzTL1CA6aGhv5rfLW3PgqA+M2ChpZQnAC8h8i4ZFkBS8X5RqkDBHA7r4hej3K9A==} - r2r-js@0.4.7: - resolution: {integrity: sha512-VZZ0JmlM41tPI5+85sj0rJ9fyZ8iOWreoNuU9Lc+aDHHwGWABYqYnA78GEyvE5j6nBx8CYpJnNmiwjOvAovfrg==} + r2r-js@0.4.8: + resolution: {integrity: sha512-DWrhvTkSZh9sH5CM7gnGJlLfAXAV1e8FBlsuISM4r/+Tz86h/0F90BRv84w4XykUmypVO52GI/KoGvzvr7jRRQ==} r2r@0.0.0-beta.0: resolution: {integrity: sha512-LwZXOYJexUffYHP+LcBP3zgF17Az9dTTM60J9q3WGBReF4qqgEvFDyzHxOukMqIR8zUz4E00rRnCW5Mdktoe6g==} @@ -7547,7 +7547,7 @@ snapshots: '@ampproject/remapping@2.3.0': dependencies: - '@jridgewell/gen-mapping': 0.3.5 + '@jridgewell/gen-mapping': 0.3.8 '@jridgewell/trace-mapping': 0.3.25 '@babel/code-frame@7.26.2': @@ -7582,15 +7582,15 @@ snapshots: dependencies: '@babel/parser': 7.26.3 '@babel/types': 7.26.3 - '@jridgewell/gen-mapping': 0.3.5 + '@jridgewell/gen-mapping': 0.3.8 '@jridgewell/trace-mapping': 0.3.25 - jsesc: 3.0.2 + jsesc: 3.1.0 '@babel/helper-compilation-targets@7.25.9': dependencies: '@babel/compat-data': 7.26.3 '@babel/helper-validator-option': 7.25.9 - browserslist: 4.24.2 + browserslist: 4.24.3 lru-cache: 5.1.1 semver: 6.3.1 @@ -14236,7 +14236,7 @@ snapshots: dependencies: commander: 1.1.1 - jsesc@3.0.2: {} + jsesc@3.1.0: {} json-buffer@3.0.1: {} @@ -15577,7 +15577,7 @@ snapshots: preact: 10.25.1 web-vitals: 4.2.4 - posthog-node@4.3.1: + posthog-node@4.3.2: dependencies: axios: 1.7.9 rusha: 0.8.14 @@ -15622,14 +15622,14 @@ snapshots: queue-microtask@1.2.3: {} - r2r-js@0.4.7: + r2r-js@0.4.8: dependencies: '@jest/globals': 29.7.0 '@rrweb/types': 2.0.0-alpha.17 axios: 1.7.9 form-data: 4.0.1 posthog-js: 1.194.3 - posthog-node: 4.3.1 + posthog-node: 4.3.2 rrweb-snapshot: 2.0.0-alpha.4 uuid: 10.0.0 transitivePeerDependencies: diff --git a/src/components/ChatDemo/answer.tsx b/src/components/ChatDemo/answer.tsx index 902c7ea..55d0f81 100644 --- a/src/components/ChatDemo/answer.tsx +++ b/src/components/ChatDemo/answer.tsx @@ -91,6 +91,7 @@ export const Answer: FC<{ } if (message.sources.kg) { + console.log('message.sources.kg', message.sources.kg); const kgLocalResult: KGSearchResult[] = JSON.parse(message.sources.kg); const entitiesArray = kgLocalResult.filter( (item: any) => item.result_type === 'entity' diff --git a/src/components/ChatDemo/result.tsx b/src/components/ChatDemo/result.tsx index b137fd9..2a52bb0 100644 --- a/src/components/ChatDemo/result.tsx +++ b/src/components/ChatDemo/result.tsx @@ -9,6 +9,7 @@ import React, { FC, useEffect, useState, useRef } from 'react'; import PdfPreviewDialog from '@/components/ChatDemo/utils/pdfPreviewDialog'; import { useUserContext } from '@/context/UserContext'; +import { extractBlocks } from '@/lib/utils'; import { Message } from '@/types'; import { Answer } from './answer'; @@ -170,6 +171,7 @@ export const Result: FC<{ sources: {}, }; + // Start with an empty assistant message const newAssistantMessage: Message = { role: 'assistant', content: '', @@ -180,13 +182,17 @@ export const Result: FC<{ searchPerformed: false, }; + // Push the user message immediately setMessages((prevMessages) => [...prevMessages, newUserMessage]); + // We'll accumulate raw text in this buffer let buffer = ''; - let inLLMResponse = false; - let fullContent = ''; - let vectorSearchSources = null; - let kgSearchResult = null; + // Flags and placeholders + let inLLMResponse = false; // Are we inside blocks? + let fullContent = ''; // Combined text for the LLM + let assistantResponse = ''; // The final text for the assistant + let vectorSearchSources: string | null = null; + let kgSearchResult: string | null = null; let searchPerformed = false; try { @@ -195,22 +201,18 @@ export const Result: FC<{ throw new Error('Failed to get authenticated client'); } + // Make sure we have a conversation let currentConversationId = selectedConversationId; - if (!currentConversationId) { try { const newConversation = await client.conversations.create(); - if (!newConversation || !newConversation.results) { throw new Error('Failed to create a new conversation'); } - currentConversationId = newConversation.results.id; - if (typeof currentConversationId !== 'string') { throw new Error('Invalid conversation ID received'); } - setSelectedConversationId(currentConversationId); } catch (error) { console.error('Error creating new conversation:', error); @@ -218,12 +220,12 @@ export const Result: FC<{ return; } } - if (!currentConversationId) { setError('No valid conversation ID. Please try again.'); return; } + // Build the config const ragGenerationConfig: GenerationConfig = { stream: true, temperature: ragTemperature ?? undefined, @@ -235,10 +237,6 @@ export const Result: FC<{ const vectorSearchSettings: ChunkSearchSettings = { indexMeasure: IndexMeasure.COSINE_DISTANCE, enabled: switches.vectorSearch?.checked ?? true, - // selectedCollectionIds: - // selectedCollectionIds.length > 0 - // ? [selectedCollectionIds].flat() - // : undefined, }; const graphSearchSettings: GraphSearchSettings = { @@ -254,25 +252,25 @@ export const Result: FC<{ graphSettings: graphSearchSettings, }; + // Call the streaming endpoint const streamResponse = mode === 'rag_agent' ? await client.retrieval.agent({ message: newUserMessage, - ragGenerationConfig: ragGenerationConfig, - searchSettings: searchSettings, + ragGenerationConfig, + searchSettings, conversationId: currentConversationId, }) : await client.retrieval.rag({ - query: query, - ragGenerationConfig: ragGenerationConfig, - searchSettings: searchSettings, + query, + ragGenerationConfig, + searchSettings, }); const reader = streamResponse.getReader(); const decoder = new TextDecoder(); - let assistantResponse = ''; - + // Continuously read chunks while (true) { if (signal.aborted) { reader.cancel(); @@ -283,65 +281,123 @@ export const Result: FC<{ if (done) { break; } + buffer += decoder.decode(value, { stream: true }); - // Handle search results - if ( - buffer.includes(CHUNK_SEARCH_STREAM_END_MARKER) || - buffer.includes(GRAPH_SEARCH_STREAM_END_MARKER) - ) { - if (buffer.includes(CHUNK_SEARCH_STREAM_MARKER)) { - vectorSearchSources = buffer - .split(CHUNK_SEARCH_STREAM_MARKER)[1] - .split(CHUNK_SEARCH_STREAM_END_MARKER)[0]; + + // + // 1) Extract any blocks + // + const chunkSearchResult = extractBlocks( + buffer, + CHUNK_SEARCH_STREAM_MARKER, + CHUNK_SEARCH_STREAM_END_MARKER + ); + buffer = chunkSearchResult.newBuffer; // leftover data + for (const rawJson of chunkSearchResult.blocks) { + try { + vectorSearchSources = rawJson; searchPerformed = true; + } catch (err) { + console.error('Failed to parse chunk_search JSON:', err, rawJson); } + // Update state so user sees search results + updateLastMessage( + fullContent, + { + vector: vectorSearchSources, + kg: kgSearchResult, + }, + true, + searchPerformed + ); + setIsSearching(false); + } - if (buffer.includes(GRAPH_SEARCH_STREAM_MARKER)) { - kgSearchResult = buffer - .split(GRAPH_SEARCH_STREAM_MARKER)[1] - .split(GRAPH_SEARCH_STREAM_END_MARKER)[0]; + // + // 2) Extract any blocks + // + const graphSearchResultBlocks = extractBlocks( + buffer, + GRAPH_SEARCH_STREAM_MARKER, + GRAPH_SEARCH_STREAM_END_MARKER + ); + buffer = graphSearchResultBlocks.newBuffer; + for (const rawJson of graphSearchResultBlocks.blocks) { + try { + kgSearchResult = rawJson; searchPerformed = true; + } catch (err) { + console.error('Failed to parse graph_search JSON:', err, rawJson); } - + // Update updateLastMessage( fullContent, - { vector: vectorSearchSources, kg: kgSearchResult }, + { + vector: vectorSearchSources, + kg: kgSearchResult, + }, true, searchPerformed ); setIsSearching(false); } - // Handle LLM response - if (buffer.includes(LLM_START_TOKEN)) { - inLLMResponse = true; - buffer = buffer.split(LLM_START_TOKEN)[1] || ''; // strip pre-stream content + // + // 3) Handle tokens for LLM text + // + // The approach below is a simplified example, just like your original code, + // but we keep leftover text in `buffer` in case the marker is partial. + // + if (!inLLMResponse) { + // See if we have the start token + const startIdx = buffer.indexOf(LLM_START_TOKEN); + if (startIdx !== -1) { + inLLMResponse = true; + // Discard anything before + buffer = buffer.slice(startIdx + LLM_START_TOKEN.length); + } } + // If we're in LLM mode, check if we found the end token if (inLLMResponse) { - const endTokenIndex = buffer.indexOf(LLM_END_TOKEN); - let chunk = ''; - - if (endTokenIndex !== -1) { - chunk = buffer.slice(0, endTokenIndex); - buffer = buffer.slice(endTokenIndex + LLM_END_TOKEN.length); + const endIdx = buffer.indexOf(LLM_END_TOKEN); + if (endIdx !== -1) { + // We have a complete chunk of LLM text + const chunk = buffer.slice(0, endIdx); + buffer = buffer.slice(endIdx + LLM_END_TOKEN.length); inLLMResponse = false; + + fullContent += chunk; + assistantResponse += chunk; + updateLastMessage( + fullContent, + { + vector: vectorSearchSources, + kg: kgSearchResult, + }, + true, + searchPerformed + ); } else { - chunk = buffer; + // No closing token yet, so the entire buffer is partial LLM text + // We append it and clear buffer for next read + fullContent += buffer; + assistantResponse += buffer; + updateLastMessage( + fullContent, + { + vector: vectorSearchSources, + kg: kgSearchResult, + }, + true, + searchPerformed + ); buffer = ''; } - - fullContent += chunk; - assistantResponse += chunk; - updateLastMessage( - fullContent, - { vector: vectorSearchSources, kg: kgSearchResult }, - true, - searchPerformed - ); } } + // After the loop completes, we have the final `assistantResponse` if (assistantResponse) { updateLastMessage( assistantResponse, @@ -349,7 +405,6 @@ export const Result: FC<{ false, searchPerformed ); - try { await client.conversations.addMessage({ id: currentConversationId, diff --git a/src/lib/utils.ts b/src/lib/utils.ts index 8558f7f..2dc66f0 100644 --- a/src/lib/utils.ts +++ b/src/lib/utils.ts @@ -101,3 +101,33 @@ export function formatFileSize(bytes: number | undefined): string { return `${size.toFixed(2)} ${units[unitIndex]}`; } + +export function extractBlocks( + buffer: string, + markerStart: string, + markerEnd: string +): { blocks: string[]; newBuffer: string } { + const blocks: string[] = []; + let startIdx = buffer.indexOf(markerStart); + + while (startIdx !== -1) { + // see if there's a matching AFTER + const endIdx = buffer.indexOf(markerEnd, startIdx + markerStart.length); + if (endIdx === -1) { + // No closing tag yet -> partial data. We'll wait for next chunk + break; + } + + const blockContent = buffer.slice(startIdx + markerStart.length, endIdx); + blocks.push(blockContent); + + // Remove the entire ... from buffer + buffer = + buffer.slice(0, startIdx) + buffer.slice(endIdx + markerEnd.length); + + // Look for the *next* + startIdx = buffer.indexOf(markerStart); + } + + return { blocks, newBuffer: buffer }; +} From d8589b7741ff05b7e77276f1ea22826d020c2462 Mon Sep 17 00:00:00 2001 From: NolanTrem <34580718+NolanTrem@users.noreply.github.com> Date: Mon, 6 Jan 2025 17:25:05 -0800 Subject: [PATCH 2/2] Fix login with token --- src/context/UserContext.tsx | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/context/UserContext.tsx b/src/context/UserContext.tsx index 5b8c3c6..0dc8ba5 100644 --- a/src/context/UserContext.tsx +++ b/src/context/UserContext.tsx @@ -174,7 +174,9 @@ export const UserProvider: React.FC<{ children: React.ReactNode }> = ({ ): Promise<{ success: boolean; userRole: 'admin' | 'user' }> => { const newClient = new r2rClient(instanceUrl); try { - const result = await newClient.loginWithToken(token); + const result = await newClient.users.loginWithToken({ + accessToken: token, + }); const userInfo = await newClient.users.me();