Feature/answer explainability

.gitignore

@@ -12,4 +12,4 @@ dummy-package/
 
 .vscode/
 
-examples/llama-index/frontend/rag-ui/
+examples/llama-index/frontend/rag-ui/

lexio/.gitignore

@@ -4,6 +4,12 @@
 # Coverage Reports
 coverage
 
+# Environment Variables
+.env.local
+
+# PDF Files
+pdfs/
+
 # Logs
 logs
 *.log

lexio/lib/components/ChatWindow/AssistantMarkdownContent.tsx

@@ -0,0 +1,111 @@
+import React from "react";
+import {ChatWindowStyles} from "./ChatWindow";
+import Markdown from "react-markdown";
+import {Light as SyntaxHighlighter} from "react-syntax-highlighter";
+import {docco} from "react-syntax-highlighter/dist/esm/styles/hljs";
+import remarkGfm from "remark-gfm";
+import {ClipboardIcon, ClipboardDocumentIcon} from "@heroicons/react/24/outline";
+
+/**
+ * Props for the AssistantMarkdownContent component.
+ * @typedef {Object} AssistantMarkdownContentProps
+ * @property {string} content - The markdown content to render.
+ * @property {ChatWindowStyles} style - Styling options for the chat window.
+ */
+interface AssistantMarkdownContentProps {
+    content: string;
+    style: ChatWindowStyles;
+}
+
+/**
+ * Renders markdown content with enhanced features such as syntax highlighting and copy functionality.
+ *
+ * @param {AssistantMarkdownContentProps} props - The props for the component.
+ * @returns {JSX.Element} The rendered markdown content.
+ *
+ * @internal
+ */
+export const AssistantMarkdownContent: React.FC<AssistantMarkdownContentProps> = ({content, style}) => {
+    return (
+        <Markdown
+            components={{
+                /**
+                 * Custom renderer for code blocks to include syntax highlighting and a copy button.
+                 *
+                 * @param {Object} props - Properties passed to the code component.
+                 * @param {Object} props.node - The AST node.
+                 * @param {string} props.className - The class name for the code block.
+                 * @param {React.ReactNode} props.children - The code content.
+                 * @returns {JSX.Element} The customized code block.
+                 */
+                code({node, className, children, ...props}) {
+                    const match = /language-(\w+)/.exec(className || '');
+                    const language = match ? match[1] : 'plaintext';
+                    const [copied, setCopied] = React.useState(false);
+
+                    /**
+                     * Handles copying the code content to the clipboard.
+                     */
+                    const handleCopy = () => {
+                        navigator.clipboard.writeText(String(children));
+                        setCopied(true);
+                        setTimeout(() => {
+                            setCopied(false);
+                        }, 3000); // Reset copied state after 3 seconds
+                    }
+
+                    // Block code (triple ```)
+                    return (
+                        <div style={{
+                            position: 'relative',
+                            marginBottom: '1rem'
+                        }} {...(props as React.HTMLProps<HTMLDivElement>)}>
+                            {match && (
+                                <div
+                                    style={{
+                                        marginLeft: '0.33rem',
+                                        fontSize: 'var(--font-size-small)',
+                                        fontWeight: 'bold',
+                                        background: '#f5f5f5',
+                                        padding: '0 8px',
+                                        borderTopLeftRadius: '5px',
+                                        borderTopRightRadius: '5px',
+                                        borderBottom: '1px solid #ccc',
+                                        display: 'inline-block',
+                                    }}
+                                >
+                                    {language}
+                                </div>
+                            )}
+                            <div className="relative">
+                                <SyntaxHighlighter
+                                    style={docco}
+                                    language={language}
+                                    PreTag="div"
+                                    wrapLongLine={true}
+                                >
+                                    {String(children).replace(/\n$/, '')}
+                                </SyntaxHighlighter>
+
+                                <div
+                                    className="absolute top-1 right-1 px-1.5 pt-0.5 text-sm rounded-md hover:shadow-md "
+                                    style={{
+                                        backgroundColor: style.messageBackgroundColor,
+                                    }}
+                                >
+                                    <button
+                                        onClick={handleCopy}
+                                    >
+                                        {copied ? <ClipboardDocumentIcon className={"w-5 h-5"} /> : <ClipboardIcon className={"w-5 h-5"} />}
+                                    </button>
+                                </div>
+                            </div>
+                        </div>
+                    );
+                },
+            }}
+            remarkPlugins={[remarkGfm]}>
+            {content}
+        </Markdown>
+    );
+}

