diff --git a/instrumentation/opentelemetry-instrumentation-genai-bedrock/tests/test_converse.py b/instrumentation/opentelemetry-instrumentation-genai-bedrock/tests/test_converse.py index 0450fa270..895388c5b 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-bedrock/tests/test_converse.py +++ b/instrumentation/opentelemetry-instrumentation-genai-bedrock/tests/test_converse.py @@ -305,95 +305,97 @@ def test_converse_with_invalid_model( def test_extract_converse_request_no_content(tracer_provider) -> None: handler = TelemetryHandler(tracer_provider=tracer_provider) - invocation = handler.inference(provider="aws.bedrock") - extract_converse_request( - { - "messages": [{"role": "user", "content": [{"text": "hello"}]}], - "system": [{"text": "system instruction"}], - "inferenceConfig": {"temperature": 0.5}, - "toolConfig": { - "tools": [{"toolSpec": {"name": "get_weather"}}], + with handler.inference(provider="aws.bedrock") as invocation: + extract_converse_request( + { + "messages": [{"role": "user", "content": [{"text": "hello"}]}], + "system": [{"text": "system instruction"}], + "inferenceConfig": {"temperature": 0.5}, + "toolConfig": { + "tools": [{"toolSpec": {"name": "get_weather"}}], + }, }, - }, - invocation, - capture_content=False, - ) - assert not invocation.input_messages - assert not invocation.system_instruction - assert invocation.tool_definitions - assert invocation.temperature == 0.5 + invocation, + capture_content=False, + ) + assert not invocation.input_messages + assert not invocation.system_instruction + assert invocation.tool_definitions + assert invocation.temperature == 0.5 def test_extract_converse_response_no_content(tracer_provider) -> None: handler = TelemetryHandler(tracer_provider=tracer_provider) - invocation = handler.inference(provider="aws.bedrock") - extract_converse_response( - { - "output": { - "message": { - "role": "assistant", - "content": [{"text": "hi"}], - } - }, - "stopReason": "end_turn", - "usage": { - "inputTokens": 5, - "outputTokens": 2, - "cacheReadInputTokens": 3, - "cacheWriteInputTokens": 7, + with handler.inference(provider="aws.bedrock") as invocation: + extract_converse_response( + { + "output": { + "message": { + "role": "assistant", + "content": [{"text": "hi"}], + } + }, + "stopReason": "end_turn", + "usage": { + "inputTokens": 5, + "outputTokens": 2, + "cacheReadInputTokens": 3, + "cacheWriteInputTokens": 7, + }, }, - }, - invocation, - capture_content=False, - ) - assert not invocation.output_messages - assert invocation.finish_reasons == ["stop"] - assert invocation.input_tokens == 5 - assert invocation.output_tokens == 2 - assert invocation.cache_read_input_tokens == 3 - assert invocation.cache_creation_input_tokens == 7 + invocation, + capture_content=False, + ) + assert not invocation.output_messages + assert invocation.finish_reasons == ["stop"] + assert invocation.input_tokens == 5 + assert invocation.output_tokens == 2 + assert invocation.cache_read_input_tokens == 3 + assert invocation.cache_creation_input_tokens == 7 def test_extract_converse_request_top_k_and_seed(tracer_provider) -> None: handler = TelemetryHandler(tracer_provider=tracer_provider) - invocation = handler.inference(provider="aws.bedrock") - extract_converse_request( - { - "inferenceConfig": {"topK": 40, "seed": 123}, - }, - invocation, - ) - assert invocation.top_k == 40.0 - assert invocation.seed == 123 + with handler.inference(provider="aws.bedrock") as invocation: + extract_converse_request( + { + "inferenceConfig": {"topK": 40, "seed": 123}, + }, + invocation, + ) + assert invocation.top_k == 40.0 + assert invocation.seed == 123 - invocation2 = handler.inference(provider="aws.bedrock") - extract_converse_request( - { - "additionalModelRequestFields": {"top_k": 250, "seed": 456}, - }, - invocation2, - ) - assert invocation2.top_k == 250.0 - assert invocation2.seed == 456 + with handler.inference(provider="aws.bedrock") as invocation2: + extract_converse_request( + { + "additionalModelRequestFields": {"top_k": 250, "seed": 456}, + }, + invocation2, + ) + assert invocation2.top_k == 250.0 + assert invocation2.seed == 456 - invocation3 = handler.inference(provider="aws.bedrock") - extract_converse_request( - { - "additionalModelRequestFields": {"inferenceConfig": {"topK": 20}}, - }, - invocation3, - ) - assert invocation3.top_k == 20.0 + with handler.inference(provider="aws.bedrock") as invocation3: + extract_converse_request( + { + "additionalModelRequestFields": { + "inferenceConfig": {"topK": 20} + }, + }, + invocation3, + ) + assert invocation3.top_k == 20.0 - invocation4 = handler.inference(provider="aws.bedrock") - extract_converse_request( - { - "inferenceConfig": {"topK": 0, "seed": 0}, - }, - invocation4, - ) - assert invocation4.top_k == 0.0 - assert invocation4.seed == 0 + with handler.inference(provider="aws.bedrock") as invocation4: + extract_converse_request( + { + "inferenceConfig": {"topK": 0, "seed": 0}, + }, + invocation4, + ) + assert invocation4.top_k == 0.0 + assert invocation4.seed == 0 def test_extract_content_block_reasoning() -> None: @@ -540,104 +542,100 @@ def test_extract_converse_request_guardrail_and_prompt_variables( tracer_provider, ) -> None: handler = TelemetryHandler(tracer_provider=tracer_provider) - invocation = handler.inference(provider="aws.bedrock") - - extract_converse_request( - { - "guardrailConfig": { - "guardrailIdentifier": "sgi5gkybzqak", - "guardrailVersion": "1", - }, - "outputConfig": {"textFormat": "json"}, - "promptVariables": { - "user_name": {"text": "Alice"}, - "language": {"text": "French"}, + with handler.inference(provider="aws.bedrock") as invocation: + extract_converse_request( + { + "guardrailConfig": { + "guardrailIdentifier": "sgi5gkybzqak", + "guardrailVersion": "1", + }, + "outputConfig": {"textFormat": "json"}, + "promptVariables": { + "user_name": {"text": "Alice"}, + "language": {"text": "French"}, + }, }, - }, - invocation, - capture_content=True, - ) + invocation, + capture_content=True, + ) - assert ( - invocation.attributes.get(AwsAttributes.AWS_BEDROCK_GUARDRAIL_ID) - == "sgi5gkybzqak" - ) - assert invocation.output_type == "json" - assert ( - invocation.attributes.get("gen_ai.prompt.variable.user_name") - == "Alice" - ) - assert ( - invocation.attributes.get("gen_ai.prompt.variable.language") - == "French" - ) + assert ( + invocation.attributes.get(AwsAttributes.AWS_BEDROCK_GUARDRAIL_ID) + == "sgi5gkybzqak" + ) + assert invocation.output_type == "json" + assert ( + invocation.attributes.get("gen_ai.prompt.variable.user_name") + == "Alice" + ) + assert ( + invocation.attributes.get("gen_ai.prompt.variable.language") + == "French" + ) def test_extract_converse_request_prompt_variables_no_content( tracer_provider, ) -> None: handler = TelemetryHandler(tracer_provider=tracer_provider) - invocation = handler.inference(provider="aws.bedrock") - - extract_converse_request( - { - "guardrailConfig": { - "guardrailIdentifier": "sgi5gkybzqak", - }, - "promptVariables": { - "user_name": {"text": "Alice"}, + with handler.inference(provider="aws.bedrock") as invocation: + extract_converse_request( + { + "guardrailConfig": { + "guardrailIdentifier": "sgi5gkybzqak", + }, + "promptVariables": { + "user_name": {"text": "Alice"}, + }, }, - }, - invocation, - capture_content=False, - ) + invocation, + capture_content=False, + ) - assert ( - invocation.attributes.get(AwsAttributes.AWS_BEDROCK_GUARDRAIL_ID) - == "sgi5gkybzqak" - ) - assert "gen_ai.prompt.variable.user_name" not in invocation.attributes + assert ( + invocation.attributes.get(AwsAttributes.AWS_BEDROCK_GUARDRAIL_ID) + == "sgi5gkybzqak" + ) + assert "gen_ai.prompt.variable.user_name" not in invocation.attributes def test_extract_converse_request_system_instruction(tracer_provider) -> None: handler = TelemetryHandler(tracer_provider=tracer_provider) - invocation = handler.inference(provider="aws.bedrock") - - extract_converse_request( - { - "system": [ - {"text": "Be concise"}, - {"text": "Answer politely"}, - ], - }, - invocation, - ) + with handler.inference(provider="aws.bedrock") as invocation: + extract_converse_request( + { + "system": [ + {"text": "Be concise"}, + {"text": "Answer politely"}, + ], + }, + invocation, + ) - assert invocation.system_instruction == [ - TextPart(content="Be concise"), - TextPart(content="Answer politely"), - ] + assert invocation.system_instruction == [ + TextPart(content="Be concise"), + TextPart(content="Answer politely"), + ] def test_extract_converse_request_system_instruction_generic( tracer_provider, ) -> None: handler = TelemetryHandler(tracer_provider=tracer_provider) - invocation = handler.inference(provider="aws.bedrock") - - extract_converse_request( - { - "system": [ - {"text": "Be concise"}, - {"guardContent": {"guardrailIdentifier": "gr-123"}}, - {"cachePoint": {"type": "default"}}, - ], - }, - invocation, - ) + with handler.inference(provider="aws.bedrock") as invocation: + extract_converse_request( + { + "system": [ + {"text": "Be concise"}, + {"guardContent": {"guardrailIdentifier": "gr-123"}}, + {"cachePoint": {"type": "default"}}, + ], + }, + invocation, + ) - assert invocation.system_instruction == [ - TextPart(content="Be concise"), - GenericPart(type="guardContent"), - GenericPart(type="cachePoint"), - ] + assert invocation.system_instruction == [ + TextPart(content="Be concise"), + GenericPart(type="guardContent"), + GenericPart(type="cachePoint"), + ] diff --git a/instrumentation/opentelemetry-instrumentation-genai-bedrock/tests/test_invoke_model.py b/instrumentation/opentelemetry-instrumentation-genai-bedrock/tests/test_invoke_model.py index 401e39a5a..c784a6afa 100644 --- a/instrumentation/opentelemetry-instrumentation-genai-bedrock/tests/test_invoke_model.py +++ b/instrumentation/opentelemetry-instrumentation-genai-bedrock/tests/test_invoke_model.py @@ -407,45 +407,45 @@ def test_invoke_model_error( def test_extract_invoke_model_response_headers(tracer_provider) -> None: handler = TelemetryHandler(tracer_provider=tracer_provider) - invocation = handler.inference(provider="aws.bedrock") - extract_invoke_model_response( - { - "ResponseMetadata": { - "HTTPHeaders": { - "X-Amzn-Bedrock-Input-Token-Count": "15", - "X-Amzn-Bedrock-Output-Token-Count": "22", + with handler.inference(provider="aws.bedrock") as invocation: + extract_invoke_model_response( + { + "ResponseMetadata": { + "HTTPHeaders": { + "X-Amzn-Bedrock-Input-Token-Count": "15", + "X-Amzn-Bedrock-Output-Token-Count": "22", + } } - } - }, - b'{"completion": "hello"}', - invocation, - ) - assert invocation.input_tokens == 15 - assert invocation.output_tokens == 22 + }, + b'{"completion": "hello"}', + invocation, + ) + assert invocation.input_tokens == 15 + assert invocation.output_tokens == 22 def test_extract_invoke_model_request_zero_values(tracer_provider) -> None: handler = TelemetryHandler(tracer_provider=tracer_provider) - invocation = handler.inference(provider="aws.bedrock") - extract_invoke_model_request( - { - "body": json.dumps( - { - "temperature": 0.0, - "top_p": 0.0, - "top_k": 0, - "max_tokens": 0, - "seed": 0, - } - ) - }, - invocation, - ) - assert invocation.temperature == 0.0 - assert invocation.top_p == 0.0 - assert invocation.top_k == 0.0 - assert invocation.max_tokens == 0 - assert invocation.seed == 0 + with handler.inference(provider="aws.bedrock") as invocation: + extract_invoke_model_request( + { + "body": json.dumps( + { + "temperature": 0.0, + "top_p": 0.0, + "top_k": 0, + "max_tokens": 0, + "seed": 0, + } + ) + }, + invocation, + ) + assert invocation.temperature == 0.0 + assert invocation.top_p == 0.0 + assert invocation.top_k == 0.0 + assert invocation.max_tokens == 0 + assert invocation.seed == 0 def test_invoke_model_anthropic_tool_call_and_result( @@ -555,20 +555,19 @@ def test_invoke_model_anthropic_tool_call_and_result( def test_extract_invoke_model_request_guardrail(tracer_provider) -> None: handler = TelemetryHandler(tracer_provider=tracer_provider) - invocation = handler.inference(provider="aws.bedrock") - - extract_invoke_model_request( - { - "guardrailIdentifier": "sgi5gkybzqak", - "body": json.dumps({"prompt": "Hello"}), - }, - invocation, - ) + with handler.inference(provider="aws.bedrock") as invocation: + extract_invoke_model_request( + { + "guardrailIdentifier": "sgi5gkybzqak", + "body": json.dumps({"prompt": "Hello"}), + }, + invocation, + ) - assert ( - invocation.attributes.get(AwsAttributes.AWS_BEDROCK_GUARDRAIL_ID) - == "sgi5gkybzqak" - ) + assert ( + invocation.attributes.get(AwsAttributes.AWS_BEDROCK_GUARDRAIL_ID) + == "sgi5gkybzqak" + ) def test_invoke_model_with_guardrail_stubber( diff --git a/util/opentelemetry-test-util-genai/src/opentelemetry/test_util_genai/fixtures.py b/util/opentelemetry-test-util-genai/src/opentelemetry/test_util_genai/fixtures.py index 6f2ea44a8..90cb343d7 100644 --- a/util/opentelemetry-test-util-genai/src/opentelemetry/test_util_genai/fixtures.py +++ b/util/opentelemetry-test-util-genai/src/opentelemetry/test_util_genai/fixtures.py @@ -49,6 +49,7 @@ import pytest +from opentelemetry.context import attach, detach, get_current from opentelemetry.sdk._logs import LoggerProvider from opentelemetry.sdk._logs.export import ( InMemoryLogRecordExporter, @@ -74,6 +75,21 @@ ) from opentelemetry.util.genai.types import ContentCapturingMode + +@pytest.fixture(autouse=True) +def _isolate_context() -> Iterator[None]: + """Isolate the active OpenTelemetry context per test. + + Restores the previous context after each test, preventing context + attached during a test from leaking into subsequent tests. + """ + token = attach(get_current()) + try: + yield + finally: + detach(token) + + # ─── In-memory exporters and providers ────────────────────────────────────── diff --git a/util/opentelemetry-util-genai/.changelog/663.added b/util/opentelemetry-util-genai/.changelog/663.added new file mode 100644 index 000000000..a57485970 --- /dev/null +++ b/util/opentelemetry-util-genai/.changelog/663.added @@ -0,0 +1 @@ +Propagate inference attributes on the active context during inference invocations. diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_context.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_context.py new file mode 100644 index 000000000..31a493e8e --- /dev/null +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_context.py @@ -0,0 +1,55 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +"""Context helpers for GenAI inference attributes.""" + +from __future__ import annotations + +from typing import Final, cast + +from opentelemetry.context import Context, get_value, set_value +from opentelemetry.util.types import AttributeValue + +INFERENCE_ATTRIBUTES_KEY: Final[str] = ( + "opentelemetry.genai.inference_attributes" +) +_INFERENCE_ATTRIBUTES_KEY = INFERENCE_ATTRIBUTES_KEY + +__all__ = [ + "INFERENCE_ATTRIBUTES_KEY", + "get_inference_attributes", + "set_inference_attributes", +] + + +def set_inference_attributes( + attributes: dict[str, AttributeValue], + context: Context | None = None, +) -> Context: + """Return a Context with the given inference attributes dictionary attached. + + Args: + attributes: The mutable inference attributes dictionary. + context: The context to attach to. Defaults to the current context. + + Returns: + A new Context containing the inference attributes dictionary. + """ + return set_value(_INFERENCE_ATTRIBUTES_KEY, attributes, context=context) + + +def get_inference_attributes( + context: Context | None = None, +) -> dict[str, AttributeValue] | None: + """Return the active inference attributes dictionary from context, if any. + + Args: + context: The context to inspect. Defaults to the current context. + + Returns: + The active inference attributes dictionary, or None if not set. + """ + attrs = get_value(_INFERENCE_ATTRIBUTES_KEY, context=context) + if isinstance(attrs, dict): + return cast("dict[str, AttributeValue]", attrs) + return None diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_inference_invocation.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_inference_invocation.py index 613417862..7c324651e 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_inference_invocation.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_inference_invocation.py @@ -8,11 +8,19 @@ from typing import Final from opentelemetry._logs import Logger, LogRecord +from opentelemetry.context import Context from opentelemetry.semconv._incubating.attributes import ( gen_ai_attributes as GenAI, ) -from opentelemetry.semconv.attributes import server_attributes +from opentelemetry.semconv.attributes import ( + error_attributes, + server_attributes, +) from opentelemetry.trace import INVALID_SPAN, Span, SpanKind, Tracer +from opentelemetry.util.genai._context import ( + get_inference_attributes, + set_inference_attributes, +) from opentelemetry.util.genai._instruments import _Instruments from opentelemetry.util.genai._invocation import ( Error, @@ -34,6 +42,16 @@ ) from opentelemetry.util.types import AttributeValue +_METRIC_SEMCONV_KEYS = ( + GenAI.GEN_AI_OPERATION_NAME, + GenAI.GEN_AI_PROVIDER_NAME, + GenAI.GEN_AI_REQUEST_MODEL, + GenAI.GEN_AI_RESPONSE_MODEL, + server_attributes.SERVER_ADDRESS, + server_attributes.SERVER_PORT, + error_attributes.ERROR_TYPE, +) + _GEN_AI_USAGE_CACHE_WRITE_INPUT_TOKENS: Final = ( "gen_ai.usage.cache_write.input_tokens" ) @@ -59,6 +77,80 @@ _GEN_AI_CONVERSATION_COMPACTED: Final = "gen_ai.conversation.compacted" _GEN_AI_PROMPT_VERSION: Final = "gen_ai.prompt.version" +_FIELD_TO_SEMCONV: Final[dict[str, str]] = { + "conversation_id": GenAI.GEN_AI_CONVERSATION_ID, + "temperature": GenAI.GEN_AI_REQUEST_TEMPERATURE, + "top_p": GenAI.GEN_AI_REQUEST_TOP_P, + "top_k": GenAI.GEN_AI_REQUEST_TOP_K, + "frequency_penalty": GenAI.GEN_AI_REQUEST_FREQUENCY_PENALTY, + "presence_penalty": GenAI.GEN_AI_REQUEST_PRESENCE_PENALTY, + "max_tokens": GenAI.GEN_AI_REQUEST_MAX_TOKENS, + "stop_sequences": GenAI.GEN_AI_REQUEST_STOP_SEQUENCES, + "seed": GenAI.GEN_AI_REQUEST_SEED, + "finish_reasons": GenAI.GEN_AI_RESPONSE_FINISH_REASONS, + "_response_model_name": GenAI.GEN_AI_RESPONSE_MODEL, + "response_model_name": GenAI.GEN_AI_RESPONSE_MODEL, + "response_id": GenAI.GEN_AI_RESPONSE_ID, + "input_tokens": GenAI.GEN_AI_USAGE_INPUT_TOKENS, + "output_tokens": GenAI.GEN_AI_USAGE_OUTPUT_TOKENS, + "thinking_tokens": GenAI.GEN_AI_USAGE_REASONING_OUTPUT_TOKENS, + "cache_write_input_tokens": _GEN_AI_USAGE_CACHE_WRITE_INPUT_TOKENS, + "cache_creation_input_tokens": _GEN_AI_USAGE_CACHE_WRITE_INPUT_TOKENS, + "cache_read_input_tokens": GenAI.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, + "text_input_tokens": _GEN_AI_USAGE_TEXT_INPUT_TOKENS, + "image_input_tokens": _GEN_AI_USAGE_IMAGE_INPUT_TOKENS, + "audio_input_tokens": _GEN_AI_USAGE_AUDIO_INPUT_TOKENS, + "text_output_tokens": _GEN_AI_USAGE_TEXT_OUTPUT_TOKENS, + "image_output_tokens": _GEN_AI_USAGE_IMAGE_OUTPUT_TOKENS, + "audio_output_tokens": _GEN_AI_USAGE_AUDIO_OUTPUT_TOKENS, + "text_cache_read_input_tokens": _GEN_AI_USAGE_TEXT_CACHE_READ_INPUT_TOKENS, + "image_cache_read_input_tokens": _GEN_AI_USAGE_IMAGE_CACHE_READ_INPUT_TOKENS, + "audio_cache_read_input_tokens": _GEN_AI_USAGE_AUDIO_CACHE_READ_INPUT_TOKENS, + "reasoning_level": _GEN_AI_REQUEST_REASONING_LEVEL, + "previous_response_id": _GEN_AI_REQUEST_PREVIOUS_RESPONSE_ID, + "conversation_compacted": _GEN_AI_CONVERSATION_COMPACTED, + "prompt_name": GenAI.GEN_AI_PROMPT_NAME, + "prompt_version": _GEN_AI_PROMPT_VERSION, + "request_choice_count": GenAI.GEN_AI_REQUEST_CHOICE_COUNT, + "output_type": GenAI.GEN_AI_OUTPUT_TYPE, + "_request_stream": GenAI.GEN_AI_REQUEST_STREAM, + "_ttfc_seconds": GenAI.GEN_AI_RESPONSE_TIME_TO_FIRST_CHUNK, +} +# Content attributes are omitted from context because capture rules, representations, +# and presence in attributes may differ between spans and events. +_OPT_IN_MESSAGE_ATTRIBUTES: Final[frozenset[str]] = frozenset( + { + "input_messages", + "output_messages", + "system_instruction", + "tool_definitions", + "prompt_variables", + } +) + + +def _filter_context_attributes( + context_attributes: Mapping[str, AttributeValue], + *, + exclude_keys: set[str] | None = None, +) -> dict[str, AttributeValue]: + filtered: dict[str, AttributeValue] = {} + for key, value in context_attributes.items(): + if exclude_keys is not None and key in exclude_keys: + continue + if ( + value == 0 + and key.startswith("gen_ai.usage.") + and key + not in ( + GenAI.GEN_AI_USAGE_INPUT_TOKENS, + GenAI.GEN_AI_USAGE_OUTPUT_TOKENS, + ) + ): + continue + filtered[key] = value + return filtered + class InferenceInvocation(GenAIInvocation): """Represents a single LLM chat/completion call. @@ -147,8 +239,42 @@ def __init__( # Rebuilt once per streaming chunk, so cache it and invalidate via # _invalidate_metric_attributes whenever an input changes. self._cached_metric_attributes: dict[str, AttributeValue] | None = None + + existing_attrs = get_inference_attributes() + if existing_attrs is not None: + self.already_started = True + self._context_attributes: dict[str, AttributeValue] = ( + existing_attrs + ) + self._context_attributes.update(self._get_start_attributes()) + else: + self.already_started = False + self._context_attributes = dict(self._get_start_attributes()) + self._start(self._get_start_attributes()) + def __setattr__(self, name: str, value: object) -> None: + super().__setattr__(name, value) + if name in _OPT_IN_MESSAGE_ATTRIBUTES: + return + # Custom attributes set on self.attributes are currently missing from + # context, which can be addressed by adding an explicit set_attribute method. + if ( + hasattr(self, "_context_attributes") + and name in _FIELD_TO_SEMCONV + and value is not None + ): + if isinstance(value, (str, bool, int, float)): + self._context_attributes[_FIELD_TO_SEMCONV[name]] = value + elif name == "stop_sequences" and self.stop_sequences is not None: + self._context_attributes[_FIELD_TO_SEMCONV[name]] = ( + self.stop_sequences + ) + elif name == "finish_reasons" and self.finish_reasons is not None: + self._context_attributes[_FIELD_TO_SEMCONV[name]] = ( + self.finish_reasons + ) + @property def cache_creation_input_tokens(self) -> int | None: """ @@ -304,6 +430,23 @@ def _get_attributes(self) -> dict[str, AttributeValue]: attrs.update({k: v for k, v in optional_attrs if v is not None}) return attrs + def _create_span_context(self) -> Context: + ctx = super()._create_span_context() + return set_inference_attributes(self._context_attributes, context=ctx) + + def _get_context_attributes(self) -> dict[str, AttributeValue]: + attrs = self._get_start_attributes() + attrs.update(self._get_attributes()) + attrs.update(self.attributes) + # Message attributes are excluded because spans and events format + # content differently and evaluate capture rules independently. + return attrs + + def _finish_already_started(self, error: Error | None = None) -> None: + # Error attributes are not recorded on inner finish to isolate errors; + # the outer invocation records them only if the error escapes unhandled. + self._context_attributes.update(self._get_context_attributes()) + def _invalidate_metric_attributes(self) -> None: """Drop the cached metric attributes so the next read rebuilds them. @@ -320,6 +463,9 @@ def _get_metric_attributes(self) -> dict[str, AttributeValue]: if self._response_model_name is not None: attrs[GenAI.GEN_AI_RESPONSE_MODEL] = self._response_model_name attrs.update(self.metric_attributes) + for key in _METRIC_SEMCONV_KEYS: + if key not in attrs and key in self._context_attributes: + attrs[key] = self._context_attributes[key] self._cached_metric_attributes = attrs return self._cached_metric_attributes @@ -330,21 +476,40 @@ def _apply_error_attributes(self, error: Error) -> None: def _get_metric_token_counts(self) -> dict[str, int]: counts: dict[str, int] = {} - if self.input_tokens is not None: - counts[GenAI.GenAiTokenTypeValues.INPUT.value] = self.input_tokens - if self.output_tokens is not None: - counts[GenAI.GenAiTokenTypeValues.OUTPUT.value] = ( - self.output_tokens + input_tokens = self.input_tokens + if input_tokens is None: + ctx_input = self._context_attributes.get( + GenAI.GEN_AI_USAGE_INPUT_TOKENS ) + if isinstance(ctx_input, int): + input_tokens = ctx_input + if input_tokens is not None: + counts[GenAI.GenAiTokenTypeValues.INPUT.value] = input_tokens + + output_tokens = self.output_tokens + if output_tokens is None: + ctx_output = self._context_attributes.get( + GenAI.GEN_AI_USAGE_OUTPUT_TOKENS + ) + if isinstance(ctx_output, int): + output_tokens = ctx_output + if output_tokens is not None: + counts[GenAI.GenAiTokenTypeValues.OUTPUT.value] = output_tokens return counts def _apply_finish(self, error: Error | None = None) -> None: if error is not None: self._apply_error_attributes(error) - attributes = self._get_attributes() + attributes = _filter_context_attributes( + self._context_attributes, + exclude_keys=set(self._get_start_attributes()), + ) + attributes.update(self._get_attributes()) attributes.update(self._get_message_attributes(for_span=True)) attributes.update(self.attributes) self.span.set_attributes(attributes) + self._context_attributes.update(self._get_context_attributes()) + self._invalidate_metric_attributes() self._record_client_metrics() log_record = self._maybe_create_event() self._call_completion_hook( @@ -366,7 +531,8 @@ def _maybe_create_event(self) -> LogRecord | None: if not should_emit_event(): return None - attributes = self._get_start_attributes() + attributes = _filter_context_attributes(self._context_attributes) + attributes.update(self._get_start_attributes()) attributes.update(self._get_attributes()) attributes.update(self._get_message_attributes(for_span=False)) attributes.update(self.attributes) @@ -490,3 +656,11 @@ def span(self) -> Span: if self._inference_invocation is not None else INVALID_SPAN ) + + @property + def already_started(self) -> bool: + return ( + self._inference_invocation.already_started + if self._inference_invocation is not None + else False + ) diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_invocation.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_invocation.py index 9829ee580..dc19b65b5 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_invocation.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/_invocation.py @@ -100,7 +100,8 @@ def __init__( self._span_name: str = span_name self._span_kind: SpanKind = span_kind self._context_token: ContextToken | None = None - self._monotonic_start_s: float + self._monotonic_start_s: float = timeit.default_timer() + self.already_started: bool = False # Streaming state, set when the invocation is handed to a stream # wrapper. ``_request_stream`` marks the request as streamed # (gen_ai.request.stream); the timing fields are populated by @@ -133,15 +134,22 @@ def _start( Args: attributes: Initial span attributes available for sampling decisions. """ + if self.already_started: + return + 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._span_context = self._create_span_context() self._monotonic_start_s = timeit.default_timer() self._context_token = attach(self._span_context) + def _create_span_context(self) -> Context: + """Create the context to attach for this invocation's span.""" + return set_span_in_context(self.span) + def _get_metric_attributes(self) -> dict[str, AttributeValue]: """Return low-cardinality attributes for metric recording.""" return self.metric_attributes @@ -152,7 +160,7 @@ def _get_metric_token_counts(self) -> dict[str, int]: # pylint: disable=no-self def record_stream_chunk(self) -> None: """Mark the request as streamed and record one output chunk arriving.""" - if self._context_token is None: + if self._context_token is None and not self.already_started: return self._request_stream = True self._on_stream_chunk(timeit.default_timer()) @@ -173,9 +181,14 @@ def _on_stream_chunk(self, chunk_at: float) -> None: self._stream_last_chunk_at = chunk_at delta = max(chunk_at - last_chunk_at, 0.0) - attributes = self._get_metric_attributes() - if self._ttfc_seconds is None: + is_first_chunk = self._ttfc_seconds is None + if is_first_chunk: self._ttfc_seconds = delta + if self.already_started: + return + + attributes = self._get_metric_attributes() + if is_first_chunk: self._instruments.time_to_first_chunk.record( delta, attributes=attributes, @@ -244,12 +257,19 @@ def _call_completion_hook( log_record=log_record, ) + def _finish_already_started(self, error: Error | None = None) -> None: + """Handle finish when the invocation was already started upstream.""" + @abstractmethod def _apply_finish(self, error: Error | None = None) -> None: """Apply finish telemetry (attributes, metrics, events).""" def _finish(self, error: Error | None = None) -> None: """Apply finish telemetry and end the span. Finishes at most once.""" + if self.already_started: + self._finish_already_started(error) + return + if self._context_token is None: return # Clear up front so a nested or repeated finish is a no-op even if diff --git a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/stream.py b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/stream.py index fc2b5a54a..7b5793425 100644 --- a/util/opentelemetry-util-genai/src/opentelemetry/util/genai/stream.py +++ b/util/opentelemetry-util-genai/src/opentelemetry/util/genai/stream.py @@ -192,7 +192,12 @@ def __next__(self) -> ChunkT: raise invocation = self._self_invocation chunk_at = timeit.default_timer() if invocation is not None else None - self._process_chunk(chunk) + # Inner stream wrappers skip chunk accumulation to avoid duplicate + # message buffering in memory. + if invocation is None or not getattr( + invocation, "already_started", False + ): + self._process_chunk(chunk) # Record after _process_chunk so response.model is on the metrics. if invocation is not None and chunk_at is not None: invocation._on_stream_chunk(chunk_at) @@ -352,7 +357,12 @@ async def __anext__(self) -> ChunkT: invocation = self._self_invocation chunk_at = timeit.default_timer() if invocation is not None else None - self._process_chunk(chunk) + # Inner stream wrappers skip chunk accumulation to avoid duplicate + # message buffering in memory. + if invocation is None or not getattr( + invocation, "already_started", False + ): + self._process_chunk(chunk) # Record after _process_chunk so response.model is on the metrics. if invocation is not None and chunk_at is not None: invocation._on_stream_chunk(chunk_at) diff --git a/util/opentelemetry-util-genai/tests/test_context.py b/util/opentelemetry-util-genai/tests/test_context.py new file mode 100644 index 000000000..539dc8fe3 --- /dev/null +++ b/util/opentelemetry-util-genai/tests/test_context.py @@ -0,0 +1,592 @@ +# Copyright The OpenTelemetry Authors +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import os +from unittest.mock import patch + +from opentelemetry.context import attach, detach +from opentelemetry.sdk._logs import LoggerProvider +from opentelemetry.sdk._logs.export import ( + InMemoryLogRecordExporter, + SimpleLogRecordProcessor, +) +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import ( + InMemorySpanExporter, +) +from opentelemetry.semconv._incubating.attributes import ( + gen_ai_attributes as GenAI, +) +from opentelemetry.semconv.attributes import ( + error_attributes, + server_attributes, +) +from opentelemetry.test.test_base import TestBase +from opentelemetry.trace.status import StatusCode +from opentelemetry.util.genai._context import ( + INFERENCE_ATTRIBUTES_KEY, + get_inference_attributes, + set_inference_attributes, +) +from opentelemetry.util.genai.handler import TelemetryHandler +from opentelemetry.util.types import AttributeValue + + +class TestInferenceContext(TestBase): + def setUp(self) -> None: + super().setUp() + self.span_exporter = InMemorySpanExporter() + self.tracer_provider.add_span_processor( + SimpleSpanProcessor(self.span_exporter) + ) + self.handler = TelemetryHandler( + tracer_provider=self.tracer_provider, + meter_provider=self.meter_provider, + ) + + def _harvest_metrics(self) -> dict[str, list[object]]: + metrics = self.get_sorted_metrics() + metrics_by_name: dict[str, list[object]] = {} + for metric in metrics or []: + points = getattr(metric.data, "data_points", None) or [] + metrics_by_name.setdefault(metric.name, []).extend(points) + return metrics_by_name + + def test_context_key_constant_value(self) -> None: + self.assertEqual( + INFERENCE_ATTRIBUTES_KEY, + "opentelemetry.genai.inference_attributes", + ) + + def test_get_inference_attributes_none_by_default(self) -> None: + self.assertIsNone(get_inference_attributes()) + + def test_set_and_get_inference_attributes(self) -> None: + attrs = {"test.attr": "value"} + ctx = set_inference_attributes(attrs) + self.assertIs(get_inference_attributes(ctx), attrs) + self.assertIsNone(get_inference_attributes()) + + token = attach(ctx) + try: + self.assertIs(get_inference_attributes(), attrs) + finally: + detach(token) + self.assertIsNone(get_inference_attributes()) + + def test_in_place_mutation_of_inference_attributes(self) -> None: + attrs: dict[str, AttributeValue] = {"initial": 1} + ctx = set_inference_attributes(attrs) + token = attach(ctx) + try: + current = get_inference_attributes() + self.assertIsNotNone(current) + assert current is not None + current.update({"updated": 2, "initial": 10}) + + after = get_inference_attributes() + assert after is not None + self.assertEqual(after["initial"], 10) + self.assertEqual(after["updated"], 2) + self.assertIs(after, attrs) + finally: + detach(token) + + def test_inference_invocation_sets_attributes_on_context(self) -> None: + self.assertIsNone(get_inference_attributes()) + + with self.handler.inference( + "openai", request_model="gpt-4o-mini" + ) as invocation: + self.assertFalse(invocation.already_started) + attrs = get_inference_attributes() + self.assertIsNotNone(attrs) + assert attrs is not None + # Start attributes are placed in context upon initialization + self.assertEqual(attrs.get(GenAI.GEN_AI_PROVIDER_NAME), "openai") + self.assertEqual( + attrs.get(GenAI.GEN_AI_REQUEST_MODEL), "gpt-4o-mini" + ) + self.assertEqual(attrs.get(GenAI.GEN_AI_OPERATION_NAME), "chat") + + # Live updates when setting typed fields + invocation.input_tokens = 42 + invocation.output_tokens = 84 + invocation.temperature = 0.7 + invocation.response_model_name = "gpt-4o-mini-2024-07-18" + self.assertEqual(attrs.get(GenAI.GEN_AI_USAGE_INPUT_TOKENS), 42) + self.assertEqual(attrs.get(GenAI.GEN_AI_USAGE_OUTPUT_TOKENS), 84) + self.assertEqual(attrs.get(GenAI.GEN_AI_REQUEST_TEMPERATURE), 0.7) + self.assertEqual( + attrs.get(GenAI.GEN_AI_RESPONSE_MODEL), + "gpt-4o-mini-2024-07-18", + ) + + self.assertIsNone(get_inference_attributes()) + + def test_inference_invocation_automatic_publish_on_finish(self) -> None: + with self.handler.inference( + "openai", request_model="gpt-4o-mini" + ) as invocation: + self.assertFalse(invocation.already_started) + invocation.input_tokens = 10 + # did not call publish_to_context() + + # Span attributes were populated automatically upon finish + spans = self.span_exporter.get_finished_spans() + self.assertEqual(len(spans), 1) + self.assertEqual( + spans[0].attributes.get(GenAI.GEN_AI_PROVIDER_NAME), "openai" + ) + self.assertEqual( + spans[0].attributes.get(GenAI.GEN_AI_USAGE_INPUT_TOKENS), 10 + ) + + def test_non_inference_invocations_do_not_set_inference_attributes( + self, + ) -> None: + self.assertIsNone(get_inference_attributes()) + + with self.handler.invoke_local_agent(agent_name="MathTutor"): + self.assertIsNone(get_inference_attributes()) + + # Nested inference invocation properly sets the inference attributes + with self.handler.inference( + "openai", request_model="gpt-4o-mini" + ) as inf_inv: + self.assertFalse(inf_inv.already_started) + self.assertIsNotNone(get_inference_attributes()) + + self.assertIsNone(get_inference_attributes()) + + def test_nested_inference_deduplication_and_enrichment(self) -> None: + log_exporter = InMemoryLogRecordExporter() + logger_provider = LoggerProvider() + logger_provider.add_log_record_processor( + SimpleLogRecordProcessor(log_exporter) + ) + handler = TelemetryHandler( + tracer_provider=self.tracer_provider, + meter_provider=self.meter_provider, + logger_provider=logger_provider, + ) + + with patch.dict( + os.environ, {"OTEL_INSTRUMENTATION_GENAI_EMIT_EVENT": "true"} + ): + with handler.inference( + "openai", request_model="gpt-4o" + ) as root_inv: + self.assertFalse(root_inv.already_started) + + with handler.inference( + "openai", + request_model="gpt-4o", + server_address="api.openai.com", + server_port=443, + ) as nested_inv: + self.assertTrue(nested_inv.already_started) + nested_inv.input_tokens = 15 + nested_inv.output_tokens = 25 + nested_inv.response_model_name = "gpt-4o-2024-08-06" + nested_inv.attributes["custom.downstream"] = "enriched" + + # While still in root context, context attributes contain downstream enrichment + attrs = get_inference_attributes() + self.assertIsNotNone(attrs) + assert attrs is not None + self.assertEqual( + attrs.get(server_attributes.SERVER_ADDRESS), + "api.openai.com", + ) + self.assertEqual(attrs.get(server_attributes.SERVER_PORT), 443) + self.assertEqual( + attrs.get("gen_ai.response.model"), + "gpt-4o-2024-08-06", + ) + self.assertEqual(attrs.get("gen_ai.usage.input_tokens"), 15) + self.assertEqual(attrs.get("gen_ai.usage.output_tokens"), 25) + self.assertEqual(attrs.get("custom.downstream"), "enriched") + + # After nested exit, span is NOT ended yet (still recording) + self.assertTrue(root_inv.span.is_recording()) + spans = self.span_exporter.get_finished_spans() + self.assertEqual(len(spans), 0) + + # After root exit, root span is ended and event is emitted + spans = self.span_exporter.get_finished_spans() + self.assertEqual(len(spans), 1) + span = spans[0] + self.assertEqual( + span.attributes.get("gen_ai.provider.name"), "openai" + ) + self.assertEqual( + span.attributes.get(server_attributes.SERVER_ADDRESS), + "api.openai.com", + ) + self.assertEqual( + span.attributes.get(server_attributes.SERVER_PORT), 443 + ) + self.assertEqual( + span.attributes.get("gen_ai.response.model"), + "gpt-4o-2024-08-06", + ) + self.assertEqual( + span.attributes.get("gen_ai.usage.input_tokens"), 15 + ) + self.assertEqual( + span.attributes.get("gen_ai.usage.output_tokens"), 25 + ) + self.assertEqual( + span.attributes.get("custom.downstream"), "enriched" + ) + + logs = log_exporter.get_finished_logs() + self.assertEqual(len(logs), 1) + event = logs[0].log_record + self.assertEqual( + event.event_name, "gen_ai.client.inference.operation.details" + ) + self.assertIsNotNone(event.attributes) + assert event.attributes is not None + self.assertEqual( + event.attributes.get(server_attributes.SERVER_ADDRESS), + "api.openai.com", + ) + self.assertEqual( + event.attributes.get(server_attributes.SERVER_PORT), 443 + ) + self.assertEqual( + event.attributes.get("gen_ai.response.model"), + "gpt-4o-2024-08-06", + ) + self.assertEqual( + event.attributes.get("gen_ai.usage.input_tokens"), 15 + ) + self.assertEqual( + event.attributes.get("gen_ai.usage.output_tokens"), 25 + ) + self.assertEqual( + event.attributes.get("custom.downstream"), "enriched" + ) + + # Check metrics - exactly 1 operation duration point for the deduplicated invocation + metrics = self._harvest_metrics() + self.assertIn("gen_ai.client.operation.duration", metrics) + duration_points = metrics["gen_ai.client.operation.duration"] + self.assertEqual(len(duration_points), 1) + duration_point = duration_points[0] + self.assertEqual( + duration_point.attributes.get( + server_attributes.SERVER_ADDRESS + ), + "api.openai.com", + ) + self.assertEqual( + duration_point.attributes.get(server_attributes.SERVER_PORT), + 443, + ) + self.assertEqual( + duration_point.attributes.get(GenAI.GEN_AI_RESPONSE_MODEL), + "gpt-4o-2024-08-06", + ) + # High-cardinality attributes must NOT leak onto metrics + self.assertNotIn("custom.downstream", duration_point.attributes) + self.assertNotIn( + GenAI.GEN_AI_CONVERSATION_ID, duration_point.attributes + ) + + # Token usage metric should be enriched from context token counts + self.assertIn("gen_ai.client.token.usage", metrics) + token_points = metrics["gen_ai.client.token.usage"] + token_by_type = { + point.attributes[GenAI.GEN_AI_TOKEN_TYPE]: point + for point in token_points + } + self.assertEqual(len(token_by_type), 2) + self.assertAlmostEqual( + token_by_type[GenAI.GenAiTokenTypeValues.INPUT.value].sum, + 15.0, + places=3, + ) + self.assertAlmostEqual( + token_by_type[GenAI.GenAiTokenTypeValues.OUTPUT.value].sum, + 25.0, + places=3, + ) + input_token_point = token_by_type[ + GenAI.GenAiTokenTypeValues.INPUT.value + ] + self.assertEqual( + input_token_point.attributes.get( + server_attributes.SERVER_ADDRESS + ), + "api.openai.com", + ) + self.assertEqual( + input_token_point.attributes.get( + server_attributes.SERVER_PORT + ), + 443, + ) + self.assertEqual( + input_token_point.attributes.get(GenAI.GEN_AI_RESPONSE_MODEL), + "gpt-4o-2024-08-06", + ) + self.assertNotIn("custom.downstream", input_token_point.attributes) + + def test_nested_inference_invocation_does_not_end_span_on_fail( + self, + ) -> None: + with self.handler.inference( + "upstream", request_model="gpt-4o" + ) as root_inv: + with self.assertRaises(ValueError) as handler_error: + with self.handler.inference( + "downstream", request_model="gpt-4o" + ) as nested_inv: + self.assertTrue(nested_inv.already_started) + raise ValueError("downstream network failure") + + self.assertEqual( + str(handler_error.exception), "downstream network failure" + ) + # Root span must NOT be ended yet + self.assertTrue(root_inv.span.is_recording()) + self.assertEqual(len(self.span_exporter.get_finished_spans()), 0) + + # Downstream error was caught by caller, so it is NOT recorded on context + attrs = get_inference_attributes() + self.assertIsNotNone(attrs) + assert attrs is not None + self.assertNotIn(error_attributes.ERROR_TYPE, attrs) + + # After root finishes normally, 1 span is ended without error attributes + spans = self.span_exporter.get_finished_spans() + self.assertEqual(len(spans), 1) + self.assertNotIn(error_attributes.ERROR_TYPE, spans[0].attributes) + + def test_inner_error_caught_by_outer_does_not_fail_root(self) -> None: + with self.handler.inference("proxy", request_model="primary-model"): + try: + with self.handler.inference( + "provider", request_model="primary-model" + ): + raise RuntimeError("primary model failed") + except RuntimeError: + pass # Model fallback / retry + + with self.handler.inference( + "provider", request_model="fallback-model" + ) as fallback_inv: + fallback_inv.output_tokens = 20 + + spans = self.span_exporter.get_finished_spans() + self.assertEqual(len(spans), 1) + root_span = spans[0] + self.assertNotIn(error_attributes.ERROR_TYPE, root_span.attributes) + self.assertNotEqual(root_span.status.status_code, StatusCode.ERROR) + + metrics = self._harvest_metrics() + duration_points = metrics.get("gen_ai.client.operation.duration", []) + self.assertEqual(len(duration_points), 1) + self.assertNotIn( + error_attributes.ERROR_TYPE, duration_points[0].attributes + ) + + def test_downstream_streaming_record_stream_chunk(self) -> None: + with self.handler.inference( + "upstream", request_model="gpt-4o" + ) as root_inv: + self.assertFalse(root_inv.already_started) + with self.handler.inference("downstream") as nested_inv: + self.assertTrue(nested_inv.already_started) + nested_inv.record_stream_chunk() + nested_inv.record_stream_chunk() + self.assertIsNotNone(nested_inv._ttfc_seconds) + + spans = self.span_exporter.get_finished_spans() + self.assertEqual(len(spans), 1) + + def test_llm_invocation_already_started(self) -> None: + from opentelemetry.util.genai._inference_invocation import ( + LLMInvocation, + ) + + inv = LLMInvocation(request_model="test") + self.assertFalse(inv.already_started) + + with self.handler.inference("upstream"): + nested_inv = LLMInvocation(request_model="nested") + self.handler.start_llm(nested_inv) + self.assertTrue(nested_inv.already_started) + attrs = get_inference_attributes() + self.assertIsNotNone(attrs) + assert attrs is not None + self.assertEqual(attrs.get(GenAI.GEN_AI_REQUEST_MODEL), "nested") + nested_inv.attributes["custom.llm"] = "val" + self.handler.stop_llm(nested_inv) + self.assertEqual(attrs.get("custom.llm"), "val") + + def test_metric_enrichment_precedence_and_error(self) -> None: + with self.assertRaises(ValueError): + with self.handler.inference( + "proxy", + request_model="gpt-4o", + server_address="proxy.internal", + server_port=8080, + ) as root_inv: + root_inv.input_tokens = 10 + with self.handler.inference( + "openai", + request_model="gpt-4o", + server_address="api.openai.com", + server_port=443, + ) as inner_inv: + inner_inv.input_tokens = 99 + inner_inv.output_tokens = 50 + inner_inv.response_model_name = "gpt-4o-2024-08-06" + raise ValueError("network reset") + + metrics = self._harvest_metrics() + duration_points = metrics["gen_ai.client.operation.duration"] + self.assertEqual(len(duration_points), 1) + point = duration_points[0] + + # Root values take precedence over downstream context + self.assertEqual( + point.attributes.get(server_attributes.SERVER_ADDRESS), + "proxy.internal", + ) + self.assertEqual( + point.attributes.get(server_attributes.SERVER_PORT), 8080 + ) + # Downstream fields not set on root are enriched from context + self.assertEqual( + point.attributes.get(GenAI.GEN_AI_RESPONSE_MODEL), + "gpt-4o-2024-08-06", + ) + self.assertEqual( + point.attributes.get(error_attributes.ERROR_TYPE), + "ValueError", + ) + + # Tokens: root input_tokens (10) takes precedence, output_tokens (50) enriched from context + token_points = metrics["gen_ai.client.token.usage"] + token_by_type = { + p.attributes[GenAI.GEN_AI_TOKEN_TYPE]: p for p in token_points + } + self.assertAlmostEqual( + token_by_type[GenAI.GenAiTokenTypeValues.INPUT.value].sum, + 10.0, + places=3, + ) + self.assertAlmostEqual( + token_by_type[GenAI.GenAiTokenTypeValues.OUTPUT.value].sum, + 50.0, + places=3, + ) + + # Root values take precedence over downstream context on span as well + spans = self.span_exporter.get_finished_spans() + self.assertEqual(len(spans), 1) + span_attrs = spans[0].attributes + self.assertEqual( + span_attrs.get(server_attributes.SERVER_ADDRESS), "proxy.internal" + ) + self.assertEqual(span_attrs.get(server_attributes.SERVER_PORT), 8080) + self.assertEqual(span_attrs.get(GenAI.GEN_AI_USAGE_INPUT_TOKENS), 10) + self.assertEqual(span_attrs.get(GenAI.GEN_AI_USAGE_OUTPUT_TOKENS), 50) + self.assertEqual( + span_attrs.get(GenAI.GEN_AI_RESPONSE_MODEL), "gpt-4o-2024-08-06" + ) + self.assertEqual( + span_attrs.get(error_attributes.ERROR_TYPE), "ValueError" + ) + + def test_root_precedence_over_nested_inference(self) -> None: + with self.handler.inference( + "proxy-provider", + request_model="proxy-model", + server_address="proxy.example.com", + server_port=8080, + ) as root: + root.temperature = 0.2 + root.input_tokens = 10 + root.attributes["custom.shared"] = "root-value" + + with self.handler.inference( + "downstream-provider", + request_model="downstream-model", + server_address="api.example.com", + server_port=443, + ) as inner: + self.assertTrue(inner.already_started) + # Overwrite shared fields downstream + inner.temperature = 0.9 + inner.input_tokens = 100 + inner.output_tokens = 50 + inner.response_id = "resp-123" + inner.attributes["custom.shared"] = "downstream-value" + inner.attributes["custom.downstream_only"] = "downstream-only" + + # While still in root context, context reflects downstream writes + attrs = get_inference_attributes() + assert attrs is not None + self.assertEqual(attrs.get(GenAI.GEN_AI_REQUEST_TEMPERATURE), 0.9) + self.assertEqual(attrs.get(GenAI.GEN_AI_USAGE_INPUT_TOKENS), 100) + + # After root finish: Option 1 root precedence applies + spans = self.span_exporter.get_finished_spans() + self.assertEqual(len(spans), 1) + root_span = spans[0] + + # Root start attributes take precedence + self.assertEqual( + root_span.attributes.get(GenAI.GEN_AI_PROVIDER_NAME), + "proxy-provider", + ) + self.assertEqual( + root_span.attributes.get(GenAI.GEN_AI_REQUEST_MODEL), + "proxy-model", + ) + self.assertEqual( + root_span.attributes.get(server_attributes.SERVER_ADDRESS), + "proxy.example.com", + ) + self.assertEqual( + root_span.attributes.get(server_attributes.SERVER_PORT), + 8080, + ) + + # Root typed attributes take precedence + self.assertEqual( + root_span.attributes.get(GenAI.GEN_AI_REQUEST_TEMPERATURE), + 0.2, + ) + self.assertEqual( + root_span.attributes.get(GenAI.GEN_AI_USAGE_INPUT_TOKENS), + 10, + ) + + # Root custom attributes take precedence + self.assertEqual( + root_span.attributes.get("custom.shared"), + "root-value", + ) + + # Downstream fields not set on root are enriched onto root span + self.assertEqual( + root_span.attributes.get(GenAI.GEN_AI_USAGE_OUTPUT_TOKENS), + 50, + ) + self.assertEqual( + root_span.attributes.get(GenAI.GEN_AI_RESPONSE_ID), + "resp-123", + ) + self.assertEqual( + root_span.attributes.get("custom.downstream_only"), + "downstream-only", + ) diff --git a/util/opentelemetry-util-genai/tests/test_stream.py b/util/opentelemetry-util-genai/tests/test_stream.py index 262db1e94..492f70a7e 100644 --- a/util/opentelemetry-util-genai/tests/test_stream.py +++ b/util/opentelemetry-util-genai/tests/test_stream.py @@ -6,7 +6,7 @@ import asyncio import inspect import timeit -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -108,6 +108,18 @@ def test_sync_stream_wrapper_processes_chunks_and_stops(): assert wrapper._self_stop_count == 1 +def test_sync_stream_wrapper_skips_chunk_accumulation_when_already_started(): + invocation = MagicMock() + invocation.already_started = True + stream = _FakeSyncStream(chunks=["chunk1", "chunk2"]) + wrapper = _TestSyncStreamWrapper(stream, invocation=invocation) + + assert next(wrapper) == "chunk1" + assert next(wrapper) == "chunk2" + assert wrapper._self_processed == [] + assert invocation._on_stream_chunk.call_count == 2 + + def test_sync_stream_wrapper_processes_iterables(): stream = _FakeSyncIterable(chunks=["chunk"]) wrapper = _TestSyncStreamWrapper(stream) @@ -300,6 +312,21 @@ async def exercise(): asyncio.run(exercise()) +def test_async_stream_wrapper_skips_chunk_accumulation_when_already_started(): + async def exercise(): + invocation = MagicMock() + invocation.already_started = True + stream = _FakeAsyncStream(chunks=["chunk1", "chunk2"]) + wrapper = _TestAsyncStreamWrapper(stream, invocation=invocation) + + assert await anext(wrapper) == "chunk1" + assert await anext(wrapper) == "chunk2" + assert wrapper._self_processed == [] + assert invocation._on_stream_chunk.call_count == 2 + + asyncio.run(exercise()) + + def test_async_stream_wrapper_processes_async_iterables(): async def exercise(): stream = _FakeAsyncIterable(chunks=["chunk"]) diff --git a/util/opentelemetry-util-genai/tests/test_utils.py b/util/opentelemetry-util-genai/tests/test_utils.py index 130265f41..d0937ac7e 100644 --- a/util/opentelemetry-util-genai/tests/test_utils.py +++ b/util/opentelemetry-util-genai/tests/test_utils.py @@ -802,25 +802,22 @@ def test_parent_child_span_relationship(self): message = _create_input_message("hi") chat_generation = _create_output_message("ok") - with self.telemetry_handler.inference( - "test-provider", request_model="parent-model" - ) as parent_invocation: - parent_invocation.input_messages = [message] + with self.telemetry_handler.workflow(name="parent-workflow"): with self.telemetry_handler.inference( "test-provider", request_model="child-model" ) as child_invocation: child_invocation.input_messages = [message] # Stop child first by exiting inner context child_invocation.output_messages = [chat_generation] - # Then stop parent by exiting outer context - parent_invocation.output_messages = [chat_generation] spans = self.span_exporter.get_finished_spans() assert len(spans) == 2 # Identify spans irrespective of export order child_span = next(s for s in spans if s.name == "chat child-model") - parent_span = next(s for s in spans if s.name == "chat parent-model") + parent_span = next( + s for s in spans if s.name == "invoke_workflow parent-workflow" + ) # Same trace assert child_span.context.trace_id == parent_span.context.trace_id