fix: loop in the schema generation

PostHog · Dec 20, 2024 · 347f271 · 347f271
1 parent 946f561
commit 347f271
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 8 deletions.
diff --git a/ee/hogai/funnels/test/test_nodes.py b/ee/hogai/funnels/test/test_nodes.py
@@ -33,6 +33,7 @@ def test_node_runs(self):
                 new_state,
                 PartialAssistantState(
                     messages=[VisualizationMessage(answer=self.schema, plan="Plan", id=new_state.messages[0].id)],
-                    intermediate_steps=None,
+                    intermediate_steps=[],
+                    plan="",
                 ),
             )
diff --git a/ee/hogai/schema_generator/nodes.py b/ee/hogai/schema_generator/nodes.py
@@ -96,7 +96,8 @@ def _run_with_prompt(
                             content=f"Oops! It looks like I’m having trouble generating this {self.INSIGHT_NAME} insight. Could you please try again?"
                         )
                     ],
-                    intermediate_steps=None,
+                    intermediate_steps=[],
+                    plan="",
                 )
 
             return PartialAssistantState(
@@ -115,7 +116,8 @@ def _run_with_prompt(
                     id=str(uuid4()),
                 )
             ],
-            intermediate_steps=None,
+            intermediate_steps=[],
+            plan="",
         )
 
     def router(self, state: AssistantState):

diff --git a/ee/hogai/schema_generator/test/test_nodes.py b/ee/hogai/schema_generator/test/test_nodes.py
@@ -54,7 +54,8 @@ def test_node_runs(self):
                 ),
                 {},
             )
-            self.assertIsNone(new_state.intermediate_steps)
+            self.assertEqual(new_state.intermediate_steps, [])
+            self.assertEqual(new_state.plan, "")
             self.assertEqual(len(new_state.messages), 1)
             self.assertEqual(new_state.messages[0].type, "ai/viz")
             self.assertEqual(new_state.messages[0].answer, self.schema)
@@ -316,7 +317,7 @@ def test_node_leaves_failover(self):
                 ),
                 {},
             )
-            self.assertIsNone(new_state.intermediate_steps)
+            self.assertEqual(new_state.intermediate_steps, [])
 
             new_state = node.run(
                 AssistantState(
@@ -328,7 +329,7 @@ def test_node_leaves_failover(self):
                 ),
                 {},
             )
-            self.assertIsNone(new_state.intermediate_steps)
+            self.assertEqual(new_state.intermediate_steps, [])
 
     def test_node_leaves_failover_after_second_unsuccessful_attempt(self):
         node = DummyGeneratorNode(self.team)
@@ -348,9 +349,10 @@ def test_node_leaves_failover_after_second_unsuccessful_attempt(self):
                 ),
                 {},
             )
-            self.assertIsNone(new_state.intermediate_steps)
+            self.assertEqual(new_state.intermediate_steps, [])
             self.assertEqual(len(new_state.messages), 1)
             self.assertIsInstance(new_state.messages[0], FailureMessage)
+            self.assertEqual(new_state.plan, "")
 
     def test_agent_reconstructs_conversation_with_failover(self):
         action = AgentAction(tool="fix", tool_input="validation error", log="exception")

diff --git a/ee/hogai/trends/test/test_nodes.py b/ee/hogai/trends/test/test_nodes.py
@@ -38,6 +38,7 @@ def test_node_runs(self):
                 new_state,
                 PartialAssistantState(
                     messages=[VisualizationMessage(answer=self.schema, plan="Plan", id=new_state.messages[0].id)],
-                    intermediate_steps=None,
+                    intermediate_steps=[],
+                    plan="",
                 ),
             )