sdk.get_code_context

4 files changed, 25 insertions, 27 deletions

diff --git a/continuedev/src/continuedev/core/autopilot.py b/continuedev/src/continuedev/core/autopilot.py
index b583d682..5c3baafd 100644
--- a/continuedev/src/continuedev/core/autopilot.py
+++ b/continuedev/src/continuedev/core/autopilot.py
@@ -157,8 +157,6 @@ class Autopilot(ContinueBaseModel):
 
         # Make sure all filepaths are relative to workspace
         workspace_path = self.continue_sdk.ide.workspace_directory
-        for rif in range_in_files:
-            rif.filepath = os.path.basename(rif.filepath)
 
         # If not adding highlighted code
         if not self._adding_highlighted_code:
@@ -170,7 +168,7 @@ class Autopilot(ContinueBaseModel):
                 # Otherwise, replace the current range with the new one
                 # This is the first range to be highlighted
                 self._highlighted_ranges = [HighlightedRangeContext(
-                    range=range_in_files[0], editing=True, pinned=False)]
+                    range=range_in_files[0], editing=True, pinned=False, display_name=os.path.basename(range_in_files[0].filepath))]
                 await self.update_subscribers()
             return
 
@@ -193,7 +191,7 @@ class Autopilot(ContinueBaseModel):
                 new_ranges.append(rif)
 
         self._highlighted_ranges = new_ranges + [HighlightedRangeContext(
-            range=rif, editing=False, pinned=False
+            range=rif, editing=False, pinned=False, display_name=os.path.basename(rif.filepath)
         ) for rif in range_in_files]
 
         self._make_sure_is_editing_range()
diff --git a/continuedev/src/continuedev/core/main.py b/continuedev/src/continuedev/core/main.py
index 4ea17f20..88690c83 100644
--- a/continuedev/src/continuedev/core/main.py
+++ b/continuedev/src/continuedev/core/main.py
@@ -205,6 +205,7 @@ class HighlightedRangeContext(ContinueBaseModel):
     range: RangeInFileWithContents
     editing: bool
     pinned: bool
+    display_name: str
 
 
 class FullState(ContinueBaseModel):
diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py
index ed670799..8649cd58 100644
--- a/continuedev/src/continuedev/core/sdk.py
+++ b/continuedev/src/continuedev/core/sdk.py
@@ -13,7 +13,7 @@ from ..libs.llm.hf_inference_api import HuggingFaceInferenceAPI
 from ..libs.llm.openai import OpenAI
 from .observation import Observation
 from ..server.ide_protocol import AbstractIdeProtocolServer
-from .main import Context, ContinueCustomException, History, Step, ChatMessage, ChatMessageRole
+from .main import Context, ContinueCustomException, HighlightedRangeContext, History, Step, ChatMessage, ChatMessageRole
 from ..steps.core.core import *
 from ..libs.llm.proxy_server import ProxyServer
 
@@ -178,6 +178,11 @@ class ContinueSDK(AbstractContinueSDK):
         else:
             return load_global_config()
 
+    def get_code_context(self, only_editing: bool = False) -> List[RangeInFileWithContents]:
+        context = list(filter(lambda x: x.editing, self.__autopilot._highlighted_ranges)
+                       ) if only_editing else self.__autopilot._highlighted_ranges
+        return [c.range for c in context]
+
     def update_default_model(self, model: str):
         config = self.config
         config.default_model = model
diff --git a/continuedev/src/continuedev/steps/main.py b/continuedev/src/continuedev/steps/main.py
index 5ccffbfe..0a20ddd7 100644
--- a/continuedev/src/continuedev/steps/main.py
+++ b/continuedev/src/continuedev/steps/main.py
@@ -97,7 +97,7 @@ class FasterEditHighlightedCodeStep(Step):
         return "Editing highlighted code"
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        range_in_files = await sdk.ide.getHighlightedCode()
+        range_in_files = await sdk.get_code_context(only_editing=True)
         if len(range_in_files) == 0:
             # Get the full contents of all open files
             files = await sdk.ide.getOpenFiles()
