-
Notifications
You must be signed in to change notification settings - Fork 7
Expand file tree
/
Copy pathCodeExplain.py
More file actions
70 lines (62 loc) · 2.76 KB
/
CodeExplain.py
File metadata and controls
70 lines (62 loc) · 2.76 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
import pynvim
from pathlib import Path
from pygments.util import ClassNotFound
from pygments.lexers import guess_lexer_for_filename, TextLexer
from langchain.llms import LlamaCpp
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
class IA():
def __init__(self,):
LLAMA_EMBEDDINGS_MODEL = str(Path.home()) +'/.codeexplain/model.bin'
MODEL_N_CTX = 1000
CALLBACKS = [StreamingStdOutCallbackHandler()]
PROMPT_TEMPLATE = """This is a piece of code written in {language}. Provide a simple walkthrough of its operation and highlight any sections of the code that could potentially lead to security risks or inefficiencies. If you don't find any security risk or inefficiencies just don't mention it.
```
{code}
```
answer:
"""
PROMPT = PromptTemplate(template=PROMPT_TEMPLATE, input_variables=["language","code"])
LLM = LlamaCpp(model_path=LLAMA_EMBEDDINGS_MODEL, n_ctx=MODEL_N_CTX, verbose=False)
self.CHAIN = LLMChain(llm=LLM, prompt=PROMPT)
def run(self, input):
explained = self.CHAIN.run(input)
lines = explained.split('\n')
return lines
@pynvim.plugin
class CodeExplain(object):
def __init__(self, nvim):
self.nvim = nvim
self.codeExplainAI = IA()
self.nvim.command("echom \"My plugin is being executed\"")
def getSelectedText(self,):
begin = self.nvim.eval("line(\"'<\")")
end = self.nvim.eval("line(\"'>\")")
lines = self.nvim.current.buffer[begin - 1:end]
return '\n'.join(lines)
def getProgrammingLanguage(self):
file_name = self.nvim.current.buffer.name
file_content = '\n'.join(self.nvim.current.buffer[:])
try:
lexer = guess_lexer_for_filename(file_name, file_content)
except ClassNotFound:
lexer = TextLexer()
return lexer.name
def createWindowBuffer(self,lines):
bufnr = self.nvim.api.create_buf(False, True)
winnr = self.nvim.api.open_win(bufnr, True, {
'relative': 'editor',
'width': 80,
'height': 10,
'row': 10,
'col': 10
})
self.nvim.api.buf_set_lines(bufnr, 0, -1, True, lines)
@pynvim.command('CodeExplain', nargs='*',range=True, sync=True)
def codeExplain(self,args,range):
selected_text = self.getSelectedText()
programming_language = self.getProgrammingLanguage() or "programming"
explained = self.codeExplainAI.run({"language": programming_language,"code":selected_text})
lines = [self.nvim.funcs.escape(line, '\"\\') for line in explained]
self.createWindowBuffer(explained)