diff options
Diffstat (limited to 'continuedev')
5 files changed, 23 insertions, 13 deletions
diff --git a/continuedev/src/continuedev/core/config.py b/continuedev/src/continuedev/core/config.py index cf723984..ed056be9 100644 --- a/continuedev/src/continuedev/core/config.py +++ b/continuedev/src/continuedev/core/config.py @@ -43,6 +43,16 @@ class ContinueConfig(BaseModel): description="Deploy a dlt pipeline to Airflow", step_name="DeployPipelineAirflowRecipe", ), + SlashCommand( + name="comment", + description="Write comments for the current file or highlighted code", + step_name="CommentCodeStep", + ), + SlashCommand( + name="pytest", + description="Write pytest unit tests for the current file", + step_name="WritePytestsRecipe", + ) ] diff --git a/continuedev/src/continuedev/core/policy.py b/continuedev/src/continuedev/core/policy.py index 37a10e36..51135f70 100644 --- a/continuedev/src/continuedev/core/policy.py +++ b/continuedev/src/continuedev/core/policy.py @@ -41,15 +41,11 @@ class DemoPolicy(Policy): after_command = " ".join(user_input.split(" ")[1:]) for slash_command in config.slash_commands: if slash_command.name == command_name[1:]: - return get_step_from_name(slash_command.step_name, slash_command.params) + params = slash_command.params + params["user_input"] = after_command + return get_step_from_name(slash_command.step_name, params) - if "/pytest" in user_input.lower(): - return WritePytestsRecipe(instructions=user_input) - if "/pytest" in observation.user_input.lower(): - return WritePytestsRecipe(instructions=observation.user_input) - elif "/comment" in observation.user_input.lower(): - return CommentCodeStep() - elif "/ask" in user_input: + if "/ask" in user_input: return AnswerQuestionChroma(question=" ".join(user_input.split(" ")[1:])) elif "/edit" in user_input: return EditFileChroma(request=" ".join(user_input.split(" ")[1:])) diff --git a/continuedev/src/continuedev/libs/util/step_name_to_steps.py b/continuedev/src/continuedev/libs/util/step_name_to_steps.py index 4023b73b..fe784964 100644 --- a/continuedev/src/continuedev/libs/util/step_name_to_steps.py +++ b/continuedev/src/continuedev/libs/util/step_name_to_steps.py @@ -7,6 +7,8 @@ from ...recipes.DDtoBQRecipe.main import DDtoBQRecipe from ...recipes.DeployPipelineAirflowRecipe.main import DeployPipelineAirflowRecipe from ...recipes.DDtoBQRecipe.main import DDtoBQRecipe from ...recipes.AddTransformRecipe.main import AddTransformRecipe +from ...recipes.WritePytestsRecipe.main import WritePytestsRecipe +from ...steps.comment_code import CommentCodeStep step_name_to_step_class = { "UserInputStep": UserInputStep, @@ -14,7 +16,9 @@ step_name_to_step_class = { "DDtoBQRecipe": DDtoBQRecipe, "DeployPipelineAirflowRecipe": DeployPipelineAirflowRecipe, "AddTransformRecipe": AddTransformRecipe, - "DDtoBQRecipe": DDtoBQRecipe + "DDtoBQRecipe": DDtoBQRecipe, + "WritePytestsRecipe": WritePytestsRecipe, + "CommentCodeStep": CommentCodeStep, } diff --git a/continuedev/src/continuedev/recipes/WritePytestsRecipe/main.py b/continuedev/src/continuedev/recipes/WritePytestsRecipe/main.py index 5994aa89..a19caf8b 100644 --- a/continuedev/src/continuedev/recipes/WritePytestsRecipe/main.py +++ b/continuedev/src/continuedev/recipes/WritePytestsRecipe/main.py @@ -7,7 +7,7 @@ import os class WritePytestsRecipe(Step): for_filepath: Union[str, None] = None - instructions: str = "Write unit tests for this file." + user_input: str = "Write unit tests for this file." async def run(self, sdk: ContinueSDK): if self.for_filepath is None: @@ -35,7 +35,7 @@ class WritePytestsRecipe(Step): Here are additional instructions: - "{self.instructions}" + "{self.user_input}" Here is a complete set of pytest unit tests:""") tests = sdk.models.gpt35.complete(prompt) diff --git a/continuedev/src/continuedev/steps/comment_code.py b/continuedev/src/continuedev/steps/comment_code.py index c538347b..aa17e62c 100644 --- a/continuedev/src/continuedev/steps/comment_code.py +++ b/continuedev/src/continuedev/steps/comment_code.py @@ -1,5 +1,5 @@ from ..core.main import ContinueSDK, Models, Step -from .main import StarCoderEditHighlightedCodeStep +from .main import EditHighlightedCodeStep class CommentCodeStep(Step): @@ -9,4 +9,4 @@ class CommentCodeStep(Step): 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")) + await sdk.run_step(EditHighlightedCodeStep(user_input="Write comprehensive comments in the canonical format for every class and function")) |