Merge pull request #157 from continuedev/less_cluttered_streaming

Less cluttered streaming

1 files changed, 68 insertions, 29 deletions

diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index 8a28490e..cdf3a78f 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -280,12 +280,19 @@ class DefaultModelEditCodeStep(Step):
         lines = []
         unfinished_line = ""
 
+        # Don't end the block until you've matched N simultaneous lines
+        # This helps avoid many tiny blocks
+        LINES_TO_MATCH_BEFORE_ENDING_BLOCK = 2
+        matched_lines_at_end_of_block = 0
+        index_of_last_matched_line = -1
+
         async def handle_generated_line(line: str):
-            nonlocal lines, current_block_start, current_line_in_file, original_lines, original_lines_below_previous_blocks, current_block_lines, offset_from_blocks
+            nonlocal lines, current_block_start, current_line_in_file, original_lines, original_lines_below_previous_blocks, current_block_lines, offset_from_blocks, matched_lines_at_end_of_block, index_of_last_matched_line, LINES_TO_MATCH_BEFORE_ENDING_BLOCK
 
             # Highlight the line to show progress
+            line_to_highlight = current_line_in_file - len(current_block_lines)
             await sdk.ide.highlightCode(RangeInFile(filepath=rif.filepath, range=Range.from_shorthand(
-                current_line_in_file, 0, current_line_in_file, 0)), "#FFFFFF22" if len(current_block_lines) == 0 else "#FFFF0022")
+                line_to_highlight, 0, line_to_highlight, 0)), "#FFFFFF22" if len(current_block_lines) == 0 else "#00FF0022")
 
             if len(current_block_lines) == 0:
                 current_block_start = current_line_in_file
@@ -295,28 +302,58 @@ class DefaultModelEditCodeStep(Step):
                     return
 
             # We are in a block currently, and checking for whether it should be ended
-            for i in range(len(original_lines_below_previous_blocks)):
-                og_line = original_lines_below_previous_blocks[i]
-                if og_line == line:
-
-                    # Insert the suggestion
-                    replacement = "\n".join(current_block_lines)
-                    await sdk.ide.showSuggestion(FileEdit(
-                        filepath=rif.filepath,
-                        range=Range.from_shorthand(
-                            current_block_start, 0, current_block_start + i, 0),
-                        replacement=replacement
-                    ))
-                    if replacement == "":
-                        current_line_in_file += 1
-
-                    # Reset current block / update variables
-                    original_lines_below_previous_blocks = original_lines_below_previous_blocks[
-                        i + 1:]
-                    offset_from_blocks += len(current_block_lines)
-                    current_block_lines = []
-                    current_block_start = -1
-                    return
+            if matched_lines_at_end_of_block == 0:
+                # Find the first matching line
+                for i in range(len(original_lines_below_previous_blocks)):
+                    og_line = original_lines_below_previous_blocks[i]
+                    # TODO: It's a bit sus to be disqualifying empty lines.
+                    # What you ideally do is find ALL matches, and then throw them out as you check the following lines
+                    if og_line == line and og_line.strip() != "":
+                        matched_lines_at_end_of_block = 1
+                        index_of_last_matched_line = i
+                        break
+            else:
+                # Check if the next line matches
+                index_of_line_to_match = index_of_last_matched_line + matched_lines_at_end_of_block
+                if len(original_lines_below_previous_blocks) > index_of_line_to_match and original_lines_below_previous_blocks[index_of_line_to_match] == line:
+                    if matched_lines_at_end_of_block >= LINES_TO_MATCH_BEFORE_ENDING_BLOCK:
+                        # We've matched the required number of lines, insert suggestion!
+
+                        # We added some lines to the block that were matched (including maybe some blank lines)
+                        # So here we will strip all matching lines from the end of current_block_lines
+                        lines_stripped = []
+                        index_of_end_of_block = index_of_line_to_match
+                        while len(current_block_lines) > 0 and current_block_lines[-1] == original_lines_below_previous_blocks[index_of_end_of_block - 1]:
+                            lines_stripped.append(current_block_lines.pop())
+                            index_of_end_of_block -= 1
+
+                        # Insert the suggestion
+                        replacement = "\n".join(current_block_lines)
+                        await sdk.ide.showSuggestion(FileEdit(
+                            filepath=rif.filepath,
+                            range=Range.from_shorthand(
+                                current_block_start + offset_from_blocks, 0, current_block_start + offset_from_blocks + index_of_end_of_block, 0),
+                            replacement=replacement
+                        ))
+                        if replacement == "":
+                            current_line_in_file += 1
+
+                        # Reset current block / update variables
+                        original_lines_below_previous_blocks = original_lines_below_previous_blocks[
+                            index_of_line_to_match + 1:]
+                        offset_from_blocks += len(current_block_lines)
+                        current_block_lines = []
+                        current_block_start = -1
+                        matched_lines_at_end_of_block = 0
+                        index_of_last_matched_line = -1
+
+                        return
+                    else:
+                        matched_lines_at_end_of_block += 1
+                else:
+                    # We matched some lines, but didn't make it to N
+                    # So this block should continue on with the matched lines as a part of it
+                    matched_lines_at_end_of_block = 0
 
             current_block_lines.append(line)
 
@@ -383,17 +420,19 @@ class DefaultModelEditCodeStep(Step):
         if len(current_block_lines) > 0:
             # We have a chance to back-track here for blank lines that are repeats of the suffix
             num_to_remove = 0
-            if repeating_file_suffix:
-                for i in range(-1, -len(current_block_lines) - 1, -1):
-                    if current_block_lines[i].strip() == "":
-                        num_to_remove += 1
+            for i in range(-1, -len(current_block_lines) - 1, -1):
+                if len(original_lines_below_previous_blocks) == 0:
+                    break
+                if current_block_lines[i] == original_lines_below_previous_blocks[-1]:
+                    num_to_remove += 1
+                    original_lines_below_previous_blocks.pop()
             current_block_lines = current_block_lines[:-
                                                       num_to_remove] if num_to_remove > 0 else current_block_lines
 
             await sdk.ide.showSuggestion(FileEdit(
                 filepath=rif.filepath,
                 range=Range.from_shorthand(
-                    current_block_start, 0, current_block_start + len(original_lines_below_previous_blocks), 0),
+                    current_block_start + offset_from_blocks, 0, current_block_start + offset_from_blocks + len(original_lines_below_previous_blocks), 0),
                 replacement="\n".join(current_block_lines)
             ))