Add unit test options to coding step

src/seer/automation/autofix/components/coding/component.py

@@ -120,6 +120,7 @@ class IsObviousOutput(BaseModel):
                     original_instruction=request.original_instruction,
                     root_cause_extra_instruction=request.root_cause_extra_instruction,
                     custom_solution=request.solution if isinstance(request.solution, str) else None,
+                    mode=request.mode,
                 ),
                 model=OpenAiProvider.model("gpt-4o-mini"),
                 response_format=IsObviousOutput,
@@ -180,7 +181,7 @@ def invoke(self, request: CodingRequest) -> CodingOutput | None:
                 RunConfig(
                     system_prompt=CodingPrompts.format_system_msg(has_tools=not is_obvious),
                     prompt=CodingPrompts.format_fix_msg(
-                        has_tools=not is_obvious, custom_solution=custom_solution
+                        has_tools=not is_obvious, custom_solution=custom_solution, mode=request.mode
                     ),
                     model=AnthropicProvider.model("claude-3-5-sonnet-v2@20241022"),
                     memory_storage_key="code",

src/seer/automation/autofix/components/coding/models.py

@@ -1,5 +1,5 @@
 import textwrap
-from typing import Annotated, Optional
+from typing import Annotated, Literal, Optional
 
 from pydantic import BaseModel, StringConstraints, field_validator
 from pydantic_xml import attr, element
@@ -25,6 +25,7 @@ class CodingRequest(BaseComponentRequest):
     summary: Optional[IssueSummary] = None
     initial_memory: list[Message] = []
     profile: Profile | None = None
+    mode: Literal["all", "fix", "test"] = "fix"
 
 
 class SnippetXml(PromptXmlModel, tag="snippet"):

src/seer/automation/autofix/components/coding/prompts.py

@@ -1,5 +1,5 @@
 import textwrap
-from typing import Optional
+from typing import Literal, Optional
 
 from seer.automation.autofix.components.coding.models import (
     FuzzyDiffChunk,
@@ -60,10 +60,14 @@ def format_custom_solution(custom_solution: str | None):
         return f"The user has provided the following solution idea: {custom_solution}"
 
     @staticmethod
-    def format_fix_msg(has_tools: bool = True, custom_solution: str | None = None):
+    def format_fix_msg(
+        has_tools: bool = True,
+        custom_solution: str | None = None,
+        mode: Literal["all", "fix", "test"] = "fix",
+    ):
         return textwrap.dedent(
             """\
-            Break down the task of fixing the issue into steps. Your list of steps should be detailed enough so that following it exactly will lead to a fully complete solution. {custom_solution_str}
+            Break down the task of {mode_str} into steps. Your list of steps should be detailed enough so that following it exactly will lead to a fully complete solution. {custom_solution_str}
 
             Enclose this plan between <plan_steps> and </plan_steps> tags. Make sure to strictly follow this format and include all necessary details within the tags. Your output must follow the format properly according to the following guidelines:
 
@@ -79,6 +83,15 @@ def format_fix_msg(has_tools: bool = True, custom_solution: str | None = None):
             {think_tools_instructions}
             - You also MUST think step-by-step before giving the final answer."""
         ).format(
+            mode_str=(
+                "fixing the issue"
+                if mode == "fix"
+                else (
+                    "writing a unit test to reproduce the issue and assert the planned solution (following test-driven development)"
+                    if mode == "test"
+                    else "writing a unit test to reproduce the issue and assert the planned solution (following test-driven development) and then fixing the issue"
+                )
+            ),
             steps_example_str=PlanStepsPromptXml.get_example().to_prompt_str(),
             use_tools_instructions=(
                 "- Make sure you use the tools provided to look through the codebase and at the files you are changing before outputting the steps."
@@ -133,6 +146,7 @@ def format_is_obvious_msg(
         original_instruction: str | None,
         root_cause_extra_instruction: str | None,
         custom_solution: str | None,
+        mode: Literal["all", "fix", "test"] = "fix",
     ):
         return (
             textwrap.dedent(
@@ -143,9 +157,18 @@ def format_is_obvious_msg(
                 {root_cause_str}{root_cause_extra_instruction}
 
                 {custom_solution_str}
-                Does the code change exist ONLY in files you can already see in your context here or do you need to look at other files?"""
+                Does the code change needed for {mode_str} exist ONLY in files you can already see in your context here or do you need to look at other files?"""
             )
             .format(
+                mode_str=(
+                    "fixing the issue"
+                    if mode == "fix"
+                    else (
+                        "writing a unit test to reproduce the issue"
+                        if mode == "test"
+                        else "writing a unit test and fixing the issue"
+                    )
+                ),
                 summary_str=format_summary(summary),
                 event_details=event_details.format_event(),
                 root_cause_str=CodingPrompts.format_root_cause(root_cause),

src/seer/automation/autofix/components/insight_sharing/component.py

@@ -23,44 +23,44 @@ def format_step_one(
 
         if step_type == "root_cause_analysis_processing":
             template = """\
+            We have a new set of notes and thoughts:
+            <NOTES>
+            {latest_thought}
+            </NOTES>
+
             We are in the process of writing a small paragraph describing the in-depth root cause of an issue in our codebase. Here it is so far:
             <PARAGRAPH_SO_FAR>
             {insights}
             </PARAGRAPH_SO_FAR>
 
+            If there is something new and useful here to extend the root cause analysis, what is the NEXT sentence in the paragraph (write it so it flows nicely; max 15 words)? {no_insight_instruction}
+            """
+        elif step_type == "plan":
+            template = """\
             We have a new set of notes and thoughts about the issue:
             <NOTES>
             {latest_thought}
             </NOTES>
 
-            If there is something new and useful here to extend the root cause analysis, what is the NEXT sentence in the paragraph (write it so it flows nicely; max 15 words)? {no_insight_instruction}
-            """
-        elif step_type == "plan":
-            template = """\
-            We are in the process of writing a small paragraph describing the code changes we need to make to fix an issue in our codebase. Here it is so far:
+            We are in the process of writing a small paragraph describing the code changes we need to make in our codebase. Here it is so far:
             <PARAGRAPH_SO_FAR>
             {insights}
             </PARAGRAPH_SO_FAR>
 
+            If there is something new and useful here to extend the description on the code changes, what is the NEXT sentence in the paragraph (write it so it flows nicely; max 15 words)? {no_insight_instruction}
+            """
+        elif step_type == "solution_processing":
+            template = """\
             We have a new set of notes and thoughts about the issue:
             <NOTES>
             {latest_thought}
             </NOTES>
 
-            If there is something new and useful here to extend the discussion on the code changes, what is the NEXT sentence in the paragraph (write it so it flows nicely; max 15 words)? {no_insight_instruction}
-            """
-        elif step_type == "solution_processing":
-            template = """\
             We are in the process of writing a small paragraph describing a solution to an issue in our codebase. Here it is so far:
             <PARAGRAPH_SO_FAR>
             {insights}
             </PARAGRAPH_SO_FAR>
 
-            We have a new set of notes and thoughts about the issue:
-            <NOTES>
-            {latest_thought}
-            </NOTES>
-
             If there is something new and useful here to extend the plan on how to fix the issue, what is the NEXT sentence in the paragraph (write it so it flows nicely; max 15 words)? {no_insight_instruction}
             """
         else:

src/seer/automation/autofix/event_manager.py

@@ -171,6 +171,7 @@ def set_selected_solution(self, payload: AutofixSolutionUpdatePayload):
             solution_step.custom_solution = (
                 payload.custom_solution if payload.custom_solution else None
             )
+            solution_step.selected_mode = payload.mode
             solution_step.solution_selected = True
             cur.delete_steps_after(solution_step, include_current=False)
             cur.clear_file_changes()
@@ -179,6 +180,7 @@ def set_selected_solution(self, payload: AutofixSolutionUpdatePayload):
 
     def send_coding_start(self):
         with self.state.update() as cur:
+            cur.clear_file_changes()
             plan_step = cur.find_step(key=self.plan_step.key)
             if not plan_step or plan_step.status != AutofixStatus.PROCESSING:
                 plan_step = cur.add_step(self.plan_step)

src/seer/automation/autofix/models.py

@@ -219,6 +219,7 @@ class SolutionStep(BaseStep):
     solution: list[SolutionTimelineEvent] = []
     custom_solution: str | None = None
     solution_selected: bool = False
+    selected_mode: Literal["all", "fix", "test"] | None = None
 
 
 class ChangesStep(BaseStep):
@@ -382,6 +383,7 @@ class AutofixSolutionUpdatePayload(BaseModel):
     type: Literal[AutofixUpdateType.SELECT_SOLUTION] = AutofixUpdateType.SELECT_SOLUTION
     custom_solution: str | None = None
     solution_selected: bool = False
+    mode: Literal["all", "fix", "test"] = "fix"
 
 
 class AutofixCreatePrUpdatePayload(BaseModel):
@@ -529,15 +531,17 @@ def get_selected_root_cause(
                     return root_cause_step.selection.custom_root_cause, None
         return None, None
 
-    def get_selected_solution(self) -> list[SolutionTimelineEvent] | str | None:
+    def get_selected_solution(
+        self,
+    ) -> tuple[list[SolutionTimelineEvent] | str | None, Literal["all", "fix", "test"] | None]:
         solution_step = self.find_step(key="solution")
         if solution_step and isinstance(solution_step, SolutionStep):
             if solution_step.solution_selected:
                 if solution_step.custom_solution:
-                    return solution_step.custom_solution
+                    return solution_step.custom_solution, solution_step.selected_mode
                 else:
-                    return solution_step.solution
-        return None
+                    return solution_step.solution, solution_step.selected_mode
+        return None, None
 
     def mark_triggered(self):
         self.last_triggered_at = datetime.datetime.now()

src/seer/automation/autofix/steps/coding_step.py

@@ -62,20 +62,23 @@ def _invoke(self, app_config: AppConfig = injected):
         self.logger.info("Executing Autofix - Coding Step")
 
         self.context.event_manager.send_coding_start()
-        if not self.request.initial_memory:
-            self.context.event_manager.add_log("Coding up a fix for this issue...")
-        else:
-            self.context.event_manager.add_log("Continuing to code...")
 
         state = self.context.state.get()
         root_cause, root_cause_extra_instruction = state.get_selected_root_cause()
         if not root_cause:
             raise ValueError("Root cause analysis must be performed before coding")
 
-        solution = state.get_selected_solution()
+        solution, coding_mode = state.get_selected_solution()
         if not solution:
             raise ValueError("Solution must be found before coding")
 
+        if not self.request.initial_memory:
+            self.context.event_manager.add_log(
+                f"Coding up a {'fix' if coding_mode == 'fix' else 'test' if coding_mode == 'test' else 'fix and a test'} for this issue..."
+            )
+        else:
+            self.context.event_manager.add_log("Continuing to code...")
+
         event_details = EventDetails.from_event(state.request.issue.events[0])
         self.context.process_event_paths(event_details)
 
@@ -93,6 +96,7 @@ def _invoke(self, app_config: AppConfig = injected):
                 summary=summary,
                 initial_memory=self.request.initial_memory,
                 profile=state.request.profile,
+                mode=coding_mode,
             )
         )
 

src/seer/automation/autofix/tasks.py

@@ -250,6 +250,7 @@ def run_autofix_coding(
         AutofixCodingStep.get_signature(
             AutofixCodingStepRequest(
                 run_id=cur.run_id,
+                mode=request.payload.mode,
             ),
             queue=app_config.CELERY_WORKER_QUEUE,
         ).apply_async()