Use gpt by default, with prefix/suffixes

5 files changed, 74 insertions, 61 deletions

diff --git a/continuedev/src/continuedev/core/policy.py b/continuedev/src/continuedev/core/policy.py
index 91ae3c83..8aea8de7 100644
--- a/continuedev/src/continuedev/core/policy.py
+++ b/continuedev/src/continuedev/core/policy.py
@@ -38,7 +38,7 @@ class DemoPolicy(Policy):
                 return EditFileChroma(request=" ".join(observation.user_input.split(" ")[1:]))
             elif "/step" in observation.user_input:
                 return ContinueStepStep(prompt=" ".join(observation.user_input.split(" ")[1:]))
-            return StarCoderEditHighlightedCodeStep(user_input=observation.user_input)
+            return EditHighlightedCodeStep(user_input=observation.user_input)
 
         state = history.get_current()
 
diff --git a/continuedev/src/continuedev/core/sdk.py b/continuedev/src/continuedev/core/sdk.py
index 76caef02..51faadf2 100644
--- a/continuedev/src/continuedev/core/sdk.py
+++ b/continuedev/src/continuedev/core/sdk.py
@@ -2,11 +2,12 @@ from abc import ABC, abstractmethod
 from typing import Coroutine, Union
 import os
 
+from ..steps.core.core import Gpt35EditCodeStep
+from ..models.main import Range
 from .abstract_sdk import AbstractContinueSDK
 from .config import ContinueConfig, load_config
 from ..models.filesystem_edit import FileEdit, FileSystemEdit, AddFile, DeleteFile, AddDirectory, DeleteDirectory
 from ..models.filesystem import RangeInFile
-from ..libs.llm import LLM
 from ..libs.llm.hf_inference_api import HuggingFaceInferenceAPI
 from ..libs.llm.openai import OpenAI
 from .observation import Observation
@@ -79,14 +80,15 @@ class ContinueSDK(AbstractContinueSDK):
         commands = commands if isinstance(commands, List) else [commands]
         return (await self.run_step(ShellCommandsStep(cmds=commands, cwd=cwd, description=description, **({'name': name} if name else {})))).text
 
-    async def edit_file(self, filename: str, prompt: str, name: str = None, description: str = None):
+    async def edit_file(self, filename: str, prompt: str, name: str = None, description: str = None, range: Range = None):
         filepath = await self._ensure_absolute_path(filename)
 
         await self.ide.setFileOpen(filepath)
         contents = await self.ide.readFile(filepath)
-        await self.run_step(EditCodeStep(
-            range_in_files=[RangeInFile.from_entire_file(filepath, contents)],
-            prompt=f'Here is the code before:\n\n{{code}}\n\nHere is the user request:\n\n{prompt}\n\nHere is the code edited to perfectly solve the user request:\n\n',
+        await self.run_step(Gpt35EditCodeStep(
+            range_in_files=[RangeInFile(filepath=filename, range=range) if range is not None else RangeInFile.from_entire_file(
+                filepath, contents)],
+            user_input=prompt,
             description=description,
             **({'name': name} if name else {})
         ))
diff --git a/continuedev/src/continuedev/recipes/CreatePipelineRecipe/steps.py b/continuedev/src/continuedev/recipes/CreatePipelineRecipe/steps.py
index c32ae923..511abd1f 100644
--- a/continuedev/src/continuedev/recipes/CreatePipelineRecipe/steps.py
+++ b/continuedev/src/continuedev/recipes/CreatePipelineRecipe/steps.py
@@ -39,7 +39,7 @@ class SetupPipelineStep(Step):
             'python3 -m venv env',
             'source env/bin/activate',
             'pip install dlt',
-            f'dlt init {source_name} duckdb\n\rY',
+            f'dlt --non-interactive init {source_name} duckdb',
             'pip install -r requirements.txt'
         ], description=dedent(f"""\
             Running the following commands:
@@ -50,10 +50,12 @@ class SetupPipelineStep(Step):
             - `pip install -r requirements.txt`: Install the Python dependencies for the pipeline"""), name="Setup Python environment")
 
         # editing the resource function to call the requested API
