diff --git a/src/airflow_wingman/static/css/wingman_chat.css b/src/airflow_wingman/static/css/wingman_chat.css
new file mode 100644
index 0000000..10080dc
--- /dev/null
+++ b/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;
+}
diff --git a/src/airflow_wingman/static/js/wingman_chat.js b/src/airflow_wingman/static/js/wingman_chat.js
new file mode 100644
index 0000000..821e5b5
--- /dev/null
+++ b/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);
+});
diff --git a/src/airflow_wingman/templates/wingman_chat.html b/src/airflow_wingman/templates/wingman_chat.html
index 56a38a6..8424866 100644
--- a/src/airflow_wingman/templates/wingman_chat.html
+++ b/src/airflow_wingman/templates/wingman_chat.html
@@ -3,6 +3,7 @@
{% block head_meta %}
{{ super() }}
+
{% endblock %}
{% block content %}
@@ -77,25 +78,7 @@
-
+
@@ -129,259 +112,5 @@
-
-
-
+
{% endblock %}
diff --git a/src/airflow_wingman/views.py b/src/airflow_wingman/views.py
index 599c4bd..c5d6233 100644
--- a/src/airflow_wingman/views.py
+++ b/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"})