refactor(genai-util): pass sampling attributes at span creation for all invocation types (#4553)

Author: lzchen

util/opentelemetry-util-genai/CHANGELOG.md

@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - 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))
+- Apply attribute for sampling on instantiation of all invocation types.
+  ([#4553](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4553))
 
 ## Version 0.4b0 (2026-05-01)
 

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

@@ -57,8 +57,7 @@ def __init__(
         request_model: str | 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,
+        agent_name: str | None = None,
     ) -> None:
         """Use handler.start_invoke_local_agent() or handler.start_invoke_remote_agent() instead of calling this directly."""
         _operation_name = GenAI.GenAiOperationNameValues.INVOKE_AGENT.value
@@ -68,17 +67,17 @@ def __init__(
             logger,
             completion_hook,
             operation_name=_operation_name,
-            span_name=_operation_name,
+            span_name=f"{_operation_name} {agent_name}"
+            if agent_name
+            else _operation_name,
             span_kind=span_kind,
-            attributes=attributes,
-            metric_attributes=metric_attributes,
         )
         self.provider = provider
         self.request_model = request_model
         self.server_address = server_address
         self.server_port = server_port
 
-        self.agent_name: str | None = None
+        self.agent_name: str | None = agent_name
         self.agent_id: str | None = None
         self.agent_description: str | None = None
         self.agent_version: str | None = None
@@ -108,7 +107,21 @@ def __init__(
         self.system_instruction: list[MessagePart] = []
         self.tool_definitions: list[ToolDefinition] | None = None
 
-        self._start()
+        self._start(self._get_base_attributes())
+
+    def _get_base_attributes(self) -> dict[str, Any]:
+        """Return sampling-relevant attributes available at span creation time."""
+        optional_attrs = (
+            (GenAI.GEN_AI_REQUEST_MODEL, self.request_model),
+            (GenAI.GEN_AI_AGENT_NAME, self.agent_name),
+            (server_attributes.SERVER_ADDRESS, self.server_address),
+            (server_attributes.SERVER_PORT, self.server_port),
+        )
+        return {
+            GenAI.GEN_AI_OPERATION_NAME: self._operation_name,
+            GenAI.GEN_AI_PROVIDER_NAME: self.provider,
+            **{k: v for k, v in optional_attrs if v is not None},
+        }
 
     def _get_common_attributes(self) -> dict[str, Any]:
         optional_attrs = (
@@ -199,10 +212,6 @@ def _apply_finish(self, error: Error | None = None) -> None:
         if error is not None:
             self._apply_error_attributes(error)
 
-        # Update span name if agent_name was set after construction
-        if self.agent_name:
-            self.span.update_name(f"{self._operation_name} {self.agent_name}")
-
         attributes: dict[str, Any] = {}
         attributes.update(self._get_common_attributes())
         attributes.update(self._get_request_attributes())

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

@@ -24,7 +24,7 @@ class EmbeddingInvocation(GenAIInvocation):
     context manager rather than constructing this directly.
     """
 
-    def __init__(  # pylint: disable=too-many-locals
+    def __init__(
         self,
         tracer: Tracer,
         metrics_recorder: InvocationMetricsRecorder,
@@ -35,12 +35,6 @@ def __init__(  # pylint: disable=too-many-locals
         request_model: str | None = None,
         server_address: str | None = None,
         server_port: int | None = None,
-        encoding_formats: list[str] | None = None,
-        input_tokens: int | None = None,
-        dimension_count: int | None = None,
-        response_model_name: str | None = None,
-        attributes: dict[str, Any] | None = None,
-        metric_attributes: dict[str, Any] | None = None,
     ) -> None:
         """Use handler.start_embedding(provider) or handler.embedding(provider) instead of calling this directly."""
         _operation_name = GenAI.GenAiOperationNameValues.EMBEDDINGS.value
@@ -54,20 +48,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  # e.g., azure.ai.openai, openai, aws.bedrock
         self.request_model = request_model
         self.server_address = server_address
         self.server_port = server_port
         # encoding_formats can be multi-value -> combinational cardinality risk.
         # Keep on spans/events only.
-        self.encoding_formats = encoding_formats
-        self.input_tokens = input_tokens
-        self.dimension_count = dimension_count
-        self.response_model_name = response_model_name
-        self._start()
+        self.encoding_formats: list[str] | None = None
+        self.input_tokens: int | None = None
+        self.dimension_count: int | None = None
+        self.response_model_name: str | None = None
+        self._start(self._get_base_attributes())
+
+    def _get_base_attributes(self) -> dict[str, Any]:
+        """Return sampling-relevant attributes available at span creation time."""
+        optional_attrs = (
+            (GenAI.GEN_AI_REQUEST_MODEL, self.request_model),
+            (GenAI.GEN_AI_PROVIDER_NAME, self.provider),
+            (server_attributes.SERVER_ADDRESS, self.server_address),
+            (server_attributes.SERVER_PORT, self.server_port),
+        )
+        return {
+            GenAI.GEN_AI_OPERATION_NAME: self._operation_name,
+            **{k: v for k, v in optional_attrs if v is not None},
+        }
 
     def _get_metric_attributes(self) -> dict[str, Any]:
         optional_attrs = (

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

@@ -47,9 +47,6 @@ def __init__(
         tool_call_id: str | None = None,
         tool_type: str | None = None,
         tool_description: str | None = None,
-        tool_result: Any = None,
-        attributes: dict[str, Any] | None = None,
-        metric_attributes: dict[str, Any] | None = None,
     ) -> None:
         """Use handler.start_tool(name) or handler.tool(name) instead of calling this directly."""
         _operation_name = GenAI.GenAiOperationNameValues.EXECUTE_TOOL.value
@@ -60,16 +57,27 @@ def __init__(
             completion_hook,
             operation_name=_operation_name,
             span_name=f"{_operation_name} {name}" if name else _operation_name,
-            attributes=attributes,
-            metric_attributes=metric_attributes,
         )
         self.name = name
         self.arguments = arguments
         self.tool_call_id = tool_call_id
         self.tool_type = tool_type
         self.tool_description = tool_description
-        self.tool_result = tool_result
-        self._start()
+        self.tool_result: Any = None
+        self._start(self._get_base_attributes())
+
+    def _get_base_attributes(self) -> dict[str, Any]:
+        """Return sampling-relevant attributes available at span creation time."""
+        optional_attrs = (
+            (GenAI.GEN_AI_TOOL_NAME, self.name),
+            (GenAI.GEN_AI_TOOL_CALL_ID, self.tool_call_id),
+            (GenAI.GEN_AI_TOOL_TYPE, self.tool_type),
+            (GenAI.GEN_AI_TOOL_DESCRIPTION, self.tool_description),
+        )
+        return {
+            GenAI.GEN_AI_OPERATION_NAME: self._operation_name,
+            **{k: v for k, v in optional_attrs if v is not None},
+        }
 
     def _get_metric_attributes(self) -> dict[str, Any]:
         attrs: dict[str, Any] = {

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

@@ -43,11 +43,6 @@ def __init__(
         logger: Logger,
         completion_hook: CompletionHook,
         name: str | None,
-        *,
-        input_messages: list[InputMessage] | None = None,
-        output_messages: list[OutputMessage] | None = None,
-        attributes: dict[str, Any] | None = None,
-        metric_attributes: dict[str, Any] | None = None,
     ) -> None:
         """Use handler.start_workflow(name) or handler.workflow(name) instead of calling this directly."""
         _operation_name = "invoke_workflow"
@@ -59,17 +54,18 @@ def __init__(
             operation_name=_operation_name,
             span_name=f"{_operation_name} {name}" if name else _operation_name,
             span_kind=SpanKind.INTERNAL,
-            attributes=attributes,
-            metric_attributes=metric_attributes,
         )
         self.name = name
-        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._start()
+        self.input_messages: list[InputMessage] = []
+        self.output_messages: list[OutputMessage] = []
+        self._start(self._get_base_attributes())
+
+    def _get_base_attributes(self) -> dict[str, Any]:
+        """Return sampling-relevant attributes available at span creation time."""
+        attrs: dict[str, Any] = {
+            GenAI.GEN_AI_OPERATION_NAME: self._operation_name,
+        }
+        return attrs
 
     def _get_messages_for_span(self) -> dict[str, Any]:
         if not is_experimental_mode() or get_content_capturing_mode() not in (

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

@@ -342,6 +342,7 @@ def start_invoke_local_agent(
         provider: str,
         *,
         request_model: str | None = None,
+        agent_name: str | None = None,
     ) -> AgentInvocation:
         """Create and start a local agent invocation (INTERNAL span kind).
 
@@ -358,6 +359,7 @@ def start_invoke_local_agent(
             provider,
             span_kind=SpanKind.INTERNAL,
             request_model=request_model,
+            agent_name=agent_name,
         )
 
     def start_invoke_remote_agent(
@@ -367,6 +369,7 @@ def start_invoke_remote_agent(
         request_model: str | None = None,
         server_address: str | None = None,
         server_port: int | None = None,
+        agent_name: str | None = None,
     ) -> AgentInvocation:
         """Create and start a remote agent invocation (CLIENT span kind).
 
@@ -383,6 +386,7 @@ def start_invoke_remote_agent(
             provider,
             span_kind=SpanKind.CLIENT,
             request_model=request_model,
+            agent_name=agent_name,
             server_address=server_address,
             server_port=server_port,
         )
@@ -392,6 +396,7 @@ def invoke_local_agent(
         provider: str,
         *,
         request_model: str | None = None,
+        agent_name: str | None = None,
     ) -> AbstractContextManager[AgentInvocation]:
         """Context manager for local agent invocations (INTERNAL span kind).
 
@@ -406,6 +411,7 @@ def invoke_local_agent(
         return self.start_invoke_local_agent(
             provider,
             request_model=request_model,
+            agent_name=agent_name,
         )._managed()
 
     def invoke_remote_agent(
@@ -415,6 +421,7 @@ def invoke_remote_agent(
         request_model: str | None = None,
         server_address: str | None = None,
         server_port: int | None = None,
+        agent_name: str | None = None,
     ) -> AbstractContextManager[AgentInvocation]:
         """Context manager for remote agent invocations (CLIENT span kind).
 
@@ -429,6 +436,7 @@ def invoke_remote_agent(
         return self.start_invoke_remote_agent(
             provider,
             request_model=request_model,
+            agent_name=agent_name,
             server_address=server_address,
             server_port=server_port,
         )._managed()