style: Format files with ruff

Author: andrewm4894

posthog/ai/otel/exporter.py

@@ -212,9 +212,7 @@ def _span_to_event(self, span: ReadableSpan) -> Optional[Dict[str, Any]]:
         trace_id = self._format_trace_id_as_uuid(span.context.trace_id)
         # Span IDs remain as hex (no dashes needed)
         span_id = format(span.context.span_id, "016x")
-        parent_span_id = (
-            format(span.parent.span_id, "016x") if span.parent else None
-        )
+        parent_span_id = format(span.parent.span_id, "016x") if span.parent else None
 
         # Check for error status
         is_error = span.status.status_code == StatusCode.ERROR if span.status else False
@@ -223,7 +221,14 @@ def _span_to_event(self, span: ReadableSpan) -> Optional[Dict[str, Any]]:
         # Model request span → $ai_generation
         if self._is_generation_span(span_name, attrs):
             return self._create_generation_event(
-                span, attrs, trace_id, span_id, parent_span_id, latency, is_error, error_message
+                span,
+                attrs,
+                trace_id,
+                span_id,
+                parent_span_id,
+                latency,
+                is_error,
+                error_message,
             )
 
         # Agent run span → skip (PostHog UI auto-creates trace wrapper from generation events)
@@ -234,13 +239,27 @@ def _span_to_event(self, span: ReadableSpan) -> Optional[Dict[str, Any]]:
         # Tool execution span → $ai_span
         if self._is_tool_span(span_name, attrs):
             return self._create_tool_span_event(
-                span, attrs, trace_id, span_id, parent_span_id, latency, is_error, error_message
+                span,
+                attrs,
+                trace_id,
+                span_id,
+                parent_span_id,
+                latency,
+                is_error,
+                error_message,
             )
 
         # Generic span that might be part of AI workflow
         if self._is_ai_related_span(span_name, attrs):
             return self._create_span_event(
-                span, attrs, trace_id, span_id, parent_span_id, latency, is_error, error_message
+                span,
+                attrs,
+                trace_id,
+                span_id,
+                parent_span_id,
+                latency,
+                is_error,
+                error_message,
             )
 
         return None
@@ -507,9 +526,7 @@ def _create_span_event(
 
         return {"name": "$ai_span", "properties": properties}
 
-    def _parse_json_attr(
-        self, value: Optional[Union[str, Any]]
-    ) -> Optional[Any]:
+    def _parse_json_attr(self, value: Optional[Union[str, Any]]) -> Optional[Any]:
         """Parse a JSON string attribute, returning the value as-is if already parsed."""
         if value is None:
             return None

posthog/ai/pydantic_ai/exporter.py

@@ -93,7 +93,9 @@ def _transform_span(self, span: ReadableSpan) -> ReadableSpan:
             normalized = self._normalize_messages(input_msgs)
             if normalized != input_msgs:
                 attrs["gen_ai.input.messages"] = (
-                    json.dumps(normalized) if isinstance(normalized, list) else normalized
+                    json.dumps(normalized)
+                    if isinstance(normalized, list)
+                    else normalized
                 )
                 modified = True
 
@@ -103,7 +105,9 @@ def _transform_span(self, span: ReadableSpan) -> ReadableSpan:
             normalized = self._normalize_messages(output_msgs)
             if normalized != output_msgs:
                 attrs["gen_ai.output.messages"] = (
-                    json.dumps(normalized) if isinstance(normalized, list) else normalized
+                    json.dumps(normalized)
+                    if isinstance(normalized, list)
+                    else normalized
                 )
                 modified = True
 
@@ -176,20 +180,24 @@ def _normalize_pydantic_message(self, msg: Dict[str, Any]) -> Dict[str, Any]:
             if part_type == "text" and "content" in part:
                 text_parts.append(str(part["content"]))
             elif part_type == "tool_call":
-                tool_calls.append({
-                    "id": part.get("id", ""),
-                    "type": "function",
-                    "function": {
-                        "name": part.get("name", ""),
-                        "arguments": part.get("arguments", "{}"),
+                tool_calls.append(
+                    {
+                        "id": part.get("id", ""),
+                        "type": "function",
+                        "function": {
+                            "name": part.get("name", ""),
+                            "arguments": part.get("arguments", "{}"),
+                        },
                     }
-                })
+                )
 
         # Build normalized message
         normalized: Dict[str, Any] = {"role": role}
 
         if text_parts:
-            normalized["content"] = "\n".join(text_parts) if len(text_parts) > 1 else text_parts[0]
+            normalized["content"] = (
+                "\n".join(text_parts) if len(text_parts) > 1 else text_parts[0]
+            )
         elif not tool_calls:
             normalized["content"] = ""
 

posthog/test/ai/otel/test_exporter.py

@@ -111,14 +111,19 @@ def test_is_agent_span(self, mock_client):
         exporter = PostHogSpanExporter(mock_client)
         assert exporter._is_agent_span("agent run", {}) is True
         assert exporter._is_agent_span("invoke_agent", {}) is True
-        assert exporter._is_agent_span("some_span", {"gen_ai.agent.name": "test"})  # truthy
+        assert exporter._is_agent_span(
+            "some_span", {"gen_ai.agent.name": "test"}
+        )  # truthy
         assert not exporter._is_agent_span("some_span", {})
 
     def test_is_tool_span(self, mock_client):
         exporter = PostHogSpanExporter(mock_client)
         assert exporter._is_tool_span("execute_tool get_weather", {}) is True
         assert exporter._is_tool_span("running tools", {}) is True
-        assert exporter._is_tool_span("some_span", {"gen_ai.tool.name": "get_weather"}) is True
+        assert (
+            exporter._is_tool_span("some_span", {"gen_ai.tool.name": "get_weather"})
+            is True
+        )
         assert exporter._is_tool_span("model call", {}) is False
 
 
@@ -135,8 +140,12 @@ def test_basic_generation_event(self, mock_client):
                 "gen_ai.system": "openai",
                 "gen_ai.usage.input_tokens": 100,
                 "gen_ai.usage.output_tokens": 50,
-                "gen_ai.input.messages": json.dumps([{"role": "user", "content": "Hello"}]),
-                "gen_ai.output.messages": json.dumps([{"role": "assistant", "content": "Hi!"}]),
+                "gen_ai.input.messages": json.dumps(
+                    [{"role": "user", "content": "Hello"}]
+                ),
+                "gen_ai.output.messages": json.dumps(
+                    [{"role": "assistant", "content": "Hi!"}]
+                ),
             },
         )
 
