Frontend capture complete response

src/airflow_wingman/static/css/wingman_chat.css

@@ -0,0 +1,85 @@
+/* Provider and model selection styling */
+.provider-section {
+    margin-bottom: 20px;
+}
+.provider-name {
+    font-size: 16px;
+    font-weight: bold;
+    margin-bottom: 10px;
+    color: #666;
+}
+.model-option {
+    margin-left: 15px;
+    margin-bottom: 8px;
+}
+.model-option label {
+    display: block;
+    cursor: pointer;
+}
+
+/* Message styling */
+.message {
+    margin-bottom: 15px;
+    max-width: 80%;
+    clear: both;
+}
+
+.message-user {
+    float: right;
+    background-color: #f0f7ff;
+    border: 1px solid #d1e6ff;
+    border-radius: 15px 15px 0 15px;
+    padding: 10px 15px;
+}
+
+.message-assistant {
+    float: left;
+    background-color: #f8f9fa;
+    border: 1px solid #e9ecef;
+    border-radius: 15px 15px 15px 0;
+    padding: 10px 15px;
+    white-space: pre-wrap;
+}
+
+#chat-messages::after {
+    content: "";
+    clear: both;
+    display: table;
+}
+
+/* Scrollbar styling */
+.panel-body::-webkit-scrollbar {
+    width: 8px;
+}
+
+.panel-body::-webkit-scrollbar-track {
+    background: #f1f1f1;
+}
+
+.panel-body::-webkit-scrollbar-thumb {
+    background: #888;
+    border-radius: 4px;
+}
+
+.panel-body::-webkit-scrollbar-thumb:hover {
+    background: #555;
+}
+
+/* Processing indicator styling */
+.processing-indicator {
+    display: none;
+    background-color: #f0f8ff;
+    padding: 8px 12px;
+    border-radius: 4px;
+    margin: 8px 0;
+    font-style: italic;
+}
+
+.processing-indicator.visible {
+    display: block;
+}
+
+.pre-formatted {
+    white-space: pre-wrap;
+    font-family: monospace;
+}

src/airflow_wingman/static/js/wingman_chat.js

