refactor(ai): address review feedback

Move DEFAULT_HOST to spans.py, extract shared make_span helper to
conftest.py, and parameterize span filter tests.

Generated-By: PostHog Code
Task-Id: 1ba1f07a-1453-4162-90a8-665958c5fe46

Author: richardsolomou

posthog/ai/otel/exporter.py

@@ -11,9 +11,7 @@
 from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 
-from .spans import is_ai_span
-
-DEFAULT_HOST = "https://us.i.posthog.com"
+from .spans import DEFAULT_HOST, is_ai_span
 
 
 class PostHogTraceExporter(SpanExporter):

posthog/ai/otel/processor.py

@@ -12,9 +12,7 @@
 from opentelemetry.sdk.trace.export import BatchSpanProcessor
 from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
 
-from .spans import is_ai_span
-
-DEFAULT_HOST = "https://us.i.posthog.com"
+from .spans import DEFAULT_HOST, is_ai_span
 
 
 class PostHogSpanProcessor(SpanProcessor):

posthog/ai/otel/spans.py

@@ -1,10 +1,12 @@
-"""Shared AI span filtering logic for OpenTelemetry integration."""
+"""Shared AI span filtering logic and constants for OpenTelemetry integration."""
 
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from opentelemetry.sdk.trace import ReadableSpan
 
+DEFAULT_HOST = "https://us.i.posthog.com"
+
 AI_SPAN_PREFIXES = ("gen_ai.", "llm.", "ai.", "traceloop.")
 
 

posthog/test/ai/otel/conftest.py

@@ -0,0 +1,8 @@
+from unittest.mock import MagicMock
+
+
+def make_span(name: str = "test", attributes: dict | None = None) -> MagicMock:
+    span = MagicMock()
+    span.name = name
+    span.attributes = attributes or {}
+    return span

posthog/test/ai/otel/test_exporter.py

@@ -1,16 +1,10 @@
 import unittest
-from unittest.mock import MagicMock, patch
+from unittest.mock import patch
 
 from opentelemetry.sdk.trace.export import SpanExportResult
 
 from posthog.ai.otel.exporter import PostHogTraceExporter
-
-
-def _make_span(name: str = "test", attributes: dict | None = None) -> MagicMock:
-    span = MagicMock()
-    span.name = name
-    span.attributes = attributes or {}
-    return span
+from posthog.test.ai.otel.conftest import make_span
 
 
 class TestPostHogTraceExporter(unittest.TestCase):
@@ -36,7 +30,7 @@ def test_exports_ai_spans(self, mock_otlp_cls):
         inner = mock_otlp_cls.return_value
         inner.export.return_value = SpanExportResult.SUCCESS
 
-        spans = [_make_span("gen_ai.chat"), _make_span("llm.call")]
+        spans = [make_span("gen_ai.chat"), make_span("llm.call")]
         result = exporter.export(spans)
 
         self.assertEqual(result, SpanExportResult.SUCCESS)
@@ -48,8 +42,8 @@ def test_filters_out_non_ai_spans(self, mock_otlp_cls):
         inner = mock_otlp_cls.return_value
         inner.export.return_value = SpanExportResult.SUCCESS
 
-        ai_span = _make_span("gen_ai.chat")
-        http_span = _make_span("http.request")
+        ai_span = make_span("gen_ai.chat")
+        http_span = make_span("http.request")
         result = exporter.export([ai_span, http_span])
 
         self.assertEqual(result, SpanExportResult.SUCCESS)
@@ -60,7 +54,7 @@ def test_returns_success_when_no_ai_spans(self, mock_otlp_cls):
         exporter = PostHogTraceExporter(api_key="phc_test")
         inner = mock_otlp_cls.return_value
 
-        result = exporter.export([_make_span("http.request"), _make_span("db.query")])
+        result = exporter.export([make_span("http.request"), make_span("db.query")])
 
         self.assertEqual(result, SpanExportResult.SUCCESS)
         inner.export.assert_not_called()
@@ -81,7 +75,7 @@ def test_exports_spans_with_ai_attributes(self, mock_otlp_cls):
         inner = mock_otlp_cls.return_value
         inner.export.return_value = SpanExportResult.SUCCESS
 
-        span = _make_span("http.request", {"gen_ai.system": "openai"})
+        span = make_span("http.request", {"gen_ai.system": "openai"})
         result = exporter.export([span])
 
         self.assertEqual(result, SpanExportResult.SUCCESS)

posthog/test/ai/otel/test_processor.py

@@ -2,13 +2,7 @@
 from unittest.mock import MagicMock, patch
 
 from posthog.ai.otel.processor import PostHogSpanProcessor
-
-
-def _make_span(name: str = "test", attributes: dict | None = None) -> MagicMock:
-    span = MagicMock()
-    span.name = name
-    span.attributes = attributes or {}
-    return span
+from posthog.test.ai.otel.conftest import make_span
 
 
 class TestPostHogSpanProcessor(unittest.TestCase):
@@ -48,7 +42,7 @@ def test_forwards_ai_spans(self, mock_batch_cls, mock_otlp_cls):
         processor = PostHogSpanProcessor(api_key="phc_test")
         inner = mock_batch_cls.return_value
 
