Merge pull request #96 from devchat-ai/fix_refactor_errors

Refactor command.yml and rewrite.py for Enhanced Readability and Error Handling

comments/main.py

@@ -5,23 +5,198 @@ import sys
 from devchat.ide.service import IDEService
 from devchat.ide.vscode_services import selected_lines, visible_lines
 from devchat.llm import chat
+from devchat.memory import FixSizeChatMemory
+
+PROMPT = """
+file: {file_path}
+```
+{selected_text}
+```
+"""
+
+PROMPT_ZH = """
+文件：{file_path}
+```
+{selected_text}
+```
+代码块中注释请使用中文描述。
+"""
+
+
+MESSAGES_FEW_SHOT = [
+    {
+        "role": "system",
+        "content": """
+You are a code assistant. Your task is to add comments for given code block.
+To add comments, you need to follow the following rules:
+1. don't change any code in code block, even space char;
+2. comment should be placed before the code line, don't append comment to the code line;
+3. keep comments already in code block;
+4. add comments for each code line;
+5. don't change string value in code block;
+
+there are a examples with correct outputs:
+input: 
+file: a1.py
+```
+# print("hello")
+
+print("Hello World") print("Hello World2")
+```
+output:
+```python
+# print("hello")
+
+# print Hello World
+print("Hello World") print("Hello World2")
+```
+In this example, "# print("hello")" is comment line, but we keep it in output.
+"print("Hello World") print("Hello World2")" is an error line, but we keep it in output,
+because output can't change any code, just insert comment lines to code block.
+
+
+here is an error example:
+```
+        // const filepath = document.uri.fsPath;
+        // const fileContent = document.getText();
+        // const posOffset = document.offsetAt(position);
+        // await outputAst(filepath, fileContent, posOffset);
+        // // await testTreesitterQuery(filepath, fileContent);
+        // logger.channel()?.info("Result:", result2);
+        // if (1) {
+        //     return [];
+        // }
+
+        let response: CodeCompleteResult | undefined = undefined;
+```
+output is :
+```typescript
+        // init response
+        let response: CodeCompleteResult | undefined = undefined;
+```
+In this example, comments lines are missed, so this output is bad.
+
+here is an error example:
+```
+    const value = "open file";
+```
+output is :
+```typescript
+    // 设置变量值为打开文件
+    const value = "打开文件";
+```
+In this example, string in code is changed, so this output is bad.
+
+Output should format as code block.
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: a1.py
+```
+    print("hello")
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```python
+    # print hello
+    print("hello")
+```
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: a1.py
+```
+    print("hell\\nworld")
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```python
+    # print hello world
+    print("hell\\nworld")
+```
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: t2.ts
+```
+                this.logEventToServer(
+                {
+                    completion_id: response!.id,
+                    type: "select",
+                    lines: response!.code.split('\\n').length,
+                    length: response!.code.length
+                });
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```typescript
+                // log event to server
+                this.logEventToServer(
+                {
+                    // completion id
+                    completion_id: response!.id,
+                    // type of event
+                    type: "select",
+                    // number of lines in response code
+                    lines: response!.code.split('\\n').length,
+                    // length of response code
+                    length: response!.code.length
+                });
+```
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: t2.ts
+```
+    // if (a == 1) {
+    //     console.log("a is 1");
+    // }
+
+    if (a == 2) {
+        console.log("a is 2");
+    }
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```typescript
+    // if (a == 1) {
+    //     console.log("a is 1");
+    // }
+
+    // check if a is 2
+    if (a == 2) {
+        // log a is 2
+        console.log("a is 2");
+    }
+```
+""",
+    },
+]
 
 
 def get_selected_code():
-    """
-    Retrieves the selected lines of code from the user's selection.
-
-    This function extracts the text selected by the user in their IDE or text editor.
-    If no text has been selected, it prints an error message to stderr and exits the
-    program with a non-zero status indicating failure.
-
-    Returns:
-        dict: A dictionary containing the key 'selectedText' with the selected text
-        as its value. If no text is selected, the program exits.
-    """
+    """Retrieves the selected lines of code from the user's selection."""
     selected_data = IDEService().get_selected_range().dict()
-
-    miss_selected_error = "Please select some text."
     if selected_data["range"]["start"] == selected_data["range"]["end"]:
         readme_path = os.path.join(os.path.dirname(__file__), "README.md")
         if os.path.exists(readme_path):
