Ah yes Uncommented Unit Test and Reproduction Fields in Root Cause Mo…

…dels
getsentry · Jan 23, 2025 · 73bfd9a · 73bfd9a
1 parent 7861f2d
commit 73bfd9a
Showing 1 changed file with 24 additions and 24 deletions.
diff --git a/src/seer/automation/autofix/components/root_cause/models.py b/src/seer/automation/autofix/components/root_cause/models.py
@@ -64,8 +64,8 @@ class RootCauseAnalysisItem(BaseModel):
     id: int = -1
     title: str
     description: str
-    # unit_test: UnitTestSnippet | None = None
-    # reproduction: str | None = None
+    unit_test: UnitTestSnippet | None = None
+    reproduction: str | None = None
     code_context: Optional[list[RootCauseRelevantContext]] = None
 
     def to_markdown_string(self) -> str:
@@ -91,25 +91,25 @@ def to_markdown_string(self) -> str:
 class RootCauseAnalysisItemPrompt(BaseModel):
     title: str
     description: str
-    # reproduction_instructions: str | None = None
-    # unit_test: UnitTestSnippetPrompt | None = None
+    reproduction_instructions: str | None = None
+    unit_test: UnitTestSnippetPrompt | None = None
     relevant_code: Optional[RootCauseAnalysisRelevantContext]
 
     @classmethod
     def from_model(cls, model: RootCauseAnalysisItem):
         return cls(
             title=model.title,
             description=model.description,
-            # reproduction_instructions=model.reproduction,
-            # unit_test=(
-            #     UnitTestSnippetPrompt(
-            #         file_path=model.unit_test.file_path,
-            #         code_snippet=model.unit_test.snippet,
-            #         description=model.unit_test.description,
-            #     )
-            #     if model.unit_test
-            #     else None
-            # ),
+            reproduction_instructions=model.reproduction,
+            unit_test=(
+                UnitTestSnippetPrompt(
+                    file_path=model.unit_test.file_path,
+                    code_snippet=model.unit_test.snippet,
+                    description=model.unit_test.description,
+                )
+                if model.unit_test
+                else None
+            ),
             relevant_code=(
                 RootCauseAnalysisRelevantContext(
                     snippets=[
@@ -131,16 +131,16 @@ def to_model(self):
         return RootCauseAnalysisItem.model_validate(
             {
                 **self.model_dump(),
-                # "reproduction": self.reproduction_instructions,
-                # "unit_test": (
-                #     {
-                #         "file_path": self.unit_test.file_path,
-                #         "snippet": self.unit_test.code_snippet,
-                #         "description": self.unit_test.description,
-                #     }
-                #     if self.unit_test
-                #     else None
-                # ),
+                "reproduction": self.reproduction_instructions,
+                "unit_test": (
+                    {
+                        "file_path": self.unit_test.file_path,
+                        "snippet": self.unit_test.code_snippet,
+                        "description": self.unit_test.description,
+                    }
+                    if self.unit_test
+                    else None
+                ),
                 "code_context": (
                     self.relevant_code.model_dump()["snippets"] if self.relevant_code else None
                 ),