improved prompt

hrshdhgd · Feb 12, 2024 · e451637 · e451637
1 parent a099592
commit e451637
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/src/codergpt/optimizer/optimizer.py b/src/codergpt/optimizer/optimizer.py
@@ -39,6 +39,7 @@ def optimize(
                     {
                         "input": f"Optimize, comment and add sphinx docstrings"
                         f" to the function '{function}' in \n\n```\n{source_code}\n```"
+                        "Return just the code block since all this will be a test file."
                         "Also explain the optimization in a systematic way as a comment."
                     }
                 )
@@ -47,6 +48,7 @@ def optimize(
                     {
                         "input": f"Optimize, comment and add sphinx docstrings"
                         f" to the class '{classname}' in \n\n```\n{source_code}\n```"
+                        "Return just the code block since all this will be a test file."
                         "Also explain the optimization in a systematic way."
                     }
                 )
@@ -56,6 +58,7 @@ def optimize(
                 {
                     "input": f"Optimize, comment and add sphinx style docstrings"
                     f" to the following code: \n\n```\n{source_code}\n```"
+                    "Return just the code block since all this will be a test file."
                     "Also explain the optimization in a systematic way."
                 }
             )