lexio/lib/components/ChatWindow/ChatWindow.tsx

@@ -199,6 +199,7 @@ const ChatWindow: React.FC<ChatWindowProps> = ({
                                 markdown={markdown}
                                 showCopy={showCopy}
                                 showFeedback={showFeedback}
+                                highlights={msg.highlights || []}
                             />
                         )}
                     </React.Fragment>
@@ -214,6 +215,7 @@ const ChatWindow: React.FC<ChatWindowProps> = ({
                         icon={assistantIcon || undefined}
                         markdown={markdown}
                         isStreaming={true}
+                        highlights={currentStream.highlights || []}
                     />
                 )}
                 <div ref={chatEndRef}/>

lexio/lib/components/ChatWindow/ChatWindowAssistantMessage.tsx

@@ -1,121 +1,17 @@
 import React, {ReactNode, useMemo, useState} from "react";
 import {ChatWindowStyles} from "./ChatWindow.tsx";
-import Markdown from "react-markdown";
-import {Light as SyntaxHighlighter} from "react-syntax-highlighter";
-import {docco} from "react-syntax-highlighter/dist/esm/styles/hljs";
-import remarkGfm from "remark-gfm";
 import "./AssistantMarkdownContent.css";
 import {ClipboardIcon, ClipboardDocumentIcon} from "@heroicons/react/24/outline";
 import {scaleFontSize} from "../../utils/scaleFontSize.tsx";
