fix: raise SchemaGenerationException so we avoid FailureMessages

Author: skoob13

ee/hogai/graph/root/tools/create_and_query_insight.py

@@ -9,7 +9,9 @@
 
 from ee.hogai.context.context import AssistantContextManager
 from ee.hogai.graph.insights_graph.graph import InsightsGraph
+from ee.hogai.graph.schema_generator.nodes import SchemaGenerationException
 from ee.hogai.tool import MaxTool, ToolMessagesArtifact
+from ee.hogai.utils.prompt import format_prompt_string
 from ee.hogai.utils.types.base import AssistantState
 
 INSIGHT_TOOL_PROMPT = """
@@ -107,14 +109,35 @@
 </system_reminder>
 """.strip()
 
-INSIGHT_TOOL_FAILURE_PROMPT = """
-The agent has encountered an error while creating an insight.
+INSIGHT_TOOL_FAILURE_SYSTEM_REMINDER_PROMPT = """
 <system_reminder>
 Inform the user that you've encountered an error during the creation of the insight. Afterwards, try to generate a new insight with a different query.
 Terminate if the error persists.
 </system_reminder>
 """.strip()
 
+INSIGHT_TOOL_HANDLED_FAILURE_PROMPT = """
+The agent has encountered the error while creating an insight.
+
+Generated output:
+```
+{{{output}}}
+```
+
+Error message:
+```
+{{{error_message}}}
+```
+
+{{{system_reminder}}}
+""".strip()
+
+
+INSIGHT_TOOL_UNHANDLED_FAILURE_PROMPT = """
+The agent has encountered an unknown error while creating an insight.
+{{{system_reminder}}}
+""".strip()
+
 
 class CreateAndQueryInsightToolArgs(BaseModel):
     tool_call_id: Annotated[str, InjectedToolCallId, SkipJsonSchema]
@@ -145,12 +168,23 @@ async def _arun_impl(self, query_description: str, tool_call_id: str) -> tuple[s
             },
             deep=True,
         )
-        dict_state = await graph.ainvoke(new_state)
+        try:
+            dict_state = await graph.ainvoke(new_state)
+        except SchemaGenerationException as e:
+            return format_prompt_string(
+                INSIGHT_TOOL_HANDLED_FAILURE_PROMPT,
+                output=e.llm_output,
+                error_message=e.validation_message,
+                system_reminder=INSIGHT_TOOL_FAILURE_SYSTEM_REMINDER_PROMPT,
+            ), None
+
         updated_state = AssistantState.model_validate(dict_state)
         maybe_viz_message, tool_call_message = updated_state.messages[-2:]
 
         if not isinstance(tool_call_message, AssistantToolCallMessage):
-            return INSIGHT_TOOL_FAILURE_PROMPT, None
+            return format_prompt_string(
+                INSIGHT_TOOL_UNHANDLED_FAILURE_PROMPT, system_reminder=INSIGHT_TOOL_FAILURE_SYSTEM_REMINDER_PROMPT
+            ), None
 
         # If the previous message is not a visualization message, the agent has requested human feedback.
         if not isinstance(maybe_viz_message, VisualizationMessage):

ee/hogai/graph/schema_generator/nodes.py

@@ -13,7 +13,7 @@
 from langchain_core.prompts import ChatPromptTemplate, HumanMessagePromptTemplate
 from langchain_core.runnables import RunnableConfig
 
-from posthog.schema import FailureMessage, VisualizationMessage
+from posthog.schema import VisualizationMessage
 
 from posthog.models.group_type_mapping import GroupTypeMapping
 
@@ -36,6 +36,15 @@
 RETRIES_ALLOWED = 2
 
 
+class SchemaGenerationException(Exception):
+    """An error occurred while generating a schema in the `SchemaGeneratorNode` node."""
+
+    def __init__(self, llm_output: str, validation_message: str):
+        super().__init__("Failed to generate schema")
+        self.llm_output = llm_output
+        self.validation_message = validation_message
+
+
 class SchemaGeneratorNode(AssistantNode, Generic[Q]):
     INSIGHT_NAME: str
     """
@@ -87,9 +96,8 @@ async def _run_with_prompt(
 
         chain = generation_prompt | merger | self._model | self._parse_output
 
-        result: SchemaGeneratorOutput[Q] | None = None
         try:
-            result = await chain.ainvoke(
+            result: SchemaGeneratorOutput[Q] = await chain.ainvoke(
                 {
                     "project_datetime": self.project_now,
                     "project_timezone": self.project_timezone,
@@ -120,18 +128,9 @@ async def _run_with_prompt(
                     query_generation_retry_count=len(intermediate_steps) + 1,
                 )
 
-        if not result:
-            # We've got no usable result after exhausting all iteration attempts - it's failure message time
-            return PartialAssistantState(
-                messages=[
-                    FailureMessage(
-                        content=f"It looks like I'm having trouble generating this {self.INSIGHT_NAME} insight."
-                    )
-                ],
-                intermediate_steps=None,
-                plan=None,
-                query_generation_retry_count=len(intermediate_steps) + 1,
-            )
+            if isinstance(e, PydanticOutputParserException):
+                raise SchemaGenerationException(e.llm_output, e.validation_message)
+            raise SchemaGenerationException(e.llm_output or "No input was provided.", str(e))
 
         # We've got a result that either passed the quality check or we've exhausted all attempts at iterating - return
         return PartialAssistantState(