feat(llma): Sessions view improvements (#40868)

Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: github-actions <41898282+github-actions[bot]@users.noreply.github.com>

products/llm_analytics/frontend/LLMAnalyticsScene.tsx

@@ -527,7 +527,14 @@ export function LLMAnalyticsScene(): JSX.Element {
     if (featureFlags[FEATURE_FLAGS.LLM_ANALYTICS_SESSIONS_VIEW]) {
         tabs.push({
             key: 'sessions',
-            label: 'Sessions',
+            label: (
+                <>
+                    Sessions{' '}
+                    <LemonTag className="ml-1" type="warning">
+                        Beta
+                    </LemonTag>
+                </>
+            ),
             content: (
                 <LLMAnalyticsSetupPrompt>
                     <LLMAnalyticsSessionsScene />

products/llm_analytics/frontend/LLMAnalyticsSessionsScene.tsx

@@ -1,10 +1,8 @@
 import { useActions, useValues } from 'kea'
-import { useState } from 'react'
 
 import { IconChevronDown, IconChevronRight } from '@posthog/icons'
 import { LemonTag } from '@posthog/lemon-ui'
 
-import api from 'lib/api'
 import { TZLabel } from 'lib/components/TZLabel'
 import { Link } from 'lib/lemon-ui/Link'
 import { Spinner } from 'lib/lemon-ui/Spinner'
@@ -13,9 +11,7 @@ import { urls } from 'scenes/urls'
 
 import { DataTable } from '~/queries/nodes/DataTable/DataTable'
 import { DataTableRow } from '~/queries/nodes/DataTable/dataTableLogic'
-import { LLMTrace, NodeKind, TraceQuery, TracesQuery } from '~/queries/schema/schema-general'
 import { isHogQLQuery } from '~/queries/utils'
-import { PropertyFilterType } from '~/types'
 
 import { LLMAnalyticsTraceEvents } from './components/LLMAnalyticsTraceEvents'
 import { useSortableColumns } from './hooks/useSortableColumns'
@@ -23,117 +19,30 @@ import { llmAnalyticsLogic } from './llmAnalyticsLogic'
 import { formatLLMCost } from './utils'
 
 export function LLMAnalyticsSessionsScene(): JSX.Element {
-    const { setDates, setShouldFilterTestAccounts, setPropertyFilters, setSessionsSort } = useActions(llmAnalyticsLogic)
-    const { sessionsQuery, dateFilter, sessionsSort } = useValues(llmAnalyticsLogic)
-    const [expandedSessionIds, setExpandedSessionIds] = useState<Set<string>>(new Set())
-    const [expandedTraceIds, setExpandedTraceIds] = useState<Set<string>>(new Set())
-    const [expandedGenerationIds, setExpandedGenerationIds] = useState<Set<string>>(new Set())
-    const [sessionTraces, setSessionTraces] = useState<Record<string, LLMTrace[]>>({})
-    const [loadingTraces, setLoadingTraces] = useState<Set<string>>(new Set())
-    const [loadingFullTraces, setLoadingFullTraces] = useState<Set<string>>(new Set())
-    const [fullTraces, setFullTraces] = useState<Record<string, LLMTrace>>({})
-
-    const handleSessionExpand = async (sessionId: string): Promise<void> => {
-        const newExpanded = new Set(expandedSessionIds)
-        if (newExpanded.has(sessionId)) {
-            newExpanded.delete(sessionId)
-            setExpandedSessionIds(newExpanded)
-        } else {
-            newExpanded.add(sessionId)
-            setExpandedSessionIds(newExpanded)
-
-            // Load traces for this session if not already loaded
-            if (!sessionTraces[sessionId] && !loadingTraces.has(sessionId)) {
-                setLoadingTraces(new Set(loadingTraces).add(sessionId))
-
-                const tracesQuerySource: TracesQuery = {
-                    kind: NodeKind.TracesQuery,
-                    dateRange: {
-                        date_from: dateFilter.dateFrom || undefined,
-                        date_to: dateFilter.dateTo || undefined,
-                    },
-                    properties: [
-                        {
-                            type: PropertyFilterType.Event,
-                            key: '$ai_session_id',
-                            operator: 'exact' as any,
-                            value: sessionId,
-                        },
-                    ],
-                }
-
-                try {
-                    const response = await api.query(tracesQuerySource)
-                    if (response.results) {
-                        setSessionTraces((prev) => ({
-                            ...prev,
-                            [sessionId]: response.results,
-                        }))
-                    }
-                } catch (error) {
-                    console.error('Error loading traces for session:', error)
-                } finally {
-                    const newLoading = new Set(loadingTraces)
-                    newLoading.delete(sessionId)
-                    setLoadingTraces(newLoading)
-                }
-            }
-        }
-    }
-
-    const handleGenerationExpand = (generationId: string): void => {
-        const newExpanded = new Set(expandedGenerationIds)
-        if (newExpanded.has(generationId)) {
-            newExpanded.delete(generationId)
-        } else {
-            newExpanded.add(generationId)
-        }
-        setExpandedGenerationIds(newExpanded)
-    }
+    const {
+        setDates,
+        setShouldFilterTestAccounts,
+        setPropertyFilters,
+        setSessionsSort,
+        toggleSessionExpanded,
+        toggleTraceExpanded,
+        toggleGenerationExpanded,
+    } = useActions(llmAnalyticsLogic)
+    const {
+        sessionsQuery,
+        dateFilter,
+        sessionsSort,
+        expandedSessionIds,
+        expandedTraceIds,
+        expandedGenerationIds,
+        sessionTraces,
+        loadingSessionTraces,
+        loadingFullTraces,
+        fullTraces,
+    } = useValues(llmAnalyticsLogic)
 
     const { renderSortableColumnTitle } = useSortableColumns(sessionsSort, setSessionsSort)
 
-    const handleTraceExpand = async (traceId: string): Promise<void> => {
-        const newExpanded = new Set(expandedTraceIds)
-        if (newExpanded.has(traceId)) {
-            newExpanded.delete(traceId)
-            setExpandedTraceIds(newExpanded)
-        } else {
-            newExpanded.add(traceId)
-            setExpandedTraceIds(newExpanded)
-
-            // Load full trace with events if not already loaded
-            if (!fullTraces[traceId] && !loadingFullTraces.has(traceId)) {
-                setLoadingFullTraces(new Set(loadingFullTraces).add(traceId))
-
-                const traceQuery: TraceQuery = {
-                    kind: NodeKind.TraceQuery,
-                    traceId,
-                    dateRange: {
-                        date_from: dateFilter.dateFrom || undefined,
-                        date_to: dateFilter.dateTo || undefined,
-                    },
-                }
-
-                try {
-                    const response = await api.query(traceQuery)
-                    if (response.results && response.results[0]) {
-                        setFullTraces((prev) => ({
-                            ...prev,
-                            [traceId]: response.results[0],
-                        }))
-                    }
-                } catch (error) {
-                    console.error('Error loading full trace:', error)
-                } finally {
-                    const newLoading = new Set(loadingFullTraces)
-                    newLoading.delete(traceId)
-                    setLoadingFullTraces(newLoading)
-                }
-            }
-        }
-    }
-
     return (
         <DataTable
             query={{
@@ -265,7 +174,7 @@ export function LLMAnalyticsSessionsScene(): JSX.Element {
 
                         const sessionId = result[0] as string
                         const traces = sessionTraces[sessionId]
-                        const isLoading = loadingTraces.has(sessionId)
+                        const isLoading = loadingSessionTraces.has(sessionId)
 
                         if (isLoading) {
                             return (
@@ -289,7 +198,7 @@ export function LLMAnalyticsSessionsScene(): JSX.Element {
                                             <div key={trace.id} className="border rounded">
                                                 <div
                                                     className="p-3 hover:bg-side-light cursor-pointer flex items-start gap-2"
-                                                    onClick={() => handleTraceExpand(trace.id)}
+                                                    onClick={() => toggleTraceExpanded(trace.id)}
                                                 >
                                                     <div className="flex-shrink-0 mt-0.5">
                                                         {isTraceExpanded ? (
@@ -343,7 +252,9 @@ export function LLMAnalyticsSessionsScene(): JSX.Element {
                                                                 trace={fullTraces[trace.id]}
                                                                 isLoading={loadingFullTraces.has(trace.id)}
                                                                 expandedEventIds={expandedGenerationIds}
-                                                                onToggleEventExpand={handleGenerationExpand}
+                                                                onToggleEventExpand={(uuid: string) =>
+                                                                    toggleGenerationExpanded(uuid, trace.id)
+                                                                }
                                                             />
                                                         </div>
                                                     </div>
@@ -361,12 +272,12 @@ export function LLMAnalyticsSessionsScene(): JSX.Element {
                         Array.isArray(result) && !!result[0] && expandedSessionIds.has(result[0] as string),
                     onRowExpand: ({ result }: DataTableRow) => {
                         if (Array.isArray(result) && result[0]) {
-                            handleSessionExpand(result[0] as string)
+                            toggleSessionExpanded(result[0] as string)
                         }
                     },
                     onRowCollapse: ({ result }: DataTableRow) => {
                         if (Array.isArray(result) && result[0]) {
-                            handleSessionExpand(result[0] as string)
+                            toggleSessionExpanded(result[0] as string)
                         }
                     },
                     noIndent: true,

products/llm_analytics/frontend/llmAnalyticsLogic.test.ts

@@ -206,4 +206,240 @@ describe('llmAnalyticsLogic', () => {
             shouldFilterTestAccounts: false,
         })
     })
+
+    describe('session expansion state', () => {
+        it('toggles session expansion state', async () => {
+            await expectLogic(logic, () => {
+                logic.actions.toggleSessionExpanded('session-123')
+            }).toMatchValues({
+                expandedSessionIds: new Set(['session-123']),
+            })
+
+            // Toggle again to collapse
+            await expectLogic(logic, () => {
+                logic.actions.toggleSessionExpanded('session-123')
+            }).toMatchValues({
+                expandedSessionIds: new Set(),
+            })
+        })
+
+        it('handles multiple expanded sessions', async () => {
+            await expectLogic(logic, () => {
+                logic.actions.toggleSessionExpanded('session-1')
+                logic.actions.toggleSessionExpanded('session-2')
+                logic.actions.toggleSessionExpanded('session-3')
+            }).toMatchValues({
+                expandedSessionIds: new Set(['session-1', 'session-2', 'session-3']),
+            })
+
+            // Collapse middle session
+            await expectLogic(logic, () => {
+                logic.actions.toggleSessionExpanded('session-2')
+            }).toMatchValues({
+                expandedSessionIds: new Set(['session-1', 'session-3']),
+            })
+        })
+
+        it('clears expanded sessions when date filter changes', async () => {
+            logic.actions.toggleSessionExpanded('session-123')
+
+            await expectLogic(logic, () => {
+                logic.actions.setDates('-7d', null)
+            }).toMatchValues({
+                expandedSessionIds: new Set(),
+                sessionTraces: {},
+            })
+        })
+
+        it('clears expanded sessions when property filters change', async () => {
+            logic.actions.toggleSessionExpanded('session-456')
+
+            await expectLogic(logic, () => {
+                logic.actions.setPropertyFilters([
+                    {
+                        type: PropertyFilterType.Event,
+                        key: 'browser',
+                        value: 'Chrome',
+                        operator: PropertyOperator.Exact,
+                    },
+                ])
+            }).toMatchValues({
+                expandedSessionIds: new Set(),
+                sessionTraces: {},
+            })
+        })
+
+        it('clears expanded sessions when test accounts filter changes', async () => {
+            logic.actions.toggleSessionExpanded('session-789')
+
+            await expectLogic(logic, () => {
+                logic.actions.setShouldFilterTestAccounts(true)
+            }).toMatchValues({
+                expandedSessionIds: new Set(),
+                sessionTraces: {},
+            })
+        })
+    })
+
+    describe('trace expansion state', () => {
+        it('toggles trace expansion state', async () => {
+            await expectLogic(logic, () => {
+                logic.actions.toggleTraceExpanded('trace-abc')
+            }).toMatchValues({
+                expandedTraceIds: new Set(['trace-abc']),
+            })
+
+            // Toggle again to collapse
+            await expectLogic(logic, () => {
+                logic.actions.toggleTraceExpanded('trace-abc')
+            }).toMatchValues({
+                expandedTraceIds: new Set(),
+            })
+        })
+
+        it('handles multiple expanded traces', async () => {
+            await expectLogic(logic, () => {
+                logic.actions.toggleTraceExpanded('trace-1')
+                logic.actions.toggleTraceExpanded('trace-2')
+            }).toMatchValues({
+                expandedTraceIds: new Set(['trace-1', 'trace-2']),
+            })
+        })
+
+        it('clears expanded traces when filters change', async () => {
+            logic.actions.toggleTraceExpanded('trace-xyz')
+
+            await expectLogic(logic, () => {
+                logic.actions.setDates('-14d', null)
+            }).toMatchValues({
+                expandedTraceIds: new Set(),
+                fullTraces: {},
+            })
+        })
+    })
+
+    describe('loading state tracking', () => {
+        it('tracks loading state for session traces', async () => {
+            logic.actions.loadSessionTraces('session-123')
+
+            expect(logic.values.loadingSessionTraces.has('session-123')).toBe(true)
+
+            logic.actions.loadSessionTracesSuccess('session-123', [])
+
+            expect(logic.values.loadingSessionTraces.has('session-123')).toBe(false)
+        })
+
+        it('clears loading state on failure', async () => {
+            logic.actions.loadSessionTraces('session-456')
+
+            expect(logic.values.loadingSessionTraces.has('session-456')).toBe(true)
+
+            logic.actions.loadSessionTracesFailure('session-456', new Error('Test error'))
+
+            expect(logic.values.loadingSessionTraces.has('session-456')).toBe(false)
+        })
+
+        it('tracks loading state for full traces', async () => {
+            logic.actions.loadFullTrace('trace-abc')
+
+            expect(logic.values.loadingFullTraces.has('trace-abc')).toBe(true)
+
+            const mockTrace = { id: 'trace-abc' } as any
+            logic.actions.loadFullTraceSuccess('trace-abc', mockTrace)
+
+            expect(logic.values.loadingFullTraces.has('trace-abc')).toBe(false)
+        })
+
+        it('handles multiple concurrent loading operations', async () => {
+            logic.actions.loadSessionTraces('session-1')
+            logic.actions.loadSessionTraces('session-2')
+            logic.actions.loadFullTrace('trace-1')
+
+            expect(logic.values.loadingSessionTraces.has('session-1')).toBe(true)
+            expect(logic.values.loadingSessionTraces.has('session-2')).toBe(true)
+            expect(logic.values.loadingFullTraces.has('trace-1')).toBe(true)
+
+            logic.actions.loadSessionTracesSuccess('session-1', [])
+
+            expect(logic.values.loadingSessionTraces.has('session-1')).toBe(false)
+            expect(logic.values.loadingSessionTraces.has('session-2')).toBe(true)
+            expect(logic.values.loadingFullTraces.has('trace-1')).toBe(true)
+        })
+    })
+
+    describe('session traces data', () => {
+        it('stores loaded session traces', async () => {
+            const mockTraces = [{ id: 'trace-1' }, { id: 'trace-2' }] as any[]
+
+            await expectLogic(logic, () => {
+                logic.actions.loadSessionTracesSuccess('session-123', mockTraces)
+            }).toMatchValues({
+                sessionTraces: {
+                    'session-123': mockTraces,
+                },
+            })
+        })
+
+        it('stores traces for multiple sessions', async () => {
+            const mockTraces1 = [{ id: 'trace-1' }] as any[]
+            const mockTraces2 = [{ id: 'trace-2' }] as any[]
+
+            logic.actions.loadSessionTracesSuccess('session-1', mockTraces1)
+            logic.actions.loadSessionTracesSuccess('session-2', mockTraces2)
+
+            expect(logic.values.sessionTraces).toEqual({
+                'session-1': mockTraces1,
+                'session-2': mockTraces2,
+            })
+        })
+
+        it('clears session traces when filters change', async () => {
+            const mockTraces = [{ id: 'trace-1' }] as any[]
+            logic.actions.loadSessionTracesSuccess('session-123', mockTraces)
+
+            await expectLogic(logic, () => {
+                logic.actions.setDates('-30d', null)
+            }).toMatchValues({
+                sessionTraces: {},
+            })
+        })
+    })
+
+    describe('full traces data', () => {
+        it('stores loaded full trace', async () => {
+            const mockTrace = { id: 'trace-abc', events: [] } as any
+
+            await expectLogic(logic, () => {
+                logic.actions.loadFullTraceSuccess('trace-abc', mockTrace)
+            }).toMatchValues({
+                fullTraces: {
+                    'trace-abc': mockTrace,
+                },
+            })
+        })
+
+        it('stores multiple full traces', async () => {
+            const mockTrace1 = { id: 'trace-1' } as any
+            const mockTrace2 = { id: 'trace-2' } as any
+
+            logic.actions.loadFullTraceSuccess('trace-1', mockTrace1)
+            logic.actions.loadFullTraceSuccess('trace-2', mockTrace2)
+
+            expect(logic.values.fullTraces).toEqual({
+                'trace-1': mockTrace1,
+                'trace-2': mockTrace2,
+            })
+        })
+
+        it('clears full traces when filters change', async () => {
+            const mockTrace = { id: 'trace-xyz' } as any
+            logic.actions.loadFullTraceSuccess('trace-xyz', mockTrace)
+
+            await expectLogic(logic, () => {
+                logic.actions.setPropertyFilters([])
+            }).toMatchValues({
+                fullTraces: {},
+            })
+        })
+    })
 })

products/llm_analytics/frontend/llmAnalyticsLogic.tsx

@@ -116,6 +116,14 @@ export const llmAnalyticsLogic = kea<llmAnalyticsLogicType>([
         toggleGenerationExpanded: (uuid: string, traceId: string) => ({ uuid, traceId }),
         setLoadedTrace: (traceId: string, trace: LLMTrace) => ({ traceId, trace }),
         clearExpandedGenerations: true,
+        toggleSessionExpanded: (sessionId: string) => ({ sessionId }),
+        toggleTraceExpanded: (traceId: string) => ({ traceId }),
+        loadSessionTraces: (sessionId: string) => ({ sessionId }),
+        loadSessionTracesSuccess: (sessionId: string, traces: LLMTrace[]) => ({ sessionId, traces }),
+        loadSessionTracesFailure: (sessionId: string, error: Error) => ({ sessionId, error }),
+        loadFullTrace: (traceId: string) => ({ traceId }),
+        loadFullTraceSuccess: (traceId: string, trace: LLMTrace) => ({ traceId, trace }),
+        loadFullTraceFailure: (traceId: string, error: Error) => ({ traceId, error }),
     }),
 
     reducers({
@@ -246,6 +254,102 @@ export const llmAnalyticsLogic = kea<llmAnalyticsLogicType>([
                 setShouldFilterTestAccounts: () => ({}),
             },
         ],
+
+        expandedSessionIds: [
+            new Set<string>() as Set<string>,
+            {
+                toggleSessionExpanded: (state, { sessionId }) => {
+                    const newSet = new Set(state)
+                    if (newSet.has(sessionId)) {
+                        newSet.delete(sessionId)
+                    } else {
+                        newSet.add(sessionId)
+                    }
+                    return newSet
+                },
+                setDates: () => new Set<string>(),
+                setPropertyFilters: () => new Set<string>(),
+                setShouldFilterTestAccounts: () => new Set<string>(),
+            },
+        ],
+
+        expandedTraceIds: [
+            new Set<string>() as Set<string>,
+            {
+                toggleTraceExpanded: (state, { traceId }) => {
+                    const newSet = new Set(state)
+                    if (newSet.has(traceId)) {
+                        newSet.delete(traceId)
+                    } else {
+                        newSet.add(traceId)
+                    }
+                    return newSet
+                },
+                setDates: () => new Set<string>(),
+                setPropertyFilters: () => new Set<string>(),
+                setShouldFilterTestAccounts: () => new Set<string>(),
+            },
+        ],
+
+        sessionTraces: [
+            {} as Record<string, LLMTrace[]>,
+            {
+                loadSessionTracesSuccess: (state, { sessionId, traces }) => ({
+                    ...state,
+                    [sessionId]: traces,
+                }),
+                setDates: () => ({}),
+                setPropertyFilters: () => ({}),
+                setShouldFilterTestAccounts: () => ({}),
+            },
+        ],
+
+        fullTraces: [
+            {} as Record<string, LLMTrace>,
+            {
+                loadFullTraceSuccess: (state, { traceId, trace }) => ({
+                    ...state,
+                    [traceId]: trace,
+                }),
+                setDates: () => ({}),
+                setPropertyFilters: () => ({}),
+                setShouldFilterTestAccounts: () => ({}),
+            },
+        ],
+
+        loadingSessionTraces: [
+            new Set<string>() as Set<string>,
+            {
+                loadSessionTraces: (state, { sessionId }) => new Set(state).add(sessionId),
+                loadSessionTracesSuccess: (state, { sessionId }) => {
+                    const newSet = new Set(state)
+                    newSet.delete(sessionId)
+                    return newSet
+                },
+                loadSessionTracesFailure: (state, { sessionId }) => {
+                    const newSet = new Set(state)
+                    newSet.delete(sessionId)
+                    return newSet
+                },
+            },
+        ],
+
+        loadingFullTraces: [
+            new Set<string>() as Set<string>,
+            {
+                loadFullTrace: (state, { traceId }) => new Set(state).add(traceId),
+                loadFullTraceSuccess: (state, { traceId }) => {
+                    const newSet = new Set(state)
+                    newSet.delete(traceId)
+                    return newSet
+                },
+                loadFullTraceFailure: (state, { traceId }) => {
+                    const newSet = new Set(state)
+                    newSet.delete(traceId)
+                    return newSet
+                },
+            },
+        ],
     }),
 
     loaders({
@@ -295,6 +399,81 @@ export const llmAnalyticsLogic = kea<llmAnalyticsLogicType>([
                 }
             }
         },
+
+        toggleSessionExpanded: async ({ sessionId }) => {
+            if (
+                values.expandedSessionIds.has(sessionId) &&
+                !values.sessionTraces[sessionId] &&
+                !values.loadingSessionTraces.has(sessionId)
+            ) {
+                actions.loadSessionTraces(sessionId)
+            }
+        },
+
+        loadSessionTraces: async ({ sessionId }) => {
+            const dateFrom = values.dateFilter.dateFrom || undefined
+            const dateTo = values.dateFilter.dateTo || undefined
+
+            const tracesQuerySource: import('~/queries/schema/schema-general').TracesQuery = {
+                kind: NodeKind.TracesQuery,
+                dateRange: {
+                    date_from: dateFrom,
+                    date_to: dateTo,
+                },
+                properties: [
+                    {
+                        type: PropertyFilterType.Event,
+                        key: '$ai_session_id',
+                        operator: 'exact' as any,
+                        value: sessionId,
+                    },
+                ],
+            }
+
+            try {
+                const response = await api.query(tracesQuerySource)
+                if (response.results) {
+                    actions.loadSessionTracesSuccess(sessionId, response.results)
+                }
+            } catch (error) {
+                console.error('Error loading traces for session:', error)
+                actions.loadSessionTracesFailure(sessionId, error as Error)
+            }
+        },
+
+        toggleTraceExpanded: async ({ traceId }) => {
+            if (
+                values.expandedTraceIds.has(traceId) &&
+                !values.fullTraces[traceId] &&
+                !values.loadingFullTraces.has(traceId)
+            ) {
+                actions.loadFullTrace(traceId)
+            }
+        },
+
+        loadFullTrace: async ({ traceId }) => {
+            const dateFrom = values.dateFilter.dateFrom || undefined
+            const dateTo = values.dateFilter.dateTo || undefined
+
+            const traceQuery: TraceQuery = {
+                kind: NodeKind.TraceQuery,
+                traceId,
+                dateRange: {
+                    date_from: dateFrom,
+                    date_to: dateTo,
+                },
+            }
+
+            try {
+                const response = await api.query(traceQuery)
+                if (response.results && response.results[0]) {
+                    actions.loadFullTraceSuccess(traceId, response.results[0])
+                }
+            } catch (error) {
+                console.error('Error loading full trace:', error)
+                actions.loadFullTraceFailure(traceId, error as Error)
+            }
+        },
     })),
 
     selectors({
@@ -922,32 +1101,22 @@ export const llmAnalyticsLogic = kea<llmAnalyticsLogicType>([
                     kind: NodeKind.HogQLQuery,
                     query: `
                 SELECT
-                    ai_session_id as session_id,
-                    countDistinctIf(ai_trace_id, notEmpty(ai_trace_id)) as traces,
-                    countIf(event_type = '$ai_span') as spans,
-                    countIf(event_type = '$ai_generation') as generations,
-                    countIf(event_type = '$ai_embedding') as embeddings,
-                    countIf(notEmpty(ai_error) OR ai_is_error = 'true') as errors,
-                    round(sum(toFloat(ai_total_cost_usd)), 4) as total_cost,
-                    round(sum(toFloat(ai_latency)), 2) as total_latency,
+                    properties.$ai_session_id as session_id,
+                    countDistinctIf(properties.$ai_trace_id, isNotNull(properties.$ai_trace_id)) as traces,
+                    countIf(event = '$ai_span') as spans,
+                    countIf(event = '$ai_generation') as generations,
+                    countIf(event = '$ai_embedding') as embeddings,
+                    countIf(isNotNull(properties.$ai_error) OR properties.$ai_is_error = 'true') as errors,
+                    round(sum(toFloat(properties.$ai_total_cost_usd)), 4) as total_cost,
+                    round(sum(toFloat(properties.$ai_latency)), 2) as total_latency,
                     min(timestamp) as first_seen,
                     max(timestamp) as last_seen
-                FROM (
-                    SELECT
-                        event as event_type,
-                        timestamp,
-                        JSONExtractString(properties, '$ai_session_id') as ai_session_id,
-                        JSONExtractRaw(properties, '$ai_trace_id') as ai_trace_id,
-                        JSONExtractRaw(properties, '$ai_total_cost_usd') as ai_total_cost_usd,
-                        JSONExtractRaw(properties, '$ai_latency') as ai_latency,
-                        JSONExtractRaw(properties, '$ai_error') as ai_error,
-                        JSONExtractString(properties, '$ai_is_error') as ai_is_error
-                    FROM events
-                    WHERE event IN ('$ai_generation', '$ai_span', '$ai_embedding', '$ai_trace')
-                        AND notEmpty(JSONExtractString(properties, '$ai_session_id'))
-                        AND {filters}
-                )
-                GROUP BY ai_session_id
+                FROM events
+                WHERE event IN ('$ai_generation', '$ai_span', '$ai_embedding', '$ai_trace')
+                    AND isNotNull(properties.$ai_session_id)
+                    AND properties.$ai_session_id != ''
+                    AND {filters}
+                GROUP BY properties.$ai_session_id
                 ORDER BY ${sessionsSort.column} ${sessionsSort.direction}
                 LIMIT 50
                     `,