From 347f271bfa56d63f2b0e9831792510a9ae31b43d Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Fri, 20 Dec 2024 17:34:01 +0100 Subject: [PATCH] fix: loop in the schema generation --- ee/hogai/funnels/test/test_nodes.py | 3 ++- ee/hogai/schema_generator/nodes.py | 6 ++++-- ee/hogai/schema_generator/test/test_nodes.py | 10 ++++++---- ee/hogai/trends/test/test_nodes.py | 3 ++- 4 files changed, 14 insertions(+), 8 deletions(-) diff --git a/ee/hogai/funnels/test/test_nodes.py b/ee/hogai/funnels/test/test_nodes.py index 4f4e9fca0e5d4..91b53d13cb9d3 100644 --- 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 index 4bed02fd462cc..3b12bbc65fe30 100644 --- 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 index b44154b93b927..3b2702b55b7e1 100644 --- 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 index 369ce8bc9b292..004ab58408d8e 100644 --- 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="", ), )