@@ -0,0 +1,285 @@
+document.addEventListener('DOMContentLoaded', function() {
+    // Add title attributes for tooltips
+    document.querySelectorAll('[data-bs-toggle="tooltip"]').forEach(function(el) {
+        el.title = el.getAttribute('title') || el.getAttribute('data-bs-original-title');
+    });
+
+    // Handle model selection and model name input
+    const modelNameInput = document.getElementById('modelName');
+    const modelRadios = document.querySelectorAll('input[name="model"]');
+
+    modelRadios.forEach(function(radio) {
+        radio.addEventListener('change', function() {
+            const provider = this.value.split(':')[0];
+            const modelName = this.getAttribute('data-model-name');
+            console.log('Selected provider:', provider);
+            console.log('Model name:', modelName);
+
+            if (provider === 'openrouter') {
+                console.log('Enabling model name input');
+                modelNameInput.disabled = false;
+                modelNameInput.value = '';
+                modelNameInput.placeholder = 'Enter model name for OpenRouter';
+            } else {
+                console.log('Disabling model name input');
+                modelNameInput.disabled = true;
+                modelNameInput.value = modelName;
+            }
+        });
+    });
+
+    // Set initial state based on default selection
+    const defaultSelected = document.querySelector('input[name="model"]:checked');
+    if (defaultSelected) {
+        const provider = defaultSelected.value.split(':')[0];
+        const modelName = defaultSelected.getAttribute('data-model-name');
+        console.log('Initial provider:', provider);
+        console.log('Initial model name:', modelName);
+
+        if (provider === 'openrouter') {
+            console.log('Initially enabling model name input');
+            modelNameInput.disabled = false;
+            modelNameInput.value = '';
+            modelNameInput.placeholder = 'Enter model name for OpenRouter';
+        } else {
+            console.log('Initially disabling model name input');
+            modelNameInput.disabled = true;
+            modelNameInput.value = modelName;
+        }
+    }
+
+    const messageInput = document.getElementById('message-input');
+    const sendButton = document.getElementById('send-button');
+    const refreshButton = document.getElementById('refresh-button');
+    const chatMessages = document.getElementById('chat-messages');
+
+    let currentMessageDiv = null;
+    let messageHistory = [];
+
+    // Create a processing indicator element
+    const processingIndicator = document.createElement('div');
+    processingIndicator.className = 'processing-indicator';
+    processingIndicator.textContent = 'Processing tool calls...';
+    chatMessages.appendChild(processingIndicator);
+
+    function clearChat() {
+        // Clear the chat messages
+        chatMessages.innerHTML = '';
+        // Add back the processing indicator
+        chatMessages.appendChild(processingIndicator);
+        // Reset message history
+        messageHistory = [];
+        // Clear the input field
+        messageInput.value = '';
+        // Enable input if it was disabled
+        messageInput.disabled = false;
+        sendButton.disabled = false;
+    }
+
+    function addMessage(content, isUser) {
+        const messageDiv = document.createElement('div');
+        messageDiv.className = `message ${isUser ? 'message-user' : 'message-assistant'}`;
+
+        // Apply pre-formatted class to preserve whitespace and newlines
+        messageDiv.classList.add('pre-formatted');
+
+        // Use innerText instead of textContent to preserve newlines
+        messageDiv.innerText = content;
+
+        chatMessages.appendChild(messageDiv);
+        chatMessages.scrollTop = chatMessages.scrollHeight;
+        return messageDiv;
+    }
+
+    function showProcessingIndicator() {
+        processingIndicator.classList.add('visible');
+        chatMessages.scrollTop = chatMessages.scrollHeight;
+    }
+
+    function hideProcessingIndicator() {
+        processingIndicator.classList.remove('visible');
+    }
+
+    async function sendMessage() {
+        const message = messageInput.value.trim();
+        if (!message) return;
+
+        // Get selected model
+        const selectedModel = document.querySelector('input[name="model"]:checked');
+        if (!selectedModel) {
+            alert('Please select a model');
+            return;
+        }
+
+        const [provider, modelId] = selectedModel.value.split(':');
+        const modelName = provider === 'openrouter' ? modelNameInput.value : modelId;
+
+        // Clear input and add user message
+        messageInput.value = '';
+        addMessage(message, true);
+
+        // Add user message to history
+        messageHistory.push({
+            role: 'user',
+            content: message
+        });
+
+        // Use full message history for the request
+        const messages = [...messageHistory];
+
+        // Create assistant message div
+        currentMessageDiv = addMessage('', false);
+
+        // Get API key
+        const apiKey = document.getElementById('api-key').value.trim();
+        if (!apiKey) {
+            alert('Please enter an API key');
+            return;
+        }
+
+        // Disable input while processing
+        messageInput.disabled = true;
+        sendButton.disabled = true;
+
+        // Get CSRF token
+        const csrfToken = document.querySelector('meta[name="csrf-token"]')?.getAttribute('content');
+        if (!csrfToken) {
+            alert('CSRF token not found. Please refresh the page.');
+            return;
+        }
+
+        // Create request data
+        const requestData = {
+            provider: provider,
+            model: modelName,
+            messages: messages,
+            api_key: apiKey,
+            stream: true,
+            temperature: 0.7
+        };
+        console.log('Sending request:', {...requestData, api_key: '***'});
+
+        try {
+            // Send request
+            const response = await fetch('/wingman/chat', {
+                method: 'POST',
+                headers: {
+                    'Content-Type': 'application/json',
+                    'X-CSRFToken': csrfToken
+                },
+                body: JSON.stringify(requestData)
+            });
+
+            if (!response.ok) {
+                const error = await response.json();
+                throw new Error(error.error || 'Failed to get response');
+            }
+
+            // Process the streaming response
+            const reader = response.body.getReader();
+            const decoder = new TextDecoder();
+            let fullResponse = '';
+
+            while (true) {
+                const { value, done } = await reader.read();
+                if (done) break;
+
+                const chunk = decoder.decode(value);
+                const lines = chunk.split('\n');
+
+                for (const line of lines) {
+                    if (line.trim() === '') continue;
+
+                    if (line.startsWith('data: ')) {
+                        const content = line.slice(6); // Remove 'data: ' prefix
+
+                        // Check for special events or end marker
+                        if (content === '[DONE]') {
+                            console.log('Stream complete');
+
+                            // Add assistant's response to history
+                            if (fullResponse) {
+                                messageHistory.push({
+                                    role: 'assistant',
+                                    content: fullResponse
+                                });
+                            }
+                            continue;
+                        }
+
+                        // Try to parse as JSON for special events
+                        try {
+                            const parsed = JSON.parse(content);
+
+                            if (parsed.event === 'tool_processing_start') {
+                                console.log('Tool processing started');
+                                showProcessingIndicator();
+                                continue;
+                            }
+
+                            if (parsed.event === 'tool_processing_complete') {
+                                console.log('Tool processing completed');
+                                hideProcessingIndicator();
+                                continue;
+                            }
+
+                            // Handle the complete response event
+                            if (parsed.event === 'complete_response') {
+                                console.log('Received complete response from backend');
+                                // Use the complete response from the backend
+                                fullResponse = parsed.content;
+
+                                // Update the display with the complete response
+                                if (!currentMessageDiv.classList.contains('pre-formatted')) {
+                                    currentMessageDiv.classList.add('pre-formatted');
+                                }
+                                currentMessageDiv.innerText = fullResponse;
+                                continue;
+                            }
+
+                            // If we have JSON that's not a special event, it might be content
+                            currentMessageDiv.textContent += JSON.stringify(parsed);
+                            fullResponse += JSON.stringify(parsed);
+                        } catch (e) {
+                            // Not JSON, handle as normal content
+                            // console.log('Received chunk:', JSON.stringify(content));
+
+                            // Add to full response
+                            fullResponse += content;
+
+                            // Create a properly formatted display
+                            if (!currentMessageDiv.classList.contains('pre-formatted')) {
+                                currentMessageDiv.classList.add('pre-formatted');
+                            }
+
+                            // Always rebuild the entire content from the full response
+                            currentMessageDiv.innerText = fullResponse;
+                        }
+                        // Scroll to bottom
+                        chatMessages.scrollTop = chatMessages.scrollHeight;
+                    }
+                }
+            }
+        } catch (error) {
+            console.error('Error:', error);
+            if (currentMessageDiv) {
+                currentMessageDiv.textContent = `Error: ${error.message}`;
+                currentMessageDiv.style.color = 'red';
+            }
+        } finally {
+            // Always re-enable input and hide indicators
+            messageInput.disabled = false;
+            sendButton.disabled = false;
+            hideProcessingIndicator();
+        }
+    }
+
+    sendButton.addEventListener('click', sendMessage);
+    messageInput.addEventListener('keypress', function(e) {
+        if (e.key === 'Enter') {
+            sendMessage();
+        }
+    });
+
+    refreshButton.addEventListener('click', clearChat);
+});

