├── autoload ├── codegeex_gen.py ├── codegeex_trans.py └── codegeex.vim └── README.md /autoload/codegeex_gen.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import json 4 | import requests 5 | import argparse 6 | import os 7 | 8 | url = "https://tianqi.aminer.cn/api/v2/multilingual_code_generate" 9 | 10 | 11 | def request_codegeex_gen(prompt, lang, apikey, apisecret): 12 | payload = { 13 | "prompt": prompt, 14 | "n": 1, 15 | "lang": lang, 16 | "apikey": apikey, 17 | "apisecret": apisecret, 18 | } 19 | headers = {'Content-Type': 'application/json'} 20 | try: 21 | response = requests.post(url, 22 | data=json.dumps(payload), 23 | headers=headers) 24 | data = json.loads(response.content.decode('utf-8')) 25 | print(''.join(data['result']['output']['code']), end='') 26 | except Exception: 27 | pass 28 | 29 | 30 | def load_key(keyfile): 31 | with open(keyfile, 'r') as f: 32 | return json.load(f) 33 | 34 | 35 | def main(): 36 | parser = argparse.ArgumentParser() 37 | parser.add_argument('--keyfile', 38 | default=os.path.join(os.environ["HOME"], 39 | ".tianqi.key")) 40 | parser.add_argument('lang') 41 | parser.add_argument('prompt') 42 | result = parser.parse_args() 43 | if not os.path.exists(result.keyfile): 44 | return 45 | key = load_key(result.keyfile) 46 | request_codegeex_gen(result.prompt, 47 | result.lang, 48 | key['apikey'], 49 | key['apisecret']) 50 | 51 | 52 | if __name__ == "__main__": 53 | main() 54 | -------------------------------------------------------------------------------- /autoload/codegeex_trans.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python3 2 | 3 | import json 4 | import requests 5 | import argparse 6 | import os 7 | 8 | url = "https://tianqi.aminer.cn/api/v2/multilingual_code_translate" 9 | 10 | 11 | def request_codegeex_trans(prompt, src_lang, dst_lang, apikey, apisecret): 12 | payload = { 13 | "prompt": prompt, 14 | "src_lang": src_lang, 15 | "dst_lang": dst_lang, 16 | "apikey": apikey, 17 | "apisecret": apisecret, 18 | } 19 | headers = {"Content-Type": "application/json"} 20 | try: 21 | response = requests.post(url, 22 | data=json.dumps(payload), 23 | headers=headers) 24 | data = json.loads(response.content.decode("utf-8")) 25 | print("".join(data["result"]["output"]["code"]), end="") 26 | except Exception: 27 | pass 28 | 29 | 30 | def load_key(keyfile): 31 | with open(keyfile, "r") as f: 32 | return json.load(f) 33 | 34 | 35 | def main(): 36 | parser = argparse.ArgumentParser() 37 | parser.add_argument("--keyfile", 38 | default=os.path.join(os.environ["HOME"], 39 | ".tianqi.key")) 40 | parser.add_argument("src_lang") 41 | parser.add_argument("dst_lang") 42 | parser.add_argument("prompt") 43 | result = parser.parse_args() 44 | if not os.path.exists(result.keyfile): 45 | return 46 | key = load_key(result.keyfile) 47 | request_codegeex_trans(result.prompt, 48 | result.src_lang, 49 | result.dst_lang, 50 | key["apikey"], 51 | key["apisecret"]) 52 | 53 | 54 | if __name__ == "__main__": 55 | main() 56 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Codegeex-Vim 2 | 3 | Codegeex-Vim is a Vim plugin that enables you to easily generate code snippets using Codegeex API. Codegeex is an AI-powered programming assistant that helps developers generate high-quality, production-ready code. With this plugin, you can quickly generate code and seamlessly integrate it into your workflow. 4 | 5 | ## Installation 6 | 7 | To install the plugin, use your preferred Vim plugin manager: 8 | 9 | - [vim-plug](https://github.com/junegunn/vim-plug): `Plug 'codegeex/codegeex-vim'` 10 | - [Vundle](https://github.com/VundleVim/Vundle.vim): `Plugin 'codegeex/codegeex-vim'` 11 | 12 | Note: Make sure Python3 is installed on your system. 13 | 14 | ## Usage 15 | 16 | The plugin provides the following commands: 17 | 18 | - `GenCode([ins])`: This command generates a code snippet based on the currently active buffer or visual selection. If the cursor is in insert mode when the command is run, the generated code will be inserted in insert mode. If `ins` is 1, the command will force insert mode. Example: `:call codegeex#GenCode(0)` or `` in normal mode. 19 | - `TransCode()`: This command translates code from the source language to a chosen destination language. Example: `:call codegeex#TransCode()` or `` in visual mode. 20 | - `SetTianqiKeyFile(keyfile)`: This command sets the path to the Tianqi API key file. By default, the plugin looks for the file at `~/.tianqi.key`. Example: `:call codegeex#SetTianqiKeyFile('/path/to/tianqi/key/file')`. 21 | 22 | ## Configuration 23 | 24 | Here's an example configuration for the plugin commands: 25 | 26 | ``` 27 | inoremap :call codegeex#GenCode(1) 28 | nnoremap :call codegeex#GenCode(0) 29 | vnoremap :call codegeex#TransCode() 30 | ``` 31 | 32 | ## License 33 | 34 | MIT License. See LICENSE for more details. 35 | -------------------------------------------------------------------------------- /autoload/codegeex.vim: -------------------------------------------------------------------------------- 1 | let s:tianqiKeyFile = expand("~/.tianqi.key") 2 | 3 | let s:file_types = { 4 | \ 'cpp': 'C++', 5 | \ 'c': 'C', 6 | \ 'csharp': 'C#', 7 | \ 'cuda': 'Cuda', 8 | \ 'objc': 'Objective-C', 9 | \ 'objcpp': 'Objective-C++', 10 | \ 'python': 'Python', 11 | \ 'java': 'Java', 12 | \ 'scala': 'Scala', 13 | \ 'tex': 'TeX', 14 | \ 'html': 'HTML', 15 | \ 'php': 'PHP', 16 | \ 'javascript': 'JavaScript', 17 | \ 'typescript': 'TypeScript', 18 | \ 'go': 'Go', 19 | \ 'sh': 'Shell', 20 | \ 'rust': 'Rust', 21 | \ 'css': 'CSS', 22 | \ 'sql': 'SQL', 23 | \ 'kotlin': 'Kotlin', 24 | \ 'pascal': 'Pascal', 25 | \ 'r': 'R', 26 | \ 'fortran': 'Fortran', 27 | \ 'lean': 'Lean', 28 | \ } 29 | 30 | function! codegeex#getLanguage() 31 | let file_type = &filetype 32 | return get(s:file_types, file_type, 'Unknown') 33 | endfunction 34 | 35 | function! codegeex#SetTianqiKeyFile(keyfile) 36 | let s:tianqiKeyFile = a:keyfile 37 | endfunction 38 | 39 | function! codegeex#GenCode(ins=0) 40 | let lang = codegeex#getLanguage() 41 | if lang == "Unknown" 42 | echo "Unsupported Language" 43 | return 44 | endif 45 | echo "Codegeex is thinking ..." 46 | let back = @a 47 | normal! "aygg 48 | let code = @a 49 | let result = system("python3 ~/.local/share/nvim/plugged/codegeex-vim/autoload/codegeex_gen.py " . lang . " " . shellescape(code) . " --keyfile " . shellescape(s:tianqiKeyFile)) 50 | let @a = result 51 | set paste 52 | call feedkeys("\a\a\") 53 | set nopaste 54 | if a:ins 55 | call feedkeys('a') 56 | endif 57 | endfunction 58 | 59 | function! codegeex#chooseDstLang(lst) 60 | let choices = [] 61 | let i = 0 62 | for l in a:lst 63 | call add(choices, i . ' ' . l) 64 | let i = i + 1 65 | endfor 66 | return inputlist(choices) 67 | endfunction 68 | 69 | 70 | function! codegeex#getVisualText() 71 | let save_reg = @a 72 | normal! gv"ay 73 | let text = @a 74 | let @a = save_reg 75 | return text 76 | endfunction 77 | 78 | 79 | function! codegeex#openWindow(tp) 80 | let buf_name = '__Codegeex_'. a:tp .'__' 81 | let index = bufnr(buf_name) 82 | let new = 0 83 | if index == -1 84 | vsplit 85 | wincmd L 86 | enew 87 | execute 'file ' . buf_name 88 | setlocal noswapfile 89 | setlocal hidden 90 | setlocal buftype=nofile 91 | execute 'setlocal filetype=' . a:tp 92 | else 93 | if index(tabpagebuflist(), index) == -1 94 | vsplit 95 | wincmd L 96 | execute 'buffer' index 97 | else 98 | call win_gotoid(win_findbuf(index)[0]) 99 | endif 100 | endif 101 | endfunction 102 | 103 | 104 | function! codegeex#handleTransResult(tp,d) 105 | call codegeex#openWindow(a:tp) 106 | let result = a:d 107 | call append(line('$'), result) 108 | call append(line('$'), '') 109 | endfunction 110 | 111 | function! codegeex#TransCode() 112 | let src_lang = codegeex#getLanguage() 113 | if src_lang == "Unknown" 114 | echo "Unsupported Language" 115 | return 116 | endif 117 | 118 | let ft_lst = [] 119 | let lst = [] 120 | let i = 0 121 | for [k,v] in items(s:file_types) 122 | call add(ft_lst, k) 123 | call add(lst, v) 124 | let i = i + 1 125 | endfor 126 | 127 | let index = codegeex#chooseDstLang(lst) 128 | if index < 0 129 | return 130 | endif 131 | 132 | let dst_lang = lst[index] 133 | 134 | if src_lang == dst_lang 135 | echom "Src language is same as dest language" 136 | return 137 | endif 138 | 139 | echo "Codegeex is thinking ..." 140 | let code = codegeex#getVisualText() 141 | let cmd = "python3 ~/.local/share/nvim/plugged/codegeex-vim/autoload/codegeex_trans.py " . src_lang . " " . dst_lang . " " . shellescape(code) 142 | call jobstart(cmd, {"on_stdout": {j,d,e->codegeex#handleTransResult(ft_lst[index], d)}, 'stadout_buffered':1}) 143 | endfunction 144 | 145 | 146 | " example config: 147 | " 148 | " inoremap :call codegeex#GenCode(1) 149 | " nnoremap :call codegeex#GenCode(0) 150 | " vnoremap :call codegeex#TransCode() 151 | --------------------------------------------------------------------------------