refactor(genai-util): pass sampling attributes at span creation time (#4538)

Author: lzchen

util/opentelemetry-util-genai/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+- Change `InferenceInvocation` init params to only accept base params
+- Pass in `attributes` on invocation `_start` so samplers have access to attributes.
+  ([#4538](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4538))
+
 ## Version 0.4b0 (2026-05-01)
 
 - Add `AgentInvocation` type with `invoke_agent` span lifecycle

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

@@ -55,7 +55,7 @@ class InferenceInvocation(GenAIInvocation):
     context manager rather than constructing this directly.
     """
 
-    def __init__(  # pylint: disable=too-many-locals
+    def __init__(
         self,
         tracer: Tracer,
         metrics_recorder: InvocationMetricsRecorder,
@@ -64,25 +64,8 @@ def __init__(  # pylint: disable=too-many-locals
         provider: str,
         *,
         request_model: str | None = None,
-        input_messages: list[InputMessage] | None = None,
-        output_messages: list[OutputMessage] | None = None,
-        system_instruction: list[MessagePart] | None = None,
-        response_model_name: str | None = None,
-        response_id: str | None = None,
-        finish_reasons: list[str] | None = None,
-        input_tokens: int | None = None,
-        output_tokens: int | None = None,
-        temperature: float | None = None,
-        top_p: float | None = None,
-        frequency_penalty: float | None = None,
-        presence_penalty: float | None = None,
-        max_tokens: int | None = None,
-        stop_sequences: list[str] | None = None,
-        seed: int | None = None,
         server_address: str | None = None,
         server_port: int | None = None,
-        attributes: dict[str, Any] | None = None,
-        metric_attributes: dict[str, Any] | None = None,
     ) -> None:
         """Use handler.start_inference(provider) or handler.inference(provider) instead of calling this directly."""
         _operation_name = GenAI.GenAiOperationNameValues.CHAT.value
@@ -96,38 +79,31 @@ def __init__(  # pylint: disable=too-many-locals
             if request_model
             else _operation_name,
             span_kind=SpanKind.CLIENT,
-            attributes=attributes,
-            metric_attributes=metric_attributes,
         )
         self.provider = provider
         self.request_model = request_model
-        self.input_messages: list[InputMessage] = (
-            [] if input_messages is None else input_messages
-        )
-        self.output_messages: list[OutputMessage] = (
-            [] if output_messages is None else output_messages
-        )
-        self.system_instruction: list[MessagePart] = (
-            [] if system_instruction is None else system_instruction
-        )
-        self.response_model_name = response_model_name
-        self.response_id = response_id
-        self.finish_reasons = finish_reasons
-        self.input_tokens = input_tokens
-        self.output_tokens = output_tokens
-        self.temperature = temperature
-        self.top_p = top_p
-        self.frequency_penalty = frequency_penalty
-        self.presence_penalty = presence_penalty
-        self.max_tokens = max_tokens
-        self.stop_sequences = stop_sequences
-        self.seed = seed
         self.server_address = server_address
         self.server_port = server_port
+
+        self.input_messages: list[InputMessage] = []
+        self.output_messages: list[OutputMessage] = []
+        self.system_instruction: list[MessagePart] = []
+        self.response_model_name: str | None = None
+        self.response_id: str | None = None
+        self.finish_reasons: list[str] | None = None
+        self.input_tokens: int | None = None
+        self.output_tokens: int | None = None
+        self.temperature: float | None = None
+        self.top_p: float | None = None
+        self.frequency_penalty: float | None = None
+        self.presence_penalty: float | None = None
+        self.max_tokens: int | None = None
+        self.stop_sequences: list[str] | None = None
+        self.seed: int | None = None
         self.cache_creation_input_tokens: int | None = None
         self.cache_read_input_tokens: int | None = None
         self.tool_definitions: list[ToolDefinition] | None = None
-        self._start()
+        self._start(self._get_base_attributes())
 
     def _get_message_attributes(self, *, for_span: bool) -> dict[str, Any]:
         return get_content_attributes(
@@ -288,33 +264,34 @@ def _start_with_handler(
         completion_hook: CompletionHook,
     ) -> None:
         """Create and start an InferenceInvocation from this data container. Called by handler.start_llm()."""
-        self._inference_invocation = InferenceInvocation(
+        inv = InferenceInvocation(
             tracer,
             metrics_recorder,
             logger,
             completion_hook,
             self.provider or "",
             request_model=self.request_model,
-            input_messages=self.input_messages,
-            output_messages=self.output_messages,
-            system_instruction=self.system_instruction,
-            response_model_name=self.response_model_name,
-            response_id=self.response_id,
-            finish_reasons=self.finish_reasons,
-            input_tokens=self.input_tokens,
-            output_tokens=self.output_tokens,
-            temperature=self.temperature,
-            top_p=self.top_p,
-            frequency_penalty=self.frequency_penalty,
-            presence_penalty=self.presence_penalty,
-            max_tokens=self.max_tokens,
-            stop_sequences=self.stop_sequences,
-            seed=self.seed,
             server_address=self.server_address,
             server_port=self.server_port,
-            attributes=self.attributes,
-            metric_attributes=self.metric_attributes,
         )
+        inv.input_messages = self.input_messages
+        inv.output_messages = self.output_messages
+        inv.system_instruction = self.system_instruction
+        inv.response_model_name = self.response_model_name
+        inv.response_id = self.response_id
+        inv.finish_reasons = self.finish_reasons
+        inv.input_tokens = self.input_tokens
+        inv.output_tokens = self.output_tokens
+        inv.temperature = self.temperature
+        inv.top_p = self.top_p
+        inv.frequency_penalty = self.frequency_penalty
+        inv.presence_penalty = self.presence_penalty
+        inv.max_tokens = self.max_tokens
+        inv.stop_sequences = self.stop_sequences
+        inv.seed = self.seed
+        inv.attributes.update(self.attributes)
+        inv.metric_attributes.update(self.metric_attributes)
+        self._inference_invocation = inv
 
     def _sync_to_invocation(self) -> None:
         inv = self._inference_invocation

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

@@ -96,11 +96,16 @@ def __init__(
         self._context_token: ContextToken | None = None
         self._monotonic_start_s: float | None = None
 
-    def _start(self) -> None:
-        """Start the invocation span and attach it to the current context."""
+    def _start(self, attributes: dict[str, Any] | None = None) -> None:
+        """Start the invocation span and attach it to the current context.
+
+        Args:
+            attributes: Initial span attributes available for sampling decisions.
+        """
         self.span = self._tracer.start_span(
             name=self._span_name,
             kind=self._span_kind,
+            attributes=attributes,
         )
         self._span_context = set_span_in_context(self.span)
         self._monotonic_start_s = timeit.default_timer()

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

@@ -33,6 +33,7 @@
 from opentelemetry.sdk.trace.export.in_memory_span_exporter import (
     InMemorySpanExporter,
 )
+from opentelemetry.sdk.trace.sampling import Decision, SamplingResult
 from opentelemetry.semconv._incubating.attributes import (
     gen_ai_attributes as GenAI,
 )
@@ -46,7 +47,10 @@
     OTEL_INSTRUMENTATION_GENAI_CAPTURE_MESSAGE_CONTENT,
     OTEL_INSTRUMENTATION_GENAI_EMIT_EVENT,
 )
-from opentelemetry.util.genai.handler import get_telemetry_handler
+from opentelemetry.util.genai.handler import (
+    TelemetryHandler,
+    get_telemetry_handler,
+)
 from opentelemetry.util.genai.types import (
     ContentCapturingMode,
     InputMessage,
@@ -357,6 +361,108 @@ def test_llm_manual_start_and_stop_creates_span(self):
             },
         )
 
+    def test_start_inference_passes_sampling_attributes_at_span_creation(self):
+        """Verify that sampling-relevant attributes are available at start_span() time."""
+        captured_attributes = {}
+
+        class AttributeCapturingSampler:  # pylint: disable=no-self-use
+            """A sampler that records the attributes passed to should_sample."""
+
+            def should_sample(
+                self,
+                parent_context,
+                trace_id,
+                name,
+                kind=None,
+                attributes=None,
+                links=None,
+            ):
+                captured_attributes.update(attributes or {})
+
+                return SamplingResult(Decision.RECORD_AND_SAMPLE, attributes)
+
+            def get_description(self):
+                return "AttributeCapturingSampler"
+
+        sampler_provider = TracerProvider(sampler=AttributeCapturingSampler())
+        sampler_provider.add_span_processor(
+            SimpleSpanProcessor(self.span_exporter)
+        )
+
+        handler = TelemetryHandler(tracer_provider=sampler_provider)
+
+        invocation = handler.start_inference(
+            "test-provider",
+            request_model="sampler-model",
+            server_address="api.example.com",
+            server_port=8080,
+        )
+        invocation.stop()
+
+        assert captured_attributes[GenAI.GEN_AI_OPERATION_NAME] == "chat"
+        assert (
+            captured_attributes[GenAI.GEN_AI_REQUEST_MODEL] == "sampler-model"
+        )
+        assert (
+            captured_attributes[GenAI.GEN_AI_PROVIDER_NAME] == "test-provider"
+        )
+        assert (
+            captured_attributes[server_attributes.SERVER_ADDRESS]
+            == "api.example.com"
+        )
+        assert captured_attributes[server_attributes.SERVER_PORT] == 8080
+
+    def test_start_inference_sampler_can_drop_span_based_on_attributes(self):
+        """Verify that a sampler can reject spans based on attributes passed at creation time."""
+
+        class ModelRejectingSampler:  # pylint: disable=no-self-use
+            """Drops spans whose gen_ai.request.model matches the reject list."""
+
+            def __init__(self, reject_models):
+                self._reject_models = reject_models
+
+            def should_sample(
+                self,
+                parent_context,
+                trace_id,
+                name,
+                kind=None,
+                attributes=None,
+                links=None,
+            ):
+                model = (attributes or {}).get(GenAI.GEN_AI_REQUEST_MODEL)
+                if model in self._reject_models:
+                    return SamplingResult(Decision.DROP)
+                return SamplingResult(Decision.RECORD_AND_SAMPLE, attributes)
+
+            def get_description(self):
+                return "ModelRejectingSampler"
+
+        sampler_provider = TracerProvider(
+            sampler=ModelRejectingSampler(reject_models={"rejected-model"})
+        )
+        sampler_provider.add_span_processor(
+            SimpleSpanProcessor(self.span_exporter)
+        )
+
+        handler = TelemetryHandler(tracer_provider=sampler_provider)
+
+        # This invocation should be dropped
+        invocation = handler.start_inference(
+            "test-provider", request_model="rejected-model"
+        )
+        invocation.stop()
+
+        # This invocation should be recorded
+        invocation = handler.start_inference(
+            "test-provider", request_model="accepted-model"
+        )
+        invocation.stop()
+
+        spans = self.span_exporter.get_finished_spans()
+        assert len(spans) == 1
+        assert spans[0].name == "chat accepted-model"
+
     def test_llm_span_finish_reasons_without_output_messages(self):
         invocation = self.telemetry_handler.start_inference(
             "test-provider", request_model="model-without-output"