@@ -178,14 +187,20 @@ def test_generation_event_with_error(self, mock_client):
         assert props["$ai_error"] == "Rate limit exceeded"
 
     def test_generation_event_privacy_mode(self, mock_client):
-        exporter = PostHogSpanExporter(mock_client, distinct_id="user_123", privacy_mode=True)
+        exporter = PostHogSpanExporter(
+            mock_client, distinct_id="user_123", privacy_mode=True
+        )
 
         span = create_mock_span(
             name="chat openai",
             attributes={
                 "gen_ai.request.model": "gpt-4",
-                "gen_ai.input.messages": json.dumps([{"role": "user", "content": "Secret data"}]),
-                "gen_ai.output.messages": json.dumps([{"role": "assistant", "content": "Response"}]),
+                "gen_ai.input.messages": json.dumps(
+                    [{"role": "user", "content": "Secret data"}]
+                ),
+                "gen_ai.output.messages": json.dumps(
+                    [{"role": "assistant", "content": "Response"}]
+                ),
             },
         )
 
@@ -224,7 +239,9 @@ class TestAgentSpanHandling:
     def test_agent_span_is_skipped(self, mock_client):
         exporter = PostHogSpanExporter(mock_client, distinct_id="user_123")
 
-        span = create_mock_span(name="agent run", attributes={"gen_ai.agent.name": "TestAgent"})
+        span = create_mock_span(
+            name="agent run", attributes={"gen_ai.agent.name": "TestAgent"}
+        )
 
         exporter.export([span])
 
@@ -250,7 +267,9 @@ def test_basic_tool_span_event(self, mock_client):
             name="execute_tool get_weather",
             attributes={
                 "gen_ai.tool.name": "get_weather",
-                "gen_ai.tool.call.arguments": json.dumps({"latitude": 37.7749, "longitude": -122.4194}),
+                "gen_ai.tool.call.arguments": json.dumps(
+                    {"latitude": 37.7749, "longitude": -122.4194}
+                ),
                 "gen_ai.tool.call.result": "Sunny, 72°F",
             },
             parent_span_id=0xABCDEF1234567890,
@@ -266,11 +285,16 @@ def test_basic_tool_span_event(self, mock_client):
         assert "$ai_trace_id" in props
         assert "$ai_span_id" in props
         assert "$ai_parent_id" in props
-        assert props["$ai_tool_arguments"] == {"latitude": 37.7749, "longitude": -122.4194}
+        assert props["$ai_tool_arguments"] == {
+            "latitude": 37.7749,
+            "longitude": -122.4194,
+        }
         assert props["$ai_tool_result"] == "Sunny, 72°F"
 
     def test_tool_span_privacy_mode(self, mock_client):
-        exporter = PostHogSpanExporter(mock_client, distinct_id="user_123", privacy_mode=True)
+        exporter = PostHogSpanExporter(
+            mock_client, distinct_id="user_123", privacy_mode=True
+        )
 
         span = create_mock_span(
             name="execute_tool get_weather",
@@ -294,7 +318,9 @@ class TestDistinctIdHandling:
     def test_distinct_id_from_constructor(self, mock_client):
         exporter = PostHogSpanExporter(mock_client, distinct_id="configured_user")
 
-        span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"})
+        span = create_mock_span(
+            name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}
+        )
 
         exporter.export([span])
 