@@ -29,54 +204,12 @@ def get_selected_code():
                 readme_text = f.read()
                 print(readme_text)
                 sys.exit(0)
-
-        print(miss_selected_error, file=sys.stderr, flush=True)
+        print("Please select some text.", file=sys.stderr, flush=True)
         sys.exit(-1)
-
     return selected_data
 
 
-def get_visible_code():
-    """
-    Retrieves visible code from the visible_lines function.
-
-    Returns:
-    visible_data: The visible code retrieved from the visible_lines function.
-    """
-    visible_data = IDEService().get_visible_range().dict()
-    return visible_data
-
-
-PROMPT = prompt = """
-Your task is:
-Add necessary line comments to the selected lines of code. Please keep in mind to ensure:
-1. Just add comments only for the selected portion of the code, \
-do not modify any thing beyond the selected portion;
-2. Add these comments above the corresponding lines of code;
-3. Output only the selected portion of the code with comments added;
-4. Maintains the same indentation as the selected portion of the code;
-5. Do not show any code beyond the selected portion;
-Following the task requirements, please ensure that the revised code segment \
-maintains the same indentation as the selected code to seamlessly integrate with \
-the existing code structure and maintain correct syntax.
-Here is the relevant context information for your reference:
-1. Selected portion of the code: {selected_text}
-2. Visible portion of the code: {visible_text}
-"""
-
-
-PROMPT_ZH = prompt = """
-你的任务是：
-使用中文给被选中的代码添加必要的注释。
-根据任务要求，仅修改被选中部分的代码。请确保：
-1. 仅修改被选中的代码，可见代码只是作为你理解的 context；
-2. 仅添加注释，代码结构保存不变；
-3. 修改后的代码用 markdown 代码块的格式输出；
-4. 修改后的代码段与所选代码保持相同的缩进，以与现有代码结构无缝集成并保持正确的语法；
-以下是你可以参考的 context 信息：
-1. 编辑器中被选中的代码：{selected_text}
-2. 当前编辑器中可见代码：{visible_text}
-"""
+memory = FixSizeChatMemory(max_size=20, messages=MESSAGES_FEW_SHOT)
 
 
 def get_prompt():
@@ -84,44 +217,62 @@ def get_prompt():
     return PROMPT_ZH if ide_language == "zh" else PROMPT
 
 
-@chat(prompt=get_prompt(), stream_out=True)
+@chat(prompt=get_prompt(), stream_out=True, memory=memory)
 # pylint: disable=unused-argument
-def add_comments(selected_text, visible_text):
-    """
-    call ai to rewrite selected code
-    """
+def add_comments(selected_text, file_path):
+    """Call AI to rewrite selected code"""
     pass  # pylint: disable=unnecessary-pass
 
 
 def extract_markdown_block(text):
-    """
-    Extracts the first Markdown code block from the given text without the language specifier.
-
-    :param text: A string containing Markdown text
-    :return: The content of the first Markdown code block, or None if not found
-    """
+    """Extracts the first Markdown code block from the given text without the language specifier."""
     pattern = r"```(?:\w+)?\s*\n(.*?)\n```"
     match = re.search(pattern, text, re.DOTALL)
-
     if match:
         block_content = match.group(1)
         return block_content
     else:
+        if text.find("```"):
+            return None
         return text
 
 
+def remove_unnecessary_escapes(code_a, code_b):
+    code_copy = code_b  # Create a copy of the original code
+    escape_chars = re.finditer(r"\\(.)", code_b)
+    remove_char_index = []
+    for match in escape_chars:
+        before = code_b[max(0, match.start() - 4) : match.start()]
+        after = code_b[match.start() + 1 : match.start() + 5]
+        substr = before + after
+        if substr in code_a:
+            remove_char_index.append(match.start())
+    remove_char_index.reverse()
+    for index in remove_char_index:
+        code_copy = code_copy[:index] + code_copy[index + 1 :]
+    return code_copy
+
+
 def main():
-    # prepare code
     selected_text = get_selected_code()
-    visible_text = get_visible_code()
+    file_path = selected_text.get("abspath", "")
+    code_text = selected_text.get("text", "")
 
-    # rewrite
-    response = add_comments(selected_text=selected_text, visible_text=visible_text)
-
-    # apply new code to editor
+    response = add_comments(selected_text=code_text, file_path=file_path)
     new_code = extract_markdown_block(response)