@@ -105,21 +105,16 @@ class FasterEditHighlightedCodeStep(Step):
             for file in files:
                 contents[file] = await sdk.ide.readFile(file)
 
-            range_in_files = [RangeInFile.from_entire_file(
+            range_in_files = [RangeInFileWithContents.from_entire_file(
                 filepath, content) for filepath, content in contents.items()]
 
-        rif_with_contents = []
-        for range_in_file in range_in_files:
-            file_contents = await sdk.ide.readRangeInFile(range_in_file)
-            rif_with_contents.append(
-                RangeInFileWithContents.from_range_in_file(range_in_file, file_contents))
-        enc_dec = MarkdownStyleEncoderDecoder(rif_with_contents)
+        enc_dec = MarkdownStyleEncoderDecoder(range_in_files)
         code_string = enc_dec.encode()
         prompt = self._prompt.format(
             code=code_string, user_input=self.user_input)
 
         rif_dict = {}
-        for rif in rif_with_contents:
+        for rif in range_in_files:
             rif_dict[rif.filepath] = rif.contents
 
         completion = await sdk.models.gpt35.complete(prompt)
@@ -193,7 +188,7 @@ class StarCoderEditHighlightedCodeStep(Step):
         return await models.gpt35.complete(f"{self._prompt_and_completion}\n\nPlease give brief a description of the changes made above using markdown bullet points:")
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        range_in_files = await sdk.ide.getHighlightedCode()
+        range_in_files = await sdk.get_code_context(only_editing=True)
         found_highlighted_code = len(range_in_files) > 0
         if not found_highlighted_code:
             # Get the full contents of all open files
@@ -202,20 +197,14 @@ class StarCoderEditHighlightedCodeStep(Step):
             for file in files:
                 contents[file] = await sdk.ide.readFile(file)
 
-            range_in_files = [RangeInFile.from_entire_file(
+            range_in_files = [RangeInFileWithContents.from_entire_file(
                 filepath, content) for filepath, content in contents.items()]
 
-        rif_with_contents = []
-        for range_in_file in range_in_files:
-            file_contents = await sdk.ide.readRangeInFile(range_in_file)
-            rif_with_contents.append(
-                RangeInFileWithContents.from_range_in_file(range_in_file, file_contents))
-
         rif_dict = {}
-        for rif in rif_with_contents:
+        for rif in range_in_files:
             rif_dict[rif.filepath] = rif.contents
 
-        for rif in rif_with_contents:
+        for rif in range_in_files:
             prompt = self._prompt.format(
                 code=rif.contents, user_request=self.user_input)
 
@@ -255,7 +244,7 @@ class EditHighlightedCodeStep(Step):
         return "Editing code"
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
-        range_in_files = await sdk.ide.getHighlightedCode()
+        range_in_files = sdk.get_code_context(only_editing=True)
         if len(range_in_files) == 0:
             # Get the full contents of all open files
             files = await sdk.ide.getOpenFiles()
@@ -263,7 +252,7 @@ class EditHighlightedCodeStep(Step):
             for file in files:
                 contents[file] = await sdk.ide.readFile(file)
 
-            range_in_files = [RangeInFile.from_entire_file(
+            range_in_files = [RangeInFileWithContents.from_entire_file(
                 filepath, content) for filepath, content in contents.items()]
 
         # If still no highlighted code, create a new file and edit there
@@ -271,7 +260,12 @@ class EditHighlightedCodeStep(Step):
             # Create a new file
             new_file_path = "new_file.txt"
             await sdk.add_file(new_file_path, "")
-            range_in_files = [RangeInFile.from_entire_file(new_file_path, "")]
+            range_in_files = [
+                RangeInFileWithContents.from_entire_file(new_file_path, "")]
+
+        range_in_files = list(map(lambda x: RangeInFile(
+            filepath=x.filepath, range=x.range
+        ), range_in_files))
 
         await sdk.run_step(DefaultModelEditCodeStep(user_input=self.user_input, range_in_files=range_in_files))