FilesExpand file tree

main

/

CodeExplain.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

70 lines (62 loc) · 2.76 KB

main

/

CodeExplain.py

Copy path

Top

File metadata and controls

• Code
• Blame

70 lines (62 loc) · 2.76 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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)