From 92c1efe6eeb4fe776e90d231ee81e13054586003 Mon Sep 17 00:00:00 2001 From: Debanjum Date: Sun, 10 Nov 2024 21:03:03 -0800 Subject: [PATCH] Fixes to render & save code context with non text based output modes - Fix to render code generated chart with images, excalidraw diagrams - Fix to save code context to chat history in image, diagram output modes - Fix bug in image markdown being wrapped twice in markdown syntax - Render newline in code references shown on chat page of web app Previously newlines weren't getting rendered. This made the code executed by Khoj hard to read in references. This changes fixes that. `dangerouslySetInnerHTML' usage is justified as rendered code snippet is being sanitized by DOMPurify before rendering. --- src/interface/web/app/common/chatFunctions.ts | 31 +-------- .../components/chatMessage/chatMessage.tsx | 66 +++++++------------ .../referencePanel/referencePanel.tsx | 12 ++-- src/khoj/routers/api_chat.py | 2 + 4 files changed, 36 insertions(+), 75 deletions(-) diff --git a/src/interface/web/app/common/chatFunctions.ts b/src/interface/web/app/common/chatFunctions.ts index 5e9ac321..aeb74ee8 100644 --- a/src/interface/web/app/common/chatFunctions.ts +++ b/src/interface/web/app/common/chatFunctions.ts @@ -139,33 +139,6 @@ export function processMessageChunk( if (onlineContext) currentMessage.onlineContext = onlineContext; if (context) currentMessage.context = context; - // Replace file links with base64 data - currentMessage.rawResponse = renderCodeGenImageInline( - currentMessage.rawResponse, - codeContext, - ); - - // Add code context files to the message - if (codeContext) { - Object.entries(codeContext).forEach(([key, value]) => { - value.results.output_files?.forEach((file) => { - if (file.filename.endsWith(".png") || file.filename.endsWith(".jpg")) { - // Don't add the image again if it's already in the message! - if (!currentMessage.rawResponse.includes(`![${file.filename}](`)) { - currentMessage.rawResponse += `\n\n![${file.filename}](data:image/png;base64,${file.b64_data})`; - } - } else if ( - file.filename.endsWith(".txt") || - file.filename.endsWith(".org") || - file.filename.endsWith(".md") - ) { - const decodedText = atob(file.b64_data); - currentMessage.rawResponse += `\n\n\`\`\`\n${decodedText}\n\`\`\``; - } - }); - }); - } - // Mark current message streaming as completed currentMessage.completed = true; } @@ -200,10 +173,10 @@ export function renderCodeGenImageInline(message: string, codeContext: CodeConte Object.values(codeContext).forEach((contextData) => { contextData.results.output_files?.forEach((file) => { const regex = new RegExp(`!?\\[.*?\\]\\(.*${file.filename}\\)`, "g"); - if (file.filename.match(/\.(png|jpg|jpeg|gif|webp)$/i)) { + if (file.filename.match(/\.(png|jpg|jpeg)$/i)) { const replacement = `![${file.filename}](data:image/${file.filename.split(".").pop()};base64,${file.b64_data})`; message = message.replace(regex, replacement); - } else if (file.filename.match(/\.(txt|org|md)$/i)) { + } else if (file.filename.match(/\.(txt|org|md|csv|json)$/i)) { // render output files generated by codegen as downloadable links const replacement = `![${file.filename}](data:text/plain;base64,${file.b64_data})`; message = message.replace(regex, replacement); diff --git a/src/interface/web/app/components/chatMessage/chatMessage.tsx b/src/interface/web/app/components/chatMessage/chatMessage.tsx index 4869b5fa..a8fb2f6f 100644 --- a/src/interface/web/app/components/chatMessage/chatMessage.tsx +++ b/src/interface/web/app/components/chatMessage/chatMessage.tsx @@ -421,6 +421,31 @@ const ChatMessage = forwardRef((props, ref) => } } + // Replace file links with base64 data + message = renderCodeGenImageInline(message, props.chatMessage.codeContext); + + // Add code context files to the message + if (props.chatMessage.codeContext) { + Object.entries(props.chatMessage.codeContext).forEach(([key, value]) => { + value.results.output_files?.forEach((file) => { + if (file.filename.endsWith(".png") || file.filename.endsWith(".jpg")) { + // Don't add the image again if it's already in the message! + if (!message.includes(`![${file.filename}](`)) { + message += `\n\n![${file.filename}](data:image/png;base64,${file.b64_data})`; + } + } else if ( + file.filename.endsWith(".txt") || + file.filename.endsWith(".org") || + file.filename.endsWith(".md") || + file.filename.endsWith(".csv") || + file.filename.endsWith(".json") + ) { + message += `\n\n## ${file.filename}\n\`\`\`\n${file.b64_data}\n\`\`\`\n`; + } + }); + }); + } + // Handle user attached images rendering let messageForClipboard = message; let messageToRender = message; @@ -446,47 +471,6 @@ const ChatMessage = forwardRef((props, ref) => messageToRender = `${userImagesInHtml}${messageToRender}`; } - if (props.chatMessage.intent && props.chatMessage.intent.type == "text-to-image") { - message = `![generated image](data:image/png;base64,${message})`; - } else if (props.chatMessage.intent && props.chatMessage.intent.type == "text-to-image2") { - message = `![generated image](${message})`; - } else if ( - props.chatMessage.intent && - props.chatMessage.intent.type == "text-to-image-v3" - ) { - message = `![generated image](data:image/webp;base64,${message})`; - } - if ( - props.chatMessage.intent && - props.chatMessage.intent.type.includes("text-to-image") && - props.chatMessage.intent["inferred-queries"]?.length > 0 - ) { - message += `\n\n${props.chatMessage.intent["inferred-queries"][0]}`; - } - - // Replace file links with base64 data - message = renderCodeGenImageInline(message, props.chatMessage.codeContext); - - // Add code context files to the message - if (props.chatMessage.codeContext) { - Object.entries(props.chatMessage.codeContext).forEach(([key, value]) => { - value.results.output_files?.forEach((file) => { - if (file.filename.endsWith(".png") || file.filename.endsWith(".jpg")) { - // Don't add the image again if it's already in the message! - if (!message.includes(`![${file.filename}](`)) { - message += `\n\n![${file.filename}](data:image/png;base64,${file.b64_data})`; - } - } else if ( - file.filename.endsWith(".txt") || - file.filename.endsWith(".org") || - file.filename.endsWith(".md") - ) { - message += `\n\n## ${file.filename}\n\`\`\`\n${file.b64_data}\n\`\`\`\n`; - } - }); - }); - } - // Set the message text setTextRendered(messageForClipboard); diff --git a/src/interface/web/app/components/referencePanel/referencePanel.tsx b/src/interface/web/app/components/referencePanel/referencePanel.tsx index 785d05d8..53016cfc 100644 --- a/src/interface/web/app/components/referencePanel/referencePanel.tsx +++ b/src/interface/web/app/components/referencePanel/referencePanel.tsx @@ -103,7 +103,7 @@ interface CodeContextReferenceCardProps { function CodeContextReferenceCard(props: CodeContextReferenceCardProps) { const fileIcon = getIconFromFilename(".py", "w-6 h-6 text-muted-foreground inline-flex mr-2"); - const snippet = DOMPurify.sanitize(props.code); + const sanitizedCodeSnippet = DOMPurify.sanitize(props.code.replace(/\n/g, "
")); const [isHovering, setIsHovering] = useState(false); return ( @@ -123,9 +123,8 @@ function CodeContextReferenceCard(props: CodeContextReferenceCardProps) {

- {snippet} -

+ dangerouslySetInnerHTML={{ __html: sanitizedCodeSnippet }} + >

@@ -136,7 +135,10 @@ function CodeContextReferenceCard(props: CodeContextReferenceCardProps) { {fileIcon} Code -

{snippet}

+

diff --git a/src/khoj/routers/api_chat.py b/src/khoj/routers/api_chat.py index 33476a97..7c79809f 100644 --- a/src/khoj/routers/api_chat.py +++ b/src/khoj/routers/api_chat.py @@ -1131,6 +1131,7 @@ async def chat( conversation_id=conversation_id, compiled_references=compiled_references, online_results=online_results, + code_results=code_results, query_images=uploaded_images, train_of_thought=train_of_thought, attached_file_context=attached_file_context, @@ -1192,6 +1193,7 @@ async def chat( conversation_id=conversation_id, compiled_references=compiled_references, online_results=online_results, + code_results=code_results, query_images=uploaded_images, train_of_thought=train_of_thought, attached_file_context=attached_file_context,