src/airflow_wingman/templates/wingman_chat.html

@@ -3,6 +3,7 @@
 {% block head_meta %}
     {{ super() }}
     <meta name="csrf-token" content="{{ csrf_token() }}">
+    <link rel="stylesheet" href="{{ url_for('wingman.static', filename='css/wingman_chat.css') }}">
 {% endblock %}
 
 {% block content %}
@@ -77,25 +78,7 @@
                     </div>
                 </div>
 
-                <style>
+
-                .provider-section {
-                    margin-bottom: 20px;
-                }
-                .provider-name {
-                    font-size: 16px;
-                    font-weight: bold;
-                    margin-bottom: 10px;
-                    color: #666;
-                }
-                .model-option {
-                    margin-left: 15px;
-                    margin-bottom: 8px;
-                }
-                .model-option label {
-                    display: block;
-                    cursor: pointer;
-                }
-                </style>
             </div>
         </div>
 
@@ -129,259 +112,5 @@
     </div>
 </div>
 
-<style>
+<script src="{{ url_for('wingman.static', filename='js/wingman_chat.js') }}"></script>
-.message {
-    margin-bottom: 15px;
-    max-width: 80%;
-    clear: both;
-}
-
-.message-user {
-    float: right;
-    background-color: #f0f7ff;
-    border: 1px solid #d1e6ff;
-    border-radius: 15px 15px 0 15px;
-    padding: 10px 15px;
-}
-
-.message-assistant {
-    float: left;
-    background-color: #f8f9fa;
-    border: 1px solid #e9ecef;
-    border-radius: 15px 15px 15px 0;
-    padding: 10px 15px;
-    white-space: pre-wrap;
-}
-
-#chat-messages::after {
-    content: "";
-    clear: both;
-    display: table;
-}
-
-.panel-body::-webkit-scrollbar {
-    width: 8px;
-}
-
-.panel-body::-webkit-scrollbar-track {
-    background: #f1f1f1;
-}
-
-.panel-body::-webkit-scrollbar-thumb {
-    background: #888;
-    border-radius: 4px;
-}
-
-.panel-body::-webkit-scrollbar-thumb:hover {
-    background: #555;
-}
-</style>
-
-<script>
-document.addEventListener('DOMContentLoaded', function() {
-    // Add title attributes for tooltips
-    document.querySelectorAll('[data-bs-toggle="tooltip"]').forEach(function(el) {
-        el.title = el.getAttribute('title') || el.getAttribute('data-bs-original-title');
-    });
-
-    // Handle model selection and model name input
-    const modelNameInput = document.getElementById('modelName');
-    const modelRadios = document.querySelectorAll('input[name="model"]');
-
-    modelRadios.forEach(function(radio) {
-        radio.addEventListener('change', function() {
-            const provider = this.value.split(':')[0];  // Get provider from value instead of data attribute
-            const modelName = this.getAttribute('data-model-name');
-            console.log('Selected provider:', provider);
-            console.log('Model name:', modelName);
-
-            if (provider === 'openrouter') {
-                console.log('Enabling model name input');
-                modelNameInput.disabled = false;
-                modelNameInput.value = '';
-                modelNameInput.placeholder = 'Enter model name for OpenRouter';
-            } else {
-                console.log('Disabling model name input');
-                modelNameInput.disabled = true;
-                modelNameInput.value = modelName;
-            }
-        });
-    });
-
-    // Set initial state based on default selection
-    const defaultSelected = document.querySelector('input[name="model"]:checked');
-    if (defaultSelected) {
-        const provider = defaultSelected.value.split(':')[0];  // Get provider from value instead of data attribute
-        const modelName = defaultSelected.getAttribute('data-model-name');
-        console.log('Initial provider:', provider);
-        console.log('Initial model name:', modelName);
-
-        if (provider === 'openrouter') {
-            console.log('Initially enabling model name input');
-            modelNameInput.disabled = false;
-            modelNameInput.value = '';
-            modelNameInput.placeholder = 'Enter model name for OpenRouter';
-        } else {
-            console.log('Initially disabling model name input');
-            modelNameInput.disabled = true;
-            modelNameInput.value = modelName;
-        }
-    }
-
-    const messageInput = document.getElementById('message-input');
-    const sendButton = document.getElementById('send-button');
-    const refreshButton = document.getElementById('refresh-button');
-    const chatMessages = document.getElementById('chat-messages');
-
-    let currentMessageDiv = null;
-    let messageHistory = [];
-
-    function clearChat() {
-        // Clear the chat messages
-        chatMessages.innerHTML = '';
-        // Reset message history
-        messageHistory = [];
-        // Clear the input field
-        messageInput.value = '';
-        // Enable input if it was disabled
-        messageInput.disabled = false;
-        sendButton.disabled = false;
-    }
-
-    function addMessage(content, isUser) {
-        const messageDiv = document.createElement('div');
-        messageDiv.className = `message ${isUser ? 'message-user' : 'message-assistant'}`;
-        messageDiv.textContent = content;
-        chatMessages.appendChild(messageDiv);
-        chatMessages.scrollTop = chatMessages.scrollHeight;
-        return messageDiv;
-    }
-
-    async function sendMessage() {
-        const message = messageInput.value.trim();
-        if (!message) return;
-
-        // Get selected model
-        const selectedModel = document.querySelector('input[name="model"]:checked');
-        if (!selectedModel) {
-            alert('Please select a model');
-            return;
-        }
-
-        const [provider, modelId] = selectedModel.value.split(':');
-        const modelName = provider === 'openrouter' ? modelNameInput.value : modelId;
-
-        // Clear input and add user message
-        messageInput.value = '';
-        addMessage(message, true);
-
-        try {
-            // Add user message to history
-            messageHistory.push({
-                role: 'user',
-                content: message
-            });
-
-            // Use full message history for the request
-            const messages = [...messageHistory];
-
-            // Create assistant message div
-            currentMessageDiv = addMessage('', false);
-
-            // Get API key
-            const apiKey = document.getElementById('api-key').value.trim();
-            if (!apiKey) {
-                alert('Please enter an API key');
-                return;
-            }
-
-            // Debug log the request
-            const requestData = {
-                provider: provider,
-                model: modelName,
-                messages: messages,
-                api_key: apiKey,
-                stream: true,
-                temperature: 0.7
-            };
-            console.log('Sending request:', {...requestData, api_key: '***'});
-
-            // Get CSRF token
-            const csrfToken = document.querySelector('meta[name="csrf-token"]')?.getAttribute('content');
-            if (!csrfToken) {
-                throw new Error('CSRF token not found. Please refresh the page.');
-            }
-
-            // Send request
-            const response = await fetch('/wingman/chat', {
-                method: 'POST',
-                headers: {
-                    'Content-Type': 'application/json',
-                    'X-CSRFToken': csrfToken
-                },
-                body: JSON.stringify({
-                    provider: provider,
-                    model: modelName,
-                    messages: messages,
-                    api_key: apiKey,
-                    stream: true,
-                    temperature: 0.7
-                })
-            });
-
-            if (!response.ok) {
-                const error = await response.json();
-                throw new Error(error.error || 'Failed to get response');
-            }
-
-            // Handle streaming response
-            const reader = response.body.getReader();
-            const decoder = new TextDecoder();
-            let fullResponse = '';
-
-            while (true) {
-                const { value, done } = await reader.read();
-                if (done) break;
-
-                const chunk = decoder.decode(value);
-                const lines = chunk.split('\n');
-
-                for (const line of lines) {
-                    if (line.startsWith('data: ')) {
-                        const content = line.slice(6);
-                        if (content) {
-                            // Use textContent to properly handle newlines
-                            console.log('Received chunk:', JSON.stringify(content)); // Debug
-                            currentMessageDiv.textContent += content;
-                            fullResponse += content;
-                            chatMessages.scrollTop = chatMessages.scrollHeight;
-                        }
-                    }
-                }
-            }
-
-            // Add assistant's response to history
-            if (fullResponse) {
-                messageHistory.push({
-                    role: 'assistant',
-                    content: fullResponse
-                });
-            }
-        } catch (error) {
-            console.error('Error:', error);
-            currentMessageDiv.textContent = `Error: ${error.message}`;
-            currentMessageDiv.style.color = 'red';
-        }
-    }
-
-    sendButton.addEventListener('click', sendMessage);
-    messageInput.addEventListener('keypress', function(e) {
-        if (e.key === 'Enter') {
-            sendMessage();
-        }
-    });
-
-    refreshButton.addEventListener('click', clearChat);
-});
-</script>
 {% endblock %}

src/airflow_wingman/views.py

@@ -116,6 +116,7 @@ class WingmanView(AppBuilderBaseView):
                 # Send SSE format for each chunk
                 for chunk in generator:
                     if chunk:
+                        complete_response += chunk
                         yield f"data: {chunk}\n\n"
 
                 # Log the complete assembled response at the end
@@ -123,6 +124,11 @@ class WingmanView(AppBuilderBaseView):
                 logger.info(complete_response)
                 logger.info("<<< COMPLETE RESPONSE END")
 
+                # Send the complete response as a special event
+                complete_event = json.dumps({"event": "complete_response", "content": complete_response})
+                yield f"data: {complete_event}\n\n"
+
+                # Signal the end of the stream
                 yield "data: [DONE]\n\n"
 
             return Response(stream_response(), mimetype="text/event-stream", headers={"Cache-Control": "no-cache", "X-Accel-Buffering": "no"})