feat(llm-observability): speed up the traces scene (#27955)

Co-authored-by: Peter Kirkham <[email protected]>

posthog/hogql/database/test/test_database.py

@@ -250,13 +250,13 @@ def test_serialize_database_warehouse_table_source_query_count(self):
         for i in range(5):
             source = ExternalDataSource.objects.create(
                 team=self.team,
-                source_id=f"source_id_{i+2}",
-                connection_id=f"connection_id_{i+2}",
+                source_id=f"source_id_{i + 2}",
+                connection_id=f"connection_id_{i + 2}",
                 status=ExternalDataSource.Status.COMPLETED,
                 source_type=ExternalDataSource.Type.STRIPE,
             )
             warehouse_table = DataWarehouseTable.objects.create(
-                name=f"table_{i+2}",
+                name=f"table_{i + 2}",
                 format="Parquet",
                 team=self.team,
                 external_data_source=source,
@@ -269,7 +269,7 @@ def test_serialize_database_warehouse_table_source_query_count(self):
             )
             ExternalDataSchema.objects.create(
                 team=self.team,
-                name=f"table_{i+2}",
+                name=f"table_{i + 2}",
                 source=source,
                 table=warehouse_table,
                 should_sync=True,

posthog/hogql_queries/ai/test/test_traces_query_runner.py

@@ -227,7 +227,6 @@ def test_field_mapping(self):
         self.assertEqual(len(response.results), 2)
 
         trace = response.results[0]
-
         self.assertTraceEqual(
             trace,
             {
@@ -244,7 +243,28 @@ def test_field_mapping(self):
             },
         )
         self.assertEqual(trace.person.distinct_id, "person1")
+        self.assertFalse(trace.events)
+
+        trace = response.results[1]
+        self.assertTraceEqual(
+            trace,
+            {
+                "id": "trace2",
+                "createdAt": datetime(2025, 1, 14, tzinfo=UTC).isoformat(),
+                "totalLatency": 1,
+                "inputTokens": 3,
+                "outputTokens": 3,
+                "inputCost": 3,
+                "outputCost": 3,
+                "totalCost": 6,
+            },
+        )
+        self.assertEqual(trace.person.distinct_id, "person2")
+        self.assertFalse(trace.events)
 
+        response = TracesQueryRunner(team=self.team, query=TracesQuery(traceId="trace1")).calculate()
+        self.assertEqual(len(response.results), 1)
+        trace = response.results[0]
         self.assertEqual(len(trace.events), 2)
         event = trace.events[0]
         self.assertIsNotNone(event.id)
@@ -286,7 +306,9 @@ def test_field_mapping(self):
             },
         )
 
-        trace = response.results[1]
+        response = TracesQueryRunner(team=self.team, query=TracesQuery(traceId="trace2")).calculate()
+        self.assertEqual(len(response.results), 1)
+        trace = response.results[0]
         self.assertTraceEqual(
             trace,
             {
@@ -385,7 +407,7 @@ def test_maps_all_fields(self):
             },
         )
 
-        response = TracesQueryRunner(team=self.team, query=TracesQuery()).calculate()
+        response = TracesQueryRunner(team=self.team, query=TracesQuery(traceId="trace1")).calculate()
         self.assertEqual(len(response.results), 1)
         self.assertEqual(response.results[0].id, "trace1")
         self.assertEqual(response.results[0].totalLatency, 10.5)
@@ -582,6 +604,7 @@ def test_model_parameters(self):
         response = TracesQueryRunner(
             team=self.team,
             query=TracesQuery(
+                traceId="trace1",
                 dateRange=DateRange(date_from="2024-12-01T00:00:00Z", date_to="2024-12-01T00:10:00Z"),
             ),
         ).calculate()
@@ -623,6 +646,7 @@ def test_full_trace(self):
         response = TracesQueryRunner(
             team=self.team,
             query=TracesQuery(
+                traceId="trace1",
                 dateRange=DateRange(date_from="2024-12-01T00:00:00Z", date_to="2024-12-01T00:10:00Z"),
             ),
         ).calculate()
