From d0e47aa18f39c1cd5b528f4c9447a0136df10514 Mon Sep 17 00:00:00 2001 From: rjmacarthy Date: Tue, 6 Feb 2024 20:44:11 +0000 Subject: [PATCH 1/2] add custom template feature initial commit --- package.json | 28 ++++++++++-- src/constants.ts | 26 ++++++++++- src/extension.ts | 74 +++++++++++++++++-------------- src/providers/sidebar.ts | 24 +++++++--- src/template-provider.ts | 12 ++++- src/webview/chat.tsx | 2 +- src/webview/code-block.tsx | 4 +- src/webview/hooks.ts | 38 +++++++++++++--- src/webview/index.module.css | 6 +++ src/webview/index.tsx | 8 +--- src/webview/main.tsx | 39 ++++++++++++++++ src/webview/suggestions.tsx | 17 ++++--- src/webview/template-settings.tsx | 71 +++++++++++++++++++++++++++++ src/webview/utils.ts | 16 ++++++- 14 files changed, 298 insertions(+), 67 deletions(-) create mode 100644 src/webview/main.tsx create mode 100644 src/webview/template-settings.tsx diff --git a/package.json b/package.json index dcb62113..7145bb5a 100644 --- a/package.json +++ b/package.json @@ -81,20 +81,30 @@ } ], "view/title": [ + { + "command": "twinny.openChat", + "group": "navigation@0", + "when": "view == twinny.sidebar && twinnyManageTemplates" + }, + { + "command": "twinny.manageTemplates", + "group": "navigation@1", + "when": "view == twinny.sidebar" + }, { "command": "twinny.templates", "when": "view == twinny.sidebar", - "group": "navigation@0" + "group": "navigation@2" }, { "command": "twinny.newChat", "when": "view == twinny.sidebar", - "group": "navigation@1" + "group": "navigation@3" }, { "command": "twinny.settings", "when": "view == twinny.sidebar", - "group": "navigation@2" + "group": "navigation@4" } ] }, @@ -154,7 +164,19 @@ "command": "twinny.templates", "shortTitle": "Edit twinny templates", "title": "Edit twinny templates", + "icon": "$(pencil)" + }, + { + "command": "twinny.manageTemplates", + "shortTitle": "Manage twinny templates", + "title": "Manage twinny templates", "icon": "$(files)" + }, + { + "command": "twinny.openChat", + "shortTitle": "Back to chat view", + "title": "Back to chat view", + "icon": "$(arrow-left)" } ], "keybindings": [ diff --git a/src/constants.ts b/src/constants.ts index 18453d79..fbef5531 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -27,7 +27,10 @@ export const MESSAGE_NAME = { twinnyClickSuggestion: 'twinny-click-suggestion', twinnyEnableModelDownload: 'twinny-enable-model-download', twinnySendLanguage: 'twinny-send-language', - twinnySendTheme : 'twinny-send-theme', + twinnySendTheme: 'twinny-send-theme', + twinnyListTemplates: 'twinny-list-templates', + twinnyManageTemplates: 'twinny-manage-templates', + twinnySetTab: 'twinny-set-tab' } export const MESSAGE_KEY = { @@ -36,10 +39,17 @@ export const MESSAGE_KEY = { selection: 'selection', chatMessage: 'chatMessage', autoScroll: 'autoScroll', + selectedTemplates: 'selectedTemplates', } export const CONTEXT_NAME = { twinnyGeneratingText: 'twinnyGeneratingText', + twinnyManageTemplates: 'twinnyManageTemplates' +} + +export const TABS = { + chat: 'chat', + templates: 'templates' } export const fimTempateFormats = { @@ -60,4 +70,16 @@ export const allBrackets = [...openingBrackets, ...closingBrackets] as const export const BRACKET_REGEX = /^[()[\]{}]+$/ export const NORMALIZE_REGEX = /\r?\n|\r/g -export const codeActionTypes = ['add-types', 'refactor', 'generate-docs', 'fix-code'] +export const CODE_ACTION_TYPES = [ + 'add-types', + 'refactor', + 'generate-docs', + 'fix-code' +] + +export const DEFAULT_TEMPLATES = [ + 'refactor', + 'add-tests', + 'add-types', + 'explain' +] diff --git a/src/extension.ts b/src/extension.ts index cf3fb70c..1b651c3c 100644 --- a/src/extension.ts +++ b/src/extension.ts @@ -4,8 +4,7 @@ import { languages, StatusBarAlignment, window, - workspace, - Uri, + workspace } from 'vscode' import * as path from 'path' import * as os from 'os' @@ -16,8 +15,15 @@ import { init } from './init' import { SidebarProvider } from './providers/sidebar' import { delayExecution, deleteTempFiles } from './utils' import { setContext } from './context' -import { EXTENSION_NAME, MESSAGE_KEY } from './constants' +import { + CONTEXT_NAME, + EXTENSION_NAME, + MESSAGE_KEY, + MESSAGE_NAME, + TABS +} from './constants' import { TemplateProvider } from './template-provider' +import { ServerMessage } from './types' export async function activate(context: ExtensionContext) { const config = workspace.getConfiguration('twinny') @@ -25,8 +31,8 @@ export async function activate(context: ExtensionContext) { const chatModel = config.get('chatModelName') as string const statusBar = window.createStatusBarItem(StatusBarAlignment.Right) const templateDir = - config.get('templateDir') as string || - path.join(os.homedir(), '.twinny/templates') as string + (config.get('templateDir') as string) || + (path.join(os.homedir(), '.twinny/templates') as string) setContext(context) try { @@ -58,48 +64,48 @@ export async function activate(context: ExtensionContext) { commands.registerCommand('twinny.disable', () => { statusBar.hide() }), - commands.registerCommand('twinny.explain', () => { - commands.executeCommand('twinny.sidebar.focus') - delayExecution(() => - sidebarProvider.chatService?.streamTemplateCompletion('explain') - ) - }), - commands.registerCommand('twinny.fixCode', () => { + commands.registerCommand('twinny.templateCompletion', (template: string) => { commands.executeCommand('twinny.sidebar.focus') delayExecution(() => - sidebarProvider.chatService?.streamTemplateCompletion('fix-code') + sidebarProvider.chatService?.streamTemplateCompletion(template) ) }), commands.registerCommand('twinny.stopGeneration', () => { completionProvider.destroyStream() sidebarProvider.destroyStream() }), - commands.registerCommand('twinny.addTypes', () => { - commands.executeCommand('twinny.sidebar.focus') - delayExecution(() => - sidebarProvider.chatService?.streamTemplateCompletion('add-types') - ) - }), - commands.registerCommand('twinny.refactor', () => { - commands.executeCommand('twinny.sidebar.focus') - delayExecution(() => - sidebarProvider.chatService?.streamTemplateCompletion('refactor') + commands.registerCommand('twinny.templates', async () => { + await vscode.commands.executeCommand( + 'vscode.openFolder', + vscode.Uri.parse(templateDir), + true ) }), - commands.registerCommand('twinny.addTests', () => { - commands.executeCommand('twinny.sidebar.focus') - delayExecution(() => - sidebarProvider.chatService?.streamTemplateCompletion('add-tests') + commands.registerCommand('twinny.manageTemplates', async () => { + commands.executeCommand( + 'setContext', + CONTEXT_NAME.twinnyManageTemplates, + true ) + sidebarProvider.view?.webview.postMessage({ + type: MESSAGE_NAME.twinnySetTab, + value: { + data: TABS.templates + } + } as ServerMessage) }), - commands.registerCommand('twinny.generateDocs', () => { - commands.executeCommand('twinny.sidebar.focus') - delayExecution(() => - sidebarProvider.chatService?.streamTemplateCompletion('generate-docs') + commands.registerCommand('twinny.openChat', () => { + commands.executeCommand( + 'setContext', + CONTEXT_NAME.twinnyManageTemplates, + false ) - }), - commands.registerCommand('twinny.templates', async () => { - await vscode.commands.executeCommand('vscode.openFolder', Uri.parse(templateDir), true); + sidebarProvider.view?.webview.postMessage({ + type: MESSAGE_NAME.twinnySetTab, + value: { + data: TABS.chat + } + } as ServerMessage) }), commands.registerCommand('twinny.settings', () => { vscode.commands.executeCommand( diff --git a/src/providers/sidebar.ts b/src/providers/sidebar.ts index 471b590c..c685b639 100644 --- a/src/providers/sidebar.ts +++ b/src/providers/sidebar.ts @@ -1,8 +1,9 @@ import * as vscode from 'vscode' import { getLanguage, getTextSelection, getTheme, openDiffView } from '../utils' -import { MESSAGE_KEY, MESSAGE_NAME } from '../constants' +import { DEFAULT_TEMPLATES, MESSAGE_KEY, MESSAGE_NAME } from '../constants' import { ChatService } from '../chat-service' import { ClientMessage, MessageType, ServerMessage } from '../types' +import { TemplateProvider } from '../template-provider' export class SidebarProvider implements vscode.WebviewViewProvider { view?: vscode.WebviewView @@ -11,6 +12,7 @@ export class SidebarProvider implements vscode.WebviewViewProvider { private _statusBar: vscode.StatusBarItem private context: vscode.ExtensionContext private _templateDir: string + private _templateProvider: TemplateProvider constructor( statusBar: vscode.StatusBarItem, @@ -20,6 +22,7 @@ export class SidebarProvider implements vscode.WebviewViewProvider { this._statusBar = statusBar this.context = context this._templateDir = templateDir + this._templateProvider = new TemplateProvider(templateDir) } public resolveWebviewView(webviewView: vscode.WebviewView) { @@ -61,9 +64,9 @@ export class SidebarProvider implements vscode.WebviewViewProvider { webviewView.webview.onDidReceiveMessage( ( - message: ClientMessage + message: ClientMessage & ClientMessage ) => { - const eventHandlers: { [k: string]: (arg0: any) => void } = { + const eventHandlers = { [MESSAGE_NAME.twinnyChatMessage]: this.streamChatCompletion, [MESSAGE_NAME.twinnyOpenDiff]: this.openDiff, [MESSAGE_NAME.twinnyClickSuggestion]: this.clickSuggestion, @@ -76,19 +79,30 @@ export class SidebarProvider implements vscode.WebviewViewProvider { this.setTwinnyWorkspaceContext, [MESSAGE_NAME.twinnySendLanguage]: this.getCurrentLanguage, [MESSAGE_NAME.twinnySendTheme]: this.getTheme, - [MESSAGE_NAME.twinnyNotification]: this.sendNotification + [MESSAGE_NAME.twinnyNotification]: this.sendNotification, + [MESSAGE_NAME.twinnyListTemplates]: this.listTemplates } eventHandlers[message.type as string](message) } ) } + public listTemplates = () => { + const templates = this._templateProvider.listTemplates() + this.view?.webview.postMessage({ + type: MESSAGE_NAME.twinnyListTemplates, + value: { + data: templates + } + } as ServerMessage) + } + public sendNotification = (data: ClientMessage) => { vscode.window.showInformationMessage(data.data as string) } public clickSuggestion = (data: ClientMessage) => { - vscode.commands.executeCommand(data.data as string) + vscode.commands.executeCommand('twinny.templateCompletion', data.data as string) } public streamChatCompletion = (data: ClientMessage) => { diff --git a/src/template-provider.ts b/src/template-provider.ts index 957dc427..85a4688f 100644 --- a/src/template-provider.ts +++ b/src/template-provider.ts @@ -34,7 +34,7 @@ export class TemplateProvider { const destPath = path.join(this._basePath) try { fs.mkdirSync(destPath, { recursive: true }) - defaultTemplates.forEach(({name, template}) => { + defaultTemplates.forEach(({ name, template }) => { const destFile = path.join(destPath, name) fs.writeFileSync(`${destFile}.hbs`, template, 'utf8') }) @@ -80,6 +80,16 @@ export class TemplateProvider { } } + public listTemplates(): string[] { + const files = fs.readdirSync(this._basePath, 'utf8') + const templates = files.filter((fileName) => fileName.endsWith('.hbs')) + const templateNames = templates + .map((fileName) => fileName.replace('.hbs', '')) + .sort((a, b) => a.localeCompare(b)) + .filter(name => name !== 'chat' && name !== 'system') + return templateNames + } + public async renderTemplate( templateName: string, data: T diff --git a/src/webview/chat.tsx b/src/webview/chat.tsx index 9710135f..e4a48508 100644 --- a/src/webview/chat.tsx +++ b/src/webview/chat.tsx @@ -22,11 +22,11 @@ import { StopIcon } from './icons' -import styles from './index.module.css' import { Suggestions } from './suggestions' import { ClientMessage, MessageType, ServerMessage } from '../types' import { Message } from './message' import { getCompletionContent } from './utils' +import styles from './index.module.css' // eslint-disable-next-line @typescript-eslint/no-explicit-any const global = globalThis as any diff --git a/src/webview/code-block.tsx b/src/webview/code-block.tsx index 4f9f7717..3133246d 100644 --- a/src/webview/code-block.tsx +++ b/src/webview/code-block.tsx @@ -3,7 +3,7 @@ import { ReactNode } from 'react' import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter' import { vscDarkPlus, vs } from 'react-syntax-highlighter/dist/esm/styles/prism' -import { MESSAGE_NAME, codeActionTypes } from '../constants' +import { MESSAGE_NAME, CODE_ACTION_TYPES } from '../constants' import styles from './index.module.css' import { LanguageType, Theme, ThemeType } from '../types' @@ -45,7 +45,7 @@ export const CodeBlock = (props: CodeBlockProps) => { language={lang} />
- {codeActionTypes.includes(completionType) && ( + {CODE_ACTION_TYPES.includes(completionType) && ( <> Accept diff --git a/src/webview/hooks.ts b/src/webview/hooks.ts index 14061ea2..840d3e18 100644 --- a/src/webview/hooks.ts +++ b/src/webview/hooks.ts @@ -1,7 +1,7 @@ import { useEffect, useState } from 'react' -import { MESSAGE_NAME } from '../constants' -import { LanguageType, ServerMessage, ThemeType } from '../types' +import { MESSAGE_KEY, MESSAGE_NAME } from '../constants' +import { ClientMessage, LanguageType, ServerMessage, ThemeType } from '../types' // eslint-disable-next-line @typescript-eslint/no-explicit-any const global = globalThis as any @@ -84,15 +84,14 @@ export const useTheme = () => { } useEffect(() => { global.vscode.postMessage({ - type: MESSAGE_NAME.twinnySendTheme, + type: MESSAGE_NAME.twinnySendTheme }) window.addEventListener('message', handler) }, []) return theme } - -export const useLanguage = (): LanguageType | undefined => { +export const useLanguage = (): LanguageType | undefined => { const [language, setLanguage] = useState() const handler = (event: MessageEvent) => { const message: ServerMessage = event.data @@ -103,9 +102,36 @@ export const useLanguage = (): LanguageType | undefined => { } useEffect(() => { global.vscode.postMessage({ - type: MESSAGE_NAME.twinnySendLanguage, + type: MESSAGE_NAME.twinnySendLanguage }) window.addEventListener('message', handler) }, []) return language } + +export const useTemplates = () => { + const [templates, setTemplates] = useState() + const handler = (event: MessageEvent) => { + const message: ServerMessage = event.data + if (message?.type === MESSAGE_NAME.twinnyListTemplates) { + setTemplates(message?.value.data) + } + return () => window.removeEventListener('message', handler) + } + + const saveTemplates = (templates: string[]) => { + global.vscode.postMessage({ + type: MESSAGE_NAME.twinnySetWorkspaceContext, + key: MESSAGE_KEY.selectedTemplates, + data: templates + } as ClientMessage) + } + + useEffect(() => { + global.vscode.postMessage({ + type: MESSAGE_NAME.twinnyListTemplates + }) + window.addEventListener('message', handler) + }, []) + return { templates, saveTemplates } +} diff --git a/src/webview/index.module.css b/src/webview/index.module.css index 1d686a8c..ede3547a 100644 --- a/src/webview/index.module.css +++ b/src/webview/index.module.css @@ -109,3 +109,9 @@ code { background: var(--vscode-inputOption-hoverBackground); cursor: pointer; } + +.templateCheckbox label { + display: flex; + align-items: center; + gap: 5px; +} diff --git a/src/webview/index.tsx b/src/webview/index.tsx index 9e43cc32..b2356a0c 100644 --- a/src/webview/index.tsx +++ b/src/webview/index.tsx @@ -1,6 +1,6 @@ import { createRoot } from 'react-dom/client' -import { Chat } from './chat' +import { Main } from './main' // eslint-disable-next-line @typescript-eslint/no-explicit-any (globalThis as any).vscode = window.acquireVsCodeApi() @@ -9,9 +9,5 @@ const container = document.querySelector('#root') if (container) { const root = createRoot(container) - root.render( - <> - - - ) + root.render(
) } diff --git a/src/webview/main.tsx b/src/webview/main.tsx new file mode 100644 index 00000000..1a733779 --- /dev/null +++ b/src/webview/main.tsx @@ -0,0 +1,39 @@ +import { useEffect, useState } from 'react' +import { Chat } from './chat' +import { TemplateSettings } from './template-settings' +import { ServerMessage } from '../types' +import { MESSAGE_NAME } from '../constants' + +interface TabComponents { + [key: string]: { component: JSX.Element } +} + +const tabs: TabComponents = { + chat: { + component: + }, + templates: { + component: + } +} + +export const Main = () => { + const [tab, setTab] = useState('chat') + + const handler = (event: MessageEvent) => { + const message: ServerMessage = event.data + if (message?.type === MESSAGE_NAME.twinnySetTab) { + setTab(message?.value.data) + } + return () => window.removeEventListener('message', handler) + } + useEffect(() => { + window.addEventListener('message', handler) + }, []) + + if (!tab) { + return null + } + + return tabs[tab].component || null +} diff --git a/src/webview/suggestions.tsx b/src/webview/suggestions.tsx index 3d636fdc..075f88d2 100644 --- a/src/webview/suggestions.tsx +++ b/src/webview/suggestions.tsx @@ -1,8 +1,10 @@ import { CodeIcon, ExplainIcon, FixCodeIcon, TestsIcon } from './icons' -import { MESSAGE_NAME } from '../constants' +import { MESSAGE_KEY, MESSAGE_NAME } from '../constants' import cn from 'classnames' import styles from './index.module.css' +import { useWorkSpaceContext } from './hooks' +import { kebabToSentence } from './utils' const SUGGESTIONS = [ { @@ -34,6 +36,8 @@ const SUGGESTIONS = [ // eslint-disable-next-line @typescript-eslint/no-explicit-any const global = globalThis as any export const Suggestions = ({ isDisabled }: { isDisabled?: boolean }) => { + const templates = useWorkSpaceContext(MESSAGE_KEY.selectedTemplates) + const handleOnClickSuggestion = (message: string) => { if (isDisabled) return @@ -45,14 +49,15 @@ export const Suggestions = ({ isDisabled }: { isDisabled?: boolean }) => { return (
- {SUGGESTIONS.map(({ name, value, icon, message }) => ( + {templates?.map((name) => (
handleOnClickSuggestion(message)} + onClick={() => handleOnClickSuggestion(name)} key={name} - className={cn(styles.suggestion, { [styles['suggestion--disabled']]: isDisabled })} + className={cn(styles.suggestion, { + [styles['suggestion--disabled']]: isDisabled + })} > -
{icon}
-
{value}
+
{kebabToSentence(name)}
))}
diff --git a/src/webview/template-settings.tsx b/src/webview/template-settings.tsx new file mode 100644 index 00000000..3b55c79b --- /dev/null +++ b/src/webview/template-settings.tsx @@ -0,0 +1,71 @@ +import { VSCodeCheckbox } from '@vscode/webview-ui-toolkit/react' +import { useTemplates, useWorkSpaceContext } from './hooks' +import { DEFAULT_TEMPLATES, MESSAGE_KEY } from '../constants' +import { useEffect, useState } from 'react' +import { kebabToSentence } from './utils' + +import styles from './index.module.css' + +export const TemplateSettings = () => { + const { templates, saveTemplates } = useTemplates() + const [selectedTemplates, setSelectedTemplates] = useState([]) + const selectedTemplatesContext = + useWorkSpaceContext(MESSAGE_KEY.selectedTemplates) || [] + + const handleTemplateClick = ( + e: React.MouseEvent + ) => { + const target = e.target as HTMLInputElement + + const template = target.value + + if (selectedTemplates.includes(template)) { + if (selectedTemplates.length === 1) { + saveTemplates([]) + setSelectedTemplates([]) + return + } + + return setSelectedTemplates((prev) => { + const newValue = prev.filter((item) => item !== template) + saveTemplates(newValue) + return newValue + }) + } + + setSelectedTemplates((prev) => { + const newValue = [...prev, template] + saveTemplates(newValue) + return newValue + }) + } + + useEffect(() => { + if (selectedTemplatesContext !== undefined) { + return setSelectedTemplates(selectedTemplatesContext) + } else { + setSelectedTemplates(DEFAULT_TEMPLATES) + } + }, [selectedTemplatesContext]) + + return ( + <> +

Set prompt templates

+ {templates && + templates.map((templateName: string) => ( +
+ +
+ ))} + + ) +} diff --git a/src/webview/utils.ts b/src/webview/utils.ts index 7a410fab..1e5e20ca 100644 --- a/src/webview/utils.ts +++ b/src/webview/utils.ts @@ -20,7 +20,9 @@ export const getLanguageMatch = ( const languageId = language.languageId.toString() const languageEntry = supportedLanguages[languageId as CodeLanguage] - return languageEntry && languageEntry.derivedFrom ? languageEntry.derivedFrom : languageId + return languageEntry && languageEntry.derivedFrom + ? languageEntry.derivedFrom + : languageId } return 'javascript' @@ -33,3 +35,15 @@ export const getCompletionContent = (message: ServerMessage) => { return message.value.completion || EMPTY_MESAGE } + +export const kebabToSentence = (kebabStr: string) => { + if (!kebabStr) { + return '' + } + + const words = kebabStr.split('-') + + words[0] = words[0].charAt(0).toUpperCase() + words[0].slice(1) + + return words.join(' ') +} From cddd8b06fb0f0db8ff74d9cf8c9f11aa120cf06f Mon Sep 17 00:00:00 2001 From: rjmacarthy Date: Tue, 6 Feb 2024 20:47:27 +0000 Subject: [PATCH 2/2] catch non array --- src/webview/utils.ts | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/webview/utils.ts b/src/webview/utils.ts index 1e5e20ca..d49ef1be 100644 --- a/src/webview/utils.ts +++ b/src/webview/utils.ts @@ -43,6 +43,10 @@ export const kebabToSentence = (kebabStr: string) => { const words = kebabStr.split('-') + if (!words.length) { + return kebabStr + } + words[0] = words[0].charAt(0).toUpperCase() + words[0].slice(1) return words.join(' ')