diff options
Diffstat (limited to 'py')
| -rw-r--r-- | py/chat.py | 41 | ||||
| -rw-r--r-- | py/complete.py | 35 | ||||
| -rw-r--r-- | py/utils.py | 37 |
3 files changed, 63 insertions, 50 deletions
@@ -16,26 +16,8 @@ prompt = f"{initial_prompt}\n{file_content}" lines = prompt.splitlines() messages = [] -def parse_messages(): - file_content = vim.eval('trim(join(getline(1, "$"), "\n"))') - lines = file_content.splitlines() - messages = [] - for line in lines: - if line.startswith(">>> system"): - messages.append({"role": "system", "content": ""}) - continue - if line.startswith(">>> user"): - messages.append({"role": "user", "content": ""}) - continue - if line.startswith("<<< assistant"): - messages.append({"role": "assistant", "content": ""}) - continue - if not messages: - continue - messages[-1]["content"] += "\n" + line - return messages - -messages = parse_messages() +chat_content = vim.eval('trim(join(getline(1, "$"), "\n"))') +messages = parse_chat_messages(chat_content) if not messages: # roles not found, put whole file content as an user prompt @@ -43,11 +25,8 @@ if not messages: vim.command("normal! G") vim.command("let &ul=&ul") # breaks undo sequence (https://vi.stackexchange.com/a/29087) vim.command("redraw") - messages = parse_messages() - -for message in messages: - # strip newlines from the content as it causes empty responses - message["content"] = message["content"].strip() + chat_content = vim.eval('trim(join(getline(1, "$"), "\n"))') + messages = parse_chat_messages(chat_content) try: if messages[-1]["content"].strip(): @@ -58,16 +37,8 @@ try: vim.command("redraw") response = openai.ChatCompletion.create(messages=messages, stream=True, **request_options) - - generating_text = False - for resp in response: - text = resp['choices'][0]['delta'].get('content', '') - if not text.strip() and not generating_text: - continue # trim newlines from the beginning - - generating_text = True - vim.command("normal! a" + text) - vim.command("redraw") + text_chunks = map(lambda resp: resp['choices'][0]['delta'].get('content', ''), response) + render_text_chunks(text_chunks) vim.command("normal! a\n\n>>> user\n\n") vim.command("redraw") diff --git a/py/complete.py b/py/complete.py index bddd8e6..96ae032 100644 --- a/py/complete.py +++ b/py/complete.py @@ -4,28 +4,35 @@ import openai plugin_root = vim.eval("s:plugin_root") vim.command(f"py3file {plugin_root}/py/utils.py") -prompt = vim.eval("prompt").strip() +engine = vim.eval("engine") +options = make_options() request_options = make_request_options() +prompt = vim.eval("prompt").strip() + openai.api_key = load_api_key() +def complete_engine(): + response = openai.Completion.create(stream=True, prompt=prompt, **request_options) + text_chunks = map(lambda resp: resp['choices'][0].get('text', ''), response) + return text_chunks + +def chat_engine(): + initial_prompt = options.get('initial_prompt', '') + chat_content = f"{initial_prompt}\n\n>>> user\n\n{prompt}".strip() + messages = parse_chat_messages(chat_content) + response = openai.ChatCompletion.create(messages=messages, stream=True, **request_options) + text_chunks = map(lambda resp: resp['choices'][0]['delta'].get('content', ''), response) + return text_chunks + +engines = {"chat": chat_engine, "complete": complete_engine} + try: if prompt: - print('Completing...') vim.command("redraw") - - response = openai.Completion.create(stream=True, prompt=prompt, **request_options) - - generating_text = False - for resp in response: - text = resp['choices'][0].get('text', '') - if not text.strip() and not generating_text: - continue # trim newlines from the beginning - - generating_text = True - vim.command("normal! a" + text) - vim.command("redraw") + text_chunks = engines[engine]() + render_text_chunks(text_chunks) except KeyboardInterrupt: vim.command("normal! a Ctrl-C...") except openai.error.Timeout: diff --git a/py/utils.py b/py/utils.py index 02ecfef..68a4ee9 100644 --- a/py/utils.py +++ b/py/utils.py @@ -11,8 +11,11 @@ def load_api_key(): pass return api_key.strip() +def make_options(): + return vim.eval("options") + def make_request_options(): - options = vim.eval("options") + options = make_options() request_options = {} request_options['model'] = options['model'] request_options['max_tokens'] = int(options['max_tokens']) @@ -20,3 +23,35 @@ def make_request_options(): request_options['request_timeout'] = float(options['request_timeout']) return request_options +def render_text_chunks(chunks): + generating_text = False + for text in chunks: + if not text.strip() and not generating_text: + continue # trim newlines from the beginning + generating_text = True + vim.command("normal! a" + text) + vim.command("redraw") + +def parse_chat_messages(chat_content): + lines = chat_content.splitlines() + messages = [] + for line in lines: + if line.startswith(">>> system"): + messages.append({"role": "system", "content": ""}) + continue + if line.startswith(">>> user"): + messages.append({"role": "user", "content": ""}) + continue + if line.startswith("<<< assistant"): + messages.append({"role": "assistant", "content": ""}) + continue + if not messages: + continue + messages[-1]["content"] += "\n" + line + + for message in messages: + # strip newlines from the content as it causes empty responses + message["content"] = message["content"].strip() + + return messages + |