-        await sdk.ide.highlightCode(RangeInFile(filepath=os.path.join(await sdk.ide.getWorkspaceDirectory(), filename), range=Range.from_shorthand(15, 0, 29, 0)), "#00ff0022")
+        resource_function_range = Range.from_shorthand(15, 0, 29, 0)
+        await sdk.ide.highlightCode(RangeInFile(filepath=os.path.join(await sdk.ide.getWorkspaceDirectory(), filename), range=resource_function_range), "#00ff0022")
 
         # sdk.set_loading_message("Writing code to call the API...")
         await sdk.edit_file(
+            range=resource_function_range,
             filename=filename,
             prompt=f'Edit the resource function to call the API described by this: {self.api_description}. Do not move or remove the exit() call in __main__.',
             name=f"Edit the resource function to call the API {AI_ASSISTED_STRING}"
@@ -144,8 +146,7 @@ class ValidatePipelineStep(Step):
 
             # print table names
             for row in rows:
-                print(row)
-        ''')
+                print(row)''')
 
         query_filename = os.path.join(workspace_dir, "query.py")
         await sdk.apply_filesystem_edit(AddFile(filepath=query_filename, content=tables_query_code), name="Add query.py file", description="Adding a file called `query.py` to the workspace that will run a test query on the DuckDB instance")
diff --git a/continuedev/src/continuedev/steps/core/core.py b/continuedev/src/continuedev/steps/core/core.py
index c6dc7c04..413bc195 100644
--- a/continuedev/src/continuedev/steps/core/core.py
+++ b/continuedev/src/continuedev/steps/core/core.py
@@ -4,7 +4,7 @@ from textwrap import dedent
 from typing import Coroutine, List, Union
 from ...libs.llm.prompt_utils import MarkdownStyleEncoderDecoder
 
-from ...models.filesystem_edit import EditDiff, FileEditWithFullContents, FileSystemEdit
+from ...models.filesystem_edit import EditDiff, FileEdit, FileEditWithFullContents, FileSystemEdit
 from ...models.filesystem import FileSystem, RangeInFile, RangeInFileWithContents
 from ...core.observation import Observation, TextObservation, TracebackObservation, UserInputObservation
 from ...core.main import Step, SequentialStep
@@ -74,26 +74,48 @@ class ShellCommandsStep(Step):
         # return None
 
 
-class EditCodeStep(Step):
-    # Might make an even more specific atomic step, which is "apply file edit"
+class Gpt35EditCodeStep(Step):
+    user_input: str
     range_in_files: List[RangeInFile]
-    prompt: str  # String with {code} somewhere
-    name: str = "Edit code"
-
-    _edit_diffs: Union[List[EditDiff], None] = None
-    _prompt: Union[str, None] = None
-    _completion: Union[str, None] = None
+    name: str = "Editing Code"
+    hide = False
+    _prompt: str = dedent("""\
+        Take the file prefix and suffix into account, but only rewrite the commit before as specified in the commit message. Here's an example:
+
+        <file_prefix>
+        a = 5
+        b = 4
+
+        <file_suffix>
+
+        def mul(a, b):
+            return a * b
+        <commit_before>
+        def sum():
+            return a + b
+        <commit_msg>
+        Make a and b parameters of sum
+        <commit_after>
+        def sum(a, b):
+            return a + b
+        <|endoftext|>
+
+        Now complete the real thing:
+
+        <file_prefix>
+        {file_prefix}
+        <file_suffix>
+        {file_suffix}
+        <commit_before>
+        {code}
+        <commit_msg>
+        {user_request}
+        <commit_after>""")
+
+    _prompt_and_completion: str = ""
 
     async def describe(self, models: Models) -> Coroutine[str, None, None]:
