From 3e4a37c3678a97d8761e0dbef6d3500d008cc95f Mon Sep 17 00:00:00 2001 From: Yaroslav Yashin Date: Sat, 8 Feb 2025 21:06:08 +0100 Subject: [PATCH] feat: add function handler for performing various text operations in Sublime Text --- plugins/function_handler.py | 122 ++++++++++++++++++++++++++++++++++++ plugins/openai_base.py | 16 ++++- 2 files changed, 136 insertions(+), 2 deletions(-) create mode 100644 plugins/function_handler.py diff --git a/plugins/function_handler.py b/plugins/function_handler.py new file mode 100644 index 0000000..e2f9489 --- /dev/null +++ b/plugins/function_handler.py @@ -0,0 +1,122 @@ +from __future__ import annotations + +import logging +from enum import Enum +from json import dumps, loads +from typing import Dict + +from sublime import Region, Window + +from .project_structure import build_folder_structure + +# FIXME: logger prints nothing from within rust context +logger = logging.getLogger(__name__) + + +# TODO: This should be deleted in favor to rust type +class Function(str, Enum): + replace_text_with_another_text = 'replace_text_with_another_text' + replace_text_for_whole_file = 'replace_text_for_whole_file' + read_region_content = 'read_region_content' + get_working_directory_content = 'get_working_directory_content' + + +class FunctionHandler: + @staticmethod + def perform_function(func_name: str, args: str, window: Window) -> str: + args_json = loads(args) + logger.debug(f'executing: {func_name}') + if func_name == Function.replace_text_with_another_text.value: + path = args_json.get('file_path') + old_content = args_json.get('old_content') + new_content = args_json.get('new_content') + + if ( + path + and isinstance(path, str) + and old_content + and isinstance(old_content, str) + and new_content + and isinstance(new_content, str) + ): + view = window.find_open_file(path) + if view: + escaped_string = ( + old_content.replace('(', r'\(') + .replace(')', r'\)') + .replace('[', r'\[') + .replace(']', r'\]') + .replace('{', r'\{') + .replace('}', r'\}') + .replace('|', r'\|') + .replace('"', r'\"') + .replace('\\', r'\\\\') + ) + region = view.find(pattern=escaped_string, start_pt=0) + logger.debug(f'region {region}') + serializable_region = { + 'a': region.begin(), + 'b': region.end(), + } + if ( + region.begin() == region.end() == -1 or region.begin() == region.end() == 0 + ): # means search found nothing + return f'Text not found: {old_content}' + else: + view.run_command( + 'replace_region', + {'region': serializable_region, 'text': new_content}, + ) + return dumps(serializable_region) + else: + return f'File under path not found: {path}' + else: + return f'Wrong attributes passed: {path}, {old_content}, {new_content}' + + elif func_name == Function.replace_text_for_whole_file.value: + path = args_json.get('file_path') + create = args_json.get('create') + content = args_json.get('content') + if path and isinstance(path, str) and content and isinstance(content, str): + if isinstance(create, bool): + window.open_file(path) + view = window.find_open_file(path) + if view: + region = Region(0, len(view)) + view.run_command( + 'replace_region', + {'region': {'a': region.begin(), 'b': region.end()}, 'text': content}, + ) + text = view.substr(region) + return dumps({'result': text}) + else: + return f'File under path not found: {path}' + else: + return f'Wrong attributes passed: {path}, {content}' + + elif func_name == Function.read_region_content.value: + path = args_json.get('file_path') + region = args_json.get('region') + if path and isinstance(path, str) and region and isinstance(region, Dict): + view = window.find_open_file(path) + if view: + a_reg: int = region.get('a') if region.get('a') != -1 else 0 # type: ignore + b_reg = region.get('b') if region.get('b') != -1 else len(view) + region_ = Region(a=a_reg, b=b_reg) + text = view.substr(region_) + return dumps({'content': f'{text}'}) + else: + return f'File under path not found: {path}' + else: + return f'Wrong attributes passed: {path}, {region}' + + elif func_name == Function.get_working_directory_content.value: + path = args_json.get('directory_path') + if path and isinstance(path, str): + folder_structure = build_folder_structure(path) + + return dumps({'content': f'{folder_structure}'}) + else: + return f'Wrong attributes passed: {path}' + else: + return f"Called function doen't exists: {func_name}" diff --git a/plugins/openai_base.py b/plugins/openai_base.py index c5e0beb..df74d8a 100644 --- a/plugins/openai_base.py +++ b/plugins/openai_base.py @@ -11,15 +11,16 @@ SublimeInputContent, # type: ignore Worker, # type: ignore ) -from sublime import Region, Settings, Sheet, View, Window, active_window +from sublime import Region, Settings, Sheet, View, Window, active_window, windows from .assistant_settings import ( CommandMode, ) from .buffer import BufferContentManager from .errors.OpenAIException import WrongUserInputException, present_error, present_error_str +from .function_handler import FunctionHandler from .image_handler import ImageValidator -from .load_model import get_cache_path, get_model_or_default +from .load_model import get_cache_path from .output_panel import SharedOutputPanelListener from .phantom_streamer import PhantomStreamer from .response_manager import ResponseManager @@ -197,6 +198,8 @@ def on_input( else ViewCapture(view).tab_handler ) + fn_handler = FunctionCapture(window).fn_handler + cls.worker.run( view.id(), assistant.output_mode, @@ -204,6 +207,7 @@ def on_input( assistant, handler, ErrorCapture.error_handler, + fn_handler, ) if assistant.output_mode == PromptMode.View: @@ -273,6 +277,14 @@ def error_handler(content: str) -> None: present_error_str('OpenAI Error', content) +class FunctionCapture: + def __init__(self, window: Window) -> None: + self.window = window + + def fn_handler(self, name: str, args: str) -> str: + return FunctionHandler.perform_function(name, args, self.window) + + class ViewCapture: def __init__(self, view: View) -> None: self.view = view