diff --git a/.sampo/changesets/ai-tool-call-spans.md b/.sampo/changesets/ai-tool-call-spans.md new file mode 100644 index 00000000..50e65baa --- /dev/null +++ b/.sampo/changesets/ai-tool-call-spans.md @@ -0,0 +1,5 @@ +--- +pypi/posthog: minor +--- + +feat(ai): emit an `$ai_span` for each tool call an LLM requests in the OpenAI, Anthropic, and Gemini wrappers. When a generation's response contains tool calls, the wrapper now captures a child span (`$ai_span_type: "tool"`) per call — carrying the tool name, arguments, and tool-call id — nested under the generation via `$ai_parent_id`, so requested tools show up in the trace tree. Generation events also now carry an `$ai_span_id`. Works for sync, async, and streaming calls; span capture is best-effort and never breaks the underlying LLM call. diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index ec369031..6c1fb2cd 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -379,6 +379,88 @@ def merge_system_prompt( return [] +def _extract_response_tool_calls(formatted_output: Any) -> List[Dict[str, Any]]: + """Pull the tool calls the model requested out of an already-formatted response. + + Every provider converter normalizes a tool call into a content item shaped + like ``{"type": "function", "id": ..., "function": {"name": ..., + "arguments": ...}}`` (see the ``format_*_output`` helpers in each + converter), so reading that shape keeps this provider agnostic. + """ + tool_calls: List[Dict[str, Any]] = [] + if not isinstance(formatted_output, list): + return tool_calls + for message in formatted_output: + if not isinstance(message, dict): + continue + content = message.get("content") + if not isinstance(content, list): + continue + for item in content: + if isinstance(item, dict) and item.get("type") == "function": + fn = item.get("function") or {} + tool_calls.append( + { + "id": item.get("id"), + "name": fn.get("name"), + "arguments": fn.get("arguments"), + } + ) + return tool_calls + + +def _capture_tool_call_spans( + ph_client: PostHogClient, + *, + formatted_output: Any, + trace_id: str, + parent_id: str, + distinct_id: Optional[str], + groups: Optional[Dict[str, Any]], + privacy_mode: bool, + include_person_profile: bool, +) -> None: + """Emit one ``$ai_span`` per tool call the LLM requested in its response. + + Each span records *which* tool the model chose (name + arguments), nested + under the generation via ``$ai_parent_id`` so it shows up in the trace tree. + The wrapper never runs the tool, so these spans carry no duration or result + — instrument the tool execution itself if you need that. Best-effort: a + failure here must never break the underlying LLM call. + """ + try: + tool_calls = _extract_response_tool_calls(formatted_output) + if not tool_calls: + return + for tool_call in tool_calls: + properties: Dict[str, Any] = { + "$ai_trace_id": trace_id, + "$ai_span_id": str(uuid.uuid4()), + "$ai_parent_id": parent_id, + "$ai_span_name": tool_call.get("name") or "tool_call", + "$ai_span_type": "tool", + "$ai_input_state": with_privacy_mode( + ph_client, privacy_mode, tool_call.get("arguments") + ), + "$ai_latency": 0, + } + tool_call_id = tool_call.get("id") + if tool_call_id is not None: + properties["$ai_tool_call_id"] = tool_call_id + if not include_person_profile: + properties["$process_person_profile"] = False + _capture_ai_event( + ph_client, + "$ai_span", + distinct_id=distinct_id, + properties=properties, + groups=groups, + ) + except Exception: + # Span capture is best-effort telemetry; never surface to the caller. + pass + + def call_llm_and_track_usage( posthog_distinct_id: Optional[str], ph_client: PostHogClient, @@ -453,19 +535,20 @@ def call_llm_and_track_usage( "$ai_input", with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages), ) + formatted_output = finalize_ai_content( + format_response(response, provider), ph_client + ) tag( "$ai_output_choices", - with_privacy_mode( - ph_client, - posthog_privacy_mode, - finalize_ai_content(format_response(response, provider), ph_client), - ), + with_privacy_mode(ph_client, posthog_privacy_mode, formatted_output), ) tag("$ai_http_status", http_status) tag("$ai_input_tokens", usage.get("input_tokens", 0)) tag("$ai_output_tokens", usage.get("output_tokens", 0)) tag("$ai_latency", latency) tag("$ai_trace_id", posthog_trace_id) + generation_span_id = str(uuid.uuid4()) + tag("$ai_span_id", generation_span_id) tag("$ai_base_url", str(base_url)) warn_if_posthog_ai_gateway(base_url) @@ -529,6 +612,16 @@ def call_llm_and_track_usage( properties=merged_properties, groups=posthog_groups, ) + _capture_tool_call_spans( + ph_client, + formatted_output=formatted_output, + trace_id=posthog_trace_id, + parent_id=generation_span_id, + distinct_id=contexts.get_context_distinct_id(), + groups=posthog_groups, + privacy_mode=posthog_privacy_mode, + include_person_profile=has_person_distinct_id, + ) if error: raise error @@ -606,19 +699,20 @@ async def call_llm_and_track_usage_async( "$ai_input", with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages), ) + formatted_output = finalize_ai_content( + format_response(response, provider), ph_client + ) tag( "$ai_output_choices", - with_privacy_mode( - ph_client, - posthog_privacy_mode, - finalize_ai_content(format_response(response, provider), ph_client), - ), + with_privacy_mode(ph_client, posthog_privacy_mode, formatted_output), ) tag("$ai_http_status", http_status) tag("$ai_input_tokens", usage.get("input_tokens", 0)) tag("$ai_output_tokens", usage.get("output_tokens", 0)) tag("$ai_latency", latency) tag("$ai_trace_id", posthog_trace_id) + generation_span_id = str(uuid.uuid4()) + tag("$ai_span_id", generation_span_id) tag("$ai_base_url", str(base_url)) warn_if_posthog_ai_gateway(base_url) @@ -682,6 +776,16 @@ async def call_llm_and_track_usage_async( properties=merged_properties, groups=posthog_groups, ) + _capture_tool_call_spans( + ph_client, + formatted_output=formatted_output, + trace_id=posthog_trace_id, + parent_id=generation_span_id, + distinct_id=contexts.get_context_distinct_id(), + groups=posthog_groups, + privacy_mode=posthog_privacy_mode, + include_person_profile=has_person_distinct_id, + ) if error: raise error @@ -728,6 +832,7 @@ def capture_streaming_event( event_data: Standardized streaming event data containing all necessary information """ trace_id = event_data.get("trace_id") or str(uuid.uuid4()) + generation_span_id = str(uuid.uuid4()) # Build base event properties event_properties = { @@ -749,6 +854,7 @@ def capture_streaming_event( "$ai_output_tokens": event_data["usage_stats"].get("output_tokens", 0), "$ai_latency": event_data["latency"], "$ai_trace_id": trace_id, + "$ai_span_id": generation_span_id, "$ai_base_url": str(event_data["base_url"]), **(event_data.get("properties") or {}), } @@ -837,3 +943,13 @@ def capture_streaming_event( properties=event_properties, groups=event_data.get("groups"), ) + _capture_tool_call_spans( + ph_client, + formatted_output=event_data["formatted_output"], + trace_id=trace_id, + parent_id=generation_span_id, + distinct_id=event_data.get("distinct_id") or trace_id, + groups=event_data.get("groups"), + privacy_mode=event_data["privacy_mode"], + include_person_profile=event_data.get("distinct_id") is not None, + ) diff --git a/posthog/test/ai/anthropic/test_anthropic.py b/posthog/test/ai/anthropic/test_anthropic.py index 62cdc24d..d5be7d0e 100644 --- a/posthog/test/ai/anthropic/test_anthropic.py +++ b/posthog/test/ai/anthropic/test_anthropic.py @@ -917,9 +917,14 @@ def test_tool_calls_in_output_choices( ) assert response == mock_anthropic_response_with_tool_calls - assert mock_client.capture.call_count == 1 + # $ai_generation plus one $ai_span for the tool call the model requested + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + call_args = next( + c[1] + for c in mock_client.capture.call_args_list + if c[1]["event"] == "$ai_generation" + ) props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -980,9 +985,14 @@ def test_tool_calls_only_no_content( ) assert response == mock_anthropic_response_tool_calls_only - assert mock_client.capture.call_count == 1 + # $ai_generation plus one $ai_span for the requested tool call + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + call_args = next( + c[1] + for c in mock_client.capture.call_args_list + if c[1]["event"] == "$ai_generation" + ) props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1049,9 +1059,14 @@ async def run_test(): response = asyncio.run(run_test()) assert response == mock_anthropic_response_with_tool_calls - assert mock_client.capture.call_count == 1 + # $ai_generation plus one $ai_span for the requested tool call + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + call_args = next( + c[1] + for c in mock_client.capture.call_args_list + if c[1]["event"] == "$ai_generation" + ) props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1117,9 +1132,14 @@ def test_streaming_with_tool_calls(mock_client, mock_anthropic_stream_with_tools list(response) # Capture happens synchronously when generator is exhausted - assert mock_client.capture.call_count == 1 + # ($ai_generation plus one $ai_span for the streamed tool call) + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + call_args = next( + c[1] + for c in mock_client.capture.call_args_list + if c[1]["event"] == "$ai_generation" + ) props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1244,9 +1264,14 @@ async def run_test(): asyncio.run(run_test()) # Capture completes before asyncio.run() returns - assert mock_client.capture.call_count == 1 + # ($ai_generation plus one $ai_span for the streamed tool call) + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + call_args = next( + c[1] + for c in mock_client.capture.call_args_list + if c[1]["event"] == "$ai_generation" + ) props = call_args["properties"] assert call_args["distinct_id"] == "test-id" diff --git a/posthog/test/ai/gemini/test_gemini.py b/posthog/test/ai/gemini/test_gemini.py index c1b82fe0..dcaf93b0 100644 --- a/posthog/test/ai/gemini/test_gemini.py +++ b/posthog/test/ai/gemini/test_gemini.py @@ -718,9 +718,14 @@ def test_function_calls_in_output_choices( ) assert response == mock_gemini_response_with_function_calls - assert mock_client.capture.call_count == 1 + # $ai_generation plus one $ai_span for the function call the model requested + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + call_args = next( + c[1] + for c in mock_client.capture.call_args_list + if c[1]["event"] == "$ai_generation" + ) props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -767,9 +772,14 @@ def test_function_calls_only_no_content( ) assert response == mock_gemini_response_function_calls_only - assert mock_client.capture.call_count == 1 + # $ai_generation plus one $ai_span for the requested function call + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + call_args = next( + c[1] + for c in mock_client.capture.call_args_list + if c[1]["event"] == "$ai_generation" + ) props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -1553,8 +1563,13 @@ def mock_streaming_response(): list(response) - assert mock_client.capture.call_count == 1 - props = mock_client.capture.call_args[1]["properties"] + # $ai_generation plus one $ai_span for the streamed function call + assert mock_client.capture.call_count == 2 + props = next( + c[1] + for c in mock_client.capture.call_args_list + if c[1]["event"] == "$ai_generation" + )["properties"] assert props["$ai_stop_reason"] == "STOP" diff --git a/posthog/test/ai/gemini/test_gemini_async.py b/posthog/test/ai/gemini/test_gemini_async.py index 53f85b26..a10ca291 100644 --- a/posthog/test/ai/gemini/test_gemini_async.py +++ b/posthog/test/ai/gemini/test_gemini_async.py @@ -586,9 +586,14 @@ async def test_async_function_calls_in_output_choices( ) assert response == mock_gemini_response_with_function_calls - assert mock_client.capture.call_count == 1 + # $ai_generation plus one $ai_span for the function call the model requested + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + call_args = next( + c[1] + for c in mock_client.capture.call_args_list + if c[1]["event"] == "$ai_generation" + ) props = call_args["properties"] assert call_args["distinct_id"] == "test-id" diff --git a/posthog/test/ai/openai/test_openai.py b/posthog/test/ai/openai/test_openai.py index e82bb902..014686a8 100644 --- a/posthog/test/ai/openai/test_openai.py +++ b/posthog/test/ai/openai/test_openai.py @@ -739,9 +739,13 @@ def test_tool_calls(mock_client, mock_openai_response_with_tool_calls): ) assert response == mock_openai_response_with_tool_calls - assert mock_client.capture.call_count == 1 + # One $ai_generation plus one $ai_span for the tool call the model requested + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + calls = mock_client.capture.call_args_list + generation_call = next(c[1] for c in calls if c[1]["event"] == "$ai_generation") + span_call = next(c[1] for c in calls if c[1]["event"] == "$ai_span") + call_args = generation_call props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -786,6 +790,20 @@ def test_tool_calls(mock_client, mock_openai_response_with_tool_calls): assert props["$ai_output_tokens"] == 15 assert props["$ai_http_status"] == 200 + # The tool call the model requested is emitted as a child $ai_span, + # nested under the generation. + span_props = span_call["properties"] + assert span_call["distinct_id"] == "test-id" + assert span_props["$ai_span_type"] == "tool" + assert span_props["$ai_span_name"] == "get_weather" + assert span_props["$ai_tool_call_id"] == "call_abc123" + assert ( + span_props["$ai_input_state"] + == '{"location": "San Francisco", "unit": "celsius"}' + ) + assert span_props["$ai_trace_id"] == props["$ai_trace_id"] + assert span_props["$ai_parent_id"] == props["$ai_span_id"] + def test_tool_calls_only_no_content(mock_client, mock_openai_response_tool_calls_only): with patch( @@ -810,9 +828,13 @@ def test_tool_calls_only_no_content(mock_client, mock_openai_response_tool_calls ) assert response == mock_openai_response_tool_calls_only - assert mock_client.capture.call_count == 1 + # One $ai_generation plus one $ai_span for the requested tool call + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + calls = mock_client.capture.call_args_list + generation_call = next(c[1] for c in calls if c[1]["event"] == "$ai_generation") + span_call = next(c[1] for c in calls if c[1]["event"] == "$ai_span") + call_args = generation_call props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -840,6 +862,14 @@ def test_tool_calls_only_no_content(mock_client, mock_openai_response_tool_calls assert props["$ai_output_tokens"] == 10 assert props["$ai_http_status"] == 200 + # Tool call emitted as a child $ai_span nested under the generation + span_props = span_call["properties"] + assert span_props["$ai_span_type"] == "tool" + assert span_props["$ai_span_name"] == "get_weather" + assert span_props["$ai_tool_call_id"] == "call_def456" + assert span_props["$ai_input_state"] == '{"location": "New York"}' + assert span_props["$ai_parent_id"] == props["$ai_span_id"] + def test_responses_api_tool_calls(mock_client, mock_responses_api_with_tool_calls): with patch( @@ -865,9 +895,13 @@ def test_responses_api_tool_calls(mock_client, mock_responses_api_with_tool_call ) assert response == mock_responses_api_with_tool_calls - assert mock_client.capture.call_count == 1 + # One $ai_generation plus one $ai_span for the requested tool call + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + calls = mock_client.capture.call_args_list + generation_call = next(c[1] for c in calls if c[1]["event"] == "$ai_generation") + span_call = next(c[1] for c in calls if c[1]["event"] == "$ai_span") + call_args = generation_call props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -897,6 +931,14 @@ def test_responses_api_tool_calls(mock_client, mock_responses_api_with_tool_call assert props["$ai_output_tokens"] == 20 assert props["$ai_http_status"] == 200 + # Tool call emitted as a child $ai_span nested under the generation + span_props = span_call["properties"] + assert span_props["$ai_span_type"] == "tool" + assert span_props["$ai_span_name"] == "get_weather" + assert span_props["$ai_tool_call_id"] == "call_xyz789" + assert span_props["$ai_input_state"] == '{"location": "Chicago"}' + assert span_props["$ai_parent_id"] == props["$ai_span_id"] + def test_streaming_with_tool_calls(mock_client, streaming_tool_call_chunks): # Mock the create method to return our chunks @@ -933,10 +975,14 @@ def test_streaming_with_tool_calls(mock_client, streaming_tool_call_chunks): assert len(chunks) == 4 assert chunks == streaming_tool_call_chunks - # Verify the capture was called with the right arguments - assert mock_client.capture.call_count == 1 + # Verify the capture was called with the right arguments: + # one $ai_generation plus one $ai_span for the streamed tool call + assert mock_client.capture.call_count == 2 - call_args = mock_client.capture.call_args[1] + calls = mock_client.capture.call_args_list + generation_call = next(c[1] for c in calls if c[1]["event"] == "$ai_generation") + span_call = next(c[1] for c in calls if c[1]["event"] == "$ai_span") + call_args = generation_call props = call_args["properties"] assert call_args["distinct_id"] == "test-id" @@ -995,6 +1041,17 @@ def test_streaming_with_tool_calls(mock_client, streaming_tool_call_chunks): assert "prompt_tokens" in props["$ai_usage"] assert "completion_tokens" in props["$ai_usage"] + # The accumulated tool call is emitted as a child $ai_span + span_props = span_call["properties"] + assert span_props["$ai_span_type"] == "tool" + assert span_props["$ai_span_name"] == "get_weather" + assert span_props["$ai_tool_call_id"] == "call_abc123" + assert ( + span_props["$ai_input_state"] + == '{"location": "San Francisco", "unit": "celsius"}' + ) + assert span_props["$ai_parent_id"] == props["$ai_span_id"] + # test responses api def test_responses_api(mock_client, mock_openai_response_with_responses_api):