remove overly complicated env check, add messages to test

Author: keith-decker

util/opentelemetry-util-genai/src/opentelemetry/util/genai/generators.py

@@ -35,7 +35,6 @@
 import json
 from contextlib import contextmanager
 from dataclasses import asdict, dataclass, field
-from enum import Enum
 from typing import Any, Dict, List, Optional, Tuple
 from uuid import UUID
 
@@ -143,41 +142,22 @@ def _collect_finish_reasons(generations: List[OutputMessage]) -> List[str]:
     return finish_reasons
 
 
-class _CaptureTarget(Enum):
-    SPAN = "span"
-    EVENT = "event"
-
-
-def _should_capture_content(mode: _CaptureTarget) -> bool:
-    # if GEN_AI stability mode is DEFAULT, do not capture message content
-    if (
+def _is_experimental_mode() -> bool:
+    return (
         _OpenTelemetrySemanticConventionStability._get_opentelemetry_stability_opt_in_mode(
             _OpenTelemetryStabilitySignalType.GEN_AI,
         )
-        == _StabilityMode.DEFAULT
-    ):
-        return False
-
-    capture_mode = get_content_capturing_mode()
-    if mode == _CaptureTarget.SPAN:
-        return capture_mode in (
-            ContentCapturingMode.SPAN_ONLY,
-            ContentCapturingMode.SPAN_AND_EVENT,
-        )
-
-    if mode == _CaptureTarget.EVENT:
-        return capture_mode in (
-            ContentCapturingMode.EVENT_ONLY,
-            ContentCapturingMode.SPAN_AND_EVENT,
-        )
-
-    return False
+        is _StabilityMode.GEN_AI_LATEST_EXPERIMENTAL
+    )
 
 
 def _maybe_set_span_input_messages(
     span: Span, messages: List[InputMessage]
 ) -> None:
-    if not _should_capture_content(_CaptureTarget.SPAN):
+    if not _is_experimental_mode() or get_content_capturing_mode() not in (
+        ContentCapturingMode.SPAN_ONLY,
+        ContentCapturingMode.SPAN_AND_EVENT,
+    ):
         return
     message_parts: List[Dict[str, Any]] = [
         asdict(message) for message in messages
@@ -189,7 +169,10 @@ def _maybe_set_span_input_messages(
 def _maybe_set_span_output_messages(
     span: Span, generations: List[OutputMessage]
 ) -> None:
-    if not _should_capture_content(_CaptureTarget.SPAN):
+    if not _is_experimental_mode() or get_content_capturing_mode() not in (
+        ContentCapturingMode.SPAN_ONLY,
+        ContentCapturingMode.SPAN_AND_EVENT,
+    ):
         return
     generation_parts: List[Dict[str, Any]] = [
         asdict(generation) for generation in generations

util/opentelemetry-util-genai/tests/test_utils.py

@@ -97,20 +97,33 @@ def test_get_content_capturing_mode_raises_exception_on_invalid_envvar(
 
 
 class TestTelemetryHandler(unittest.TestCase):
-    def setUp(self):
-        self.span_exporter = InMemorySpanExporter()
+    @classmethod
+    def setUpClass(cls):
+        cls.span_exporter = InMemorySpanExporter()
         tracer_provider = TracerProvider()
         tracer_provider.add_span_processor(
-            SimpleSpanProcessor(self.span_exporter)
+            SimpleSpanProcessor(cls.span_exporter)
         )
-        trace.set_tracer_provider(tracer_provider)
+        try:
+            trace.set_tracer_provider(tracer_provider)
+        except Exception:  # pragma: no cover - defensive
+            pass
+
+    def setUp(self):
+        self.span_exporter = self.__class__.span_exporter
+        self.span_exporter.clear()
         self.telemetry_handler = get_telemetry_handler()
 
     def tearDown(self):
+        # Clear spans and reset the singleton telemetry handler so each test starts clean
         self.span_exporter.clear()
-        # Reset to default tracer provider
-        trace.set_tracer_provider(trace.NoOpTracerProvider())
+        if hasattr(get_telemetry_handler, "_default_handler"):
+            delattr(get_telemetry_handler, "_default_handler")
 
+    @patch_env_vars(
+        stability_mode="gen_ai_latest_experimental",
+        content_capturing="SPAN_ONLY",
+    )
     def test_llm_start_and_stop_creates_span(self):  # pylint: disable=no-self-use
         run_id = uuid4()
         message = InputMessage(
@@ -147,4 +160,11 @@ def test_llm_start_and_stop_creates_span(self):  # pylint: disable=no-self-use
         assert invocation.attributes.get("custom_attr") == "value"
         assert invocation.attributes.get("extra") == "info"
 
-        # TODO: check messages
+        # Check messages captured on span
+        input_messages_json = span_attrs.get("gen_ai.input.messages")
+        output_messages_json = span_attrs.get("gen_ai.output.messages")
+        assert input_messages_json is not None
+        assert output_messages_json is not None
+
+        assert isinstance(input_messages_json, str)
+        assert isinstance(output_messages_json, str)