Merge pull request #85 from madox2/custom-roles

Custom roles

4 files changed, 95 insertions, 6 deletions

diff --git a/py/chat.py b/py/chat.py
index ff70904..67f2d5d 100644
--- a/py/chat.py
+++ b/py/chat.py
@@ -4,10 +4,14 @@ import vim
 plugin_root = vim.eval("s:plugin_root")
 vim.command(f"py3file {plugin_root}/py/utils.py")
 
+prompt, role_options = parse_prompt_and_role(vim.eval("l:prompt"))
 config = normalize_config(vim.eval("l:config"))
-config_options = config['options']
+config_options = {
+    **config['options'],
+    **role_options['options_default'],
+    **role_options['options_chat'],
+}
 config_ui = config['ui']
-prompt = vim.eval("l:prompt").strip()
 
 def initialize_chat_window():
     lines = vim.eval('getline(1, "$")')
@@ -38,7 +42,7 @@ def initialize_chat_window():
         vim.command("normal! i\n>>> user\n\n")
 
     if prompt:
-        vim.command("normal! a" + prompt)
+        vim.command("normal! i" + prompt)
         vim_break_undo_sequence()
         vim.command("redraw")
 
@@ -72,7 +76,7 @@ try:
             **openai_options
         }
         printDebug("[chat] request: {}", request)
-        url = config_options['endpoint_url']
+        url = options['endpoint_url']
         response = openai_request(url, request, http_options)
         def map_chunk(resp):
             printDebug("[chat] response: {}", resp)
diff --git a/py/complete.py b/py/complete.py
index debe275..f340e96 100644
--- a/py/complete.py
+++ b/py/complete.py
@@ -6,11 +6,16 @@ vim.command(f"py3file {plugin_root}/py/utils.py")
 
 config = normalize_config(vim.eval("l:config"))
 engine = config['engine']
-config_options = config['options']
+
+prompt, role_options = parse_prompt_and_role(vim.eval("l:prompt"))
+config_options = {
+    **config['options'],
+    **role_options['options_default'],
+    **role_options['options_complete'],
+}
 openai_options = make_openai_options(config_options)
 http_options = make_http_options(config_options)
 
-prompt = vim.eval("l:prompt").strip()
 is_selection = vim.eval("l:is_selection")
 
 def complete_engine(prompt):
diff --git a/py/roles.py b/py/roles.py
new file mode 100644
index 0000000..b7d19e1
--- /dev/null
+++ b/py/roles.py
@@ -0,0 +1,23 @@
+import vim
+
+# import utils
+plugin_root = vim.eval("s:plugin_root")
+vim.command(f"py3file {plugin_root}/py/utils.py")
+
+roles_config_path = os.path.expanduser(vim.eval("g:vim_ai_roles_config_file"))
+if not os.path.exists(roles_config_path):
+    raise Exception(f"Role config file does not exist: {roles_config_path}")
+
+roles = configparser.ConfigParser()
+roles.read(roles_config_path)
+
+enhance_roles_with_custom_function(roles)
+
+role_names = [name for name in roles.sections() if not '.' in name]
+
+role_list = [f'"{name}"' for name in role_names]
+role_list = ", ".join(role_list)
+
+role_list = f"[{role_list}]"
+
+vim.command(f'let l:role_list = {role_list}')
diff --git a/py/utils.py b/py/utils.py
index 19de7c4..963377d 100644
--- a/py/utils.py
+++ b/py/utils.py
@@ -11,6 +11,7 @@ import re
 from urllib.error import URLError
 from urllib.error import HTTPError
 import traceback
+import configparser
 
 is_debugging = vim.eval("g:vim_ai_debug") == "1"
 debug_log_file = vim.eval("g:vim_ai_debug_log_file")
@@ -261,3 +262,59 @@ def handle_completion_error(error):
 def clear_echo_message():
     # https://neovim.discourse.group/t/how-to-clear-the-echo-message-in-the-command-line/268/3
     vim.command("call feedkeys(':','nx')")
+
+def enhance_roles_with_custom_function(roles):
+    if vim.eval("exists('g:vim_ai_roles_config_function')") == '1':
+        roles_config_function = vim.eval("g:vim_ai_roles_config_function")
+        if not vim.eval("exists('*" + roles_config_function + "')"):
+            raise Exception(f"Role config function does not exist: {roles_config_function}")
+        else:
+            roles.update(vim.eval(roles_config_function + "()"))
+
+def load_role_config(role):
+    roles_config_path = os.path.expanduser(vim.eval("g:vim_ai_roles_config_file"))
+    if not os.path.exists(roles_config_path):
+        raise Exception(f"Role config file does not exist: {roles_config_path}")
+
+    roles = configparser.ConfigParser()
+    roles.read(roles_config_path)
+
+    enhance_roles_with_custom_function(roles)
+
+    if not role in roles:
+        raise Exception(f"Role `{role}` not found")
+
+    options = roles[f"{role}.options"] if f"{role}.options" in roles else {}
+    options_complete =roles[f"{role}.options-complete"] if f"{role}.options-complete" in roles else {}
+    options_chat = roles[f"{role}.options-chat"] if f"{role}.options-chat" in roles else {}
+
+    return {
+        'role': dict(roles[role]),
+        'options': {
+            'options_default': dict(options),
+            'options_complete': dict(options_complete),
+            'options_chat': dict(options_chat),
+        },
+    }
+
+empty_role_options = {
+    'options_default': {},
+    'options_complete': {},
+    'options_chat': {},
+}
+
+def parse_prompt_and_role(raw_prompt):
+    prompt = raw_prompt.strip()
+    role = re.split(' |:', prompt)[0]
+    if not role.startswith('/'):
+        # does not require role
+        return (prompt, empty_role_options)
+
+    prompt = prompt[len(role):].strip()
+    role = role[1:]
+
+    config = load_role_config(role)
+    if 'prompt' in config['role'] and config['role']['prompt']:
+        delim = '' if prompt.startswith(':') else ':\n'
+        prompt = config['role']['prompt'] + delim + prompt
+    return (prompt, config['options'])