-    IDEService().diff_apply("", new_code)
 
+    if not new_code:
+        ide_lang = IDEService().ide_language()
+        error_msg = (
+            "\n\nThe output of the LLM is incomplete and cannot perform code operations.\n\n"
+            if ide_lang != "zh"
+            else "\n\n大模型输出不完整，不能进行代码操作。\n\n"
+        )
+        print(error_msg)
+        sys.exit(0)
+
+    new_code = remove_unnecessary_escapes(code_text, new_code)
+    IDEService().diff_apply("", new_code)
     sys.exit(0)
 
 

docstring/main.py

@@ -5,6 +5,144 @@ import sys
 from devchat.ide.service import IDEService
 from devchat.ide.vscode_services import selected_lines, visible_lines
 from devchat.llm import chat
+from devchat.memory import FixSizeChatMemory
+
+PROMPT = prompt = """
+file: {file_path}
+```
+{selected_text}
+```
+"""
+
+PROMPT_ZH = prompt = """
+文件: {file_path}
+```
+{selected_text}
+```
+
+输出内容使用中文，我的母语为中文。
+"""
+
+
+def get_prompt():
+    ide_language = IDEService().ide_language()
+    return PROMPT_ZH if ide_language == "zh" else PROMPT
+
+
+MESSAGES_A = [
+    {
+        "role": "system",
+        "content": """
+Your task is:
+Write a documentation comment to the selected code. Please pay attention to using \
+standard comment format, such as method comments, please explain parameters and return values. \
+And just add the documents for the selected portion of the code.
+Output documentation comment is format as code block.\
+
+You must follow the following rules:
+1. Output documentation comment in ```comment <documentation comments without code lines> ``` \
+    format.
+2. Different languages have different comment symbols, please choose the correct comment symbol \
+    according to the file name.
+3. You must output ... to indicate the remaining code, output all code block can make more errors.
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: a1.py
+```
+    def print_hello():
+        print("hello")
+        print("world")
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```comment
+    def print_hello():
+        \"\"\"
+        print hello
+
+        Parameters:
+        None
+
+        Returns:
+        None
+        \"\"\"
+        print("hello")
+        ...
+```
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: t1.java
+```
+    public static void main(String[] args) {
+        System.out.println("Hello, World!");
+    }
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```comment
+    /**
+     * The main method is the entry point of the program.
+     * This method prints "Hello, World!" to the console.
+     *
+     * @param args command line arguments (not used in this program)
+     */
+    public static void main(String[] args) {
+        ...
+```
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: t1.py
+```
+def content_to_json(content):
+    try:
+        content_no_block = _try_remove_markdown_block_flag(content)
+        response_obj = json.loads(content_no_block)
+        return response_obj
+    except json.JSONDecodeError as err:
+        raise RetryException(err) from err
+    except Exception as err:
+        raise err
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```comment
+def content_to_json(content):
+    \"\"\"
+    Convert the given content to a JSON object.
+
+    Parameters:
+    content (str): The content to convert.
+
+    Returns:
+    dict: The JSON object.
+
+    Raises:
+    RetryException: If the content cannot be decoded to JSON.
+    \"\"\"
+    try:
+        ...
+```
+""",
+    },
+]
 
 
 def get_selected_code():
@@ -36,49 +174,12 @@ def get_selected_code():
     return selected_data
 
 
-def get_visible_code():
-    """
-    Retrieves visible code from the visible_lines function.
-
-    Returns:
-    visible_data: The visible code retrieved from the visible_lines function.
-    """
-    visible_data = IDEService().get_visible_range().dict()
-    return visible_data
+memory = FixSizeChatMemory(max_size=20, messages=MESSAGES_A)
 
 
-PROMPT = prompt = """
-Your task is:
-Add a documentation comment to the selected code. Please pay attention to using \
-standard comment format, such as method comments, please explain parameters and return values. \
-And just add the documents for the selected portion of the code.
-Following the task requirements, modify only the selected portion of the code. \
-Please ensure that the revised code segment maintains the same indentation as the \
-selected code to seamlessly integrate with the existing code structure and maintain \
-correct syntax. Keep all other information as it is. \
-Here is the relevant context information for your reference:
-1. Selected portion of the code: {selected_text}
-"""
-
-
-PROMPT_ZH = prompt = """
-你的任务是：
-使用中文给被选中的代码添加头部文档注释。请注意使用规范的注释格式，如方法注释请解释参数以及返回值。
-根据任务要求，仅修改被选中部分的代码。请确保修改后的代码段与所选代码保持相同的缩进，\
-以与现有代码结构无缝集成并保持正确的语法。保留所有其他信息不变。
-以下是你可以参考的 context 信息：
-1. 编辑器中被选中的代码：{selected_text}
-"""
-
-
-def get_prompt():
-    ide_language = IDEService().ide_language()
-    return PROMPT_ZH if ide_language == "zh" else PROMPT
-
-
-@chat(prompt=get_prompt(), stream_out=True)
+@chat(prompt=get_prompt(), stream_out=True, memory=memory)
 # pylint: disable=unused-argument
