MtoRagSystem/public/assets/js/base.js

document.addEventListener('DOMContentLoaded', () => {
    const chatEl = document.getElementById('chat');
    const promptEl = document.getElementById('prompt');
    const sendBtn = document.getElementById('send');
    const abortBtn = document.getElementById('abort');
    const clearBtn = document.getElementById('clear');

    let abort = false;

    marked.setOptions({ breaks: true });

    function renderMarkdown(text) {
        return DOMPurify.sanitize(marked.parse(text));
    }

    function enhanceChatLinks(container = chatEl) {
        if (!container) {
            return;
        }

        container.querySelectorAll('a[href]').forEach((link) => {
            link.setAttribute('target', '_blank');
            link.setAttribute('rel', 'noopener noreferrer');
        });
    }

    function addMessage(role, html = '', extra = '') {
        const msg = document.createElement('div');
        msg.className = 'message ' + role;

        const bubble = document.createElement('div');
        bubble.className = 'bubble ' + extra;
        bubble.innerHTML = html;

        msg.appendChild(bubble);
        chatEl.appendChild(msg);

        enhanceChatLinks(bubble);

        chatEl.scrollTop = chatEl.scrollHeight;

        return bubble;
    }

    function addLoader() {
        return addMessage('assistant', 'AI is thinking…', 'loader');
    }

    function cleanupEmptyBlocks(container) {
        container.querySelectorAll('p, div, li, blockquote').forEach((el) => {
            const html = el.innerHTML
                .replace(/<br\s*\/?>/gi, '')
                .replace(/&nbsp;/gi, '')
                .trim();

            const text = (el.textContent || '').trim();

            if (html === '' || text === '') {
                el.remove();
            }
        });
    }

    function stripAllThinkContent(container) {
        const blockSelector = 'p, div, li, blockquote';
        const thinkSpans = Array.from(container.querySelectorAll('.think'));

        if (thinkSpans.length === 0) {
            return;
        }

        const handledBlocks = new Set();

        thinkSpans.forEach((span) => {
            const block = span.closest(blockSelector) || span.parentElement;

            if (!block || handledBlocks.has(block)) {
                return;
            }

            handledBlocks.add(block);

            const thinksInBlock = Array.from(block.querySelectorAll('.think'));
            const lastThinkInBlock = thinksInBlock[thinksInBlock.length - 1];

            if (!lastThinkInBlock) {
                return;
            }

            let node = block.firstChild;

            while (node) {
                const next = node.nextSibling;
                node.remove();

                if (node === lastThinkInBlock) {
                    break;
                }

                node = next;
            }

            while (
                block.firstChild &&
                (
                    (block.firstChild.nodeType === Node.TEXT_NODE &&
                        block.firstChild.textContent.trim() === '') ||
                    (block.firstChild.nodeType === Node.ELEMENT_NODE &&
                        block.firstChild.tagName === 'BR')
                )
                ) {
                block.firstChild.remove();
            }
        });

        cleanupEmptyBlocks(container);
    }

    function hasNonThinkContent(container) {
        const clone = container.cloneNode(true);

        stripAllThinkContent(clone);
        cleanupEmptyBlocks(clone);

        if ((clone.textContent || '').trim() !== '') {
            return true;
        }

        return clone.querySelector(
            'img, table, pre, code, ul, ol, h1, h2, h3, h4, h5, h6, a'
        ) !== null;
    }

    function cleanupThinkSpans(container) {
        if (!container) {
            return;
        }

        const thinkSpans = Array.from(container.querySelectorAll('.think'));

        if (thinkSpans.length === 0) {
            return;
        }

        if (hasNonThinkContent(container)) {
            stripAllThinkContent(container);
            return;
        }

        if (thinkSpans.length <= 1) {
            return;
        }

        const blockSelector = 'p, div, li, blockquote';
        const lastThink = thinkSpans[thinkSpans.length - 1];
        const lastBlock = lastThink.closest(blockSelector) || lastThink.parentElement;

        thinkSpans.slice(0, -1).forEach((span) => {
            const block = span.closest(blockSelector) || span.parentElement;

            if (block && block !== lastBlock) {
                block.remove();
                return;
            }

            if (block === lastBlock) {
                span.remove();
            }
        });

        if (lastBlock && lastBlock.contains(lastThink)) {
            let node = lastBlock.firstChild;

            while (node && node !== lastThink) {
                const next = node.nextSibling;
                node.remove();
                node = next;
            }

            while (
                lastThink.nextSibling &&
                (
                    (lastThink.nextSibling.nodeType === Node.TEXT_NODE &&
                        lastThink.nextSibling.textContent.trim() === '') ||
                    (lastThink.nextSibling.nodeType === Node.ELEMENT_NODE &&
                        lastThink.nextSibling.tagName === 'BR')
                )
                ) {
                lastThink.nextSibling.remove();
            }
        }

        cleanupEmptyBlocks(container);
    }

    function renderBubbleContent(bubble, raw) {
        bubble.innerHTML = renderMarkdown(raw);
        cleanupThinkSpans(bubble);
        enhanceChatLinks(bubble);
        chatEl.scrollTop = chatEl.scrollHeight;
    }

    async function loadHistory() {
        try {
            const res = await fetch('/history');
            if (!res.ok) return;
            const messages = await res.json();
            messages.forEach(m => addMessage(m.role, renderMarkdown(m.text)));
            enhanceChatLinks(chatEl);
        } catch {}
    }

    loadHistory();

    promptEl.addEventListener('keydown', (e) => {
        if (e.key === 'Enter' && !e.shiftKey) {
            e.preventDefault();
            sendBtn.click();
        }
    });

    sendBtn.addEventListener('click', async () => {
        const prompt = promptEl.value.trim();
        if (!prompt) return;

        addMessage('user', renderMarkdown(prompt));
        promptEl.value = '';

        const bubble = addLoader();
        let raw = '';
        let firstChunk = true;
        let renderTimer = null;

        function scheduleRender() {
            if (renderTimer) return;

            renderTimer = setTimeout(() => {
                renderBubbleContent(bubble, raw);
                renderTimer = null;
            }, 100);
        }

        abort = false;
        sendBtn.disabled = true;
        abortBtn.disabled = false;
        clearBtn.disabled = true;

        try {
            document.getElementById('ai-cloud')?.classList.remove('d-none');

            const res = await fetch('/ask-sse', {
                method: 'POST',
                headers: { 'Content-Type': 'application/json' },
                body: JSON.stringify({ prompt })
            });

            const reader = res.body.getReader();
            const decoder = new TextDecoder();
            let sseBuffer = '';

            while (!abort) {
                const { value, done } = await reader.read();
                if (done) break;

                const chunk = decoder.decode(value, { stream: true });
                sseBuffer += chunk;

                const events = sseBuffer.split('\n\n');
                sseBuffer = events.pop() || '';

                events.forEach(event => {
                    if (!event.trim()) return;

                    const dataLines = event
                        .split('\n')
                        .filter(line => line.startsWith('data: '))
                        .map(line => line.slice(6));

                    if (dataLines.length === 0) return;

                    const text = dataLines.join('\n');

                    if (text === '[DONE]') {
                        if (renderTimer) {
                            clearTimeout(renderTimer);
                            renderTimer = null;
                        }

                        renderBubbleContent(bubble, raw);
                        abort = true;
                        return;
                    }

                    if (firstChunk) {
                        bubble.classList.remove('loader');
                        bubble.innerHTML = '';
                        firstChunk = false;
                        document.getElementById('ai-cloud')?.classList.add('d-none');
                    }

                    raw += text;
                    scheduleRender();
                });
            }
        } catch {
            bubble.innerHTML += '<br><em>Error occurred.</em>';
            enhanceChatLinks(bubble);
        } finally {
            if (renderTimer) {
                clearTimeout(renderTimer);
            }

            sendBtn.disabled = false;
            abortBtn.disabled = true;
            clearBtn.disabled = false;
        }
    });

    abortBtn.addEventListener('click', () => {
        abort = true;
        addMessage('assistant', '<em>[aborted]</em>');
    });

    clearBtn.addEventListener('click', async () => {
        await fetch('/history/delete', { method: 'POST' });
        chatEl.innerHTML = '';
        addMessage('assistant', '<em>History cleared.</em>');
    });
});