@@ -667,26 +691,27 @@ def test_properties_filter_with_multiple_events_in_group(self):
         response = TracesQueryRunner(
             team=self.team,
             query=TracesQuery(
+                traceId="trace1",
                 properties=[EventPropertyFilter(key="foo", value="bar", operator=PropertyOperator.EXACT)],
                 dateRange=DateRange(date_from="2024-12-01T00:00:00Z", date_to="2024-12-01T00:10:00Z"),
             ),
         ).calculate()
         self.assertEqual(len(response.results), 1)
-        self.assertEqual(len(response.results[0].events), 2)
 
         response = TracesQueryRunner(
             team=self.team,
             query=TracesQuery(
+                traceId="trace1",
                 properties=[EventPropertyFilter(key="foo", value="baz", operator=PropertyOperator.EXACT)],
                 dateRange=DateRange(date_from="2024-12-01T00:00:00Z", date_to="2024-12-01T00:10:00Z"),
             ),
         ).calculate()
         self.assertEqual(len(response.results), 1)
-        self.assertEqual(len(response.results[0].events), 2)
 
         response = TracesQueryRunner(
             team=self.team,
             query=TracesQuery(
+                traceId="trace1",
                 properties=[EventPropertyFilter(key="foo", value="barz", operator=PropertyOperator.EXACT)],
                 dateRange=DateRange(date_from="2024-12-01T00:00:00Z", date_to="2024-12-01T00:10:00Z"),
             ),
@@ -728,7 +753,6 @@ def test_trace_property_filter_for_event_group(self):
             ),
         ).calculate()
         self.assertEqual(len(response.results), 1)
-        self.assertEqual(len(response.results[0].events), 2)
 
         response = TracesQueryRunner(
             team=self.team,
@@ -738,7 +762,6 @@ def test_trace_property_filter_for_event_group(self):
             ),
         ).calculate()
         self.assertEqual(len(response.results), 1)
-        self.assertEqual(len(response.results[0].events), 2)
 
         # Shouldn't return anything because there isn't such trace
         response = TracesQueryRunner(
@@ -807,3 +830,54 @@ def test_trace_name_fallback(self):
         ).calculate()
         self.assertEqual(len(response.results), 1)
         self.assertEqual(response.results[0].traceName, "bar")
+
+    def test_returns_metrics_and_feedback_events(self):
+        _create_person(distinct_ids=["person1"], team=self.team)
+        _create_ai_generation_event(
+            distinct_id="person1",
+            trace_id="trace1",
+            team=self.team,
+            timestamp=datetime(2024, 12, 1, 0, 0),
+        )
+        _create_event(
+            distinct_id="person1",
+            team=self.team,
+            timestamp=datetime(2024, 12, 1, 0, 1),
+            event="$ai_metric",
+            properties={
+                "$ai_trace_id": "trace1",
+            },
+        )
+        _create_event(
+            distinct_id="person1",
+            team=self.team,
+            timestamp=datetime(2024, 12, 1, 0, 2),
+            event="$ai_feedback",
+            properties={
+                "$ai_trace_id": "trace1",
+            },
+        )
+
+        response = TracesQueryRunner(
+            team=self.team,
+            query=TracesQuery(
+                traceId="trace1",
+                dateRange=DateRange(date_from="2024-12-01T00:00:00Z", date_to="2024-12-01T00:10:00Z"),
+            ),
+        ).calculate()
+        self.assertEqual(len(response.results), 1)
+        self.assertEqual(len(response.results[0].events), 3)
+        self.assertEqual(response.results[0].events[0].event, "$ai_generation")
+        self.assertEqual(response.results[0].events[1].event, "$ai_metric")
+        self.assertEqual(response.results[0].events[2].event, "$ai_feedback")
+
+        response = TracesQueryRunner(
+            team=self.team,
+            query=TracesQuery(
+                dateRange=DateRange(date_from="2024-12-01T00:00:00Z", date_to="2024-12-01T00:10:00Z"),
+            ),
+        ).calculate()
+        self.assertEqual(len(response.results), 1)
+        self.assertEqual(len(response.results[0].events), 2)
+        self.assertEqual(response.results[0].events[0].event, "$ai_metric")
+        self.assertEqual(response.results[0].events[1].event, "$ai_feedback")

