cleaner highlighting (from server side)

3 files changed, 63 insertions, 11 deletions

diff --git a/continuedev/src/continuedev/server/session_manager.py b/continuedev/src/continuedev/server/session_manager.py
index 0dbfaf38..ebea08a5 100644
--- a/continuedev/src/continuedev/server/session_manager.py
+++ b/continuedev/src/continuedev/server/session_manager.py
@@ -28,6 +28,7 @@ class DemoAutopilot(Autopilot):
     cumulative_edit_string = ""
 
     def handle_manual_edits(self, edits: List[FileEditWithFullContents]):
+        return
         for edit in edits:
             self.cumulative_edit_string += edit.fileEdit.replacement
             self._manual_edits_buffer.append(edit)
diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index 57689f19..4288ffd2 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -3,8 +3,10 @@ import os
 import subprocess
 from textwrap import dedent
 from typing import Coroutine, List, Union
-from ...libs.llm.prompt_utils import MarkdownStyleEncoderDecoder
 
+from ...models.main import Range
+from ...libs.util.calculate_diff import calculate_diff2, apply_edit_to_str
+from ...libs.llm.prompt_utils import MarkdownStyleEncoderDecoder
 from ...models.filesystem_edit import EditDiff, FileEdit, FileEditWithFullContents, FileSystemEdit
 from ...models.filesystem import FileSystem, RangeInFile, RangeInFileWithContents
 from ...core.observation import Observation, TextObservation, TracebackObservation, UserInputObservation
@@ -149,7 +151,11 @@ class Gpt35EditCodeStep(Step):
     _prompt_and_completion: str = ""
 
     async def describe(self, models: Models) -> Coroutine[str, None, None]:
-        return models.gpt35.complete(f"{self._prompt_and_completion}\n\nPlease give brief a description of the changes made above using markdown bullet points:")
+        description = models.gpt35.complete(
+            f"{self._prompt_and_completion}\n\nPlease give brief a description of the changes made above using markdown bullet points. Be concise and only mention changes made to the commit before, not prefix or suffix:")
+        self.name = models.gpt35.complete(
+            f"Write a short title for this description: {description}")
+        return description
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
         rif_with_contents = []
@@ -174,11 +180,40 @@ class Gpt35EditCodeStep(Step):
 
             self._prompt_and_completion += prompt + completion
 
-            await sdk.ide.applyFileSystemEdit(
-                FileEdit(filepath=rif.filepath, range=rif.range, replacement=completion))
-            await sdk.ide.saveFile(rif.filepath)
+            # Calculate diff, open file, apply edits, and highlight changed lines
+            edits = calculate_diff2(
+                rif.filepath, rif.contents, completion.removesuffix("\n"))
+
             await sdk.ide.setFileOpen(rif.filepath)
 
+            lines_to_highlight = set()
+            for edit in edits:
+                edit.range.start.line += rif.range.start.line
+                edit.range.start.character += rif.range.start.character
+                edit.range.end.line += rif.range.start.line
+                edit.range.end.character += rif.range.start.character if edit.range.end.line == 0 else 0
+
+                for line in range(edit.range.start.line, edit.range.end.line + 1):
+                    lines_to_highlight.add(line)
+
+                await sdk.ide.applyFileSystemEdit(edit)
+
+            current_start = None
+            last_line = None
+            for line in sorted(list(lines_to_highlight)):
+                if current_start is None:
+                    current_start = line
+                elif line != last_line + 1:
+                    await sdk.ide.highlightCode(RangeInFile(filepath=edit.filepath, range=Range.from_shorthand(current_start, 0, last_line, 0)))
+                    current_start = line
+
+                last_line = line
+
+            if current_start is not None:
+                await sdk.ide.highlightCode(RangeInFile(filepath=edit.filepath, range=Range.from_shorthand(current_start, 0, last_line, 0)))
+
+            await sdk.ide.saveFile(rif.filepath)
+
 
 class EditFileStep(Step):
     filepath: str
diff --git a/continuedev/src/continuedev/steps/main.py b/continuedev/src/continuedev/steps/main.py
index 24335b4f..9634c726 100644
--- a/continuedev/src/continuedev/steps/main.py
+++ b/continuedev/src/continuedev/steps/main.py
@@ -16,6 +16,7 @@ from ..core.sdk import ContinueSDK, Models
 from ..core.observation import Observation
 import subprocess
 from .core.core import Gpt35EditCodeStep
+from ..libs.util.calculate_diff import calculate_diff2
 
 
 class SetupContinueWorkspaceStep(Step):
@@ -216,7 +217,8 @@ class StarCoderEditHighlightedCodeStep(Step):
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
         range_in_files = await sdk.ide.getHighlightedCode()
-        if len(range_in_files) == 0:
+        found_highlighted_code = len(range_in_files) > 0
+        if not found_highlighted_code:
             # Get the full contents of all open files
             files = await sdk.ide.getOpenFiles()
             contents = {}
@@ -239,15 +241,29 @@ class StarCoderEditHighlightedCodeStep(Step):
         for rif in rif_with_contents:
             prompt = self._prompt.format(
                 code=rif.contents, user_request=self.user_input)
-            completion = str(sdk.models.starcoder.complete(prompt))
+
+            if found_highlighted_code:
+                full_file_contents = await sdk.ide.readFile(rif.filepath)
+                segs = full_file_contents.split(rif.contents)
+                prompt = f"<file_prefix>{segs[0]}<file_suffix>{segs[1]}" + prompt
+
+            completion = str((await sdk.models.starcoder()).complete(prompt))
             eot_token = "<|endoftext|>"
-            if completion.endswith(eot_token):
-                completion = completion[:completion.rindex(eot_token)]
+            completion = completion.removesuffix(eot_token)
+
+            if found_highlighted_code:
+                rif.contents = segs[0] + rif.contents + segs[1]
+                completion = segs[0] + completion + segs[1]
 
             self._prompt_and_completion += prompt + completion
 
-            await sdk.ide.applyFileSystemEdit(
-                FileEdit(filepath=rif.filepath, range=rif.range, replacement=completion))
+            edits = calculate_diff2(
+                rif.filepath, rif.contents, completion.removesuffix("\n"))
+            for edit in edits:
+                await sdk.ide.applyFileSystemEdit(edit)
+
+            # await sdk.ide.applyFileSystemEdit(
+            #     FileEdit(filepath=rif.filepath, range=rif.range, replacement=completion))
             await sdk.ide.saveFile(rif.filepath)
             await sdk.ide.setFileOpen(rif.filepath)