-def add_docstring(selected_text, visible_text):
+def add_docstring(selected_text, file_path):
     """
     call ai to rewrite selected code
     """
@@ -99,23 +200,125 @@ def extract_markdown_block(text):
         block_content = match.group(1)
         return block_content
     else:
+        # whether exist ```language?
+        if text.find("```"):
+            return None
         return text
 
 
+def get_indent_level(text):
+    """
+    Returns the indentation level of the given text.
+
+    :param text: A string containing the text to be analyzed
+    :return: The indentation level of the text, as an integer
+    """
+    indent_level = 0
+    for char in text:
+        if char == " ":
+            indent_level += 1
+        elif char == "\t":
+            indent_level += 4
+        else:
+            break
+    return indent_level
+
+
+def offset_indent_level(text, indent_level):
+    """
+    Offsets the indentation level of the given text by the specified amount.
+
+    :param text: A string containing the text to be modified
+    :param indent_level: The amount by which to offset the indentation level
+    :return: The modified text with the indentation level offset
+    """
+    current_indent = get_indent_level(text)
+    offset_indent = indent_level - current_indent
+    if offset_indent > 0:
+        lines = text.splitlines()
+        new_lines = []
+        for line in lines:
+            new_lines.append(" " * offset_indent + line)
+        text = "\n".join(new_lines)
+    return text
+
+
+def merge_code(selected_text, docstring):
+    user_selected_lines = selected_text.split("\n")
+    docstring_lines = docstring.split("\n")
+
+    user_selected_trim_lines = [line.replace(" ", "").strip() for line in user_selected_lines]
+    docstring_trim_lines = [line.replace(" ", "").strip() for line in docstring_lines]
+
+    # match user_selected_trim_line == docstring_trim_line
+    #    and index_selected_line != index_docstring_line
+    has_match = False
+    for index, line in enumerate(user_selected_trim_lines):
+        for index_doc, line_doc in enumerate(docstring_trim_lines):
+            if line_doc == "..." and has_match:
+                line_doc = line
+                break
+            if line == line_doc:
+                has_match = True
+                break
+        if line != line_doc or index == index_doc:
+            continue
+        return "\n".join(docstring_lines[:index_doc] + user_selected_lines[index:])
+
+    # match with part of code
+    for index, line in enumerate(user_selected_trim_lines):
+        for index_doc, line_doc in enumerate(docstring_trim_lines):
+            if line_doc == "...":
+                break
+            if (
+                line.strip().find(line_doc.strip()) != -1
+                or line_doc.strip().find(line.strip()) != -1
+            ):
+                break
+        if (
+            line.strip().find(line_doc.strip()) == -1 and line_doc.strip().find(line.strip()) == -1
+        ) or index == index_doc:
+            continue
+        return "\n".join(docstring_lines[:index_doc] + user_selected_lines[index:])
+    return docstring + "\n" + selected_text
+
+
 def main():
-    # prepare code
+    # Prepare code
     selected_text = get_selected_code()
-    visible_text = get_visible_code()
 
-    # rewrite
-    response = add_docstring(selected_text=selected_text, visible_text=visible_text)
+    # Rewrite
+    response = add_docstring(
+        selected_text=selected_text.get("text", ""), file_path=selected_text.get("abspath", "")
+    )
 
-    # apply new code to editor
+    # Get indent level
+    indent = get_indent_level(selected_text.get("text", ""))
+
+    # Apply new code to editor
     new_code = extract_markdown_block(response)
