From b2fad558b4e465322171df61386d42e94a7545df Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Mon, 27 Jan 2025 12:48:47 +0100 Subject: [PATCH 01/23] feat: taxonomy --- frontend/src/lib/taxonomy.tsx | 20 ++++++++++++++++++++ posthog/taxonomy/taxonomy.py | 23 +++++++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/frontend/src/lib/taxonomy.tsx b/frontend/src/lib/taxonomy.tsx index e2d792af0b90a..299bb17286b9d 100644 --- a/frontend/src/lib/taxonomy.tsx +++ b/frontend/src/lib/taxonomy.tsx @@ -184,6 +184,11 @@ export const CORE_FILTER_DEFINITIONS_BY_GROUP = { description: 'A generative AI trace. Usually a trace tracks a single user interaction and contains one or more AI generation calls', }, + $ai_span: { + label: 'AI Span', + description: + 'A generative AI span. Usually a span tracks a unit of work for a trace of generative AI models (LLMs)', + }, $ai_metric: { label: 'AI Metric', description: 'An evaluation metric for a trace of generative AI models (LLMs)', @@ -1490,6 +1495,21 @@ export const CORE_FILTER_DEFINITIONS_BY_GROUP = { description: 'The text provided by the user for feedback on the LLM trace', examples: ['"The response was helpful, but it did not use the provided context."'], }, + $ai_parent_id: { + label: 'AI Parent ID (LLM)', + description: 'The parent span ID of a span or generation, used to group a trace into a tree view', + examples: ['bdf42359-9364-4db7-8958-c001f28c9255'], + }, + $ai_generation_id: { + label: 'AI Generation ID (LLM)', + description: 'The unique identifier for a LLM generation', + examples: ['bdf42359-9364-4db7-8958-c001f28c9255'], + }, + $ai_span_id: { + label: 'AI Span ID (LLM)', + description: 'The unique identifier for a LLM span', + examples: ['bdf42359-9364-4db7-8958-c001f28c9255'], + }, }, numerical_event_properties: {}, // Same as event properties, see assignment below person_properties: {}, // Currently person properties are the same as event properties, see assignment below diff --git a/posthog/taxonomy/taxonomy.py b/posthog/taxonomy/taxonomy.py index 02e9ccf13d128..d14edf95e181a 100644 --- a/posthog/taxonomy/taxonomy.py +++ b/posthog/taxonomy/taxonomy.py @@ -196,6 +196,14 @@ class CoreFilterDefinition(TypedDict): "label": "AI Feedback (LLM)", "description": "User-provided feedback for a trace of a generative AI model (LLM).", }, + "$ai_trace": { + "label": "AI Trace (LLM)", + "description": "A generative AI trace. Usually a trace tracks a single user interaction and contains one or more AI generation calls", + }, + "$ai_span": { + "label": "AI Span (LLM)", + "description": "A generative AI span. Usually a span tracks a unit of work for a trace of generative AI models (LLMs)", + }, "Application Opened": { "label": "Application Opened", "description": "When a user opens the mobile app either for the first time or from the foreground.", @@ -1377,6 +1385,21 @@ class CoreFilterDefinition(TypedDict): "description": "The text provided by the user for feedback on the LLM trace", "examples": ['"The response was helpful, but it did not use the provided context."'], }, + "$ai_parent_id": { + "label": "AI Parent ID (LLM)", + "description": "The parent span ID of a span or generation, used to group a trace into a tree view", + "examples": ["bdf42359-9364-4db7-8958-c001f28c9255"], + }, + "$ai_generation_id": { + "label": "AI Generation ID (LLM)", + "description": "The unique identifier for a LLM generation", + "examples": ["bdf42359-9364-4db7-8958-c001f28c9255"], + }, + "$ai_span_id": { + "label": "AI Span ID (LLM)", + "description": "The unique identifier for a LLM span", + "examples": ["bdf42359-9364-4db7-8958-c001f28c9255"], + }, }, "numerical_event_properties": {}, "person_properties": {}, From 3a6a7f23fb4c2fa758eaaf7e96ca645e0b1ad859 Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Mon, 27 Jan 2025 14:26:21 +0100 Subject: [PATCH 02/23] feat: output spans --- .../ai/test/test_traces_query_runner.py | 55 ++++++++++++++++++- .../hogql_queries/ai/traces_query_runner.py | 2 +- 2 files changed, 54 insertions(+), 3 deletions(-) diff --git a/posthog/hogql_queries/ai/test/test_traces_query_runner.py b/posthog/hogql_queries/ai/test/test_traces_query_runner.py index 53b44eb690077..ebe8658bfbb3c 100644 --- a/posthog/hogql_queries/ai/test/test_traces_query_runner.py +++ b/posthog/hogql_queries/ai/test/test_traces_query_runner.py @@ -123,6 +123,41 @@ def _create_ai_trace_event( ) +def _create_ai_span_event( + *, + trace_id: str, + input_state: Any, + output_state: Any, + span_id: str | None = None, + parent_id: str | None = None, + span_name: str | None = None, + team: Team | None = None, + distinct_id: str | None = None, + properties: dict[str, Any] | None = None, + timestamp: datetime | None = None, + event_uuid: str | UUID | None = None, +): + props = { + "$ai_trace_id": trace_id, + "$ai_span_name": span_name, + "$ai_input_state": input_state, + "$ai_output_state": output_state, + "$ai_span_id": span_id or str(uuid.uuid4()), + "$ai_parent_id": parent_id or trace_id, + } + if properties: + props.update(properties) + + _create_event( + event="$ai_span", + distinct_id=distinct_id, + properties=props, + team=team, + timestamp=timestamp, + event_uuid=str(event_uuid) if event_uuid else None, + ) + + class TestTracesQueryRunner(ClickhouseTestMixin, BaseTest): def setUp(self): super().setUp() @@ -556,11 +591,19 @@ def test_model_parameters(self): def test_full_trace(self): _create_person(distinct_ids=["person1"], team=self.team, properties={"foo": "bar"}) + _create_ai_span_event( + trace_id="trace1", + span_name="runnable", + input_state={"messages": [{"role": "user", "content": "Foo"}]}, + output_state={"messages": [{"role": "user", "content": "Foo"}, {"role": "assistant", "content": "Bar"}]}, + team=self.team, + timestamp=datetime(2024, 12, 1, 0, 9), + ) _create_ai_generation_event( distinct_id="person1", trace_id="trace1", team=self.team, - timestamp=datetime(2024, 12, 1, 0, 0), + timestamp=datetime(2024, 12, 1, 0, 9, 30), ) _create_ai_generation_event( distinct_id="person1", @@ -591,10 +634,18 @@ def test_full_trace(self): response.results[0].outputState, {"messages": [{"role": "user", "content": "Foo"}, {"role": "assistant", "content": "Bar"}]}, ) - self.assertEqual(len(response.results[0].events), 2) + self.assertEqual(len(response.results[0].events), 3) + + self.assertEqual(response.results[0].events[0].event, "$ai_span") self.assertEqual(response.results[0].events[0].properties["$ai_trace_id"], "trace1") + self.assertEqual(response.results[0].events[0].properties["$ai_span_name"], "runnable") + + self.assertEqual(response.results[0].events[1].event, "$ai_generation") self.assertEqual(response.results[0].events[1].properties["$ai_trace_id"], "trace1") + self.assertEqual(response.results[0].events[2].event, "$ai_generation") + self.assertEqual(response.results[0].events[2].properties["$ai_trace_id"], "trace1") + @snapshot_clickhouse_queries def test_properties_filter_with_multiple_events_in_group(self): _create_person(distinct_ids=["person1"], team=self.team) diff --git a/posthog/hogql_queries/ai/traces_query_runner.py b/posthog/hogql_queries/ai/traces_query_runner.py index 0cccb0ebe6943..7784ba13bd34d 100644 --- a/posthog/hogql_queries/ai/traces_query_runner.py +++ b/posthog/hogql_queries/ai/traces_query_runner.py @@ -222,7 +222,7 @@ def _get_event_query(self) -> ast.SelectQuery: arraySort(x -> x.3, groupArray(tuple(uuid, event, timestamp, properties))) as events, {filter_conditions} FROM events - WHERE event IN ('$ai_generation', '$ai_metric', '$ai_feedback') AND {common_conditions} + WHERE event IN ('$ai_span', '$ai_generation', '$ai_metric', '$ai_feedback') AND {common_conditions} GROUP BY id ) AS generations LEFT JOIN ( From fd250f74cde464f3e0666d6a22243b3fbb7b9f13 Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Mon, 27 Jan 2025 17:52:36 +0100 Subject: [PATCH 03/23] feat: frontend tree --- .../frontend/LLMObservabilityTraceScene.tsx | 79 +++++++++++++------ .../llmObservabilityTraceDataLogic.ts | 48 +++++++++++ 2 files changed, 103 insertions(+), 24 deletions(-) diff --git a/products/llm_observability/frontend/LLMObservabilityTraceScene.tsx b/products/llm_observability/frontend/LLMObservabilityTraceScene.tsx index 78d1e50cdc69f..fb3b6cde9bc84 100644 --- a/products/llm_observability/frontend/LLMObservabilityTraceScene.tsx +++ b/products/llm_observability/frontend/LLMObservabilityTraceScene.tsx @@ -20,7 +20,7 @@ import { ConversationMessagesDisplay } from './ConversationDisplay/ConversationM import { MetadataHeader } from './ConversationDisplay/MetadataHeader' import { ParametersHeader } from './ConversationDisplay/ParametersHeader' import { LLMInputOutput } from './LLMInputOutput' -import { llmObservabilityTraceDataLogic } from './llmObservabilityTraceDataLogic' +import { llmObservabilityTraceDataLogic, TraceTreeNode } from './llmObservabilityTraceDataLogic' import { llmObservabilityTraceLogic } from './llmObservabilityTraceLogic' import { formatLLMCost, formatLLMLatency, formatLLMUsage, isLLMTraceEvent, removeMilliseconds } from './utils' @@ -44,7 +44,7 @@ export function LLMObservabilityTraceScene(): JSX.Element { function TraceSceneWrapper(): JSX.Element { const { eventId } = useValues(llmObservabilityTraceLogic) - const { trace, showableEvents, event, responseLoading, responseError, feedbackEvents, metricEvents } = + const { tree, trace, event, responseLoading, responseError, feedbackEvents, metricEvents } = useValues(llmObservabilityTraceDataLogic) return ( @@ -63,7 +63,7 @@ function TraceSceneWrapper(): JSX.Element { feedbackEvents={feedbackEvents as LLMTraceEvent[]} />
- +
@@ -145,35 +145,28 @@ function TraceMetadata({ function TraceSidebar({ trace, eventId, - events, + tree, }: { trace: LLMTrace eventId?: string | null - events: LLMTraceEvent[] + tree: TraceTreeNode[] }): JSX.Element { return ( - ) } function NestingGroup({ level = 0, children }: { level?: number; children: React.ReactNode }): JSX.Element { - const listEl = ( -
    {children}
- ) + const listEl =
    {children}
if (!level) { return listEl @@ -190,68 +190,65 @@ function NestingGroup({ level = 0, children }: { level?: number; children: React ) } -const TraceNode = React.memo( - ({ - topLevelTrace, - item, - isSelected, - }: { - topLevelTrace: LLMTrace - item: LLMTrace | LLMTraceEvent - isSelected: boolean - }): JSX.Element => { - const totalCost = 'properties' in item ? item.properties.$ai_total_cost_usd : item.totalCost - const latency = 'properties' in item ? item.properties.$ai_latency : item.totalLatency - const usage = formatLLMUsage(item) +const TraceLeaf = React.memo(function TraceNode({ + topLevelTrace, + item, + isSelected, +}: { + topLevelTrace: LLMTrace + item: LLMTrace | LLMTraceEvent + isSelected: boolean +}): JSX.Element { + const totalCost = 'properties' in item ? item.properties.$ai_total_cost_usd : item.totalCost + const latency = 'properties' in item ? item.properties.$ai_latency : item.totalLatency + const usage = formatLLMUsage(item) - const children = [ - isLLMTraceEvent(item) && item.properties.$ai_is_error && ( - - Error - - ), - latency >= 0.01 && ( - - {formatLLMLatency(latency)} - - ), - (usage != null || totalCost != null) && ( - - {usage} - {usage != null && totalCost != null && {' / '}} - {totalCost != null && formatLLMCost(totalCost)} - - ), - ] - const hasChildren = children.find((child) => !!child) + const children = [ + isLLMTraceEvent(item) && item.properties.$ai_is_error && ( + + Error + + ), + latency >= 0.01 && ( + + {formatLLMLatency(latency)} + + ), + (usage != null || totalCost != null) && ( + + {usage} + {usage != null && totalCost != null && {' / '}} + {totalCost != null && formatLLMCost(totalCost)} + + ), + ] + const hasChildren = children.some((child) => !!child) - return ( -
  • - -
    - - {formatLLMEventTitle(item)} -
    - {hasChildren && ( -
    {children}
    - )} - -
    • - ) - } -) -TraceNode.displayName = 'TraceNode' + return ( +
  • + +
    + + {formatLLMEventTitle(item)} +
    + {hasChildren && ( +
    {children}
    + )} + +
    • + ) +}) -function RecursiveTreeDisplay({ +function TreeNode({ tree, trace, selectedEventId, @@ -264,16 +261,16 @@ function RecursiveTreeDisplay({ {tree.map(({ event, children }) => ( - -
  • - {children && ( - - )} -
    • + {children && ( +
  • + +
    • + )}     ))}     @@ -304,7 +301,7 @@ function EventContentDisplay({
    {isObject(output) ? ( @@ -320,7 +317,7 @@ function EventContentDisplay({ function EventContent({ event }: { event: LLMTrace | LLMTraceEvent | null }): JSX.Element { return ( -
    +
    {!event ? ( ) : ( @@ -334,7 +331,7 @@ function EventContent({ event }: { event: LLMTrace | LLMTraceEvent | null }): JS
    {isLLMTraceEvent(event) ? ( () // Map all events with parents to their parent IDs - events.forEach((event) => { + for (const event of events) { // Exclude all metric and feedback events. if (FEEDBACK_EVENTS.has(event.event)) { - return + continue } const eventId = event.properties.$ai_generation_id ?? event.properties.$ai_span_id ?? event.id @@ -115,7 +115,7 @@ export function restoreTree(events: LLMTraceEvent[], traceId: string): TraceTree childrenMap.set(parentId, [eventId]) } } - }) + } function traverse(spanId: any): TraceTreeNode { const children = childrenMap.get(spanId) From 77dcdc8afa0a23249f64f3cebffd1f188d4884bc Mon Sep 17 00:00:00 2001 From: Georgiy Tarasov Date: Tue, 28 Jan 2025 19:18:02 +0100 Subject: [PATCH 20/23] fix: overflow --- .../frontend/LLMObservabilityTraceScene.tsx | 26 +++++++++---------- 1 file changed, 12 insertions(+), 14 deletions(-) diff --git a/products/llm_observability/frontend/LLMObservabilityTraceScene.tsx b/products/llm_observability/frontend/LLMObservabilityTraceScene.tsx index 4b8e7c929870a..d295bf928a0a9 100644 --- a/products/llm_observability/frontend/LLMObservabilityTraceScene.tsx +++ b/products/llm_observability/frontend/LLMObservabilityTraceScene.tsx @@ -160,37 +160,37 @@ function TraceSidebar({ tree: TraceTreeNode[] }): JSX.Element { return ( -