├── .eslintrc.json ├── .gitignore ├── .vscode ├── extensions.json ├── launch.json ├── settings.json └── tasks.json ├── .vscodeignore ├── CHANGELOG.md ├── README.md ├── images └── worldspider.png ├── language-configuration.json ├── package-lock.json ├── package.json ├── src ├── callModel.ts ├── document.ts ├── extension.ts ├── logging.ts ├── test │ ├── runTest.ts │ └── suite │ │ ├── extension.test.ts │ │ └── index.ts └── utils.ts ├── syntaxes └── worldspider.tmLanguage.json ├── tsconfig.json └── vsc-extension-quickstart.md /.eslintrc.json: -------------------------------------------------------------------------------- 1 | { 2 | "root": true, 3 | "parser": "@typescript-eslint/parser", 4 | "parserOptions": { 5 | "ecmaVersion": 6, 6 | "sourceType": "module" 7 | }, 8 | "plugins": [ 9 | "@typescript-eslint" 10 | ], 11 | "rules": { 12 | "@typescript-eslint/naming-convention": "warn", 13 | "@typescript-eslint/semi": "warn", 14 | "curly": "warn", 15 | "eqeqeq": "warn", 16 | "no-throw-literal": "warn", 17 | "semi": "off" 18 | }, 19 | "ignorePatterns": [ 20 | "out", 21 | "dist", 22 | "**/*.d.ts" 23 | ] 24 | } 25 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | out 2 | dist 3 | node_modules 4 | .vscode-test/ 5 | *.vsix 6 | -------------------------------------------------------------------------------- /.vscode/extensions.json: -------------------------------------------------------------------------------- 1 | { 2 | // See http://go.microsoft.com/fwlink/?LinkId=827846 3 | // for the documentation about the extensions.json format 4 | "recommendations": [ 5 | "dbaeumer.vscode-eslint" 6 | ] 7 | } 8 | -------------------------------------------------------------------------------- /.vscode/launch.json: -------------------------------------------------------------------------------- 1 | // A launch configuration that compiles the extension and then opens it inside a new window 2 | // Use IntelliSense to learn about possible attributes. 3 | // Hover to view descriptions of existing attributes. 4 | // For more information, visit: https://go.microsoft.com/fwlink/?linkid=830387 5 | { 6 | "version": "0.2.0", 7 | "configurations": [ 8 | { 9 | "name": "Run Extension", 10 | "type": "extensionHost", 11 | "request": "launch", 12 | "args": [ 13 | "--extensionDevelopmentPath=${workspaceFolder}" 14 | ], 15 | "outFiles": [ 16 | "${workspaceFolder}/out/**/*.js" 17 | ], 18 | "preLaunchTask": "${defaultBuildTask}" 19 | }, 20 | { 21 | "name": "Extension Tests", 22 | "type": "extensionHost", 23 | "request": "launch", 24 | "args": [ 25 | "--extensionDevelopmentPath=${workspaceFolder}", 26 | "--extensionTestsPath=${workspaceFolder}/out/test/suite/index" 27 | ], 28 | "outFiles": [ 29 | "${workspaceFolder}/out/test/**/*.js" 30 | ], 31 | "preLaunchTask": "${defaultBuildTask}" 32 | } 33 | ] 34 | } 35 | -------------------------------------------------------------------------------- /.vscode/settings.json: -------------------------------------------------------------------------------- 1 | // Place your settings in this file to overwrite default and user settings. 2 | { 3 | "files.exclude": { 4 | "out": false // set this to true to hide the "out" folder with the compiled JS files 5 | }, 6 | "search.exclude": { 7 | "out": true // set this to false to include "out" folder in search results 8 | }, 9 | // Turn off tsc task auto detection since we have the necessary tasks as npm scripts 10 | "typescript.tsc.autoDetect": "off" 11 | } -------------------------------------------------------------------------------- /.vscode/tasks.json: -------------------------------------------------------------------------------- 1 | // See https://go.microsoft.com/fwlink/?LinkId=733558 2 | // for the documentation about the tasks.json format 3 | { 4 | "version": "2.0.0", 5 | "tasks": [ 6 | { 7 | "type": "npm", 8 | "script": "watch", 9 | "problemMatcher": "$tsc-watch", 10 | "isBackground": true, 11 | "presentation": { 12 | "reveal": "never" 13 | }, 14 | "group": { 15 | "kind": "build", 16 | "isDefault": true 17 | } 18 | } 19 | ] 20 | } 21 | -------------------------------------------------------------------------------- /.vscodeignore: -------------------------------------------------------------------------------- 1 | .vscode/** 2 | .vscode-test/** 3 | src/** 4 | .gitignore 5 | .yarnrc 6 | vsc-extension-quickstart.md 7 | **/tsconfig.json 8 | **/.eslintrc.json 9 | **/*.map 10 | **/*.ts 11 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | # Change Log 2 | 3 | All notable changes to the "worldspider" extension will be documented in this file. 4 | 5 | Check [Keep a Changelog](http://keepachangelog.com/) for recommendations on how to structure this file. 6 | 7 | ## [0.0.6] - 2023-03-23 8 | 9 | ### Added 10 | 11 | - Comments 12 | - Alternate base paths for openAI API 13 | 14 | ## [0.0.5] - 2023-03-17 15 | 16 | ### Added 17 | 18 | - Support for chat models like gpt-3.5-turbo and gpt-4 19 | - Generated text is highlighted in the editor (can be disabled with `worldspider.highlightModelText` setting) 20 | - Added customizable model operations 21 | 22 | ## [0.0.4] 23 | 24 | ### Added 25 | 26 | - Token hover shows counterfactual logprobs if available. 27 | 28 | ## [0.0.3] - 2023-01-25 29 | 30 | ### Added 31 | 32 | - Added `worldspider.copyCompletions` command to copy model response to the clipboard. 33 | - Added `worldspider.copyCompletionsText` command to copy model response text to the clipboard. 34 | - Added `worldspider.prevCompletions` command to show the previous batch of completions. 35 | - Added `worldspider.nextCompletions` command to show the next batch of completions. 36 | - Added `worldspider.scorePrompt` command to score prompt logprobs without generating completions. 37 | - Added `worldspider.generation.logprobs` setting to retrieve log probabilities. 38 | - Added `worldspider.generation.echo` setting to control whether prompt logprobs are returned by the model. 39 | - Hovering over text now shows token probabilities if the information is available. 40 | - Model responses are logged to `worldspider` output channel. 41 | 42 | ## [0.0.2] - 2023-01-23 43 | 44 | ### Added 45 | 46 | - Added `worldspider.getInfills` command to generate infills from the selected text. 47 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # worldspider 2 | 3 | ## Requirements 4 | 5 | vscode 1.74.0 or later. 6 | 7 | You will need an OpenAI API key. 8 | 9 | ## Instructions 10 | 11 | ### Installing from vsix 12 | 13 | Install using the command `code --install-extension worldspider-0.0.N.vsix`. 14 | 15 | ### Testing from source 16 | 17 | 1. Clone this repository. 18 | 2. Run `npm install` to install dependencies. 19 | 3. Press `F5` to open a new window with the extension loaded. 20 | 21 | ### Required settings 22 | 23 | After installing, set `worldspider.apiKey` to your OpenAI API key. 24 | 25 | ### Generating text 26 | 27 | To generate continuations from the caret, press `Alt+D`. To generate infills, select the text you want to replace and press `Alt+Shift+D`. 28 | 29 | Use arrow keys to scroll through completions and hit enter to append the selected completion to the prompt, or click a completion to append it. 30 | 31 | If you close the completions dropdown, you can reopen it with `Alt+S`. 32 | 33 | To show the previous batch of completions, press `Alt+Left`. To show the next batch of completions, press `Alt+Right`. 34 | 35 | ### Token logprobs 36 | 37 | To view token logprobs, hover over model-generated text. If `worldspider.generation.echo` is set to `true`, you can also hover over the prompt to view token logprobs. 38 | 39 | You can also highlight text and press `Alt+P` to view token logprobs without generating completions. 40 | 41 | ## Keybindings 42 | 43 | You can change any keybindings by going to `File > Preferences > Keyboard Shortcuts` and searching for `worldspider`. 44 | 45 | ## Customizations 46 | 47 | If you find the completion dropdown too small, you can change the size of the dropdown by setting `editor.suggestLineHeight` to larger number of lines (16 is default). 48 | 49 | If you have Copilot enabled at the same time as Worldspider, you will also get inline Copilot suggestions, which does not interfere with Worldspider's functionality but may be confusing. 50 | -------------------------------------------------------------------------------- /images/worldspider.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/socketteer/worldspider/9d26999740f93f00102ede9376f1195c81101582/images/worldspider.png -------------------------------------------------------------------------------- /language-configuration.json: -------------------------------------------------------------------------------- 1 | {"comments": 2 | { 3 | "blockComment": [ "<|BEGINCOMMENT|>", "<|ENDCOMMENT|>" ] 4 | } 5 | } -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "worldspider", 3 | "displayName": "worldspider", 4 | "description": "loom in vscode", 5 | "icon": "images/worldspider.png", 6 | "version": "0.0.6", 7 | "engines": { 8 | "vscode": "^1.74.0" 9 | }, 10 | "publisher": "janus", 11 | "categories": [ 12 | "Other" 13 | ], 14 | "activationEvents": [ 15 | "*" 16 | ], 17 | "main": "./out/extension.js", 18 | "contributes": { 19 | "commands": [ 20 | { 21 | "command": "worldspider.showCompletions", 22 | "title": "Show completions" 23 | }, 24 | { 25 | "command": "worldspider.getContinuations", 26 | "title": "Call model" 27 | }, 28 | { 29 | "command": "worldspider.getInfills", 30 | "title": "Call model with infill" 31 | }, 32 | { 33 | "command": "worldspider.scorePrompt", 34 | "title": "Get prompt logprobs without generating completions" 35 | }, 36 | { 37 | "command": "worldspider.copyCompletions", 38 | "title": "Copy completions to clipboard" 39 | }, 40 | { 41 | "command": "worldspider.copyCompletionsText", 42 | "title": "Copy completions text to clipboard" 43 | }, 44 | { 45 | "command": "worldspider.prevCompletions", 46 | "title": "Show previous completions" 47 | }, 48 | { 49 | "command": "worldspider.nextCompletions", 50 | "title": "Show next completions" 51 | }, 52 | { 53 | "command": "worldspider.customOp", 54 | "title": "Custom operation" 55 | } 56 | ], 57 | "keybindings": [ 58 | { 59 | "command": "worldspider.showCompletions", 60 | "key": "alt+s", 61 | "mac": "alt+s" 62 | }, 63 | { 64 | "command": "worldspider.getContinuations", 65 | "key": "alt+d", 66 | "mac": "alt+d" 67 | }, 68 | { 69 | "command": "worldspider.scorePrompt", 70 | "key": "alt+p", 71 | "mac": "alt+p" 72 | }, 73 | { 74 | "command": "worldspider.getInfills", 75 | "key": "alt+shift+d", 76 | "mac": "alt+shift+d" 77 | }, 78 | { 79 | "command": "worldspider.copyCompletions", 80 | "key": "alt+c", 81 | "mac": "alt+c" 82 | }, 83 | { 84 | "command": "worldspider.copyCompletionsText", 85 | "key": "alt+shift+c", 86 | "mac": "alt+shift+c" 87 | }, 88 | { 89 | "command": "worldspider.prevCompletions", 90 | "key": "alt+left", 91 | "mac": "alt+left" 92 | }, 93 | { 94 | "command": "worldspider.nextCompletions", 95 | "key": "alt+right", 96 | "mac": "alt+right" 97 | }, 98 | { 99 | "command": "worldspider.customOp", 100 | "key": "alt+shift+o", 101 | "mac": "alt+shift+o" 102 | } 103 | ], 104 | "configuration": { 105 | "title": "worldspider", 106 | "properties": { 107 | "worldspider.apiKey": { 108 | "type": "string", 109 | "default": "", 110 | "description": "OpenAI API key" 111 | }, 112 | "worldspider.basePath": { 113 | "type": "string", 114 | "default": "https://api.openai.com/v1", 115 | "description": "API base path" 116 | }, 117 | "worldspider.generation.model": { 118 | "type": "string", 119 | "default": "code-davinci-002", 120 | "description": "OpenAI model" 121 | }, 122 | "worldspider.generation.maxTokens": { 123 | "type": "number", 124 | "default": 64, 125 | "description": "Maximum number of tokens to generate" 126 | }, 127 | "worldspider.generation.numCompletions": { 128 | "type": "number", 129 | "default": 5, 130 | "description": "Number of completions to generate" 131 | }, 132 | "worldspider.generation.temperature": { 133 | "type": "number", 134 | "default": 1, 135 | "description": "Temperature" 136 | }, 137 | "worldspider.generation.prefixLength": { 138 | "type": "number", 139 | "default": 10000, 140 | "description": "Maximum length of prompt prefix" 141 | }, 142 | "worldspider.generation.suffixLength": { 143 | "type": "number", 144 | "default": 5000, 145 | "description": "Maximum length of prompt suffix" 146 | }, 147 | "worldspider.generation.topP": { 148 | "type": "number", 149 | "default": 1, 150 | "description": "Top p" 151 | }, 152 | "worldspider.generation.frequencyPenalty": { 153 | "type": "number", 154 | "default": 0, 155 | "description": "Frequency penalty" 156 | }, 157 | "worldspider.generation.presencePenalty": { 158 | "type": "number", 159 | "default": 0, 160 | "description": "Presence penalty" 161 | }, 162 | "worldspider.generation.stopSequences": { 163 | "type": "array", 164 | "default": [], 165 | "description": "Stop sequences" 166 | }, 167 | "worldspider.generation.logprobs": { 168 | "type": "number", 169 | "default": 0, 170 | "description": "Number of counterfactual top logprobs per token to return" 171 | }, 172 | "worldspider.generation.echo": { 173 | "type": "boolean", 174 | "default": false, 175 | "description": "Model returns prompt logprobs" 176 | }, 177 | "worldspider.filterDuplicates": { 178 | "type": "boolean", 179 | "default": true, 180 | "description": "Don't show duplicate completions" 181 | }, 182 | "worldspider.log": { 183 | "type": "boolean", 184 | "default": false, 185 | "description": "Log model responses" 186 | }, 187 | "worldspider.savePath": { 188 | "type": "string", 189 | "default": "history.json", 190 | "description": "Path to save model response history" 191 | }, 192 | "worldspider.highlightModelText": { 193 | "type": "boolean", 194 | "default": true, 195 | "description": "Highlight model-generated text" 196 | }, 197 | "worldspider.generation.customInstruction": { 198 | "type": "string", 199 | "default": "This GPT-4 powered, AGI-complete text mutator rewrites anything the user sends in the style of David Foster Wallace.", 200 | "description": "Instruction for the custom operation" 201 | }, 202 | "worldspider.generation.customOpModel": { 203 | "type": "string", 204 | "default": "gpt-4", 205 | "description": "Model for the custom operation" 206 | } 207 | } 208 | }, 209 | "languages": [{ 210 | "id": "worldspider", 211 | "extensions": [".wspdr"], 212 | "configuration": "./language-configuration.json" 213 | }], 214 | "grammars": [ 215 | { 216 | "language": "worldspider", 217 | "scopeName": "source.worldspider", 218 | "path": "./syntaxes/worldspider.tmLanguage.json" 219 | } 220 | ] 221 | }, 222 | "scripts": { 223 | "vscode:prepublish": "npm run compile", 224 | "compile": "tsc -p ./", 225 | "watch": "tsc -watch -p ./", 226 | "pretest": "npm run compile && npm run lint", 227 | "lint": "eslint src --ext ts", 228 | "test": "node ./out/test/runTest.js" 229 | }, 230 | "devDependencies": { 231 | "@types/glob": "^8.0.0", 232 | "@types/mocha": "^10.0.1", 233 | "@types/node": "16.x", 234 | "@types/vscode": "^1.74.0", 235 | "@typescript-eslint/eslint-plugin": "^5.45.0", 236 | "@typescript-eslint/parser": "^5.45.0", 237 | "@vscode/test-electron": "^2.2.0", 238 | "eslint": "^8.28.0", 239 | "glob": "^8.0.3", 240 | "mocha": "^10.1.0", 241 | "typescript": "^4.9.3" 242 | }, 243 | "dependencies": { 244 | "openai": "^3.1.0" 245 | } 246 | } 247 | -------------------------------------------------------------------------------- /src/callModel.ts: -------------------------------------------------------------------------------- 1 | const { Configuration, OpenAIApi } = require("openai"); 2 | import * as vscode from 'vscode'; 3 | import { saveModelResponse } from './logging'; 4 | 5 | export interface LogprobsObject { 6 | tokens: string[]; 7 | token_logprobs: number[]; 8 | top_logprobs: any[] | null; 9 | text_offset: number[]; 10 | } 11 | 12 | export interface ModelCompletion { 13 | text: string; 14 | index: number; 15 | logprobs: LogprobsObject | null; 16 | finish_reason: string; 17 | } 18 | 19 | export interface ModelResponse { 20 | choices: ModelCompletion[]; 21 | id: string; 22 | model: string; 23 | object: string; 24 | created: number; 25 | prompt: string; 26 | suffix: string | null; 27 | echo: boolean; 28 | usage: { 29 | prompt_tokens: number; 30 | completion_tokens: number; 31 | total_tokens: number; 32 | } 33 | } 34 | 35 | function removePromptsFromCompletions (completions: ModelCompletion[], prompt: string) { 36 | return completions.map((completion) => { 37 | const text = completion.text; 38 | const textWithoutPrompt = text.replace(prompt, ''); 39 | return { 40 | ...completion, 41 | text: textWithoutPrompt, 42 | }; 43 | }); 44 | } 45 | 46 | export async function getModelResponse(prompt: string, suffix: string | null, scorePromptOnly: boolean = false): Promise { 47 | const workbenchConfig = vscode.workspace.getConfiguration('worldspider'); 48 | const model = workbenchConfig.get('generation.model'); 49 | const numCompletions = scorePromptOnly ? 1 : workbenchConfig.get('generation.numCompletions'); 50 | const temperature = workbenchConfig.get('generation.temperature'); 51 | const maxTokens = scorePromptOnly ? 0 : workbenchConfig.get('generation.maxTokens'); 52 | const topP = workbenchConfig.get('generation.topP'); 53 | const frequencyPenalty = workbenchConfig.get('generation.frequencyPenalty'); 54 | const presencePenalty = workbenchConfig.get('generation.presencePenalty'); 55 | const logprobs = workbenchConfig.get('generation.logprobs'); 56 | const echo = scorePromptOnly ? true : workbenchConfig.get('generation.echo'); 57 | const log = workbenchConfig.get('log'); 58 | const apiKey = workbenchConfig.get('apiKey'); 59 | const basePath = workbenchConfig.get('basePath'); 60 | 61 | const chat = model === 'gpt-3.5-turbo' || model === 'gpt-4'; 62 | 63 | if(!apiKey) { 64 | vscode.window.showErrorMessage("No API key set. Please set the 'worldspider.apiKey' setting."); 65 | return null; 66 | } 67 | 68 | const configuration = new Configuration({ 69 | apiKey: apiKey, 70 | basePath: basePath, 71 | }); 72 | const openai = new OpenAIApi(configuration); 73 | 74 | const request = { 75 | model: model, 76 | // prompt: prompt, 77 | max_tokens: maxTokens, 78 | top_p: topP, 79 | frequency_penalty: frequencyPenalty, 80 | presence_penalty: presencePenalty, 81 | n: numCompletions, 82 | temperature: temperature, 83 | ...(!chat && {logprobs: logprobs}), 84 | ...(!chat && {echo: echo}), 85 | ...(!chat && {suffix: suffix}), 86 | ...(chat && { messages: [{ role: "assistant", content: prompt }] }), 87 | ...(!chat && { prompt }), 88 | }; 89 | 90 | try { 91 | const response = chat ? await openai.createChatCompletion(request) : await openai.createCompletion(request); 92 | 93 | // if chat, change the format of the response from choices[i].message.content to choices[i].text 94 | 95 | const choices = response.data.choices.map((choice: any) => { 96 | if(chat) { 97 | return { 98 | ...choice, 99 | text: " " + choice.message.content, // TODO remove this space and find less hacky way to fix the issue 100 | }; 101 | } else { 102 | return choice; 103 | } 104 | }); 105 | 106 | const responseObject = { 107 | ...response.data, 108 | choices: echo ? removePromptsFromCompletions(choices, prompt) : choices, 109 | prompt: prompt, 110 | suffix: suffix, 111 | echo: echo, 112 | }; 113 | if(log) { 114 | saveModelResponse(responseObject); 115 | } 116 | return responseObject; 117 | } catch (error: unknown) { 118 | console.log(error); 119 | let errorMessage = (error as Error).message; 120 | vscode.window.showErrorMessage(`model call failed: ${errorMessage}`); 121 | return null; 122 | } 123 | } 124 | 125 | 126 | export async function modelOperation(prefix:string, suffix: string, selectedText: string ) { 127 | const workbenchConfig = vscode.workspace.getConfiguration('worldspider'); 128 | 129 | const apiKey = workbenchConfig.get('apiKey'); 130 | const instruction = workbenchConfig.get('generation.customInstruction'); 131 | const model = workbenchConfig.get('generation.customOpModel'); 132 | 133 | const messages = [ 134 | { role: "system", content: instruction }, 135 | { role: "user", content: selectedText } 136 | ]; 137 | 138 | const request = { 139 | model: model, 140 | max_tokens: 200, 141 | n: 3, 142 | temperature: 1.2, 143 | messages: messages, 144 | }; 145 | 146 | const configuration = new Configuration({ 147 | apiKey: apiKey, 148 | }); 149 | const openai = new OpenAIApi(configuration); 150 | 151 | // console.log(request); 152 | 153 | try { 154 | const response = await openai.createChatCompletion(request); 155 | 156 | const choices = response.data.choices.map((choice: any) => { 157 | return { 158 | ...choice, 159 | text: choice.message.content, 160 | }; 161 | }); 162 | 163 | console.log(choices); 164 | 165 | const responseObject = { 166 | ...response.data, 167 | choices: choices, 168 | prefix: prefix, 169 | suffix: suffix, 170 | }; 171 | 172 | return responseObject; 173 | } 174 | catch (error) { 175 | console.log(error); 176 | vscode.window.showErrorMessage("model call failed"); 177 | return null; 178 | } 179 | } -------------------------------------------------------------------------------- /src/document.ts: -------------------------------------------------------------------------------- 1 | import * as vscode from 'vscode'; 2 | 3 | 4 | export function getCurrentContext() { 5 | const workbenchConfig = vscode.workspace.getConfiguration('worldspider'); 6 | const prefixLength = workbenchConfig.get('generation.prefixLength'); 7 | const suffixLength = workbenchConfig.get('generation.suffixLength'); 8 | let editor = vscode.window.activeTextEditor; 9 | if(!editor) { 10 | return {prefix: '', suffix: ''}; 11 | } 12 | let context = editor.document.getText(); 13 | let currentContextPrefix = context.slice(0, editor.document.offsetAt(editor.selection.start)).slice(-prefixLength); 14 | let currentContextSuffix = context.slice(editor.document.offsetAt(editor.selection.end)).slice(0, suffixLength); 15 | return {prefix: currentContextPrefix, suffix: currentContextSuffix}; 16 | } 17 | 18 | export function getSelectionInfo() { 19 | // returns currently selected text and absolute position in the document of the selection 20 | let editor = vscode.window.activeTextEditor; 21 | if(!editor) { 22 | return {prefix: '', suffix: ''}; 23 | } 24 | let selectedText = editor.document.getText(editor.selection); 25 | let selectedStart = editor.selection.start; 26 | let selectedEnd = editor.selection.end; 27 | let selectedStartOffset = editor.document.offsetAt(selectedStart); 28 | let selectedEndOffset = editor.document.offsetAt(selectedEnd); 29 | return {selectedText, selectedStart, selectedEnd, selectedStartOffset, selectedEndOffset}; 30 | } -------------------------------------------------------------------------------- /src/extension.ts: -------------------------------------------------------------------------------- 1 | // The module 'vscode' contains the VS Code extensibility API 2 | // Import the module and reference it with the alias vscode in your code below 3 | import * as vscode from 'vscode'; 4 | 5 | import { getModelResponse, ModelResponse, ModelCompletion, modelOperation } from './callModel'; 6 | import { getCurrentContext, getSelectionInfo } from './document'; 7 | // import { probToColor } from './utils'; 8 | 9 | interface InsertedHistoryItem { 10 | startPosition: number; 11 | endPosition: number; 12 | prefix: string; 13 | echo: boolean; 14 | modelCompletion: ModelCompletion; 15 | } 16 | 17 | 18 | export function removeMetaTokensFromPrompt(prompt: string): string { 19 | // Strip comments, i.e. the tokens and everything between them 20 | // Also strip a trailing newline after the comment, 21 | // so a commented line doesn't leave an extra newline in the prompt 22 | prompt = prompt.replace(/<\|BEGINCOMMENT\|>.*?<\|ENDCOMMENT\|>\n?/g, ""); 23 | 24 | // Strip everything before the <|start|> token. 25 | // If there is a newline directly after the start token, also strip it 26 | prompt = prompt.replace(/^.*<\|start\|>\n?/, ""); 27 | 28 | // Strip everything after the <|end|> token 29 | prompt = prompt.replace(/<\|end\|>.*$/, ""); 30 | 31 | // Strip any other kind of token 32 | prompt = prompt.replace(/<\|.*?\|>/g, ""); 33 | 34 | return prompt; 35 | } 36 | 37 | 38 | // This method is called when your extension is activated 39 | // Your extension is activated the very first time the command is executed 40 | export function activate(context: vscode.ExtensionContext) { 41 | 42 | // Use the console to output diagnostic information (console.log) and errors (console.error) 43 | // This line of code will only be executed once when your extension is activated 44 | vscode.window.showInformationMessage('worldspider is active'); 45 | 46 | console.log('Congratulations, your extension "worldspider" is now active!'); 47 | 48 | let wsconsole = vscode.window.createOutputChannel("worldspider"); 49 | 50 | let history: ModelResponse[] = []; 51 | let currentHistoryIndex = 0; 52 | let insertedHistory: InsertedHistoryItem[] = []; 53 | 54 | function afterInsertText(position: number, completionIndex: number, text: string) { 55 | const prompt = history[currentHistoryIndex].prompt; 56 | insertedHistory.push({ 57 | startPosition: history[currentHistoryIndex].echo ? position - prompt.length : position, 58 | endPosition: position + text.length, 59 | prefix: prompt, 60 | echo: history[currentHistoryIndex].echo, 61 | modelCompletion: history[currentHistoryIndex].choices[completionIndex], 62 | }); 63 | const workbenchConfig = vscode.workspace.getConfiguration('worldspider'); 64 | 65 | if (!workbenchConfig.get('highlightModelText')) { 66 | return; 67 | } 68 | // create decoration for the inserted text 69 | const decoration = vscode.window.createTextEditorDecorationType({ 70 | backgroundColor: 'rgba(0, 0, 0, 0.1)', 71 | borderRadius: '3px', 72 | border: '1px solid rgba(0, 0, 0, 0.4)', 73 | rangeBehavior: vscode.DecorationRangeBehavior.ClosedClosed, 74 | }); 75 | const editor = vscode.window.activeTextEditor; 76 | if (!editor) { 77 | return; 78 | } 79 | const startPosition = editor.document.positionAt(position); 80 | const endPosition = editor.document.positionAt(position + text.length); 81 | const range = new vscode.Range(startPosition, endPosition); 82 | editor.setDecorations(decoration, [range]); 83 | } 84 | 85 | let afterInsertTextDisposable = vscode.commands.registerCommand('worldspider.afterInsertText', afterInsertText); 86 | 87 | let registerCompletionItemsProvider = vscode.languages.registerCompletionItemProvider( 88 | '*', 89 | { 90 | provideCompletionItems(document, position, token, context) { 91 | console.log(history[currentHistoryIndex]); 92 | const completionItems = history[currentHistoryIndex].choices.map((response: { text: string; }, i: any) => { 93 | const text = response.text; 94 | const completionItem = new vscode.CompletionItem(text); 95 | completionItem.insertText = text; 96 | completionItem.keepWhitespace = true; 97 | completionItem.documentation = text; 98 | completionItem.kind = vscode.CompletionItemKind.Text; 99 | completionItem.range = new vscode.Range(position, position); // to prevent trying to overwrite the previous word 100 | 101 | const absolutePosition = document.offsetAt(position); 102 | // when text is inserted, call afterInsertText 103 | completionItem.command = { 104 | command: 'worldspider.afterInsertText', 105 | title: 'afterInsertText', 106 | arguments: [absolutePosition, i, text] 107 | }; 108 | 109 | return completionItem; 110 | }); 111 | return completionItems; 112 | } 113 | } 114 | );//, ...triggers); 115 | 116 | let registerHoverProvider = vscode.languages.registerHoverProvider( 117 | '*', 118 | { 119 | provideHover(document, position, token) { 120 | const absolutePosition = document.offsetAt(position); 121 | 122 | // Take the most recent element of InsertedHistory 123 | // where startPosition <= absolutePosition <= endPosition 124 | const insertedHistoryElement = insertedHistory.filter((element: InsertedHistoryItem) => { 125 | return element.startPosition <= absolutePosition && element.endPosition >= absolutePosition; 126 | }).pop(); 127 | 128 | if (!insertedHistoryElement) { 129 | return; 130 | } 131 | 132 | const modelCompletion = insertedHistoryElement.modelCompletion; 133 | const logprobs = modelCompletion.logprobs; 134 | 135 | if (!logprobs) { 136 | return; 137 | } 138 | 139 | const localOffsets = logprobs.text_offset.map((offset: number) => { 140 | return insertedHistoryElement.echo ? offset : offset - insertedHistoryElement.prefix.length; 141 | }); 142 | 143 | // get offset between absolutePosition and startPosition 144 | const offset = absolutePosition - insertedHistoryElement.startPosition; 145 | 146 | // find the first index in localOffsets where offset >= localOffset 147 | const index = localOffsets.findIndex((localOffset: number) => { 148 | return offset <= localOffset; 149 | }) - 1; 150 | 151 | const hoverToken = logprobs.tokens[index]; 152 | 153 | // check if word at hovered position contains the token 154 | // or if the token contains the hovered word 155 | const hoverWord = document.getText(document.getWordRangeAtPosition(position)); 156 | if (!hoverToken.includes(hoverWord) && !hoverWord.includes(hoverToken)) { 157 | return; 158 | } 159 | 160 | const hoverTokenLogprob = logprobs.token_logprobs[index]; 161 | const hoverTokenProb = Math.exp(hoverTokenLogprob) * 100; 162 | const hoverTokenAlternatives = logprobs.top_logprobs ? logprobs.top_logprobs[index] : undefined; 163 | // const color = probToColor(hoverTokenProb); 164 | 165 | // const markDownString = new vscode.MarkdownString(`${hoverToken}: ${hoverTokenLogprob.toPrecision(4)} | ${hoverTokenProb.toPrecision(4)}%`); 166 | // const markDownString = new vscode.MarkdownString(`'${hoverToken}': ${hoverTokenLogprob.toPrecision(4)} | ${hoverTokenProb.toPrecision(4)}%`); 167 | const markDownString = new vscode.MarkdownString(); 168 | markDownString.isTrusted = true; 169 | markDownString.supportHtml = true; 170 | 171 | 172 | if(hoverTokenAlternatives) { 173 | // append markdown table header 174 | markDownString.appendMarkdown(`\n\n| Token | Logprob | Prob`); 175 | markDownString.appendMarkdown(`\n| --- | --- | --- |`); 176 | 177 | const sortedTokenAlternatives = []; 178 | 179 | for (const key in hoverTokenAlternatives) { 180 | sortedTokenAlternatives.push([key, hoverTokenAlternatives[key]]); 181 | } 182 | 183 | sortedTokenAlternatives.sort((a, b) => { 184 | return b[1] - a[1]; 185 | }); 186 | 187 | sortedTokenAlternatives.forEach((tokenAlternative: any) => { 188 | const token = tokenAlternative[0]; 189 | const logprob = tokenAlternative[1]; 190 | const prob = Math.exp(logprob) * 100; 191 | // append row of markdown table 192 | markDownString.appendMarkdown(`\n| '${token}' | ${logprob.toPrecision(4)} | ${prob.toPrecision(4)}% |`); 193 | // markDownString.appendMarkdown(`
'${token}': ${logprob.toPrecision(4)} | ${prob.toPrecision(4)}%
`); 194 | }); 195 | } 196 | 197 | markDownString.appendMarkdown(`\n\n| Sampled | Logprob | Prob`); 198 | markDownString.appendMarkdown(`\n| --- | --- | --- |`); 199 | 200 | markDownString.appendMarkdown(`\n| '${hoverToken}' | ${hoverTokenLogprob.toPrecision(4)} | ${hoverTokenProb.toPrecision(4)}% |`); 201 | 202 | return new vscode.Hover(markDownString); 203 | } 204 | } 205 | ); 206 | 207 | 208 | let showCompletions = vscode.commands.registerCommand('worldspider.showCompletions', () => { 209 | vscode.commands.executeCommand('editor.action.triggerSuggest'); 210 | }); 211 | 212 | 213 | async function handleGetModelCompletions(callModelFunction: Function) { 214 | vscode.window.withProgress({ 215 | location: vscode.ProgressLocation.Window, 216 | cancellable: false, 217 | title: 'Generating...' 218 | }, async (progress) => { 219 | progress.report({ increment: 0 }); 220 | const modelResponse = await callModelFunction(); 221 | progress.report({ increment: 100 }); 222 | // append to history 223 | if (modelResponse) { 224 | wsconsole.appendLine(JSON.stringify(modelResponse)); 225 | history.push(modelResponse); 226 | currentHistoryIndex = history.length - 1; 227 | vscode.commands.executeCommand('editor.action.triggerSuggest'); 228 | } 229 | }); 230 | } 231 | 232 | let getContinuations = vscode.commands.registerCommand('worldspider.getContinuations', () => { 233 | let { prefix, suffix } = getCurrentContext(); 234 | prefix = removeMetaTokensFromPrompt(prefix); 235 | suffix = removeMetaTokensFromPrompt(suffix); 236 | handleGetModelCompletions(() => getModelResponse(prefix, null)); 237 | }); 238 | 239 | let getInfills = vscode.commands.registerCommand('worldspider.getInfills', () => { 240 | let { prefix, suffix } = getCurrentContext(); 241 | prefix = removeMetaTokensFromPrompt(prefix); 242 | suffix = removeMetaTokensFromPrompt(suffix); 243 | handleGetModelCompletions(() => getModelResponse(prefix, suffix)); 244 | }); 245 | 246 | let customOp = vscode.commands.registerCommand('worldspider.customOp', () => { 247 | let { prefix, suffix } = getCurrentContext(); 248 | prefix = removeMetaTokensFromPrompt(prefix); 249 | suffix = removeMetaTokensFromPrompt(suffix); 250 | const { selectedText } = getSelectionInfo(); 251 | if(!selectedText) { 252 | vscode.window.showInformationMessage(`No text selected`); 253 | return; 254 | } 255 | handleGetModelCompletions(() => modelOperation(prefix, suffix, selectedText)); 256 | }); 257 | 258 | let scorePrompt = vscode.commands.registerCommand('worldspider.scorePrompt', () => { 259 | const { selectedText, selectedStartOffset, selectedEndOffset } = getSelectionInfo(); 260 | if (selectedText) { 261 | vscode.window.withProgress({ 262 | location: vscode.ProgressLocation.Window, 263 | cancellable: false, 264 | title: 'Getting prompt logprobs...' 265 | }, async (progress) => { 266 | progress.report({ increment: 0 }); 267 | const modelResponse = await getModelResponse(selectedText, '', true); 268 | progress.report({ increment: 100 }); 269 | vscode.window.showInformationMessage(`Added prompt logprob information`); 270 | if(modelResponse) { 271 | wsconsole.appendLine(JSON.stringify(modelResponse)); 272 | insertedHistory.push({ 273 | startPosition: selectedStartOffset, 274 | endPosition: selectedEndOffset, 275 | prefix: selectedText, 276 | echo: true, 277 | modelCompletion: modelResponse.choices[0], 278 | }); 279 | } 280 | }); 281 | } 282 | }); 283 | 284 | 285 | 286 | let copyCompletions = vscode.commands.registerCommand('worldspider.copyCompletions', () => { 287 | vscode.env.clipboard.writeText(JSON.stringify(history[currentHistoryIndex])); 288 | vscode.window.showInformationMessage('Copied model completions to clipboard'); 289 | }); 290 | 291 | let copyCompletionsText = vscode.commands.registerCommand('worldspider.copyCompletionsText', () => { 292 | const completionsTextList = history[currentHistoryIndex].choices.map((choice: { text: string; }) => choice.text); 293 | const completionsText = JSON.stringify(completionsTextList); 294 | vscode.env.clipboard.writeText(completionsText); 295 | vscode.window.showInformationMessage('Copied completions text to clipboard'); 296 | }); 297 | 298 | let prevCompletions = vscode.commands.registerCommand('worldspider.prevCompletions', () => { 299 | if (currentHistoryIndex > 0) { 300 | currentHistoryIndex--; 301 | vscode.commands.executeCommand('editor.action.triggerSuggest'); 302 | } 303 | }); 304 | 305 | let nextCompletions = vscode.commands.registerCommand('worldspider.nextCompletions', () => { 306 | if (currentHistoryIndex < history.length - 1) { 307 | currentHistoryIndex++; 308 | vscode.commands.executeCommand('editor.action.triggerSuggest'); 309 | } 310 | }); 311 | 312 | 313 | context.subscriptions.push(registerCompletionItemsProvider); 314 | context.subscriptions.push(showCompletions); 315 | context.subscriptions.push(getContinuations); 316 | context.subscriptions.push(getInfills); 317 | context.subscriptions.push(customOp); 318 | context.subscriptions.push(copyCompletions); 319 | context.subscriptions.push(copyCompletionsText); 320 | context.subscriptions.push(prevCompletions); 321 | context.subscriptions.push(nextCompletions); 322 | context.subscriptions.push(afterInsertTextDisposable); 323 | context.subscriptions.push(scorePrompt); 324 | 325 | } 326 | 327 | // This method is called when your extension is deactivated 328 | export function deactivate() {} 329 | -------------------------------------------------------------------------------- /src/logging.ts: -------------------------------------------------------------------------------- 1 | import * as vscode from 'vscode'; 2 | 3 | export function saveModelResponse(modelResponse: any) { 4 | var path = require('path'); 5 | const workbenchConfig = vscode.workspace.getConfiguration('worldspider'); 6 | const folders = vscode.workspace.workspaceFolders; 7 | if(!folders) { 8 | return; 9 | } 10 | const rootPath = folders[0].uri.fsPath; 11 | const savePath = workbenchConfig.get('savePath'); 12 | const fs = require('fs'); 13 | const filePath = path.join(rootPath, savePath); 14 | // check if file exists 15 | if (!fs.existsSync(filePath)) { 16 | fs.writeFileSync(filePath, '[]'); 17 | } 18 | fs.readFile(filePath, 'utf8', function (err: any, data: any) { 19 | if (err) {throw err;} 20 | var json = JSON.parse(data); 21 | json.push(modelResponse); 22 | fs.writeFile(filePath, JSON.stringify(json), function (err: any) { 23 | if (err) {throw err;} 24 | // console.log('Saved to ' + filePath); 25 | }); 26 | } 27 | ); 28 | } -------------------------------------------------------------------------------- /src/test/runTest.ts: -------------------------------------------------------------------------------- 1 | import * as path from 'path'; 2 | 3 | import { runTests } from '@vscode/test-electron'; 4 | 5 | async function main() { 6 | try { 7 | // The folder containing the Extension Manifest package.json 8 | // Passed to `--extensionDevelopmentPath` 9 | const extensionDevelopmentPath = path.resolve(__dirname, '../../'); 10 | 11 | // The path to test runner 12 | // Passed to --extensionTestsPath 13 | const extensionTestsPath = path.resolve(__dirname, './suite/index'); 14 | 15 | // Download VS Code, unzip it and run the integration test 16 | await runTests({ extensionDevelopmentPath, extensionTestsPath }); 17 | } catch (err) { 18 | console.error('Failed to run tests'); 19 | process.exit(1); 20 | } 21 | } 22 | 23 | main(); 24 | -------------------------------------------------------------------------------- /src/test/suite/extension.test.ts: -------------------------------------------------------------------------------- 1 | import * as assert from 'assert'; 2 | 3 | // You can import and use all API from the 'vscode' module 4 | // as well as import your extension to test it 5 | import * as vscode from 'vscode'; 6 | import * as worldspider from '../../extension'; 7 | 8 | suite('Extension Test Suite', () => { 9 | vscode.window.showInformationMessage('Start all tests.'); 10 | 11 | test('Prompt Strip Test: Blockcomment', () => { 12 | 13 | assert.strictEqual( 14 | worldspider.removeMetaTokensFromPrompt( 15 | "<|BEGINCOMMENT|>Some Comment<|ENDCOMMENT|>The Prompt" 16 | ), 17 | "The Prompt" 18 | ); 19 | }); 20 | 21 | test('Prompt Strip Test: Start and End Token', () => { 22 | assert.strictEqual( 23 | worldspider.removeMetaTokensFromPrompt( 24 | "Preamble<|start|>\nThe Prompt<|end|>Addendum" 25 | ), 26 | "The Prompt" 27 | ); 28 | }); 29 | 30 | test('Prompt Strip Test: Arbitrary Token', () => { 31 | assert.strictEqual( 32 | worldspider.removeMetaTokensFromPrompt( 33 | "The <|some_token|>Prompt" 34 | ), 35 | "The Prompt" 36 | ); 37 | }); 38 | 39 | }); 40 | -------------------------------------------------------------------------------- /src/test/suite/index.ts: -------------------------------------------------------------------------------- 1 | import * as path from 'path'; 2 | import * as Mocha from 'mocha'; 3 | import * as glob from 'glob'; 4 | 5 | export function run(): Promise { 6 | // Create the mocha test 7 | const mocha = new Mocha({ 8 | ui: 'tdd', 9 | color: true 10 | }); 11 | 12 | const testsRoot = path.resolve(__dirname, '..'); 13 | 14 | return new Promise((c, e) => { 15 | glob('**/**.test.js', { cwd: testsRoot }, (err, files) => { 16 | if (err) { 17 | return e(err); 18 | } 19 | 20 | // Add files to the test suite 21 | files.forEach(f => mocha.addFile(path.resolve(testsRoot, f))); 22 | 23 | try { 24 | // Run the mocha test 25 | mocha.run(failures => { 26 | if (failures > 0) { 27 | e(new Error(`${failures} tests failed.`)); 28 | } else { 29 | c(); 30 | } 31 | }); 32 | } catch (err) { 33 | console.error(err); 34 | e(err); 35 | } 36 | }); 37 | }); 38 | } 39 | -------------------------------------------------------------------------------- /src/utils.ts: -------------------------------------------------------------------------------- 1 | 2 | export function probToColor ( prob : number ) { 3 | // returns hex color string 4 | // prob is a number between 0 and 100 5 | const hue = ( prob * 1.2 ) / 360 ; 6 | return `hsl(${hue}, 100%, 50%)` ; 7 | } -------------------------------------------------------------------------------- /syntaxes/worldspider.tmLanguage.json: -------------------------------------------------------------------------------- 1 | { 2 | "scopeName": "source.worldspider", 3 | "patterns": [{ "include": "#expression" }], 4 | "repository": { 5 | "expression": { 6 | "patterns": [ 7 | { "include": "#blockcomment" }, 8 | { "include": "#token" } 9 | ] 10 | }, 11 | "blockcomment": { 12 | "begin": "<\\|BEGINCOMMENT\\|>", 13 | "end": "<\\|ENDCOMMENT\\|>", 14 | "name": "comment.block" 15 | }, 16 | "token": { 17 | "match": "<\\|\\w+\\|>", 18 | "name": "keyword.control" 19 | } 20 | } 21 | } -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "module": "commonjs", 4 | "target": "ES2020", 5 | "outDir": "out", 6 | "lib": [ 7 | "ES2020" 8 | ], 9 | "sourceMap": true, 10 | "rootDir": "src", 11 | "strict": true /* enable all strict type-checking options */ 12 | /* Additional Checks */ 13 | // "noImplicitReturns": true, /* Report error when not all code paths in function return a value. */ 14 | // "noFallthroughCasesInSwitch": true, /* Report errors for fallthrough cases in switch statement. */ 15 | // "noUnusedParameters": true, /* Report errors on unused parameters. */ 16 | } 17 | } 18 | -------------------------------------------------------------------------------- /vsc-extension-quickstart.md: -------------------------------------------------------------------------------- 1 | # Welcome to your VS Code Extension 2 | 3 | ## What's in the folder 4 | 5 | * This folder contains all of the files necessary for your extension. 6 | * `package.json` - this is the manifest file in which you declare your extension and command. 7 | * The sample plugin registers a command and defines its title and command name. With this information VS Code can show the command in the command palette. It doesn’t yet need to load the plugin. 8 | * `src/extension.ts` - this is the main file where you will provide the implementation of your command. 9 | * The file exports one function, `activate`, which is called the very first time your extension is activated (in this case by executing the command). Inside the `activate` function we call `registerCommand`. 10 | * We pass the function containing the implementation of the command as the second parameter to `registerCommand`. 11 | 12 | ## Get up and running straight away 13 | 14 | * Press `F5` to open a new window with your extension loaded. 15 | * Run your command from the command palette by pressing (`Ctrl+Shift+P` or `Cmd+Shift+P` on Mac) and typing `Hello World`. 16 | * Set breakpoints in your code inside `src/extension.ts` to debug your extension. 17 | * Find output from your extension in the debug console. 18 | 19 | ## Make changes 20 | 21 | * You can relaunch the extension from the debug toolbar after changing code in `src/extension.ts`. 22 | * You can also reload (`Ctrl+R` or `Cmd+R` on Mac) the VS Code window with your extension to load your changes. 23 | 24 | ## Explore the API 25 | 26 | * You can open the full set of our API when you open the file `node_modules/@types/vscode/index.d.ts`. 27 | 28 | ## Run tests 29 | 30 | * Open the debug viewlet (`Ctrl+Shift+D` or `Cmd+Shift+D` on Mac) and from the launch configuration dropdown pick `Extension Tests`. 31 | * Press `F5` to run the tests in a new window with your extension loaded. 32 | * See the output of the test result in the debug console. 33 | * Make changes to `src/test/suite/extension.test.ts` or create new test files inside the `test/suite` folder. 34 | * The provided test runner will only consider files matching the name pattern `**.test.ts`. 35 | * You can create folders inside the `test` folder to structure your tests any way you want. 36 | 37 | ## Go further 38 | 39 | * [Follow UX guidelines](https://code.visualstudio.com/api/ux-guidelines/overview) to create extensions that seamlessly integrate with VS Code's native interface and patterns. 40 | * Reduce the extension size and improve the startup time by [bundling your extension](https://code.visualstudio.com/api/working-with-extensions/bundling-extension). 41 | * [Publish your extension](https://code.visualstudio.com/api/working-with-extensions/publishing-extension) on the VS Code extension marketplace. 42 | * Automate builds by setting up [Continuous Integration](https://code.visualstudio.com/api/working-with-extensions/continuous-integration). 43 | --------------------------------------------------------------------------------