Skip to content

Commit 15d2802

Browse files
lzchenCopilot
andcommitted
refactor(genai-util): pass sampling attributes at span creation for all invocation types
Extend the pattern from InferenceInvocation to AgentInvocation, WorkflowInvocation, EmbeddingInvocation, and ToolInvocation so that all invocation types pass base attributes to start_span() for head-based sampling decisions. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Assisted-by: Claude Opus 4.6
1 parent b4a9084 commit 15d2802

10 files changed

Lines changed: 331 additions & 42 deletions

File tree

util/opentelemetry-util-genai/CHANGELOG.md

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
1010
- Change `InferenceInvocation` init params to only accept base params
1111
- Pass in `attributes` on invocation `_start` so samplers have access to attributes.
1212
([#4538](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4538))
13+
- Apply attribute for sampling on instantiation of all invocation types.
14+
([#4553](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4553))
1315

1416
## Version 0.4b0 (2026-05-01)
1517

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

Lines changed: 20 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -68,8 +68,7 @@ def __init__(
6868
request_model: str | None = None,
6969
server_address: str | None = None,
7070
server_port: int | None = None,
71-
attributes: dict[str, Any] | None = None,
72-
metric_attributes: dict[str, Any] | None = None,
71+
agent_name: str | None = None,
7372
) -> None:
7473
"""Use handler.start_invoke_local_agent() or handler.start_invoke_remote_agent() instead of calling this directly."""
7574
_operation_name = GenAI.GenAiOperationNameValues.INVOKE_AGENT.value
@@ -79,17 +78,17 @@ def __init__(
7978
logger,
8079
completion_hook,
8180
operation_name=_operation_name,
82-
span_name=_operation_name,
81+
span_name=f"{_operation_name} {agent_name}"
82+
if agent_name
83+
else _operation_name,
8384
span_kind=span_kind,
84-
attributes=attributes,
85-
metric_attributes=metric_attributes,
8685
)
8786
self.provider = provider
8887
self.request_model = request_model
8988
self.server_address = server_address
9089
self.server_port = server_port
9190

92-
self.agent_name: str | None = None
91+
self.agent_name: str | None = agent_name
9392
self.agent_id: str | None = None
9493
self.agent_description: str | None = None
9594
self.agent_version: str | None = None
@@ -119,7 +118,21 @@ def __init__(
119118
self.system_instruction: list[MessagePart] = []
120119
self.tool_definitions: list[ToolDefinition] | None = None
121120

122-
self._start()
121+
self._start(self._get_base_attributes())
122+
123+
def _get_base_attributes(self) -> dict[str, Any]:
124+
"""Return sampling-relevant attributes available at span creation time."""
125+
optional_attrs = (
126+
(GenAI.GEN_AI_REQUEST_MODEL, self.request_model),
127+
(GenAI.GEN_AI_AGENT_NAME, self.agent_name),
128+
(server_attributes.SERVER_ADDRESS, self.server_address),
129+
(server_attributes.SERVER_PORT, self.server_port),
130+
)
131+
return {
132+
GenAI.GEN_AI_OPERATION_NAME: self._operation_name,
133+
GenAI.GEN_AI_PROVIDER_NAME: self.provider,
134+
**{k: v for k, v in optional_attrs if v is not None},
135+
}
123136

124137
def _get_common_attributes(self) -> dict[str, Any]:
125138
optional_attrs = (

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

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ class EmbeddingInvocation(GenAIInvocation):
3535
context manager rather than constructing this directly.
3636
"""
3737

38-
def __init__( # pylint: disable=too-many-locals
38+
def __init__(
3939
self,
4040
tracer: Tracer,
4141
metrics_recorder: InvocationMetricsRecorder,
@@ -46,12 +46,6 @@ def __init__( # pylint: disable=too-many-locals
4646
request_model: str | None = None,
4747
server_address: str | None = None,
4848
server_port: int | None = None,
49-
encoding_formats: list[str] | None = None,
50-
input_tokens: int | None = None,
51-
dimension_count: int | None = None,
52-
response_model_name: str | None = None,
53-
attributes: dict[str, Any] | None = None,
54-
metric_attributes: dict[str, Any] | None = None,
5549
) -> None:
5650
"""Use handler.start_embedding(provider) or handler.embedding(provider) instead of calling this directly."""
5751
_operation_name = GenAI.GenAiOperationNameValues.EMBEDDINGS.value
@@ -65,20 +59,31 @@ def __init__( # pylint: disable=too-many-locals
6559
if request_model
6660
else _operation_name,
6761
span_kind=SpanKind.CLIENT,
68-
attributes=attributes,
69-
metric_attributes=metric_attributes,
7062
)
7163
self.provider = provider # e.g., azure.ai.openai, openai, aws.bedrock
7264
self.request_model = request_model
7365
self.server_address = server_address
7466
self.server_port = server_port
7567
# encoding_formats can be multi-value -> combinational cardinality risk.
7668
# Keep on spans/events only.
77-
self.encoding_formats = encoding_formats
78-
self.input_tokens = input_tokens
79-
self.dimension_count = dimension_count
80-
self.response_model_name = response_model_name
81-
self._start()
69+
self.encoding_formats: list[str] | None = None
70+
self.input_tokens: int | None = None
71+
self.dimension_count: int | None = None
72+
self.response_model_name: str | None = None
73+
self._start(self._get_base_attributes())
74+
75+
def _get_base_attributes(self) -> dict[str, Any]:
76+
"""Return sampling-relevant attributes available at span creation time."""
77+
optional_attrs = (
78+
(GenAI.GEN_AI_REQUEST_MODEL, self.request_model),
79+
(GenAI.GEN_AI_PROVIDER_NAME, self.provider),
80+
(server_attributes.SERVER_ADDRESS, self.server_address),
81+
(server_attributes.SERVER_PORT, self.server_port),
82+
)
83+
return {
84+
GenAI.GEN_AI_OPERATION_NAME: self._operation_name,
85+
**{k: v for k, v in optional_attrs if v is not None},
86+
}
8287

8388
def _get_metric_attributes(self) -> dict[str, Any]:
8489
optional_attrs = (

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

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -58,9 +58,6 @@ def __init__(
5858
tool_call_id: str | None = None,
5959
tool_type: str | None = None,
6060
tool_description: str | None = None,
61-
tool_result: Any = None,
62-
attributes: dict[str, Any] | None = None,
63-
metric_attributes: dict[str, Any] | None = None,
6461
) -> None:
6562
"""Use handler.start_tool(name) or handler.tool(name) instead of calling this directly."""
6663
_operation_name = GenAI.GenAiOperationNameValues.EXECUTE_TOOL.value
@@ -71,16 +68,27 @@ def __init__(
7168
completion_hook,
7269
operation_name=_operation_name,
7370
span_name=f"{_operation_name} {name}" if name else _operation_name,
74-
attributes=attributes,
75-
metric_attributes=metric_attributes,
7671
)
7772
self.name = name
7873
self.arguments = arguments
7974
self.tool_call_id = tool_call_id
8075
self.tool_type = tool_type
8176
self.tool_description = tool_description
82-
self.tool_result = tool_result
83-
self._start()
77+
self.tool_result: Any = None
78+
self._start(self._get_base_attributes())
79+
80+
def _get_base_attributes(self) -> dict[str, Any]:
81+
"""Return sampling-relevant attributes available at span creation time."""
82+
optional_attrs = (
83+
(GenAI.GEN_AI_TOOL_NAME, self.name),
84+
(GenAI.GEN_AI_TOOL_CALL_ID, self.tool_call_id),
85+
(GenAI.GEN_AI_TOOL_TYPE, self.tool_type),
86+
(GenAI.GEN_AI_TOOL_DESCRIPTION, self.tool_description),
87+
)
88+
return {
89+
GenAI.GEN_AI_OPERATION_NAME: self._operation_name,
90+
**{k: v for k, v in optional_attrs if v is not None},
91+
}
8492

8593
def _get_metric_attributes(self) -> dict[str, Any]:
8694
attrs: dict[str, Any] = {

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

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -54,11 +54,6 @@ def __init__(
5454
logger: Logger,
5555
completion_hook: CompletionHook,
5656
name: str | None,
57-
*,
58-
input_messages: list[InputMessage] | None = None,
59-
output_messages: list[OutputMessage] | None = None,
60-
attributes: dict[str, Any] | None = None,
61-
metric_attributes: dict[str, Any] | None = None,
6257
) -> None:
6358
"""Use handler.start_workflow(name) or handler.workflow(name) instead of calling this directly."""
6459
_operation_name = "invoke_workflow"
@@ -70,17 +65,18 @@ def __init__(
7065
operation_name=_operation_name,
7166
span_name=f"{_operation_name} {name}" if name else _operation_name,
7267
span_kind=SpanKind.INTERNAL,
73-
attributes=attributes,
74-
metric_attributes=metric_attributes,
7568
)
7669
self.name = name
77-
self.input_messages: list[InputMessage] = (
78-
[] if input_messages is None else input_messages
79-
)
80-
self.output_messages: list[OutputMessage] = (
81-
[] if output_messages is None else output_messages
82-
)
83-
self._start()
70+
self.input_messages: list[InputMessage] = []
71+
self.output_messages: list[OutputMessage] = []
72+
self._start(self._get_base_attributes())
73+
74+
def _get_base_attributes(self) -> dict[str, Any]:
75+
"""Return sampling-relevant attributes available at span creation time."""
76+
attrs: dict[str, Any] = {
77+
GenAI.GEN_AI_OPERATION_NAME: self._operation_name,
78+
}
79+
return attrs
8480

8581
def _get_messages_for_span(self) -> dict[str, Any]:
8682
if not is_experimental_mode() or get_content_capturing_mode() not in (

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -353,6 +353,7 @@ def start_invoke_local_agent(
353353
provider: str,
354354
*,
355355
request_model: str | None = None,
356+
agent_name: str | None = None,
356357
) -> AgentInvocation:
357358
"""Create and start a local agent invocation (INTERNAL span kind).
358359
@@ -369,6 +370,7 @@ def start_invoke_local_agent(
369370
provider,
370371
span_kind=SpanKind.INTERNAL,
371372
request_model=request_model,
373+
agent_name=agent_name,
372374
)
373375

374376
def start_invoke_remote_agent(
@@ -378,6 +380,7 @@ def start_invoke_remote_agent(
378380
request_model: str | None = None,
379381
server_address: str | None = None,
380382
server_port: int | None = None,
383+
agent_name: str | None = None,
381384
) -> AgentInvocation:
382385
"""Create and start a remote agent invocation (CLIENT span kind).
383386
@@ -394,6 +397,7 @@ def start_invoke_remote_agent(
394397
provider,
395398
span_kind=SpanKind.CLIENT,
396399
request_model=request_model,
400+
agent_name=agent_name,
397401
server_address=server_address,
398402
server_port=server_port,
399403
)
@@ -403,6 +407,7 @@ def invoke_local_agent(
403407
provider: str,
404408
*,
405409
request_model: str | None = None,
410+
agent_name: str | None = None,
406411
) -> AbstractContextManager[AgentInvocation]:
407412
"""Context manager for local agent invocations (INTERNAL span kind).
408413
@@ -417,6 +422,7 @@ def invoke_local_agent(
417422
return self.start_invoke_local_agent(
418423
provider,
419424
request_model=request_model,
425+
agent_name=agent_name,
420426
)._managed()
421427

422428
def invoke_remote_agent(
@@ -426,6 +432,7 @@ def invoke_remote_agent(
426432
request_model: str | None = None,
427433
server_address: str | None = None,
428434
server_port: int | None = None,
435+
agent_name: str | None = None,
429436
) -> AbstractContextManager[AgentInvocation]:
430437
"""Context manager for remote agent invocations (CLIENT span kind).
431438
@@ -440,6 +447,7 @@ def invoke_remote_agent(
440447
return self.start_invoke_remote_agent(
441448
provider,
442449
request_model=request_model,
450+
agent_name=agent_name,
443451
server_address=server_address,
444452
server_port=server_port,
445453
)._managed()

0 commit comments

Comments
 (0)