├── .gitignore ├── assets └── code_generation.gif ├── .npmignore ├── src ├── constants │ └── index.ts ├── utils │ ├── getDocumentLanguage.ts │ └── getCodeCompletions.ts └── index.ts ├── tsconfig.json ├── .eslintrc.js ├── esbuild.js ├── LICENSE ├── README.md └── package.json /.gitignore: -------------------------------------------------------------------------------- 1 | node_modules 2 | lib 3 | package-lock.json 4 | yarn.lock 5 | poc 6 | -------------------------------------------------------------------------------- /assets/code_generation.gif: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/LanceZhu/coc-codegeex/HEAD/assets/code_generation.gif -------------------------------------------------------------------------------- /.npmignore: -------------------------------------------------------------------------------- 1 | src 2 | node_modules 3 | tsconfig.json 4 | *.map 5 | .tags 6 | .DS_Store 7 | webpack.config.js 8 | esbuild.js 9 | pack e-lock.json 10 | yarn.lock 11 | yarn-error.log 12 | .github 13 | .eslintrc.js 14 | .prettierrc 15 | assets 16 | -------------------------------------------------------------------------------- /src/constants/index.ts: -------------------------------------------------------------------------------- 1 | export const languageList = [ 2 | 'C++', 3 | 'C', 4 | 'C#', 5 | 'Cuda', 6 | 'Objective-C', 7 | 'Objective-C++', 8 | 'Python', 9 | 'Java', 10 | 'TeX', 11 | 'HTML', 12 | 'PHP', 13 | 'JavaScript', 14 | 'TypeScript', 15 | 'Go', 16 | 'Shell', 17 | 'Rust', 18 | 'CSS', 19 | 'SQL', 20 | 'R', 21 | ]; 22 | -------------------------------------------------------------------------------- /tsconfig.json: -------------------------------------------------------------------------------- 1 | { 2 | "compilerOptions": { 3 | "target": "es2017", 4 | "lib": ["es2017", "es2018", "dom"], 5 | "module": "commonjs", 6 | "declaration": false, 7 | "sourceMap": true, 8 | "outDir": "lib", 9 | "strict": true, 10 | "moduleResolution": "node", 11 | "noImplicitAny": false, 12 | "esModuleInterop": true, 13 | }, 14 | "include": ["src"] 15 | } 16 | -------------------------------------------------------------------------------- /.eslintrc.js: -------------------------------------------------------------------------------- 1 | module.exports = { 2 | env: { 3 | node: true, 4 | }, 5 | parser: '@typescript-eslint/parser', 6 | extends: ['plugin:@typescript-eslint/recommended', 'plugin:prettier/recommended'], 7 | rules: { 8 | '@typescript-eslint/ban-ts-comment': 'off', 9 | '@typescript-eslint/no-explicit-any': 'off', 10 | '@typescript-eslint/no-non-null-assertion': 'off', 11 | '@typescript-eslint/no-namespace': 'off', 12 | '@typescript-eslint/no-empty-function': 'off', 13 | '@typescript-eslint/explicit-function-return-type': 'off', 14 | '@typescript-eslint/explicit-module-boundary-types': 'off', 15 | }, 16 | }; 17 | -------------------------------------------------------------------------------- /esbuild.js: -------------------------------------------------------------------------------- 1 | /* eslint-disable @typescript-eslint/no-var-requires */ 2 | async function start(watch) { 3 | await require('esbuild').build({ 4 | entryPoints: ['src/index.ts'], 5 | bundle: true, 6 | watch, 7 | minify: process.env.NODE_ENV === 'production', 8 | sourcemap: process.env.NODE_ENV === 'development', 9 | mainFields: ['module', 'main'], 10 | external: ['coc.nvim'], 11 | platform: 'node', 12 | // target: 'node10.12', 13 | target: 'node14.20', 14 | outfile: 'lib/index.js', 15 | }); 16 | } 17 | 18 | let watch = false; 19 | if (process.argv.length > 2 && process.argv[2] === '--watch') { 20 | console.log('watching...'); 21 | watch = { 22 | onRebuild(error) { 23 | if (error) { 24 | console.error('watch build failed:', error); 25 | } else { 26 | console.log('watch build succeeded'); 27 | } 28 | }, 29 | }; 30 | } 31 | 32 | start(watch).catch((e) => { 33 | console.error(e); 34 | }); 35 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2022 Lance Zhu 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # coc-codegeex 2 | 3 | [codegeex][codegeex] in vim using [coc.nvim][coc.nvim] 4 | 5 | ![code_generation](./assets/code_generation.gif) 6 | 7 | ## Install 8 | 9 | first install [coc.nvim][coc.nvim] 10 | 11 | and in vim install plugin: coc-codegeex `:CocInstall coc-codegeex` 12 | 13 | ## Usage 14 | 15 | ### setup 16 | 17 | use `:CocConfig` open configuration file, add `apiKey` and `apiSecret`. You can apply this in [tianqi][tianqi] 18 | 19 | ``` 20 | { 21 | "codegeex": { 22 | "apiKey": "your_api_key", 23 | "apiSecret": "your_api_secret", 24 | } 25 | } 26 | ``` 27 | 28 | ### stealth mode 29 | 30 | when you stop type, you will see the completion list. use `Enter` to confirm selection, the selected code will be inserted. 31 | 32 | ### translation mode 33 | 34 | set keymap in .vimrc 35 | 36 | `vmap (coc-codegeex-translate-keymap)` 37 | 38 | in vim visual mode, select some code and use `ctrl + l` to trigger translation. 39 | 40 | ## License 41 | 42 | MIT 43 | 44 | --- 45 | 46 | > This extension is built with [create-coc-extension](https://github.com/fannheyward/create-coc-extension) 47 | 48 | [codegeex]: https://codegeex.ai/ 49 | [coc.nvim]: https://github.com/neoclide/coc.nvim 50 | [tianqi]: https://tianqi.aminer.cn/open/ 51 | -------------------------------------------------------------------------------- /package.json: -------------------------------------------------------------------------------- 1 | { 2 | "name": "coc-codegeex", 3 | "version": "0.0.1", 4 | "description": "codegeex in vim using coc", 5 | "author": "Lance Zhu ", 6 | "license": "MIT", 7 | "main": "lib/index.js", 8 | "keywords": [ 9 | "coc.nvim" 10 | ], 11 | "engines": { 12 | "coc": "^0.0.80" 13 | }, 14 | "scripts": { 15 | "lint": "eslint src --ext ts", 16 | "lint:fix": "eslint src --ext ts --fix", 17 | "prettier:fix": "prettier --write './**/*'", 18 | "clean": "rimraf lib", 19 | "watch": "node esbuild.js --watch", 20 | "build": "node esbuild.js", 21 | "prepare": "node esbuild.js" 22 | }, 23 | "prettier": { 24 | "singleQuote": true, 25 | "printWidth": 120, 26 | "semi": true 27 | }, 28 | "devDependencies": { 29 | "@typescript-eslint/eslint-plugin": "^5.13.0", 30 | "@typescript-eslint/parser": "^5.13.0", 31 | "coc.nvim": "^0.0.80", 32 | "esbuild": "^0.14.24", 33 | "eslint": "^8.10.0", 34 | "eslint-config-prettier": "^8.5.0", 35 | "eslint-plugin-prettier": "^4.0.0", 36 | "prettier": "^2.5.1", 37 | "rimraf": "^3.0.2", 38 | "typescript": "^4.6.2" 39 | }, 40 | "activationEvents": [ 41 | "*" 42 | ], 43 | "contributes": { 44 | "configuration": { 45 | "type": "object", 46 | "title": "coc-codegeex configuration", 47 | "properties": { 48 | "coc-codegeex.apiKey": { 49 | "type": "string", 50 | "default": "", 51 | "description": "apiKey for codegeex" 52 | }, 53 | "coc-codegeex.apiSecret": { 54 | "type": "string", 55 | "default": "", 56 | "description": "apiSecret for codegeex" 57 | } 58 | } 59 | } 60 | }, 61 | "dependencies": { 62 | "got": "^12.5.3" 63 | } 64 | } 65 | -------------------------------------------------------------------------------- /src/utils/getDocumentLanguage.ts: -------------------------------------------------------------------------------- 1 | export function getDocumentLanguage(documentLanguageId: string): string { 2 | let lang = ''; 3 | switch (documentLanguageId) { 4 | case 'cpp': 5 | lang = 'C++'; 6 | break; 7 | case 'c': 8 | lang = 'C'; 9 | break; 10 | case 'csharp': 11 | lang = 'C#'; 12 | break; 13 | case 'cuda-cpp': 14 | lang = 'Cuda'; 15 | break; 16 | case 'objective-c': 17 | lang = 'Objective-C'; 18 | break; 19 | case 'objective-cpp': 20 | lang = 'Objective-C++'; 21 | break; 22 | case 'python': 23 | lang = 'Python'; 24 | break; 25 | case 'java': 26 | lang = 'Java'; 27 | break; 28 | case 'tex': 29 | lang = 'TeX'; 30 | break; 31 | case 'html': 32 | lang = 'HTML'; 33 | break; 34 | case 'php': 35 | lang = 'PHP'; 36 | break; 37 | case 'javascript': 38 | case 'javascriptreact': 39 | lang = 'JavaScript'; 40 | break; 41 | case 'typescript': 42 | case 'typescriptreact': 43 | lang = 'TypeScript'; 44 | break; 45 | case 'go': 46 | lang = 'Go'; 47 | break; 48 | case 'shellscript': 49 | lang = 'Shell'; 50 | break; 51 | case 'rust': 52 | lang = 'Rust'; 53 | break; 54 | case 'css': 55 | case 'less': 56 | case 'sass': 57 | case 'scss': 58 | lang = 'CSS'; 59 | break; 60 | case 'sql': 61 | lang = 'SQL'; 62 | break; 63 | // case 'html': 64 | // lang = 'Kotlin' 65 | // break 66 | // case 'html': 67 | // lang = 'Pascal' 68 | // break 69 | case 'r': 70 | lang = 'R'; 71 | break; 72 | // case 'html': 73 | // lang = 'Fortran' 74 | // break 75 | // case 'html': 76 | // lang = 'Lean' 77 | // break 78 | default: 79 | lang = ''; 80 | } 81 | return lang; 82 | } 83 | -------------------------------------------------------------------------------- /src/utils/getCodeCompletions.ts: -------------------------------------------------------------------------------- 1 | import got from 'got'; 2 | 3 | export type CodeCompletions = { 4 | completions: Array; 5 | elapse: string; 6 | }; 7 | 8 | type HTTPAPIResponse = { 9 | message: string; 10 | status: number; 11 | result: any; 12 | }; 13 | 14 | export async function getCodeCompletions( 15 | prompt: string, 16 | num: number, 17 | lang: string, 18 | apiKey: string, 19 | apiSecret: string 20 | ): Promise { 21 | const API_URL = `https://tianqi.aminer.cn/api/v2/multilingual_code_generate`; 22 | const payload = { 23 | prompt: prompt, 24 | n: num, 25 | lang, 26 | apikey: apiKey, 27 | apisecret: apiSecret, 28 | }; 29 | const res: HTTPAPIResponse = await got 30 | .post(API_URL, { 31 | json: payload, 32 | }) 33 | .json(); 34 | // console.log('res:', JSON.stringify(res)); 35 | if (res.status !== 0) { 36 | return { 37 | completions: [], 38 | elapse: '0s', 39 | }; 40 | } 41 | let completions = res.result.output.code; 42 | completions = completions.filter((el: any) => el.trim() !== ''); 43 | 44 | return { 45 | completions, 46 | elapse: `${res.result.process_time.toFixed(2)}s`, 47 | }; 48 | } 49 | 50 | export async function getCodeTranslation( 51 | prompt: string, 52 | srcLang: string, 53 | dstLang: string, 54 | apiKey: string, 55 | apiSecret: string 56 | ): Promise { 57 | const API_URL = `https://tianqi.aminer.cn/api/v2/multilingual_code_translate`; 58 | const payload = { 59 | prompt: prompt, 60 | src_lang: srcLang, 61 | dst_lang: dstLang, 62 | apikey: apiKey, 63 | apisecret: apiSecret, 64 | }; 65 | const res: HTTPAPIResponse = await got 66 | .post(API_URL, { 67 | json: payload, 68 | }) 69 | .json(); 70 | if (res.status !== 0) { 71 | return ''; 72 | } 73 | const translation = res.result.output.code; 74 | return translation; 75 | } 76 | -------------------------------------------------------------------------------- /src/index.ts: -------------------------------------------------------------------------------- 1 | import { 2 | CompleteOption, 3 | CompleteResult, 4 | ExtensionContext, 5 | sources, 6 | window, 7 | workspace, 8 | Range, 9 | VimCompleteItem, 10 | WorkspaceConfiguration, 11 | } from 'coc.nvim'; 12 | 13 | import { getCodeCompletions, getCodeTranslation } from './utils/getCodeCompletions'; 14 | import { getDocumentLanguage } from './utils/getDocumentLanguage'; 15 | import { languageList } from './constants/index'; 16 | 17 | const SOURCE_NAME = 'coc-codegeex'; 18 | 19 | export async function activate(context: ExtensionContext): Promise { 20 | const config = workspace.getConfiguration('codegeex'); 21 | const statusBarItem = window.createStatusBarItem(0, { progress: true }); 22 | statusBarItem.text = 'coc-codegeex is generating completions...'; 23 | 24 | context.subscriptions.push( 25 | // source 26 | sources.createSource({ 27 | name: 'coc-codegeex completion source', // unique id 28 | // @ts-ignore 29 | triggerCharacters: [], 30 | doComplete: async (option: CompleteOption) => { 31 | statusBarItem.show(); 32 | const items = await getCompletionItems(option, config); 33 | statusBarItem.hide(); 34 | return items; 35 | }, 36 | onCompleteDone: async (item: VimCompleteItem) => { 37 | // @ts-ignore 38 | if (item.source !== SOURCE_NAME) { 39 | return; 40 | } 41 | 42 | const lines = item.user_data?.split('\n'); 43 | const lnum = (await workspace.nvim.call('line', ['.'])) as number; 44 | if (lines != null && lines[1] != null) { 45 | const appendLines = lines.slice(1); 46 | await workspace.nvim.call('append', [lnum, appendLines]); 47 | await workspace.nvim.call('setpos', [ 48 | '.', 49 | [0, lnum + appendLines.length, appendLines.slice(-1)[0].length + 1], 50 | ]); 51 | } 52 | if (item.user_data?.endsWith('\n')) { 53 | await workspace.nvim.call('append', [lnum, ['']]); 54 | await workspace.nvim.call('setpos', ['.', [0, lnum + 1, 1]]); 55 | } 56 | }, 57 | }), 58 | 59 | // keymap 60 | workspace.registerKeymap( 61 | ['v'], 62 | 'codegeex-translate-keymap', 63 | async () => { 64 | const document = await workspace.document; 65 | const documentLanguageId = document.textDocument.languageId; 66 | const srcLang = getDocumentLanguage(documentLanguageId); 67 | if (languageList.indexOf(srcLang) === -1) { 68 | window.showMessage(`current language: ${srcLang} is not supported.`); 69 | return; 70 | } 71 | const targetLanguageIdx = await window.showQuickpick(languageList, 'Target Language'); 72 | const dstLang = languageList[targetLanguageIdx]; 73 | 74 | const { nvim } = workspace; 75 | const start = await nvim.call('getpos', [`'<`]); 76 | const end = await nvim.call('getpos', [`'>`]); 77 | const [startRow, startCol] = start.slice(1, 3); 78 | const [endRow, endCol] = end.slice(1, 3); 79 | const text = document.textDocument.getText(Range.create(startRow - 1, startCol - 1, endRow - 1, endCol - 1)); 80 | 81 | statusBarItem.show(); 82 | const translation = await getCodeTranslation(text, srcLang, dstLang, config.apiKey, config.apiSecret); 83 | statusBarItem.hide(); 84 | const outputChannelName = 'coc-codegeex translation'; 85 | const outputChannel = window.createOutputChannel(outputChannelName); 86 | const translationLines = translation[0].split('\n'); 87 | for (const line of translationLines) { 88 | outputChannel.appendLine(line); 89 | } 90 | window.showOutputChannel(outputChannelName, true); 91 | }, 92 | { sync: false } 93 | ) 94 | ); 95 | } 96 | 97 | async function getCompletionItems( 98 | option: CompleteOption, 99 | config: WorkspaceConfiguration 100 | ): Promise { 101 | const num = 3; 102 | const document = await workspace.document; 103 | const documentLanguageId = document.textDocument.languageId; 104 | const lang = getDocumentLanguage(documentLanguageId); 105 | const { linenr, colnr, line } = option; 106 | const maxLines = 100; 107 | const startLine = Math.max(linenr - maxLines, 0); 108 | const text = document.textDocument.getText(Range.create(startLine, 0, linenr - 1, colnr - 1)); 109 | const prompt = text; 110 | try { 111 | const codeCompletions = await getCodeCompletions(prompt, num, lang, config.apiKey, config.apiSecret); 112 | const { completions, elapse } = codeCompletions; 113 | const completionItems = completions.map((comp) => { 114 | return { 115 | word: comp.split('\n')[0], 116 | // word: comp, 117 | menu: '[coc-codegeex]', 118 | filterText: `${line}`, 119 | user_data: comp, 120 | source: SOURCE_NAME, 121 | documentation: [ 122 | { 123 | filetype: 'markdown', 124 | content: elapse, 125 | }, 126 | ], 127 | }; 128 | }); 129 | 130 | if (completionItems && completionItems.length > 0) { 131 | // @ts-ignore 132 | completionItems[0].preSelect = true; 133 | } 134 | 135 | return { 136 | items: completionItems, 137 | priority: 1000, 138 | isIncomplete: true, 139 | startcol: colnr, 140 | }; 141 | } catch (e) { 142 | console.error(e); 143 | return null; 144 | } 145 | } 146 | --------------------------------------------------------------------------------