diff --git a/lua/avante/selection.lua b/lua/avante/selection.lua
index 387ffa6..2036488 100644
--- a/lua/avante/selection.lua
+++ b/lua/avante/selection.lua
@@ -165,11 +165,19 @@ function Selection:create_editing_input()
       full_response = full_response .. chunk
       local response_lines_ = vim.split(full_response, "\n")
       local response_lines = {}
-      for i, line in ipairs(response_lines_) do
-        if string.match(line, "^```") and (i == 1 or i == #response_lines_) then goto continue end
-        if string.match(line, "^```$") then goto continue end
-        table.insert(response_lines, line)
-        ::continue::
+      local in_code_block = false
+      for _, line in ipairs(response_lines_) do
+        if line:match("^<code>") then
+          in_code_block = true
+          line = line:gsub("^<code>", "")
+          if line ~= "" then table.insert(response_lines, line) end
+        elseif line:match("</code>") then
+          in_code_block = false
+          line = line:gsub("</code>.*$", "")
+          if line ~= "" then table.insert(response_lines, line) end
+        elseif in_code_block then
+          table.insert(response_lines, line)
+        end
       end
       if #response_lines == 1 then
         local first_line = response_lines[1]
diff --git a/lua/avante/templates/editing.avanterules b/lua/avante/templates/editing.avanterules
index f0f90cf..723ae13 100644
--- a/lua/avante/templates/editing.avanterules
+++ b/lua/avante/templates/editing.avanterules
@@ -2,7 +2,7 @@
 {% block extra_prompt %}
 Your task is to modify the provided code according to the user's request. Follow these instructions precisely:
 
-1. Return *ONLY* the complete modified code.
+1. The code you return must be wrapped in <code></code>, and cannot contain any other code.
 
 2. *DO NOT* include three backticks: {%raw%}```{%endraw%} in your suggestion! Treat the suggested code AS IS.
 
@@ -19,4 +19,39 @@ Your task is to modify the provided code according to the user's request. Follow
 8. *ONLY* return the new code snippets to be updated, *DO NOT* return the entire file content.
 
 Remember that Your response SHOULD CONTAIN ONLY THE MODIFIED CODE to be used as DIRECT REPLACEMENT to the original file.
+
+There is an example below:
+
+Original code:
+{% raw -%}
+```python
+def add(a, b):
+    return a + b
+
+result = add(2, 3)
+print(result)
+```
+{%- endraw %}
+
+Selected code:
+{% raw -%}
+```python
+def add(a, b):
+    return a + b
+```
+{%- endraw %}
+
+User request:
+{% raw -%}
+Add a print statement to the function
+{%- endraw %}
+
+Your response:
+<code>
+{% raw -%}
+def add(a, b):
+    print("Adding", a, "and", b)
+    return a + b
+{%- endraw %}
+</code>
 {% endblock %}