Merge pull request #85 from madox2/custom-roles

Custom roles

1 files changed, 57 insertions, 0 deletions

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'])