populate options in chat

3 files changed, 60 insertions, 38 deletions

diff --git a/py/chat.py b/py/chat.py
index 5714e63..2983701 100644
--- a/py/chat.py
+++ b/py/chat.py
@@ -4,41 +4,64 @@ import openai
 plugin_root = vim.eval("s:plugin_root")
 vim.command(f"py3file {plugin_root}/py/utils.py")
 
+config_options = vim.eval("options")
+config_ui = vim.eval("ui")
+
 openai.api_key = load_api_key()
 
-lines = vim.eval('getline(1, "$")')
-contains_user_prompt = '>>> user' in lines
-if not contains_user_prompt:
-    # user role not found, put whole file content as an user prompt
-    vim.command("normal! ggO>>> user\n")
-    vim.command("normal! G")
-    vim.command("let &ul=&ul") # breaks undo sequence (https://vi.stackexchange.com/a/29087)
-    vim.command("redraw")
-
-options_chat = {}
-lines = vim.eval('getline(1, "$")')
-contains_chat_options = '[chat-options]' in lines
-if contains_chat_options:
-    # parse options that are defined in the chat header
-    options_index = lines.index('[chat-options]')
-    for line in lines[options_index + 1:]:
-        if line.startswith('#'):
-            # ignore comments
-            continue
-        if line.strip() == '':
-            # stop at the end of the region
-            break
-        (key, value) = line.split('=')
-        options_chat[key.strip()] = value.strip()
-
-options = make_options(options_chat)
+def initialize_chat_window():
+    lines = vim.eval('getline(1, "$")')
+    contains_user_prompt = '>>> user' in lines
+    if not contains_user_prompt:
+        # user role not found, put whole file content as an user prompt
+        vim.command("normal! gg")
+        populates_options = config_ui['populate_options'] == '1'
+        if populates_options:
+            vim.command("normal! O[chat-options]")
+            vim.command("normal! o")
+            for key, value in config_options.items():
+                if key == 'initial_prompt':
+                    value = "\\n".join(value)
+                vim.command("normal! i" + key + "=" + value + "\n")
+        vim.command("normal! " + ("o" if populates_options else "O"))
+        vim.command("normal! i>>> user\n")
+        vim.command("normal! G")
+        vim_break_undo_sequence()
+        vim.command("redraw")
+
+def parse_chat_header_options():
+    options = {}
+    lines = vim.eval('getline(1, "$")')
+    contains_chat_options = '[chat-options]' in lines
+    if contains_chat_options:
+        # parse options that are defined in the chat header
+        options_index = lines.index('[chat-options]')
+        for line in lines[options_index + 1:]:
+            if line.startswith('#'):
+                # ignore comments
+                continue
+            if line == '':
+                # stop at the end of the region
+                break
+            (key, value) = line.strip().split('=')
+            if key == 'initial_prompt':
+                value = value.split('\\n')
+            options[key] = value
+    return options
+
+initialize_chat_window()
+
+chat_options = parse_chat_header_options()
+options = {**config_options, **chat_options}
 request_options = make_request_options(options)
 
-initial_prompt = options.get('initial_prompt', [])
-initial_prompt = '\n'.join(initial_prompt)
-file_content = vim.eval('trim(join(getline(1, "$"), "\n"))')
-chat_content = f"{initial_prompt}\n{file_content}"
-messages = parse_chat_messages(chat_content)
+initial_prompt = '\n'.join(options.get('initial_prompt', []))
+initial_messages = parse_chat_messages(initial_prompt)
+
+chat_content = vim.eval('trim(join(getline(1, "$"), "\n"))')
+chat_messages = parse_chat_messages(chat_content)
+
+messages = initial_messages + chat_messages
 
 try:
     if messages[-1]["content"].strip():
diff --git a/py/complete.py b/py/complete.py
index cdf8da9..b0b0dc5 100644
--- a/py/complete.py
+++ b/py/complete.py
@@ -5,8 +5,8 @@ plugin_root = vim.eval("s:plugin_root")
 vim.command(f"py3file {plugin_root}/py/utils.py")
 
 engine = vim.eval("engine")
-options = make_options()
-request_options = make_request_options(options)
+config_options = vim.eval("options")
+request_options = make_request_options(config_options)
 
 prompt = vim.eval("prompt").strip()
 
@@ -18,7 +18,7 @@ def complete_engine(prompt):
     return text_chunks
 
 def chat_engine(prompt):
-    initial_prompt = options.get('initial_prompt', [])
+    initial_prompt = config_options.get('initial_prompt', [])
     initial_prompt = '\n'.join(initial_prompt)
     chat_content = f"{initial_prompt}\n\n>>> user\n\n{prompt}".strip()
     messages = parse_chat_messages(chat_content)
diff --git a/py/utils.py b/py/utils.py
index bcc5309..dee7d42 100644
--- a/py/utils.py
+++ b/py/utils.py
@@ -11,10 +11,6 @@ def load_api_key():
         pass
     return api_key.strip()
 
-def make_options(options_custom = {}):
-    options_config = vim.eval("options")
-    return {**options_config, **options_custom}
-
 def make_request_options(options):
     request_options = {}
     request_options['model'] = options['model']
@@ -55,3 +51,6 @@ def parse_chat_messages(chat_content):
 
     return messages
 
+def vim_break_undo_sequence():
+    # breaks undo sequence (https://vi.stackexchange.com/a/29087)
+    vim.command("let &ul=&ul")