Merge pull request #95 from rjmacarthy/feature/custom-templates

Add custom templates

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": [

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'
+]

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,17 +15,24 @@ 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')
   const fimModel = config.get('fimModelName') as string
   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<string>)
     }),
-    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<string>)
     }),
     commands.registerCommand('twinny.settings', () => {
       vscode.commands.executeCommand(

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<string | boolean | MessageType[]>
+        message: ClientMessage<string | boolean> & ClientMessage<MessageType[]>
       ) => {
-        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<string[]>)
+  }
+
   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<MessageType[]>) => {

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<T extends DefaultTemplate>(
     templateName: string,
     data: T

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

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}
       />
       <div className={styles.codeOptions}>
-        {codeActionTypes.includes(completionType) && (
+        {CODE_ACTION_TYPES.includes(completionType) && (
           <>
             <VSCodeButton onClick={handleAccept}>Accept</VSCodeButton>
           </>