fix: describe insights better

PostHog · Feb 11, 2025 · 42ad064 · 42ad064
1 parent e1f5064
commit 42ad064
Show file tree

Hide file tree

Showing 5 changed files with 38 additions and 22 deletions.
diff --git a/ee/hogai/query_executor/test/test_nodes.py b/ee/hogai/query_executor/test/test_nodes.py
@@ -32,7 +32,7 @@ def test_node_runs(self, mock_process_query_dict):
                         tool_calls=[
                             {
                                 "id": "tool1",
-                                "name": "retrieve_data_for_question",
+                                "name": "create_and_query_insight",
                                 "args": {"query_kind": "trends", "query_description": "test query"},
                             }
                         ],

diff --git a/ee/hogai/root/nodes.py b/ee/hogai/root/nodes.py
@@ -28,10 +28,9 @@
 
 
 # Lower casing matters here. Do not change it.
-class retrieve_data_for_question(BaseModel):
+class create_and_query_insight(BaseModel):
     """
-    Retrieve results for a specific data question by creating a query, or iterate on a previous query.
-
+    Retrieve results for a specific data question by creating a query or iterate on a previous query.
     This tool only retrieves data for a single insight at a time.
     The `trends` insight type is the only insight that can display multiple trends insights in one request.
     All other insight types strictly return data for a single insight.
@@ -41,8 +40,8 @@ class retrieve_data_for_question(BaseModel):
     query_kind: Literal["trends", "funnel", "retention"] = Field(description=ROOT_INSIGHT_DESCRIPTION_PROMPT)
 
 
-RootToolCall = retrieve_data_for_question
-root_tools_parser = PydanticToolsParser(tools=[retrieve_data_for_question])
+RootToolCall = create_and_query_insight
+root_tools_parser = PydanticToolsParser(tools=[create_and_query_insight])
 
 
 class RootNode(AssistantNode):
@@ -87,7 +86,7 @@ def _model(self):
         # conversational context we're using a temperature of 0, for near determinism (https://arxiv.org/html/2405.00492v1)
         return ChatOpenAI(
             model="gpt-4o", temperature=0.0, streaming=True, stream_usage=True, parallel_tool_calls=False
-        ).bind_tools([retrieve_data_for_question], strict=True)
+        ).bind_tools([create_and_query_insight], strict=True)
 
     def _construct_messages(self, state: AssistantState) -> list[BaseMessage]:
         # `assistant` messages must be contiguous with the respective `tool` messages.

diff --git a/ee/hogai/root/prompts.py b/ee/hogai/root/prompts.py
@@ -58,18 +58,33 @@
 
 A trends insight visualizes events over time using time series. They're useful for finding patterns in historical data.
 
+The trends insights have the following features:
+- The insight can show multiple trends in one request.
+- Custom formulas can calculate derived metrics, like `A/B*100` to calculate a ratio.
+- Filter and break down data using multiple properties.
+- Compare with the previous period and sample data.
+- Apply various aggregation types, like sum, average, etc., and chart types.
+- And more.
+
 Examples of use cases include:
 - How the product's most important metrics change over time.
 - Long-term patterns, or cycles in product's usage.
 - The usage of different features side-by-side.
 - How the properties of events vary using aggregation (sum, average, etc).
 - Users can also visualize the same data points in a variety of ways.
-- Custom formulas can calculate derived metrics, like `A/B*100` to calculate a ratio.
 
 ## `funnel`
 
 A funnel insight visualizes a sequence of events that users go through in a product. They use percentages as the primary aggregation type. Funnels use two or more series, so the conversation history should mention at least two events.
 
+The funnel insights have the following features:
+- Various visualization types (steps, time-to-convert, historical trends).
+- Filter data and apply exclusion steps.
+- Break down data using a single property.
+- Specify conversion windows, details of conversion calculation, attribution settings.
+- Sample data.
+- And more.
+
 Examples of use cases include:
 - Conversion rates.
 - Drop off steps.
@@ -80,6 +95,8 @@
 
 A retention insight visualizes how many users return to the product after performing some action. They're useful for understanding user engagement and retention.
 
+The retention insights have the following features: filter data, sample data, and more.
+
 Examples of use cases include:
 - How many users come back and perform an action after their first visit.
 - How many users come back to perform action X after performing action Y.

diff --git a/ee/hogai/root/test/test_nodes.py b/ee/hogai/root/test/test_nodes.py
@@ -48,7 +48,7 @@ def test_node_handles_insight_tool_call(self, insight_type):
                     tool_calls=[
                         {
                             "id": "xyz",
-                            "name": "retrieve_data_for_question",
+                            "name": "create_and_query_insight",
                             "args": {"query_description": "Foobar", "query_kind": insight_type},
                         }
                     ],
@@ -69,7 +69,7 @@ def test_node_handles_insight_tool_call(self, insight_type):
                 assistant_message.tool_calls[0],
                 AssistantToolCall(
                     id="xyz",
-                    name="retrieve_data_for_question",
+                    name="create_and_query_insight",
                     args={"query_description": "Foobar", "query_kind": insight_type},
                 ),
             )
@@ -90,7 +90,7 @@ def test_node_handles_insight_tool_call_without_message(self, insight_type):
                     tool_calls=[
                         {
                             "id": "xyz",
-                            "name": "retrieve_data_for_question",
+                            "name": "create_and_query_insight",
                             "args": {"query_description": "Foobar", "query_kind": insight_type},
                         }
                     ],
@@ -111,7 +111,7 @@ def test_node_handles_insight_tool_call_without_message(self, insight_type):
                 assistant_message.tool_calls[0],
                 AssistantToolCall(
                     id="xyz",
-                    name="retrieve_data_for_question",
+                    name="create_and_query_insight",
                     args={"query_description": "Foobar", "query_kind": insight_type},
                 ),
             )
@@ -148,7 +148,7 @@ def test_node_reconstructs_conversation_with_tool_calls(self):
                     tool_calls=[
                         {
                             "id": "xyz",
-                            "name": "retrieve_data_for_question",
+                            "name": "create_and_query_insight",
                             "args": {},
                         }
                     ],
@@ -167,7 +167,7 @@ def test_node_reconstructs_conversation_with_tool_calls(self):
                     tool_calls=[
                         {
                             "id": "xyz",
-                            "name": "retrieve_data_for_question",
+                            "name": "create_and_query_insight",
                             "args": {},
                         }
                     ],
@@ -225,7 +225,7 @@ def test_run_validation_error(self):
                     tool_calls=[
                         AssistantToolCall(
                             id="xyz",
-                            name="retrieve_data_for_question",
+                            name="create_and_query_insight",
                             args={"invalid_field": "should fail validation"},
                         )
                     ],
@@ -249,7 +249,7 @@ def test_run_valid_tool_call(self):
                     tool_calls=[
                         AssistantToolCall(
                             id="xyz",
-                            name="retrieve_data_for_question",
+                            name="create_and_query_insight",
                             args={"query_kind": "trends", "query_description": "test query"},
                         )
                     ],
@@ -272,12 +272,12 @@ def test_run_multiple_tool_calls_raises(self):
                     tool_calls=[
                         AssistantToolCall(
                             id="xyz1",
-                            name="retrieve_data_for_question",
+                            name="create_and_query_insight",
                             args={"query_kind": "trends", "query_description": "test query 1"},
                         ),
                         AssistantToolCall(
                             id="xyz2",
-                            name="retrieve_data_for_question",
+                            name="create_and_query_insight",
                             args={"query_kind": "funnel", "query_description": "test query 2"},
                         ),
                     ],

diff --git a/ee/hogai/test/test_assistant.py b/ee/hogai/test/test_assistant.py
@@ -455,7 +455,7 @@ def test_full_trends_flow(self, memory_collector_mock, root_mock, planner_mock,
                         tool_calls=[
                             {
                                 "id": "xyz",
-                                "name": "retrieve_data_for_question",
+                                "name": "create_and_query_insight",
                                 "args": {"query_description": "Foobar", "query_kind": "trends"},
                             }
                         ],
@@ -493,7 +493,7 @@ def test_full_trends_flow(self, memory_collector_mock, root_mock, planner_mock,
                     tool_calls=[
                         {
                             "id": "xyz",
-                            "name": "retrieve_data_for_question",
+                            "name": "create_and_query_insight",
                             "args": {"query_description": "Foobar", "query_kind": "trends"},
                         }
                     ],
@@ -531,7 +531,7 @@ def test_full_funnel_flow(self, memory_collector_mock, root_mock, planner_mock,
                         tool_calls=[
                             {
                                 "id": "xyz",
-                                "name": "retrieve_data_for_question",
+                                "name": "create_and_query_insight",
                                 "args": {"query_description": "Foobar", "query_kind": "funnel"},
                             }
                         ],
@@ -574,7 +574,7 @@ def test_full_funnel_flow(self, memory_collector_mock, root_mock, planner_mock,
                     tool_calls=[
                         {
                             "id": "xyz",
-                            "name": "retrieve_data_for_question",
+                            "name": "create_and_query_insight",
                             "args": {"query_description": "Foobar", "query_kind": "funnel"},
                         }
                     ],