Enhancements to Root Cause Analysis with Error Handling and Model Upd…

…ates

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

@@ -1,6 +1,7 @@
 import logging
 
 from langfuse.decorators import observe
+from sentry_sdk import capture_exception
 from sentry_sdk.ai.monitoring import ai_track
 
 from seer.automation.agent.agent import AgentConfig, RunConfig
@@ -121,11 +122,26 @@ def invoke(
                     run_name="Root Cause Extraction & Formatting",
                 )
 
-                # Assign the ids to be the numerical indices of the causes and relevant code context
-                cause_model = formatted_response.parsed.cause.to_model()
-                cause_model.id = 0
+                try:
+                    # Attempt to create the model with validation
+                    cause_model = formatted_response.parsed.cause.to_model()
+                    cause_model.id = 0
+                except ValidationError as e:
+                    logger.error(f"Validation error in root cause analysis: {str(e)}")
+                    capture_exception(e)  # Log to Sentry for monitoring
+
+                    # Create a fallback model with required fields
+                    cause_model = RootCauseAnalysisItem(
+                        id=0,
+                        title="Error During Analysis",
+                        description=(
+                            f"An error occurred while processing the root cause analysis: {str(e)}\n\n"
+                            "The system encountered an issue while validating the analysis data structure."
+                        )
+                    )
+
                 if cause_model.code_context:
-                    for i, item in enumerate(cause_model.code_context):
+                    for j, snippet in enumerate(cause_model.code_context):
                         item.id = i
                         # Find line range for the snippet
                         if item.snippet and item.snippet.file_path and item.snippet.snippet:

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

@@ -1,6 +1,6 @@
 from typing import Annotated, Optional
 
-from pydantic import BaseModel, StringConstraints, field_validator
+from pydantic import BaseModel, ConfigDict, StringConstraints, field_validator
 from pydantic_xml import attr
 
 from seer.automation.agent.models import Message
@@ -64,10 +64,17 @@ class RootCauseAnalysisItem(BaseModel):
     id: int = -1
     title: str
     description: str
-    # unit_test: UnitTestSnippet | None = None
-    # reproduction: str | None = None
+    unit_test: Optional[UnitTestSnippet] = None
+    reproduction: Optional[str] = None
     code_context: Optional[list[RootCauseRelevantContext]] = None
 
+    model_config = ConfigDict(
+        extra='ignore',  # Ignore extra fields during validation
+        validate_assignment=True,  # Validate assignments to fields
+        arbitrary_types_allowed=True  # Allow arbitrary types in validation
+    )
+
+
     def to_markdown_string(self) -> str:
         markdown = f"# {self.title}\n\n"
         markdown += f"## Description\n{self.description}\n\n" if self.description else ""
@@ -91,25 +98,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: Optional[str] = None
+    unit_test: Optional[UnitTestSnippetPrompt] = 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=[
@@ -128,24 +135,25 @@ def from_model(cls, model: RootCauseAnalysisItem):
         )
 
     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": (
+        model_data = {
+            'title': self.title,
+            'description': self.description,
+            'code_context': (
                     self.relevant_code.model_dump()["snippets"] if self.relevant_code else None
-                ),
-            }
-        )
+            )
+        }
+
+        if self.reproduction_instructions:
+            model_data['reproduction'] = self.reproduction_instructions
+
+        if self.unit_test:
+            model_data['unit_test'] = UnitTestSnippet(
+                file_path=self.unit_test.file_path,
+                snippet=self.unit_test.code_snippet,
+                description=self.unit_test.description
+            )
+
+        return RootCauseAnalysisItem.model_validate(model_data)
 
 
 class MultipleRootCauseAnalysisOutputPrompt(BaseModel):