posthog/hogql_queries/ai/traces_query_runner.py

@@ -77,6 +77,7 @@ def calculate(self):
                 placeholders={
                     "common_conditions": self._get_where_clause(),
                     "filter_conditions": self._get_properties_filter(),
+                    "return_full_trace": ast.Constant(value=1 if self.query.traceId is not None else 0),
                 },
                 team=self.team,
                 query_type=NodeKind.TRACES_QUERY,
@@ -101,7 +102,7 @@ def get_cache_payload(self):
         return {
             **super().get_cache_payload(),
             # When the response schema changes, increment this version to invalidate the cache.
-            "schema_version": 1,
+            "schema_version": 2,
         }
 
     @cached_property
@@ -200,7 +201,7 @@ def _get_event_query(self) -> ast.SelectQuery:
                 generations.input_cost AS input_cost,
                 generations.output_cost AS output_cost,
                 generations.total_cost AS total_cost,
-                generations.events AS events,
+                if({return_full_trace}, generations.events, arrayFilter(x -> x.2 IN ('$ai_metric', '$ai_feedback'), generations.events)) as events,
                 traces.input_state AS input_state,
                 traces.output_state AS output_state,
                 traces.trace_name AS trace_name,

posthog/temporal/data_imports/pipelines/rest_source/__init__.py

@@ -250,7 +250,7 @@ def create_resources(
         include_from_parent: list[str] = endpoint_resource.get("include_from_parent", [])
         if not resolved_param and include_from_parent:
             raise ValueError(
-                f"Resource {resource_name} has include_from_parent but is not " "dependent on another resource"
+                f"Resource {resource_name} has include_from_parent but is not dependent on another resource"
             )
 
         (

products/llm_observability/frontend/LLMObservabilityScene.tsx

@@ -23,14 +23,13 @@ export const scene: SceneExport = {
 }
 
 const Filters = (): JSX.Element => {
-    const {
-        dateFilter: { dateTo, dateFrom },
-        shouldFilterTestAccounts,
-        generationsQuery,
-        propertyFilters,
-    } = useValues(llmObservabilityLogic)
+    const { dashboardDateFilter, dateFilter, shouldFilterTestAccounts, generationsQuery, propertyFilters, activeTab } =
+        useValues(llmObservabilityLogic)
     const { setDates, setShouldFilterTestAccounts, setPropertyFilters } = useActions(llmObservabilityLogic)
 
+    const dateFrom = activeTab === 'dashboard' ? dashboardDateFilter.dateFrom : dateFilter.dateFrom
+    const dateTo = activeTab === 'dashboard' ? dashboardDateFilter.dateTo : dateFilter.dateTo
+
     return (
         <div className="flex gap-x-4 gap-y-2 items-center flex-wrap py-4 -mt-4 mb-4 border-b">
             <DateFilter dateFrom={dateFrom} dateTo={dateTo} onChange={setDates} />

products/llm_observability/frontend/LLMObservabilityTraceScene.tsx

@@ -38,13 +38,10 @@ export const scene: SceneExport = {
 }
 
 export function LLMObservabilityTraceScene(): JSX.Element {
-    const { traceId, query, cachedTraceResponse } = useValues(llmObservabilityTraceLogic)
+    const { traceId, query } = useValues(llmObservabilityTraceLogic)
 
     return (
-        <BindLogic
-            logic={llmObservabilityTraceDataLogic}
-            props={{ traceId, query, cachedResults: cachedTraceResponse }}
-        >
+        <BindLogic logic={llmObservabilityTraceDataLogic} props={{ traceId, query }}>
             <TraceSceneWrapper />
         </BindLogic>
     )