ZenML Assistant
+ `; + + chatMessages.appendChild(messageDiv, chatMessages.firstChild); + currentAssistantMessage = ''; + } + + currentAssistantMessage += text; + + const html = marked.parse(currentAssistantMessage); + const sanitizeHtml = DOMPurify.sanitize(html); + + requestAnimationFrame(() => { + messageDiv.innerHTML = ` +ZenML Assistant
+ ${sanitizeHtml} + `; + chatMessages.scrollTop = chatMessages.scrollHeight; + }); + } + } + + let isInputDisabled = false; + + function disableInput() { + isInputDisabled = true; + document.getElementById('sendMessage').disabled = true; + } + + function enableInput() { + isInputDisabled = false; + document.getElementById('sendMessage').disabled = false; + } + + const providerDropdown = document.getElementById('provider-dropdown'); + const modelDropdown = document.getElementById('model-dropdown'); + + providerDropdown.addEventListener('change', event => { + const selectedProvider = event.target.value; + vscode.postMessage({ command: 'updateProvider', provider: selectedProvider }); + }); + + modelDropdown.addEventListener('change', event => { + const selectedModel = event.target.value; + vscode.postMessage({ command: 'updateModel', model: selectedModel }); + }); + + window.addEventListener('message', event => { + const message = event.data; + switch (message.command) { + case 'updateChatLog': { + const sanitizeHtml = DOMPurify.sanitize(message.chatLogHtml); + document.getElementById('chatMessages').innerHTML = sanitizeHtml; + addCopyButtonsToAssistantMessages(); // For the potential view refresh command + break; + } + case 'receiveMessage': { + if (message.text === 'disableInput') { + disableInput(); + } else if (message.text === 'enableInput') { + enableInput(); + addCopyButtonToLastAssistantMessage(); + } else { + appendToChat(message.text, 'assistant'); + } + break; + } + case 'showInfo': { + vscode.window.showInformationMessage({ + command: 'showInfoToExtension', + text: message.text, + }); + break; + } + case 'updateModelList': { + updateModelDropdown(message.models); + break; + } + case 'updateModel': { + localStorage.setItem('selectedModel', message.text); + restoreState(); + break; + } + case 'hideLoader': { + hideLoader(); + break; + } + } + }); + + function hideLoader() { + loader.classList.remove('loader'); + } + + function updateModelDropdown(models) { + modelDropdown.innerHTML = models + .map(model => ``) + .join(''); + } + + function addCopyButtonsToAssistantMessages() { + const assistantMessages = document.querySelectorAll('.assistant'); + assistantMessages.forEach(addCopyButtonToMessage); + } + + function addCopyButtonToLastAssistantMessage() { + const lastAssistantMessage = document.querySelector('.assistant:last-child'); + if (lastAssistantMessage) { + addCopyButtonToMessage(lastAssistantMessage); + } + } + + function addCopyButtonToMessage(messageDiv) { + if (!messageDiv.querySelector('.copy-button')) { + const copyButton = document.createElement('button'); + copyButton.className = 'copy-button'; + copyButton.textContent = 'Copy'; + copyButton.addEventListener('click', () => { + // Find all text content within the message div, excluding the "ZenML Assistant" header and the copy button + const content = Array.from(messageDiv.childNodes) + .filter( + node => + node.nodeType === Node.TEXT_NODE || + (node.nodeType === Node.ELEMENT_NODE && + !node.classList.contains('font-semibold') && + !node.classList.contains('copy-button')) + ) + .map(node => node.textContent) + .join('') + .trim(); + + if (content) { + navigator.clipboard.writeText(content).then(() => { + vscode.postMessage({ command: 'showInfo', text: 'Message copied to clipboard' }); + }); + } else { + vscode.postMessage({ command: 'showInfo', text: 'No content to copy' }); + } + }); + messageDiv.appendChild(copyButton); + } + } + + // Add event listeners to save state when dropdowns or checkboxes change + document.querySelector('#provider-dropdown').addEventListener('change', saveState); + document.querySelector('#model-dropdown').addEventListener('change', saveState); + document.querySelectorAll('#tree-view input[type="checkbox"]').forEach(checkbox => { + checkbox.addEventListener('change', saveState); + }); + + function sendSampleMessage() { + const provider = document.querySelector('#provider-dropdown').value; + const model = document.querySelector('#model-dropdown').value; + let context = [model]; + let buttonValue = this.value; + let message; + + switch (buttonValue) { + case 'aboutChat': { + message = 'What can this chat do?'; + break; + } + case 'summarizeStats': { + message = 'Generate a summary of my stats.'; + context.push('serverContext'); + context.push('environmentContext'); + context.push('pipelineContext'); + context.push('stackContext'); + context.push('stackComponentsContext'); + break; + } + case 'summarizeLogs': { + message = 'Generate a summary of my logs.'; + context.push('logsContext'); + break; + } + default: { + break; + } + } + + if (message) { + vscode.postMessage({ + command: 'sendMessage', + text: message, + context: context, + provider: provider, + model: model, + }); + + saveState(); + } + } + + document.querySelectorAll('.sampleQuestions').forEach(button => { + button.addEventListener('click', sendSampleMessage); + }); + + const prevPageButton = document.getElementById('prevPage'); + const nextPageButton = document.getElementById('nextPage'); + + if (prevPageButton) { + prevPageButton.addEventListener('click', () => { + vscode.postMessage({ command: 'prevPage' }); + }); + } + + if (nextPageButton) { + nextPageButton.addEventListener('click', () => { + vscode.postMessage({ command: 'nextPage' }); + }); + } + + document.addEventListener('DOMContentLoaded', event => { + const contextButton = document.getElementById('contextButton'); + const optionsDropdown = document.getElementById('optionsDropdown'); + + let isContextDisplayed = localStorage.getItem('displayContext') === 'true'; + if (isContextDisplayed) { + optionsDropdown.classList.remove('hidden'); + } else { + optionsDropdown.classList.add('hidden'); + } + + contextButton.addEventListener('click', event => { + event.stopPropagation(); + optionsDropdown.classList.toggle('hidden'); + isContextDisplayed = !isContextDisplayed; + localStorage.setItem('displayContext', String(isContextDisplayed)); + }); + + document.addEventListener('click', event => { + if (!contextButton.contains(event.target) && !optionsDropdown.contains(event.target)) { + optionsDropdown.classList.add('hidden'); + isContextDisplayed = false; + localStorage.setItem('displayContext', String(isContextDisplayed)); + } + }); + + const pipelineRunsDropdown = document.querySelectorAll('div.tree-item-children')[0]; + let isContextPipelineRunsDisplayed = + localStorage.getItem('displayContextPipelineRuns') === 'true'; + if (isContextPipelineRunsDisplayed === true && isContextDisplayed === true) { + pipelineRunsDropdown.classList.add('open'); + } else { + pipelineRunsDropdown.classList.remove('open'); + } + + const treeItemsWithChildren = Array.from( + document.querySelectorAll('div.tree-item-wrapper') + ).filter( + treeItemsWithChildren => + treeItemsWithChildren.querySelector('div.tree-item-children') !== null + ); + + treeItemsWithChildren.forEach(wrapper => { + wrapper.querySelector('div.tree-item-content').addEventListener('click', function (event) { + const checkboxEl = this.querySelector('input[type="checkbox"]'); + const childrenEl = this.parentNode.querySelector('div.tree-item-children'); + const chevronEl = this.querySelector('span.tree-item-icon'); + const prevPageButton = document.getElementById('prevPage'); + const nextPageButton = document.getElementById('nextPage'); + + if ( + event.target !== checkboxEl && + event.target !== nextPageButton && + event.target !== prevPageButton + ) { + event.stopPropagation(); + childrenEl.classList.toggle('open'); + isContextPipelineRunsDisplayed = !isContextPipelineRunsDisplayed; + localStorage.setItem( + 'displayContextPipelineRuns', + String(isContextPipelineRunsDisplayed) + ); + + if (childrenEl.classList.contains('open')) { + chevronEl.innerHTML = + ''; + } else { + chevronEl.innerHTML = + ''; + } + } + }); + }); + + const textarea = document.getElementById('messageInput'); + const loader = document.getElementById('loader'); + isInputDisabled = false; + + function showLoader() { + loader.classList.add('loader'); + } + + textarea.addEventListener('keydown', e => { + if (e.key === 'Enter') { + if (e.shiftKey) { + // Insert a new line when Shift+Enter is pressed + e.preventDefault(); + textarea.value += '\n'; + textarea.scrollTop = textarea.scrollHeight; + } else { + e.preventDefault(); + if (!isInputDisabled) { + const form = document.getElementById('chatForm'); + const event = new CustomEvent('submit', { + bubbles: true, + cancelable: true, + }); + form.dispatchEvent(event); + showLoader(); + } + } + } + }); + + // Restore dropdown state + restoreState(); + }); +})(); diff --git a/resources/chat-view/styles.css b/resources/chat-view/styles.css new file mode 100644 index 00000000..70f467dc --- /dev/null +++ b/resources/chat-view/styles.css @@ -0,0 +1,237 @@ +/* resources/chat-view/styles.css */ +@tailwind base; +@tailwind components; +@tailwind utilities; + +body { + background-color: var(--vscode-editor-background); + color: var(--vscode-editor-foreground); +} + +ul { + color: var(--vscode-editor-foreground); +} + +form { + display: flex; +} + +#chatForm { + background: var(--vscode-input-background); + border: 1px solid var(--vscode-input-border); + border-radius: 0.375rem 0.375rem 0 0; +} + +#chatMessages { + color: var(--vscode-editor-foreground); +} + +#messageInput { + color: var(--vscode-input-foreground); + background-color: var(--vscode-input-background); +} + +.rounded-lg { + background-color: var(--vscode-editor-background); + color: var(--vscode-editor-foreground); +} + +.rounded-lg:hover { + background-color: var(--vscode-button-hoverBackground); +} + +.hide { + display: none; +} + +.tree-view { + background-color: var(--vscode-editor-background); + color: var(--vscode-editor-foreground); + width: 256px; + max-height: 300px; + overflow: auto; + border: 2px var(--vscode-editor-foreground) solid; + border-radius: 0.375rem; + box-shadow: + 0 4px 6px -1px rgba(0, 0, 0, 0.1), + 0 2px 4px -1px rgba(0, 0, 0, 0.06); +} + +.tree-item { + display: flex; + align-items: center; + cursor: pointer; +} + +.tree-item:hover { + background-color: var(--vscode-list-background); +} + +.tree-item-content { + display: flex; + align-items: center; + flex-grow: 1; + padding: 2px 8px; +} + +.tree-item-icon { + width: 16px; + height: 16px; + display: inline-flex; + align-items: center; + justify-content: center; +} + +.tree-item-name { + font-size: 14px; + color: var(--vscode-editor-foreground); + margin-left: 4px; + flex-grow: 1; +} + +.tree-item-children { + display: none; + width: 100%; +} + +.tree-item-children.open { + display: block; +} + +.tree-item-wrapper { + display: flex; + flex-direction: column; + width: 100%; +} + +.tree-item-checkbox { + width: 16px; + height: 16px; + margin-left: 8px; +} + +.chatbar-options { + display: flex; + justify-content: flex-start; + align-items: center; + padding: 0.5rem; + background-color: var(--vscode-input-background); + border: 1px solid var(--vscode-input-border); + border-top: none; + border-radius: 0 0 0.375rem 0.375rem; +} + +.model-dropdown { + font-size: 0.75rem; + padding: 0.25rem; + background-color: var(--vscode-dropdown-background); + color: var(--vscode-dropdown-foreground); + border: 1px solid var(--vscode-dropdown-border); + border-radius: 0.25rem; + margin-right: 0.5rem; +} + +.context-button { + font-size: 0.75rem; + padding: 0.25rem; + background-color: var(--vscode-dropdown-background); + color: var(--vscode-dropdown-foreground); + border: none; + border-radius: 0.25rem; + cursor: pointer; + margin-right: 0.5rem; +} + +.context-button:hover { + background-color: var(--vscode-button-hoverBackground); +} + +.context-dropdown { + position: absolute; + left: 0; + right: 0; + bottom: 100%; + margin-bottom: 0.25rem; + z-index: 10; +} + +#chatMessages h1, +#chatMessages h2, +#chatMessages h3, +#chatMessages h4, +#chatMessages h5, +#chatMessages h6 { + font-weight: bold; +} + +#chatMessages pre { + padding: 1em; + border-radius: 5px; +} + +#chatMessages code { + padding: 0.2em 0.4em; + border-radius: 3px; +} + +#chatMessages pre code { + display: block; + background-color: var(--vscode-textCodeBlock-background); + color: var(--vscode-editor-foreground); + padding: 1em; + border-radius: 5px; +} + +#chatMessages a { + color: var(--vscode-textLink-foreground); + text-decoration: none; +} + +#chatMessages strong { + font-weight: bold; +} + +#chatMessages em { + font-style: italic; +} + +.copy-button { + position: absolute; + bottom: 8px; + right: 8px; + background-color: var(--vscode-button-background); + color: var(--vscode-button-foreground); + border: none; + border-radius: 4px; + padding: 4px 8px; + font-size: 12px; + cursor: pointer; +} + +.copy-button:hover { + background-color: var(--vscode-button-hoverBackground); +} + +.assistant { + background-color: var(--vscode-input-background); + position: relative; +} + +.expand span { + color: var(--vscode-textLink-foreground); +} + +.loader { + width: 50px; + aspect-ratio: 1; + border-radius: 50%; + border: 8px solid; + border-color: var(--vscode-input-background) var(--vscode-input-foreground); + animation: l1 1s infinite; +} + +@keyframes l1 { + to { + transform: rotate(0.5turn); + } +} diff --git a/resources/chat.png b/resources/chat.png new file mode 100644 index 00000000..9d8e7fc1 Binary files /dev/null and b/resources/chat.png differ diff --git a/src/commands/chat/cmds.ts b/src/commands/chat/cmds.ts new file mode 100644 index 00000000..45bee052 --- /dev/null +++ b/src/commands/chat/cmds.ts @@ -0,0 +1,47 @@ +// Copyright(c) ZenML GmbH 2024. All Rights Reserved. +// Licensed under the Apache License, Version 2.0(the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at: +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +// or implied.See the License for the specific language governing +// permissions and limitations under the License. +// ... existing code ... +import * as vscode from 'vscode'; +import { ChatDataProvider } from '../../views/chatView/ChatDataProvider'; + +const openChat = (context: vscode.ExtensionContext) => { + const panel = vscode.window.createWebviewPanel('zenmlChat', 'ZenML Chat', vscode.ViewColumn.One, { + enableScripts: true, + retainContextWhenHidden: true, + }); + + const chatDataProvider = new ChatDataProvider(context); + + // Create a WebviewViewResolveContext + const webviewViewResolveContext: vscode.WebviewViewResolveContextUsing the chat feature will incur charges from your model provider as normal.
+ + + + `; + } + + dispose(): void { + this._disposables.forEach(disposable => disposable.dispose()); + } +} + +function getNonce() { + let text = ''; + const possible = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789'; + for (let i = 0; i < 32; i++) { + text += possible.charAt(Math.floor(Math.random() * possible.length)); + } + return text; +} diff --git a/src/views/activityBar/common/PaginatedDataProvider.ts b/src/views/activityBar/common/PaginatedDataProvider.ts index 00732048..041e1b71 100644 --- a/src/views/activityBar/common/PaginatedDataProvider.ts +++ b/src/views/activityBar/common/PaginatedDataProvider.ts @@ -25,7 +25,7 @@ export class PaginatedDataProvider implements TreeDataProvider${formattedCode}`
+ : '' + formattedCode + '${roleName}
+ ${content} +ZenML Assistant
+ ${marked.parse(streamingMessage.content)} + for inline code and for bolding.
+ `;
+ if (!tokenjs) {
+ throw new Error('TokenJS not initialized');
+ }
+
+ console.log(`getChatResponse called with provider: ${provider}, model: ${model}`);
+
+ const fullMessages = [
+ { role: 'system', content: template },
+ { role: 'system', content: await addContext(context) },
+ ...messages,
+ ];
+
+ const stream = await tokenjs.chat.completions.create({
+ stream: true,
+ provider: provider.toLowerCase(),
+ model: model,
+ messages: fullMessages.map(msg => ({
+ role: msg.role as 'system' | 'user' | 'assistant',
+ content: msg.content,
+ })),
+ });
+
+ let buffer = '';
+ const delay = (ms: number) => new Promise(resolve => setTimeout(resolve, ms));
+
+ try {
+ for await (const part of stream) {
+ if (part.choices[0]?.delta?.content) {
+ buffer += part.choices[0].delta.content;
+ }
+
+ // Stream from the buffer continuously
+ while (buffer.length > 0) {
+ yield buffer[0];
+ buffer = buffer.slice(1);
+ await delay(5);
+ }
+ }
+
+ // Continue emptying the buffer after the stream ends
+ while (buffer.length > 0) {
+ yield buffer[0];
+ buffer = buffer.slice(1);
+ await delay(5);
+ }
+ } catch (streamError) {
+ console.error('Streaming error in getChatResponse:', streamError);
+ throw new Error(
+ `Streaming error with ${provider} API: ${streamError instanceof Error ? streamError.message : String(streamError)}`
+ );
+ }
+}
diff --git a/tailwind.config.js b/tailwind.config.js
new file mode 100644
index 00000000..d9d0be6d
--- /dev/null
+++ b/tailwind.config.js
@@ -0,0 +1,8 @@
+/** @type {import('tailwindcss').Config} */
+module.exports = {
+ content: ['./resources/chat-view/**/*.{html,js}'],
+ theme: {
+ extend: {},
+ },
+ plugins: [],
+};