diff options
| -rw-r--r-- | continuedev/src/continuedev/core/policy.py | 5 | ||||
| -rw-r--r-- | continuedev/src/continuedev/recipes/WritePytestsRecipe/main.py | 18 | ||||
| -rw-r--r-- | continuedev/src/continuedev/steps/comment_code.py | 12 | ||||
| -rw-r--r-- | continuedev/src/continuedev/steps/main.py | 1 |
4 files changed, 30 insertions, 6 deletions
diff --git a/continuedev/src/continuedev/core/policy.py b/continuedev/src/continuedev/core/policy.py index f7fcf21a..dc71cfa9 100644 --- a/continuedev/src/continuedev/core/policy.py +++ b/continuedev/src/continuedev/core/policy.py @@ -8,6 +8,7 @@ from ..steps.main import EditHighlightedCodeStep, SolveTracebackStep, RunCodeSte from ..recipes.WritePytestsRecipe.main import WritePytestsRecipe # from ..libs.steps.chroma import AnswerQuestionChroma, EditFileChroma from ..recipes.ContinueRecipeRecipe.main import ContinueStepStep +from ..steps.comment_code import CommentCodeStep class DemoPolicy(Policy): @@ -21,10 +22,12 @@ class DemoPolicy(Policy): observation = history.get_current().observation if observation is not None and isinstance(observation, UserInputObservation): # This could be defined with ObservationTypePolicy. Ergonomics not right though. - if " test" in observation.user_input.lower(): + if "/pytest" in observation.user_input.lower(): return WritePytestsRecipe(instructions=observation.user_input) elif "/dlt" in observation.user_input.lower() or " dlt" in observation.user_input.lower(): return CreatePipelineRecipe() + elif "/comment" in observation.user_input.lower(): + return CommentCodeStep() # elif "/ask" in observation.user_input: # return AnswerQuestionChroma(question=" ".join(observation.user_input.split(" ")[1:])) # elif "/edit" in observation.user_input: diff --git a/continuedev/src/continuedev/recipes/WritePytestsRecipe/main.py b/continuedev/src/continuedev/recipes/WritePytestsRecipe/main.py index a35055f3..82876a08 100644 --- a/continuedev/src/continuedev/recipes/WritePytestsRecipe/main.py +++ b/continuedev/src/continuedev/recipes/WritePytestsRecipe/main.py @@ -1,14 +1,20 @@ from textwrap import dedent +from typing import Union from ...models.filesystem_edit import AddDirectory, AddFile from ...core.main import Step, ContinueSDK import os class WritePytestsRecipe(Step): - for_filepath: str + for_filepath: Union[str, None] = None + instructions: str = "Write unit tests for this file." async def run(self, sdk: ContinueSDK): - filename, dirname = os.path.split(self.for_filepath) + if self.for_filepath is None: + self.for_filepath = (await sdk.ide.getOpenFiles())[0] + + filename = os.path.basename(self.for_filepath) + dirname = os.path.dirname(self.for_filepath) path_dir = os.path.join(dirname, "tests") if not os.path.exists(path_dir): @@ -16,7 +22,7 @@ class WritePytestsRecipe(Step): path = os.path.join(path_dir, f"test_{filename}") if os.path.exists(path): - return + return None for_file_contents = await sdk.ide.readFile(self.for_filepath) @@ -31,7 +37,9 @@ class WritePytestsRecipe(Step): "{self.instructions}" - Here is a complete set of pytest unit tests: - """) + Here is a complete set of pytest unit tests:""") tests = (await sdk.models.gpt35()).complete(prompt) + await sdk.apply_filesystem_edit(AddFile(filepath=path, content=tests)) + + return None diff --git a/continuedev/src/continuedev/steps/comment_code.py b/continuedev/src/continuedev/steps/comment_code.py new file mode 100644 index 00000000..c538347b --- /dev/null +++ b/continuedev/src/continuedev/steps/comment_code.py @@ -0,0 +1,12 @@ +from ..core.main import ContinueSDK, Models, Step +from .main import StarCoderEditHighlightedCodeStep + + +class CommentCodeStep(Step): + hide: bool = True + + async def describe(self, models: Models): + return "Writing comments" + + async def run(self, sdk: ContinueSDK): + await sdk.run_step(StarCoderEditHighlightedCodeStep(user_input="Wrote comprehensive comments in the canonical format for every class and function")) diff --git a/continuedev/src/continuedev/steps/main.py b/continuedev/src/continuedev/steps/main.py index 8588ec92..a0872564 100644 --- a/continuedev/src/continuedev/steps/main.py +++ b/continuedev/src/continuedev/steps/main.py @@ -215,6 +215,7 @@ class FasterEditHighlightedCodeStep(Step): class StarCoderEditHighlightedCodeStep(Step): user_input: str + name: str = "Editing code" hide = False _prompt: str = "<commit_before>{code}<commit_msg>{user_request}<commit_after>" |