-    IDEService().diff_apply("", new_code)
+    if not new_code:
+        language = IDEService().ide_language()
+        print_message(language)
+        sys.exit(0)
+
+    # Offset indent level
+    new_code = offset_indent_level(new_code, indent)
+
+    # Merge code
+    docstring_code = merge_code(selected_text.get("text", ""), new_code)
+    # Apply diff
+    IDEService().diff_apply("", docstring_code)
 
     sys.exit(0)
 
 
+def print_message(language):
+    if language == "zh":
+        print("\n\n大模型输出不完整，不能进行代码操作。\n\n")
+    else:
+        print("\n\nThe output of the LLM is incomplete and cannot perform code operations.\n\n")
+
+
 if __name__ == "__main__":
     main()

refactor/names/command.yml

@@ -1,4 +1,4 @@
 description: Automatically rename poorly-named local variables for improved readability. Select the code segment and execute this command to optimize variable names.
 steps:
-  - run: $devchat_python $command_path/../rewrite.py "Refine internal variable and function names within the code to achieve concise and meaningful identifiers that comply with English naming conventions."
+  - run: $devchat_python $command_path/main.py
 

refactor/names/main.py

@@ -0,0 +1,251 @@
+import os
+import re
+import sys
+
+from devchat.ide.service import IDEService
+from devchat.ide.vscode_services import selected_lines, visible_lines
+from devchat.llm import chat
+from devchat.memory import FixSizeChatMemory
+
+PROMPT = prompt = """
+file: {file_path}
+```
+{selected_text}
+```
+"""
+
+
+PROMPT_ZH = prompt = """
+文件: {file_path}
+```
+{selected_text}
+```
+
+输出内容使用中文，我的母语为中文。
+"""
+
+
+def get_prompt():
+    ide_language = IDEService().ide_language()
+    return PROMPT_ZH if ide_language == "zh" else PROMPT
+
+
+MESSAGES_A = [
+    {
+        "role": "system",
+        "content": """
+Your task is:
+Refine internal variable and function names within the code to achieve concise and \
+    meaningful identifiers that comply with English naming conventions.
+
+Rules:
+1. Don't rename a call or global variable. for example, xx() is function call, xx \
+    is a bad name, but you MUST not rename it .
+2. You can rename a local variable or parameter variable name.
+3. Current function's name can be renamed. Always this is a new function.
+
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: a1.py
+```
+    def print_hello():
+        a = "hello world"
+        print(a)
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```python
+    def print_hello():
+        msg = "hello world"
+        print(msg)
+```
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: t1.py
+```
+    def print_hello(a: str):
+        print(a)
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```python
+    def print_hello(msg: str):
+        print(msg)
+```
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: t1.py
+```
+    def some():
+        print("hello")
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```python
+    def output_hello():
+        print("hello")
+```
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+file: t1.py
+```
+    def print_hello():
+        print("hello")
+```
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```python
+    def print_hello():
+        output("hello")
+```
+""",
+    },
+    {
+        "role": "user",
+        "content": """
+Your response is error, you changed call name.
+print is a function call, if you rename it, this will make a compile error.
+""",
+    },
+    {
+        "role": "assistant",
+        "content": """
+```python
+    def print_hello():
+        print("hello")
+```
+""",
+    },
+]
+
+
+def get_selected_code():
+    """
+    Retrieves the selected lines of code from the user's selection.
+
+    This function extracts the text selected by the user in their IDE or text editor.
+    If no text has been selected, it prints an error message to stderr and exits the
+    program with a non-zero status indicating failure.
+
+    Returns:
+        dict: A dictionary containing the key 'selectedText' with the selected text
+        as its value. If no text is selected, the program exits.
+    """
+    selected_data = IDEService().get_selected_range().dict()
+
+    miss_selected_error = "Please select some text."
+    if selected_data["range"]["start"] == selected_data["range"]["end"]:
+        readme_path = os.path.join(os.path.dirname(__file__), "README.md")
+        if os.path.exists(readme_path):
+            with open(readme_path, "r", encoding="utf-8") as f:
+                readme_text = f.read()
+                print(readme_text)
+                sys.exit(0)
+
+        print(miss_selected_error, file=sys.stderr, flush=True)
+        sys.exit(-1)
+
+    return selected_data
+
+
+memory = FixSizeChatMemory(max_size=20, messages=MESSAGES_A)
+
+
+@chat(prompt=get_prompt(), stream_out=True, memory=memory)
+# pylint: disable=unused-argument
+def reanme_variable(selected_text, file_path):
+    """
+    call ai to rewrite selected code
+    """
+    pass  # pylint: disable=unnecessary-pass
+
+
+def extract_markdown_block(text):
+    """
+    Extracts the first Markdown code block from the given text without the language specifier.
+
+    :param text: A string containing Markdown text
+    :return: The content of the first Markdown code block, or None if not found
+    """
+    pattern = r"```(?:\w+)?\s*\n(.*?)\n```"
+    match = re.search(pattern, text, re.DOTALL)
+
+    if match:
+        block_content = match.group(1)
+        return block_content
+    else:
+        # whether exist ```language?
+        if text.find("```"):
+            return None
+        return text
+
+
+def remove_unnecessary_escapes(code_a, code_b):
+    code_copy = code_b  # Create a copy of the original code
+    escape_chars = re.finditer(r"\\(.)", code_b)
+
+    remove_char_index = []
+    for match in escape_chars:
+        before = code_b[max(0, match.start() - 4) : match.start()]
+        after = code_b[match.start() + 1 : match.start() + 5]
+        substr = before + after
+        if substr in code_a:
+            remove_char_index.append(match.start())
+
+    # visit remove_char_index in reverse order
+    remove_char_index.reverse()
+    for index in remove_char_index:
+        code_copy = code_copy[:index] + code_copy[index + 1 :]
+    return code_copy
+
+
+def main():
+    # prepare code
+    selected_text = get_selected_code()
+    selected_code = selected_text.get("text", "")
+    selected_file = selected_text.get("abspath", "")
+
+    # rewrite
+    response = reanme_variable(selected_text=selected_code, file_path=selected_file)
+
+    # apply new code to editor
+    new_code = extract_markdown_block(response)
+    if not new_code:
+        if IDEService().ide_language() == "zh":
+            print("\n\n大模型输出不完整，不能进行代码操作。\n\n")
+        else:
+            print("\n\nThe output of the LLM is incomplete and cannot perform code operations.\n\n")
+        sys.exit(0)
+
+    new_code = remove_unnecessary_escapes(selected_text.get("text", ""), new_code)
+    IDEService().diff_apply("", new_code)
+
+    sys.exit(0)
+
+
+if __name__ == "__main__":
+    main()