-        ai_span = _make_span("gen_ai.chat")
+        ai_span = make_span("gen_ai.chat")
         processor.on_end(ai_span)
         inner.on_end.assert_called_once_with(ai_span)
 
@@ -58,7 +52,7 @@ def test_drops_non_ai_spans(self, mock_batch_cls, mock_otlp_cls):
         processor = PostHogSpanProcessor(api_key="phc_test")
         inner = mock_batch_cls.return_value
 
-        processor.on_end(_make_span("http.request"))
+        processor.on_end(make_span("http.request"))
         inner.on_end.assert_not_called()
 
     @patch("posthog.ai.otel.processor.OTLPSpanExporter")
@@ -67,7 +61,7 @@ def test_forwards_span_with_ai_attributes(self, mock_batch_cls, mock_otlp_cls):
         processor = PostHogSpanProcessor(api_key="phc_test")
         inner = mock_batch_cls.return_value
 
-        span = _make_span("http.request", {"gen_ai.system": "openai"})
+        span = make_span("http.request", {"gen_ai.system": "openai"})
         processor.on_end(span)
         inner.on_end.assert_called_once_with(span)
 

posthog/test/ai/otel/test_spans.py

@@ -1,58 +1,52 @@
 import unittest
-from unittest.mock import MagicMock
-
-from posthog.ai.otel.spans import is_ai_span
 
+from parameterized import parameterized
 
-def _make_span(name: str = "test", attributes: dict | None = None) -> MagicMock:
-    span = MagicMock()
-    span.name = name
-    span.attributes = attributes or {}
-    return span
+from posthog.ai.otel.spans import is_ai_span
+from posthog.test.ai.otel.conftest import make_span
 
 
 class TestIsAISpan(unittest.TestCase):
-    def test_matches_gen_ai_name_prefix(self):
-        self.assertTrue(is_ai_span(_make_span("gen_ai.chat")))
-
-    def test_matches_llm_name_prefix(self):
-        self.assertTrue(is_ai_span(_make_span("llm.call")))
-
-    def test_matches_ai_name_prefix(self):
-        self.assertTrue(is_ai_span(_make_span("ai.completion")))
-
-    def test_matches_traceloop_name_prefix(self):
-        self.assertTrue(is_ai_span(_make_span("traceloop.workflow")))
-
-    def test_rejects_non_ai_name(self):
-        self.assertFalse(is_ai_span(_make_span("http.request")))
-        self.assertFalse(is_ai_span(_make_span("db.query")))
-        self.assertFalse(is_ai_span(_make_span("my_function")))
-
-    def test_matches_gen_ai_attribute_key(self):
-        span = _make_span("http.request", {"gen_ai.system": "openai"})
-        self.assertTrue(is_ai_span(span))
-
-    def test_matches_llm_attribute_key(self):
-        span = _make_span("http.request", {"llm.model": "gpt-4"})
-        self.assertTrue(is_ai_span(span))
-
-    def test_matches_ai_attribute_key(self):
-        span = _make_span("http.request", {"ai.provider": "anthropic"})
-        self.assertTrue(is_ai_span(span))
-
-    def test_matches_traceloop_attribute_key(self):
-        span = _make_span("http.request", {"traceloop.entity.name": "chain"})
-        self.assertTrue(is_ai_span(span))
+    @parameterized.expand(
+        [
+            ("gen_ai", "gen_ai.chat"),
+            ("llm", "llm.call"),
+            ("ai", "ai.completion"),
+            ("traceloop", "traceloop.workflow"),
+        ]
+    )
+    def test_matches_ai_name_prefix(self, _name, span_name):
+        self.assertTrue(is_ai_span(make_span(span_name)))
+
+    @parameterized.expand(
+        [
+            ("gen_ai", {"gen_ai.system": "openai"}),
+            ("llm", {"llm.model": "gpt-4"}),
+            ("ai", {"ai.provider": "anthropic"}),
+            ("traceloop", {"traceloop.entity.name": "chain"}),
+        ]
+    )
+    def test_matches_ai_attribute_key(self, _name, attrs):
+        self.assertTrue(is_ai_span(make_span("http.request", attrs)))
+
+    @parameterized.expand(
+        [
+            ("http", "http.request"),
+            ("db", "db.query"),
+            ("custom", "my_function"),
+        ]
+    )
+    def test_rejects_non_ai_name(self, _name, span_name):
+        self.assertFalse(is_ai_span(make_span(span_name)))
 
     def test_rejects_non_ai_attributes(self):
-        span = _make_span("http.request", {"http.method": "GET", "http.url": "/"})
+        span = make_span("http.request", {"http.method": "GET", "http.url": "/"})
         self.assertFalse(is_ai_span(span))
 
     def test_empty_span(self):
-        self.assertFalse(is_ai_span(_make_span("test", {})))
+        self.assertFalse(is_ai_span(make_span("test", {})))
 
     def test_none_attributes(self):
-        span = _make_span("test")
+        span = make_span("test")
         span.attributes = None
         self.assertFalse(is_ai_span(span))