tutorial steps: comment, pytest

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>"