diff options
Diffstat (limited to '')
| -rw-r--r-- | py/chat.py | 3 | ||||
| -rw-r--r-- | py/complete.py | 3 | ||||
| -rw-r--r-- | py/config.py | 105 | ||||
| -rw-r--r-- | py/utils.py | 108 | ||||
| -rw-r--r-- | pytest.ini | 4 |
5 files changed, 139 insertions, 84 deletions
@@ -4,7 +4,8 @@ import vim plugin_root = vim.eval("s:plugin_root") vim.command(f"py3file {plugin_root}/py/utils.py") -prompt, config = load_config_and_prompt('chat') +prompt = make_prompt(vim.eval("l:prompt"), vim.eval("l:role_prompt")) +config = make_config(vim.eval("l:config")) config_options = config['options'] config_ui = config['ui'] diff --git a/py/complete.py b/py/complete.py index 8d85581..9b63b0b 100644 --- a/py/complete.py +++ b/py/complete.py @@ -4,7 +4,8 @@ import vim plugin_root = vim.eval("s:plugin_root") vim.command(f"py3file {plugin_root}/py/utils.py") -prompt, config = load_config_and_prompt('complete') +prompt = make_prompt(vim.eval("l:prompt"), vim.eval("l:role_prompt")) +config = make_config(vim.eval("l:config")) config_options = config['options'] config_ui = config['ui'] diff --git a/py/config.py b/py/config.py new file mode 100644 index 0000000..03eb3ca --- /dev/null +++ b/py/config.py @@ -0,0 +1,105 @@ +import vim +import re +import os +import configparser + +def merge_deep_recursive(target, source = {}): + source = source.copy() + for key, value in source.items(): + if isinstance(value, dict): + target_child = target.setdefault(key, {}) + merge_deep_recursive(target_child, value) + else: + target[key] = value + return target + +def merge_deep(objects): + result = {} + for o in objects: + merge_deep_recursive(result, o) + return result + +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) + roles = dict(roles) + + enhance_roles_with_custom_function(roles) + + if not role in roles: + raise Exception(f"Role `{role}` not found") + + options = roles.get(f"{role}.options", {}) + options_complete = roles.get(f"{role}.options-complete", {}) + options_chat = roles.get(f"{role}.options-chat", {}) + + ui = roles.get(f"{role}.ui", {}) + ui_complete = roles.get(f"{role}.ui-complete", {}) + ui_chat = roles.get(f"{role}.ui-chat", {}) + + return { + 'role': dict(roles[role]), + 'config_default': { + 'options': dict(options), + 'ui': dict(ui), + }, + 'config_complete': { + 'options': dict(options_complete), + 'ui': dict(ui_complete), + }, + 'config_chat': { + 'options': dict(options_chat), + 'ui': dict(ui_chat), + }, + } + +def parse_role_names(prompt): + chunks = re.split(r'[ :]+', prompt) + roles = [] + for chunk in chunks: + if not chunk.startswith("/"): + break + roles.append(chunk) + return [raw_role[1:] for raw_role in roles] + +def parse_prompt_and_role_config(instruction, command_type): + instruction = instruction.strip() + roles = parse_role_names(instruction) + if not roles: + # does not require role + return ('', {}) + + last_role = roles[-1] + role_configs = merge_deep([load_role_config(role) for role in roles]) + config = merge_deep([role_configs['config_default'], role_configs['config_' + command_type]]) + role_prompt = role_configs['role'].get('prompt', '') + return role_prompt, config + +def make_config(input_var, output_var): + input_options = vim.eval(input_var) + config_default = input_options['config_default'] + config_extension = input_options['config_extension'] + instruction = input_options['instruction'] + command_type = input_options['command_type'] + + role_prompt, role_config = parse_prompt_and_role_config(instruction, command_type) + + final_config = merge_deep([config_default, config_extension, role_config]) + + output = {} + output['config'] = final_config + output['role_prompt'] = role_prompt + vim.command(f'let {output_var}={output}') + return output diff --git a/py/utils.py b/py/utils.py index d1855a7..fc888ab 100644 --- a/py/utils.py +++ b/py/utils.py @@ -44,21 +44,36 @@ def load_api_key(config_token_file_path): return (api_key, org_id) -def load_config_and_prompt(command_type): - prompt, role_options = parse_prompt_and_role(vim.eval("l:prompt")) - config = vim.eval("l:config") - config['options'] = { - **normalize_options(config['options']), - **normalize_options(role_options['options_default']), - **normalize_options(role_options['options_' + command_type]), - } - return prompt, config +def strip_roles(prompt): + chunks = re.split(r'[ :]+', prompt) + roles = [] + for chunk in chunks: + if not chunk.startswith("/"): + break + roles.append(chunk) + if not roles: + return prompt + last_role = roles[-1] + return prompt[prompt.index(last_role) + len(last_role):].strip() -def normalize_options(options): +def make_prompt(raw_prompt, role_prompt): + prompt = raw_prompt.strip() + prompt = strip_roles(prompt) + + if not role_prompt: + return prompt + + delim = '' if prompt.startswith(':') else ':\n' + prompt = role_prompt + delim + prompt + + return prompt + +def make_config(config): + options = config['options'] # initial prompt can be both a string and a list of strings, normalize it to list if 'initial_prompt' in options and isinstance(options['initial_prompt'], str): options['initial_prompt'] = options['initial_prompt'].split('\n') - return options + return config def make_openai_options(options): max_tokens = int(options['max_tokens']) @@ -302,77 +317,6 @@ def enhance_roles_with_custom_function(roles): 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_roles(prompt): - chunks = re.split(r'[ :]+', prompt) - roles = [] - for chunk in chunks: - if not chunk.startswith("/"): - break - roles.append(chunk) - return [raw_role[1:] for raw_role in roles] - -def merge_role_configs(configs): - merged_options = empty_role_options - merged_role = {} - for config in configs: - options = config['options'] - merged_options = { - 'options_default': { **merged_options['options_default'], **options['options_default'] }, - 'options_complete': { **merged_options['options_complete'], **options['options_complete'] }, - 'options_chat': { **merged_options['options_chat'], **options['options_chat'] }, - } - merged_role ={ **merged_role, **config['role'] } - return { 'role': merged_role, 'options': merged_options } - -def parse_prompt_and_role(raw_prompt): - prompt = raw_prompt.strip() - roles = parse_roles(prompt) - if not roles: - # does not require role - return (prompt, empty_role_options) - - last_role = roles[-1] - prompt = prompt[prompt.index(last_role) + len(last_role):].strip() - - role_configs = [load_role_config(role) for role in roles] - config = merge_role_configs(role_configs) - 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']) - def make_chat_text_chunks(messages, config_options): openai_options = make_openai_options(config_options) http_options = make_http_options(config_options) diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..cc50b10 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,4 @@ +[pytest] +pythonpath = + ./py + ./tests/mocks |