app/frontend/templates/pages/chat.html

{% extends "base.html" %}

{% block title %}
    <title>The Ultimate RAG</title>
{% endblock %}

{% block content %}
    <div class="chat-page">
        <div class="container py-4">
            <div id="chat-messages" class="chat-messages">
                <!-- Messages will be dynamically added by JavaScript -->
            </div>

            <form id="chat-form" class="input-group mt-4" enctype="multipart/form-data">
                <input type="text" class="form-control" name="prompt" placeholder="Ask your question here" id="queryInput">
                <label class="btn btn-outline-secondary btn-primary">
                    📎<input type="file" id="fileInput" name="files" multiple hidden>
                </label>
                <button type="button" class="btn text-white" id="searchButton">Send</button>
            </form>
        </div>
    </div>
{% endblock %}

{% block body_scripts %}
    <script src="https://cdn.jsdelivr.net/npm/marked/marked.min.js"></script>
    <script>
        const initialChatId = "{{ chat_id }}";
        const initialHistory = {{ history | tojson | safe }};
        let conversationId = initialChatId || null;

        // Initialize chat history
        if (initialHistory && Array.isArray(initialHistory)) {
            initialHistory.forEach(msg => {
                addMessageToChat(msg.role, msg.content);
            });
        }

        // Main chat function with WebSocket
        document.getElementById('searchButton').addEventListener('click', async function() {
            const query = document.getElementById('queryInput').value.trim();
            if (!query) return alert('Please enter a question');

            addMessageToChat('user', escapeHTML(query));
            document.getElementById('queryInput').value = '';
            const loadingId = addMessageToChat('assistant', 'Thinking...', true);

            try {
                const formData = new FormData();
                const fileInput = document.getElementById('fileInput');
                const files = fileInput.files;
                for (let i = 0; i < files.length; i++) {
                    formData.append('files', files[i]);
                }
                formData.append('prompt', query);
                if (conversationId) formData.append('chat_id', conversationId);

                const response = await fetch('/message_with_docs', {
                    method: 'POST',
                    body: formData
                });

                if (!response.ok) throw new Error(`HTTP error: ${response.status}`);

                const data = await response.json();
                const taskId = data.resp_task_id;
                console.log(taskId)
                if (!taskId) throw new Error('No task ID received');

                // Connect to WebSocket for streaming
                streamResponse(taskId, loadingId);

            } catch (error) {
                removeMessage(loadingId);
                addMessageToChat('assistant', `Error: ${error.message}`, false, 'error');
                console.error('Error:', error);
            }
        });

        // WebSocket streaming function
        function streamResponse(taskId, loadingId) {
            const ws = new WebSocket(`wss://${window.location.host}/ws/response/${taskId}`);
            let fullMessage = '';

            ws.onopen = () => {
                console.log(`WebSocket connected for task: ${taskId}`);
            };

            ws.onmessage = async (event) => {
                try {
                    const data = JSON.parse(event.data);
                    if (data.status === 'completed') {
                        ws.close();
                        try {
                            const replaceResp = await fetch('/replace_message', {
                                method: 'POST',
                                headers: { 'Content-Type': 'application/json' },
                                body: JSON.stringify({
                                    message: fullMessage,
                                    chat_id: conversationId
                                })
                            });
                            const json = await replaceResp.json();
                            const updatedMessage = json.updated_message;

                            removeMessage(loadingId);
                            addMessageToChat('assistant', updatedMessage);
                        } catch (error) {
                            removeMessage(loadingId);
                            addMessageToChat('assistant', `Error replacing message: ${error.message}`, false, 'error');
                        }
                    } else if (data.status === 'failed') {
                        ws.close();
                        removeMessage(loadingId);
                        addMessageToChat('assistant', `Error: ${data.error}`, false, 'error');
                    } else {
                        fullMessage += event.data;
                        updateMessageContent(loadingId, marked.parse(fullMessage));
                    }
                } catch {
                    fullMessage += event.data;
                    updateMessageContent(loadingId, marked.parse(fullMessage));
                }
            };

            ws.onclose = () => {
                console.log(`WebSocket closed for task: ${taskId}`);
                if (!fullMessage) {
                    removeMessage(loadingId);
                    addMessageToChat('assistant', 'No response received', false, 'error');
                }
            };

            ws.onerror = (error) => {
                console.error(`WebSocket error for task ${taskId}:`, error);
                removeMessage(loadingId);
                addMessageToChat('assistant', 'WebSocket error occurred', false, 'error');
                ws.close();
            };
        }


        // Update message content
        function updateMessageContent(messageId, newContent) {
            const element = document.getElementById(messageId);
            if (element) {
                const contentDiv = element.querySelector('.message-content');
                if (contentDiv) contentDiv.innerHTML = newContent;
                document.getElementById('chat-messages').scrollTop = document.getElementById('chat-messages').scrollHeight;
            }
        }

        // Add message to chat
        function addMessageToChat(role, content, isTemporary = false, className = '') {
            const chatMessages = document.getElementById('chat-messages');
            const messageId = 'msg-' + Date.now();

            const messageDiv = document.createElement('div');
            messageDiv.className = `message ${role}-message ${className}`;
            messageDiv.id = messageId;

            messageDiv.innerHTML = `
                <div class="message-header">${role === 'user' ? 'You' : 'Assistant'}</div>
                <div class="message-content">${marked.parse(content)}</div>
            `;

            chatMessages.appendChild(messageDiv);
            chatMessages.scrollTop = chatMessages.scrollHeight;

            return messageId;
        }

        // Remove message
        function removeMessage(messageId) {
            const element = document.getElementById(messageId);
            if (element) element.remove();
        }

        // Escape HTML
        function escapeHTML(str) {
            const div = document.createElement('div');
            div.textContent = str;
            return div.innerHTML;
        }

        // New chat handler
        document.querySelector('form[action="/new_chat"]').addEventListener('submit', function(e) {
            e.preventDefault();
            conversationId = null;
            document.getElementById('chat-messages').innerHTML = '';
            this.submit();
        });
    </script>
{% endblock %}