-        if self._edit_diffs is None:
-            return "Editing files: " + ", ".join(map(lambda rif: rif.filepath, self.range_in_files))
-        elif len(self._edit_diffs) == 0:
-            return "No edits made"
-        else:
-            return (await models.gpt35()).complete(dedent(f"""{self._prompt}{self._completion}
-
-                Maximally concise summary of changes in bullet points (can use markdown):
-            """))
+        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]:
         rif_with_contents = []
@@ -101,28 +123,27 @@ class EditCodeStep(Step):
             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)
-        code_string = enc_dec.encode()
-        prompt = self.prompt.format(code=code_string)
-
-        completion = (await sdk.models.gpt35()).complete(prompt)
 
-        # Temporarily doing this to generate description.
-        self._prompt = prompt
-        self._completion = completion
+        rif_dict = {}
+        for rif in rif_with_contents:
+            rif_dict[rif.filepath] = rif.contents
 
-        file_edits = enc_dec.decode(completion)
+        for rif in rif_with_contents:
+            full_file_contents = await sdk.ide.readFile(rif.filepath)
+            segs = full_file_contents.split(rif.contents)
+            prompt = self._prompt.format(
+                code=rif.contents, user_request=self.user_input, file_prefix=segs[0], file_suffix=segs[1])
 
-        self._edit_diffs = []
-        for file_edit in file_edits:
-            diff = await sdk.apply_filesystem_edit(file_edit)
-            self._edit_diffs.append(diff)
+            completion = str((await sdk.models.gpt35()).complete(prompt))
+            eot_token = "<|endoftext|>"
+            completion = completion.removesuffix(eot_token)
 
-        for filepath in set([file_edit.filepath for file_edit in file_edits]):
-            await sdk.ide.saveFile(filepath)
-            await sdk.ide.setFileOpen(filepath)
+            self._prompt_and_completion += prompt + completion
 
-        return None
+            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)
 
 
 class EditFileStep(Step):
@@ -135,10 +156,10 @@ class EditFileStep(Step):
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
         file_contents = await sdk.ide.readFile(self.filepath)
-        await sdk.run_step(EditCodeStep(
+        await sdk.run_step(Gpt35EditCodeStep(
             range_in_files=[RangeInFile.from_entire_file(
                 self.filepath, file_contents)],
-            prompt=self.prompt
+            user_input=self.prompt
         ))
 
 
diff --git a/continuedev/src/continuedev/steps/main.py b/continuedev/src/continuedev/steps/main.py
index da0fc8d2..69c98bd4 100644
--- a/continuedev/src/continuedev/steps/main.py
+++ b/continuedev/src/continuedev/steps/main.py
@@ -15,7 +15,7 @@ from ..core.main import Step
 from ..core.sdk import ContinueSDK, Models
 from ..core.observation import Observation
 import subprocess
-from .core.core import EditCodeStep
+from .core.core import Gpt35EditCodeStep
 
 
 class SetupContinueWorkspaceStep(Step):
@@ -255,19 +255,9 @@ class StarCoderEditHighlightedCodeStep(Step):
 class EditHighlightedCodeStep(Step):
     user_input: str
     hide = True
-    _prompt: str = dedent("""Below is the code before changes:
-
-{code}
-
-This is the user request:
-
-{user_input}
-
-This is the code after being changed to perfectly satisfy the user request:
-    """)
 
     async def describe(self, models: Models) -> Coroutine[str, None, None]:
-        return "Editing highlighted code"
+        return "Editing code"
 
     async def run(self, sdk: ContinueSDK) -> Coroutine[Observation, None, None]:
         range_in_files = await sdk.ide.getHighlightedCode()
@@ -281,8 +271,7 @@ This is the code after being changed to perfectly satisfy the user request:
             range_in_files = [RangeInFile.from_entire_file(
                 filepath, content) for filepath, content in contents.items()]
 
-        await sdk.run_step(EditCodeStep(
-            range_in_files=range_in_files, prompt=self._prompt.format(code="{code}", user_input=self.user_input)))
+        await sdk.run_step(Gpt35EditCodeStep(user_input=self.user_input, range_in_files=range_in_files))
 
 
 class FindCodeStep(Step):