-import { MessageFeedback } from "./MessageFeedback.tsx";
-/**
- * Props for the AssistantMarkdownContent component.
- * @typedef {Object} AssistantMarkdownContentProps
- * @property {string} content - The markdown content to render.
- * @property {ChatWindowStyles} style - Styling options for the chat window.
- */
-interface AssistantMarkdownContentProps {
-    content: string;
-    style: ChatWindowStyles;
-}
-
-/**
- * Renders markdown content with enhanced features such as syntax highlighting and copy functionality.
- *
- * @param {AssistantMarkdownContentProps} props - The props for the component.
- * @returns {JSX.Element} The rendered markdown content.
- *
- * @internal
- */
-const AssistantMarkdownContent: React.FC<AssistantMarkdownContentProps> = ({content, style}) => {
-    return (
-        <Markdown
-            components={{
-                /**
-                 * Custom renderer for code blocks to include syntax highlighting and a copy button.
-                 *
-                 * @param {Object} props - Properties passed to the code component.
-                 * @param {Object} props.node - The AST node.
-                 * @param {string} props.className - The class name for the code block.
-                 * @param {React.ReactNode} props.children - The code content.
-                 * @returns {JSX.Element} The customized code block.
-                 */
-                code({node, className, children, ...props}) {
-                    const match = /language-(\w+)/.exec(className || '');
-                    const language = match ? match[1] : 'plaintext';
-                    const [copied, setCopied] = React.useState(false);
-
-                    /**
-                     * Handles copying the code content to the clipboard.
-                     */
-                    const handleCopy = () => {
-                        navigator.clipboard.writeText(String(children));
-                        setCopied(true);
-                        setTimeout(() => {
-                            setCopied(false);
-                        }, 3000); // Reset copied state after 3 seconds
-                    }
-
-                    // Block code (triple ```)
-                    return (
-                        <div style={{
-                            position: 'relative',
-                            marginBottom: '1rem'
-                        }} {...(props as React.HTMLProps<HTMLDivElement>)}>
-                            {match && (
-                                <div
-                                    style={{
-                                        marginLeft: '0.33rem',
-                                        fontSize: 'var(--font-size-small)',
-                                        fontWeight: 'bold',
-                                        background: '#f5f5f5',
-                                        padding: '0 8px',
-                                        borderTopLeftRadius: '5px',
-                                        borderTopRightRadius: '5px',
-                                        borderBottom: '1px solid #ccc',
-                                        display: 'inline-block',
-                                    }}
-                                >
-                                    {language}
-                                </div>
-                            )}
-                            <div className="relative">
-                                <SyntaxHighlighter
-                                    style={docco}
-                                    language={language}
-                                    PreTag="div"
-                                    wrapLongLine={true}
-                                >
-                                    {String(children).replace(/\n$/, '')}
-                                </SyntaxHighlighter>
-
-                                <div
-                                    className="absolute top-1 right-1 px-1.5 pt-0.5 text-sm rounded-md hover:shadow-md "
-                                    style={{
-                                        backgroundColor: style.messageBackgroundColor,
-                                    }}
-                                >
-                                    <button
-                                        onClick={handleCopy}
-                                    >
-                                        {copied ? <ClipboardDocumentIcon className={"w-5 h-5"} /> : <ClipboardIcon className={"w-5 h-5"} />}
-                                    </button>
-                                </div>
-                            </div>
-                        </div>
-                    );
-                },
-            }}
-            remarkPlugins={[remarkGfm]}>
-            {content}
-        </Markdown>
-    );
-}
+import {MessageFeedback} from "./MessageFeedback.tsx";
+import {MessageHighlight} from '../../types';
+import {renderHighlightedContent} from './MessageHighlighting';
 
 /**
  * Props for the ChatWindowAssistantMessage component.
  * @typedef {Object} ChatWindowAssistantMessageProps
  * @property {string} message - The assistant message content.
+ * @property {string | null} messageId - ID of the message, null for streaming messages
  * @property {ChatWindowStyles} style - Styling options for the chat window / assistant message
  * @property {string} roleLabel - Label indicating the role of the assistant.
  * @property {boolean} showRoleIndicator - Flag to show or hide the role indicator.
@@ -124,6 +20,7 @@ const AssistantMarkdownContent: React.FC<AssistantMarkdownContentProps> = ({cont
  * @property {boolean} [isStreaming] - Flag indicating if the message is being streamed.
  * @property {boolean} [showCopy] - Flag to show or hide the copy button.
  * @property {boolean} [showFeedback] - Flag to show or hide the feedback buttons.
+ * @property {MessageHighlight[]} [highlights] - Highlights to highlight in the message, used for completed messages
  */
 interface ChatWindowAssistantMessageProps {
     message: string;
@@ -136,6 +33,7 @@ interface ChatWindowAssistantMessageProps {
     isStreaming?: boolean;
     showCopy?: boolean;
     showFeedback?: boolean;
+    highlights?: MessageHighlight[];
 }
 
 /**
@@ -155,13 +53,16 @@ const ChatWindowAssistantMessage: React.FC<ChatWindowAssistantMessageProps> = ({
                                                                                    isStreaming,
                                                                                    showCopy,
                                                                                    showFeedback,
+                                                                                   highlights
                                                                                }) => {
     const [copied, setCopied] = useState(false);
     // show icon if showRoleIndicator is true and icon is not null
     const showIcon = showRoleIndicator && !!icon;
     // show label if showRoleIndicator is true and roleLabel is not null and icon is not shown
     const showLabel = showRoleIndicator && !!roleLabel && !showIcon;
 
+    const highlightedContentRef = React.useRef<HTMLDivElement>(null);
+
     /**
      * Handles copying the message content to the clipboard.
      */
@@ -210,16 +111,15 @@ const ChatWindowAssistantMessage: React.FC<ChatWindowAssistantMessageProps> = ({
                     borderRadius: style.messageBorderRadius,
                     fontSize: style.messageFontSize,
                 }}>
-                    {markdown ? (
-                        <div className={"assistant-markdown-content"}
-                             style={assistantMarkdownContentStyling}>
-                            <AssistantMarkdownContent content={message} style={style}/>
-                        </div>
-                    ) : (
-                        <div className="inline whitespace-pre-wrap">
-                            {message}
-                        </div>
-                    )}
+                    <div ref={highlightedContentRef}>
+                        {renderHighlightedContent(
+                            message,
+                            highlights,
+                            markdown || false,
+                            style,
+                            assistantMarkdownContentStyling
+                        )}
+                    </div>
                 </div>
                 {messageId && !isStreaming && (
                     <div