@@ -326,12 +352,17 @@ def test_distinct_id_fallback_to_trace_id(self, mock_client):
 
         exporter.export([span])
 
-        assert mock_client.capture.call_args[1]["distinct_id"] == "abcdef1234567890abcdef1234567890"
+        assert (
+            mock_client.capture.call_args[1]["distinct_id"]
+            == "abcdef1234567890abcdef1234567890"
+        )
 
     def test_process_person_profile_false_when_no_distinct_id(self, mock_client):
         exporter = PostHogSpanExporter(mock_client)
 
-        span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"})
+        span = create_mock_span(
+            name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}
+        )
 
         exporter.export([span])
 
@@ -349,7 +380,9 @@ def test_additional_properties_included(self, mock_client):
             properties={"$ai_session_id": "session_abc", "custom_prop": "value"},
         )
 
-        span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"})
+        span = create_mock_span(
+            name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}
+        )
 
         exporter.export([span])
 
@@ -364,7 +397,9 @@ def test_groups_included(self, mock_client):
             groups={"company": "posthog", "team": "product"},
         )
 
-        span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"})
+        span = create_mock_span(
+            name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}
+        )
 
         exporter.export([span])
 
@@ -380,7 +415,9 @@ class TestExportResult:
     def test_export_returns_success(self, mock_client):
         exporter = PostHogSpanExporter(mock_client, distinct_id="user_123")
 
-        span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"})
+        span = create_mock_span(
+            name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}
+        )
 
         result = exporter.export([span])
 
@@ -390,7 +427,9 @@ def test_export_handles_exceptions_gracefully(self, mock_client):
         mock_client.capture.side_effect = Exception("Network error")
         exporter = PostHogSpanExporter(mock_client, distinct_id="user_123")
 
-        span = create_mock_span(name="chat openai", attributes={"gen_ai.request.model": "gpt-4"})
+        span = create_mock_span(
+            name="chat openai", attributes={"gen_ai.request.model": "gpt-4"}
+        )
 
         result = exporter.export([span])
 

posthog/test/ai/pydantic_ai/test_exporter.py

@@ -67,7 +67,10 @@ def test_normalize_simple_text_message(self, mock_client):
         exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123")
 
         pydantic_format = [
-            {"parts": [{"content": "Hello, how are you?", "type": "text"}], "role": "user"}
+            {
+                "parts": [{"content": "Hello, how are you?", "type": "text"}],
+                "role": "user",
+            }
         ]
         result = exporter._normalize_messages(pydantic_format)
 
@@ -168,7 +171,9 @@ def test_normalize_preserves_finish_reason(self, mock_client):
         ]
         result = exporter._normalize_messages(pydantic_format)
 
-        assert result == [{"content": "Done!", "role": "assistant", "finish_reason": "stop"}]
+        assert result == [
+            {"content": "Done!", "role": "assistant", "finish_reason": "stop"}
+        ]
 
     def test_normalize_already_openai_format(self, mock_client):
         exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123")
@@ -301,7 +306,12 @@ def test_export_normalizes_and_captures(self, mock_client):
         exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123")
 
         pydantic_input = json.dumps(
-            [{"parts": [{"content": "What's the weather?", "type": "text"}], "role": "user"}]
+            [
+                {
+                    "parts": [{"content": "What's the weather?", "type": "text"}],
+                    "role": "user",
+                }
+            ]
         )
         pydantic_output = json.dumps(
             [
@@ -329,7 +339,9 @@ def test_export_normalizes_and_captures(self, mock_client):
         mock_client.capture.assert_called_once()
 
         props = mock_client.capture.call_args[1]["properties"]
-        assert props["$ai_input"] == [{"content": "What's the weather?", "role": "user"}]
+        assert props["$ai_input"] == [
+            {"content": "What's the weather?", "role": "user"}
+        ]
         assert props["$ai_output_choices"] == [
             {"content": "It's sunny!", "role": "assistant", "finish_reason": "stop"}
         ]
@@ -440,7 +452,10 @@ def test_does_not_overwrite_existing_genai_attributes(self, mock_client):
         transformed = exporter._transform_span(span)
 
         # Should preserve existing GenAI attribute, not overwrite
-        assert transformed.attributes["gen_ai.tool.call.arguments"] == '{"existing": "value"}'
+        assert (
+            transformed.attributes["gen_ai.tool.call.arguments"]
+            == '{"existing": "value"}'
+        )
 
     def test_tool_span_export_with_mapped_attributes(self, mock_client):
         exporter = PydanticAISpanExporter(mock_client, distinct_id="user_123")