Refactor Root Cause Analysis Models

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

@@ -1,7 +1,7 @@
 from typing import Annotated, Optional
 
 from pydantic import BaseModel, StringConstraints, field_validator
-from pydantic_xml import attr
+from pydantic_xml import attr 
 
 from seer.automation.agent.models import Message
 from seer.automation.autofix.utils import remove_code_backticks
@@ -31,41 +31,19 @@ class RootCauseRelevantCodeSnippet(BaseModel):
     start_line: int | None = None
     end_line: int | None = None
 
-
 class RootCauseRelevantContext(BaseModel):
     id: int
     title: str
     description: str
     snippet: Optional[RootCauseRelevantCodeSnippet]
-
-
+
 class RootCauseAnalysisRelevantContext(BaseModel):
     snippets: list[RootCauseRelevantContext]
 
-
-class UnitTestSnippetPrompt(BaseModel):
-    file_path: str
-    code_snippet: str
-    description: str
-
-    @field_validator("code_snippet")
-    @classmethod
-    def clean_code_snippet(cls, v: str) -> str:
-        return remove_code_backticks(v)
-
-
-class UnitTestSnippet(BaseModel):
-    file_path: str
-    snippet: str
-    description: str
-
-
 class RootCauseAnalysisItem(BaseModel):
     id: int = -1
     title: str
     description: str
-    # unit_test: UnitTestSnippet | None = None
-    # reproduction: str | None = None
     code_context: Optional[list[RootCauseRelevantContext]] = None
 
     def to_markdown_string(self) -> str:
@@ -81,35 +59,21 @@ def to_markdown_string(self) -> str:
                     markdown += f"**File:** {context.snippet.file_path}\n"
                     if context.snippet.repo_name:
                         markdown += f"**Repository:** {context.snippet.repo_name}\n"
-                    markdown += "```\n"
+                    markdown += f"```\n"
                     markdown += f"{context.snippet.snippet}\n"
                     markdown += "```\n\n"
-
         return markdown.strip()
 
-
 class RootCauseAnalysisItemPrompt(BaseModel):
     title: str
     description: str
-    # 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
-            # ),
             relevant_code=(
                 RootCauseAnalysisRelevantContext(
                     snippets=[
@@ -131,16 +95,6 @@ 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
-                # ),
                 "code_context": (
                     self.relevant_code.model_dump()["snippets"] if self.relevant_code else None
                 ),