refactor/rewrite.py

@@ -36,17 +36,6 @@ def get_selected_code():
     return selected_data
 
 
-def get_visible_code():
-    """
-    Retrieves visible code from the visible_lines function.
-
-    Returns:
-    visible_data: The visible code retrieved from the visible_lines function.
-    """
-    visible_data = IDEService().get_visible_range().dict()
-    return visible_data
-
-
 REWRITE_PROMPT = prompt = """
 Your task is:
 {question}
@@ -57,13 +46,12 @@ correct syntax. Just refactor the selected code. Keep all other information as i
 Here is the relevant context \
 information for your reference:
 1.  selected code info: {selected_text}
-2.  current visible code info: {visible_text}
 """
 
 
 @chat(prompt=REWRITE_PROMPT, stream_out=True)
 # pylint: disable=unused-argument
-def ai_rewrite(question, selected_text, visible_text):
+def ai_rewrite(question, selected_text):
     """
     call ai to rewrite selected code
     """
@@ -84,6 +72,9 @@ def extract_markdown_block(text):
         block_content = match.group(1)
         return block_content
     else:
+        # whether exist ```language?
+        if text.find("```"):
+            return None
         return text
 
 
@@ -126,13 +117,19 @@ def main():
     question = sys.argv[1]
     # prepare code
     selected_text = get_selected_code()
-    visible_text = get_visible_code()
 
     # rewrite
-    response = ai_rewrite(question=question, selected_text=selected_text, visible_text=visible_text)
+    response = ai_rewrite(question=question, selected_text=selected_text)
 
     # apply new code to editor
     new_code = extract_markdown_block(response)
+    if not new_code:
+        if IDEService().ide_language() == "zh":
+            print("\n\n大模型输出不完整，不能进行代码操作。\n\n")
+        else:
+            print("\n\nThe output of the LLM is incomplete and cannot perform code operations.\n\n")
+        sys.exit(0)
+
     IDEService().diff_apply("", new_code)
 
     sys.exit(0)