From c0dccb33a118438cc2c9a77bf134c430d943d44f Mon Sep 17 00:00:00 2001 From: Manasjyoti Sharma Date: Mon, 20 Apr 2026 20:30:03 +0530 Subject: [PATCH 1/2] fix: harden ST-6 framework instrumentation lifecycle and safety flows --- .../instrumentation/langchain/__init__.py | 66 ++++---- .../instrumentation/litellm/__init__.py | 150 ++++++++++++++++-- .../litellm/streaming_safety.py | 11 +- .../tests/test_logger_integration.py | 123 ++++++++++++++ .../tests/test_safety_hooks.py | 71 +++++++-- .../llamaindex/dispatcher_wrapper.py | 106 +++++++++++++ .../instrumentation/llamaindex/safety.py | 28 +++- .../instrumentation/llamaindex/span_utils.py | 32 +++- .../tests/test_safety_hooks.py | 21 +++ .../openai/shared/chat_wrappers.py | 84 +++++----- .../instrumentation/openai/utils.py | 14 ++ .../instrumentation/openai/v0/__init__.py | 18 ++- .../instrumentation/openai/v1/__init__.py | 50 +++--- .../openai/v1/instrumentation_runtime.py | 16 +- .../openai/v1/safety_registration.py | 46 +++--- .../tests/traces/test_chat.py | 57 +++++++ 16 files changed, 723 insertions(+), 170 deletions(-) diff --git a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py index 1a2defc272..423935c015 100644 --- a/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py +++ b/packages/opentelemetry-instrumentation-langchain/opentelemetry/instrumentation/langchain/__init__.py @@ -95,9 +95,9 @@ def _instrument(self, **kwargs): tracer, duration_histogram, token_histogram ) wrap_function_wrapper( - module="langchain_core.callbacks", - name="BaseCallbackManager.__init__", - wrapper=_BaseCallbackManagerInitWrapper(traceloopCallbackHandler), + "langchain_core.callbacks", + "BaseCallbackManager.__init__", + _BaseCallbackManagerInitWrapper(traceloopCallbackHandler), ) instrument_safety_wrappers() @@ -110,66 +110,66 @@ def _wrap_openai_functions_for_tracing(self, traceloopCallbackHandler): if is_package_available("langchain_community"): # Wrap langchain_community.llms.openai.BaseOpenAI wrap_function_wrapper( - module="langchain_community.llms.openai", - name="BaseOpenAI._generate", - wrapper=openai_tracing_wrapper, + "langchain_community.llms.openai", + "BaseOpenAI._generate", + openai_tracing_wrapper, ) wrap_function_wrapper( - module="langchain_community.llms.openai", - name="BaseOpenAI._agenerate", - wrapper=openai_tracing_wrapper, + "langchain_community.llms.openai", + "BaseOpenAI._agenerate", + openai_tracing_wrapper, ) wrap_function_wrapper( - module="langchain_community.llms.openai", - name="BaseOpenAI._stream", - wrapper=openai_tracing_wrapper, + "langchain_community.llms.openai", + "BaseOpenAI._stream", + openai_tracing_wrapper, ) wrap_function_wrapper( - module="langchain_community.llms.openai", - name="BaseOpenAI._astream", - wrapper=openai_tracing_wrapper, + "langchain_community.llms.openai", + "BaseOpenAI._astream", + openai_tracing_wrapper, ) if is_package_available("langchain_openai"): # Wrap langchain_openai.llms.base.BaseOpenAI wrap_function_wrapper( - module="langchain_openai.llms.base", - name="BaseOpenAI._generate", - wrapper=openai_tracing_wrapper, + "langchain_openai.llms.base", + "BaseOpenAI._generate", + openai_tracing_wrapper, ) wrap_function_wrapper( - module="langchain_openai.llms.base", - name="BaseOpenAI._agenerate", - wrapper=openai_tracing_wrapper, + "langchain_openai.llms.base", + "BaseOpenAI._agenerate", + openai_tracing_wrapper, ) wrap_function_wrapper( - module="langchain_openai.llms.base", - name="BaseOpenAI._stream", - wrapper=openai_tracing_wrapper, + "langchain_openai.llms.base", + "BaseOpenAI._stream", + openai_tracing_wrapper, ) wrap_function_wrapper( - module="langchain_openai.llms.base", - name="BaseOpenAI._astream", - wrapper=openai_tracing_wrapper, + "langchain_openai.llms.base", + "BaseOpenAI._astream", + openai_tracing_wrapper, ) # langchain_openai.chat_models.base.BaseOpenAI wrap_function_wrapper( - module="langchain_openai.chat_models.base", - name="BaseChatOpenAI._generate", - wrapper=openai_tracing_wrapper, + "langchain_openai.chat_models.base", + "BaseChatOpenAI._generate", + openai_tracing_wrapper, ) wrap_function_wrapper( - module="langchain_openai.chat_models.base", - name="BaseChatOpenAI._agenerate", - wrapper=openai_tracing_wrapper, + "langchain_openai.chat_models.base", + "BaseChatOpenAI._agenerate", + openai_tracing_wrapper, ) # Doesn't work :( diff --git a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py index 4b33ab2b99..ec70deae5b 100644 --- a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py +++ b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/__init__.py @@ -3,6 +3,7 @@ import asyncio # FR: async safety import inspect import logging +import threading from typing import Collection from opentelemetry import context as context_api @@ -43,6 +44,107 @@ _FR_SPAN_ROLE_KEY = "fortifyroot.span.role" _FR_SPAN_ROLE_VALUE = "safety_wrapper" +# Marker set on the safety_wrapper span when LiteLLM's native OTel callback +# will emit a separate ``litellm_request`` span. The FR backend's +# LLMUsageExtractor uses this marker to skip the safety_wrapper for dedup +# WITHOUT needing to see the child span in the same OTLP batch — important +# because ``disable_batch=True`` (SimpleSpanProcessor) sends each span in its +# own ResourceSpans and the existing parent→child correlation in +# ``buildSafetyWrapperDedupeSet`` cannot observe siblings across batches. +_FR_HAS_NATIVE_OTEL_CHILD_KEY = "fortifyroot.span.has_native_otel_child" +_FR_COMPLETION_SAFETY_MARKER = "_fortifyroot_completion_safety_applied" +_FR_COMPLETION_SAFETY_FALLBACK_MARKERS: set[tuple[int, type]] = set() +_FR_COMPLETION_SAFETY_MARKERS_LOCK = threading.Lock() +_FR_COMPLETION_SAFETY_MARKERS_MAX = 4096 + + +def _native_otel_callback_active() -> bool: + """Return True iff LiteLLM's native ``OpenTelemetry`` callback will + emit a separate ``litellm_request`` span. + + LiteLLM's callback does not create that primary span when an ambient + parent span already exists unless ``USE_OTEL_LITELLM_REQUEST_SPAN`` is + enabled. FR always attaches its safety span as the ambient parent, so a + registered callback alone is not enough to mark the safety span for + backend dedup. + """ + try: + import litellm + # LiteLLM's OpenTelemetry callback class lives at a known path. + from litellm.integrations.opentelemetry import OpenTelemetry as _LiteLLMNativeOTel # noqa: N814 + from litellm.secret_managers.main import get_secret_bool + except ImportError: + return False + callbacks = getattr(litellm, "callbacks", None) or [] + try: + has_native_callback = any( + isinstance(cb, _LiteLLMNativeOTel) for cb in callbacks + ) + if not has_native_callback: + return False + return bool(get_secret_bool("USE_OTEL_LITELLM_REQUEST_SPAN", False)) + except Exception: + return False + + +def _completion_safety_already_applied(response) -> bool: + try: + if bool(getattr(response, _FR_COMPLETION_SAFETY_MARKER, False)): + return True + except Exception: + pass + + key = (id(response), type(response)) + with _FR_COMPLETION_SAFETY_MARKERS_LOCK: + return key in _FR_COMPLETION_SAFETY_FALLBACK_MARKERS + + +def _mark_completion_safety_applied(response) -> None: + try: + setattr(response, _FR_COMPLETION_SAFETY_MARKER, True) + return + except Exception: + pass + + key = (id(response), type(response)) + with _FR_COMPLETION_SAFETY_MARKERS_LOCK: + if len(_FR_COMPLETION_SAFETY_FALLBACK_MARKERS) >= _FR_COMPLETION_SAFETY_MARKERS_MAX: + _FR_COMPLETION_SAFETY_FALLBACK_MARKERS.clear() + _FR_COMPLETION_SAFETY_FALLBACK_MARKERS.add(key) + + +def _ensure_completion_safety_applied(span, response, request_type, span_name) -> None: + """Apply completion safety once for a concrete LiteLLM response object. + + LiteLLM async callbacks run through a global worker and may execute after + the wrapper has returned, after the event loop has changed, or during + process teardown. Running completion safety during FR span finalization + makes the customer-visible response deterministic. The callback path still + calls this helper first when it wins the race, preserving the + "FR logger before native OTel" ordering while preventing duplicate finding + events and ended-span writes when the worker runs late. + """ + if response is None or _completion_safety_already_applied(response): + return + try: + apply_completion_safety(span, response, request_type, span_name) + finally: + _mark_completion_safety_applied(response) + + +async def _ensure_completion_safety_applied_async( + span, response, request_type, span_name +) -> None: + if response is None or _completion_safety_already_applied(response): + return + try: + await asyncio.to_thread( + apply_completion_safety, span, response, request_type, span_name + ) + finally: + _mark_completion_safety_applied(response) + + _WRAPPED_METHODS = [ ("litellm", "completion", False, False), ("litellm", "acompletion", True, False), @@ -75,7 +177,9 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): request_type = _request_type(kwargs, is_tc) span = trace.get_current_span() try: - apply_completion_safety(span, response_obj, request_type, span_name) + _ensure_completion_safety_applied( + span, response_obj, request_type, span_name + ) except Exception: pass @@ -87,8 +191,8 @@ async def async_log_success_event(self, kwargs, response_obj, start_time, end_ti request_type = _request_type(kwargs, is_tc) span = trace.get_current_span() try: - await asyncio.to_thread( - apply_completion_safety, span, response_obj, request_type, span_name + await _ensure_completion_safety_applied_async( + span, response_obj, request_type, span_name ) except Exception: pass @@ -180,15 +284,22 @@ def _invoke_completion(tracer, wrapped, args, kwargs, *, is_text_completion=Fals span_name = _span_name(kwargs, is_text_completion) request_type = _request_type(kwargs, is_text_completion) + span_attrs = { + GenAIAttributes.GEN_AI_SYSTEM: "litellm", + SpanAttributes.LLM_REQUEST_TYPE: request_type, + SpanAttributes.LLM_IS_STREAMING: bool(kwargs.get("stream")), + _FR_SPAN_ROLE_KEY: _FR_SPAN_ROLE_VALUE, + } + if _native_otel_callback_active(): + # Hint to the FR backend that this safety_wrapper WILL have a + # sibling litellm_request child emitted by LiteLLM's native OTel + # callback — enables single-pass dedup in proc_llm_extractor.go + # without needing to see the child in the same OTLP batch. + span_attrs[_FR_HAS_NATIVE_OTEL_CHILD_KEY] = True span = tracer.start_span( _FR_SAFETY_SPAN_NAME, kind=SpanKind.CLIENT, - attributes={ - GenAIAttributes.GEN_AI_SYSTEM: "litellm", - SpanAttributes.LLM_REQUEST_TYPE: request_type, - SpanAttributes.LLM_IS_STREAMING: bool(kwargs.get("stream")), - _FR_SPAN_ROLE_KEY: _FR_SPAN_ROLE_VALUE, - }, + attributes=span_attrs, ) # Attach FR span as ambient OTel context so LiteLLM's native OTel callback # creates its litellm_request span as a child of this span. @@ -245,15 +356,18 @@ async def _invoke_acompletion( span_name = _span_name(kwargs, is_text_completion) request_type = _request_type(kwargs, is_text_completion) + span_attrs = { + GenAIAttributes.GEN_AI_SYSTEM: "litellm", + SpanAttributes.LLM_REQUEST_TYPE: request_type, + SpanAttributes.LLM_IS_STREAMING: bool(kwargs.get("stream")), + _FR_SPAN_ROLE_KEY: _FR_SPAN_ROLE_VALUE, + } + if _native_otel_callback_active(): + span_attrs[_FR_HAS_NATIVE_OTEL_CHILD_KEY] = True span = tracer.start_span( _FR_SAFETY_SPAN_NAME, kind=SpanKind.CLIENT, - attributes={ - GenAIAttributes.GEN_AI_SYSTEM: "litellm", - SpanAttributes.LLM_REQUEST_TYPE: request_type, - SpanAttributes.LLM_IS_STREAMING: bool(kwargs.get("stream")), - _FR_SPAN_ROLE_KEY: _FR_SPAN_ROLE_VALUE, - }, + attributes=span_attrs, ) ctx = set_span_in_context(span) ctx = context_api.set_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, True, ctx) @@ -433,8 +547,9 @@ async def _finalize_awaitable_response( async def _async_finalize_response(span, response, request_type, span_name): # FR: async safety - """Async variant of _finalize_response. Completion safety handled by logger.""" # FR: async safety + """Async variant of _finalize_response with deterministic completion safety.""" # FR: async safety try: # FR: async safety + _ensure_completion_safety_applied(span, response, request_type, span_name) # FR: async safety _set_response_attributes(span, response) # FR: async safety span.set_status(Status(StatusCode.OK)) # FR: async safety return response # FR: async safety @@ -443,8 +558,9 @@ async def _async_finalize_response(span, response, request_type, span_name): # def _finalize_response(span, response, request_type, span_name): - """Finalize a non-streaming response. Completion safety handled by logger.""" + """Finalize a non-streaming response with deterministic completion safety.""" try: + _ensure_completion_safety_applied(span, response, request_type, span_name) _set_response_attributes(span, response) span.set_status(Status(StatusCode.OK)) return response diff --git a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py index 0e37ea4a2e..fbe1d30390 100644 --- a/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py +++ b/packages/opentelemetry-instrumentation-litellm/opentelemetry/instrumentation/litellm/streaming_safety.py @@ -17,14 +17,21 @@ def is_sync_streaming_response(kwargs, response) -> bool: """Check if this is a streaming response.""" if kwargs.get("stream"): - return not inspect.iscoroutine(response) and not inspect.isasyncgen(response) + return ( + not inspect.iscoroutine(response) + and not inspect.isasyncgen(response) + ) return False def is_async_streaming_response(kwargs, response) -> bool: """Check if this is an async streaming response.""" if kwargs.get("stream"): - return inspect.iscoroutine(response) or inspect.isasyncgen(response) + return ( + inspect.iscoroutine(response) + or inspect.isasyncgen(response) + or hasattr(response, "__aiter__") + ) return False diff --git a/packages/opentelemetry-instrumentation-litellm/tests/test_logger_integration.py b/packages/opentelemetry-instrumentation-litellm/tests/test_logger_integration.py index 21f1213621..355ff2d25c 100644 --- a/packages/opentelemetry-instrumentation-litellm/tests/test_logger_integration.py +++ b/packages/opentelemetry-instrumentation-litellm/tests/test_logger_integration.py @@ -145,6 +145,55 @@ async def test_async_logger_masks_non_streaming_completion(): assert response.choices[0].message.content == "[SECRET.token]" +@pytest.mark.asyncio +async def test_async_logger_skips_response_already_processed_by_wrapper(): + """Late LiteLLM worker callbacks must not re-run safety on a response + already handled by FR finalization.""" + from opentelemetry.instrumentation.litellm import _invoke_acompletion + + exp, tracer = _make_tracer() + calls = [] + + def handler(context): + calls.append(context.text) + if context.text == "token-abc": + return _completion_handler("[SECRET.token]", context) + return None + + register_completion_safety_handler(handler) + + async def wrapped(*args, **kwargs): + return SimpleNamespace( + model="gpt-4o-mini", + choices=[ + SimpleNamespace( + message=SimpleNamespace(content="token-abc"), + text="token-abc", + finish_reason="stop", + ) + ], + ) + + response = await _invoke_acompletion( + tracer, + wrapped, + (), + {"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + ) + assert response.choices[0].message.content == "[SECRET.token]" + assert calls == ["token-abc"] + + # Simulate LiteLLM's async logging worker flushing after FR already ended + # the span. The marker should make this a no-op. + logger = _FortifyRootCompletionLogger() + await logger.async_log_success_event({}, response, None, None) + + assert calls == ["token-abc"] + spans = exp.get_finished_spans() + assert len(spans) == 1 + assert len(spans[0].events) == 1 + + # --------------------------------------------------------------------------- # Logger uses current OTel span for finding emission # --------------------------------------------------------------------------- @@ -274,3 +323,77 @@ def wrapped(*args, **kwargs): assert spans[0].name == "fortifyroot.litellm.safety" assert spans[0].attributes["fortifyroot.span.role"] == "safety_wrapper" assert spans[0].attributes["gen_ai.system"] == "litellm" + + +def test_native_otel_marker_requires_litellm_request_span_flag(monkeypatch): + """Do not mark the safety span unless LiteLLM will emit litellm_request.""" + from opentelemetry.instrumentation.litellm import _invoke_completion + import litellm + import litellm.integrations.opentelemetry as native_otel + + class DummyOpenTelemetry: + pass + + original_callbacks = list(getattr(litellm, "callbacks", [])) + monkeypatch.setattr(native_otel, "OpenTelemetry", DummyOpenTelemetry) + monkeypatch.delenv("USE_OTEL_LITELLM_REQUEST_SPAN", raising=False) + + try: + litellm.callbacks = [DummyOpenTelemetry()] + exp, tracer = _make_tracer() + + def wrapped(*args, **kwargs): + return SimpleNamespace( + model="gpt-4o", + usage=None, + choices=[SimpleNamespace(message=SimpleNamespace(content="reply"), text="reply")], + ) + + _invoke_completion( + tracer, + wrapped, + (), + {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + ) + + span = exp.get_finished_spans()[0] + assert "fortifyroot.span.has_native_otel_child" not in span.attributes + finally: + litellm.callbacks = original_callbacks + + +def test_native_otel_marker_set_when_litellm_request_span_enabled(monkeypatch): + """Mark the safety span when LiteLLM is configured to emit litellm_request.""" + from opentelemetry.instrumentation.litellm import _invoke_completion + import litellm + import litellm.integrations.opentelemetry as native_otel + + class DummyOpenTelemetry: + pass + + original_callbacks = list(getattr(litellm, "callbacks", [])) + monkeypatch.setattr(native_otel, "OpenTelemetry", DummyOpenTelemetry) + monkeypatch.setenv("USE_OTEL_LITELLM_REQUEST_SPAN", "true") + + try: + litellm.callbacks = [DummyOpenTelemetry()] + exp, tracer = _make_tracer() + + def wrapped(*args, **kwargs): + return SimpleNamespace( + model="gpt-4o", + usage=None, + choices=[SimpleNamespace(message=SimpleNamespace(content="reply"), text="reply")], + ) + + _invoke_completion( + tracer, + wrapped, + (), + {"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]}, + ) + + span = exp.get_finished_spans()[0] + assert span.attributes["fortifyroot.span.has_native_otel_child"] is True + finally: + litellm.callbacks = original_callbacks diff --git a/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py b/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py index 336da9b068..5436b9e458 100644 --- a/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py +++ b/packages/opentelemetry-instrumentation-litellm/tests/test_safety_hooks.py @@ -107,29 +107,36 @@ def _completion_result(masked_text, context): def test_sync_completion_masks_prompt_and_sets_span_attributes(): - """Prompt is masked by _invoke_completion; completion safety is handled by the - logger (not by _finalize_response). Span attributes reflect the response as - returned from the LLM (pre-masked by the logger in production).""" + """Prompt and completion safety are deterministic in the wrapper. + + The LiteLLM callback may still win the race in production and mask first, + but finalization must also handle a raw response so async/global-worker + teardown cannot leak unmasked completion text to the caller. + """ exporter, tracer = _test_tracer() register_prompt_safety_handler( lambda context: _prompt_result("[PII.email]", context) if context.location == SafetyLocation.PROMPT and context.text == "secret" else None ) + register_completion_safety_handler( + lambda context: _completion_result("[SECRET.token]", context) + if context.location == SafetyLocation.COMPLETION and context.text == "token-abc" + else None + ) messages = [{"role": "user", "content": "secret"}] def wrapped(*args, **kwargs): assert kwargs["messages"][0]["content"] == "[PII.email]" assert context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY) is True - # Simulate logger having already masked the response (as in production). return ModelResponse( model="gpt-4o-mini", usage={"prompt_tokens": 3, "completion_tokens": 5, "total_tokens": 8}, choices=[ { "finish_reason": "stop", - "message": {"role": "assistant", "content": "[SECRET.token]"}, + "message": {"role": "assistant", "content": "token-abc"}, } ], ) @@ -153,8 +160,8 @@ def wrapped(*args, **kwargs): assert span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.content"] == "[PII.email]" assert span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"] == "[SECRET.token]" assert span.attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 8 - # One prompt-safety finding event (completion safety handled by logger). - assert len(span.events) == 1 + # One prompt-safety finding and one completion-safety finding. + assert len(span.events) == 2 assert span.attributes["fortifyroot.span.role"] == "safety_wrapper" @@ -165,15 +172,19 @@ def test_sync_text_completion_masks_text_choices(): if context.location == SafetyLocation.PROMPT and context.text == "secret" else None ) + register_completion_safety_handler( + lambda context: _completion_result("[SECRET.token]", context) + if context.location == SafetyLocation.COMPLETION and context.text == "token-abc" + else None + ) def wrapped(*args, **kwargs): assert args[0] == "[PII.email]" assert kwargs["model"] == "gpt-3.5-turbo-instruct" assert context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY) is True - # Response pre-masked (as logger would do in production). return TextCompletionResponse( model="gpt-3.5-turbo-instruct", - choices=[{"text": "[SECRET.token]"}], + choices=[{"text": "token-abc"}], ) response = _invoke_completion( @@ -194,6 +205,7 @@ def wrapped(*args, **kwargs): assert span.attributes[GenAIAttributes.GEN_AI_REQUEST_MODEL] == "gpt-3.5-turbo-instruct" assert span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.content"] == "[PII.email]" assert span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"] == "[SECRET.token]" + assert len(span.events) == 2 @pytest.mark.asyncio @@ -204,6 +216,11 @@ async def test_async_completion_masks_prompt_and_response(): if context.location == SafetyLocation.PROMPT and context.text == "secret" else None ) + register_completion_safety_handler( + lambda context: _completion_result("[SECRET.token]", context) + if context.location == SafetyLocation.COMPLETION and context.text == "token-abc" + else None + ) async def wrapped(*args, **kwargs): assert kwargs["messages"][0]["content"] == "[PII.email]" @@ -212,7 +229,7 @@ async def wrapped(*args, **kwargs): choices=[ { "finish_reason": "stop", - "message": {"role": "assistant", "content": "[SECRET.token]"}, + "message": {"role": "assistant", "content": "token-abc"}, } ], ) @@ -230,6 +247,7 @@ async def wrapped(*args, **kwargs): span = spans[0] assert span.name == _FR_SPAN_NAME assert span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"] == "[SECRET.token]" + assert len(span.events) == 2 @pytest.mark.asyncio @@ -240,13 +258,18 @@ async def test_async_text_completion_masks_prompt_and_response(): if context.location == SafetyLocation.PROMPT and context.text == "secret" else None ) + register_completion_safety_handler( + lambda context: _completion_result("[SECRET.token]", context) + if context.location == SafetyLocation.COMPLETION and context.text == "token-abc" + else None + ) async def wrapped(*args, **kwargs): assert args[0] == "[PII.email]" assert kwargs["model"] == "gpt-3.5-turbo-instruct" return TextCompletionResponse( model="gpt-3.5-turbo-instruct", - choices=[{"text": "[SECRET.token]"}], + choices=[{"text": "token-abc"}], ) response = await _invoke_acompletion( @@ -262,6 +285,7 @@ async def wrapped(*args, **kwargs): assert span.name == _FR_SPAN_NAME assert span.attributes[f"{SpanAttributes.LLM_PROMPTS}.0.content"] == "[PII.email]" assert span.attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"] == "[SECRET.token]" + assert len(span.events) == 2 @pytest.mark.asyncio @@ -272,6 +296,11 @@ async def test_sync_wrapper_handles_awaitable_response(): if context.location == SafetyLocation.PROMPT and context.text == "secret" else None ) + register_completion_safety_handler( + lambda context: _completion_result("[SECRET.token]", context) + if context.location == SafetyLocation.COMPLETION and context.text == "token-abc" + else None + ) async def response_coro(): assert context_api.get_value(SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY) is True @@ -280,7 +309,7 @@ async def response_coro(): choices=[ { "finish_reason": "stop", - "message": {"role": "assistant", "content": "[SECRET.token]"}, + "message": {"role": "assistant", "content": "token-abc"}, } ], ) @@ -303,6 +332,7 @@ def wrapped(*args, **kwargs): spans = exporter.get_finished_spans() assert len(spans) == 1 assert spans[0].attributes[f"{SpanAttributes.LLM_COMPLETIONS}.0.content"] == "[SECRET.token]" + assert len(spans[0].events) == 2 @pytest.mark.asyncio @@ -940,8 +970,8 @@ def __next__(self): assert is_sync_streaming_response({"stream": False}, CustomStreamWrapper()) is False -def test_is_async_streaming_response_detects_coroutines(): - """Coroutines and async generators are detected when stream=True.""" +def test_is_async_streaming_response_detects_async_response_types(): + """Coroutines, async generators, and async iterators are detected when stream=True.""" async def _coro(): pass @@ -961,6 +991,18 @@ async def _agen(): import asyncio asyncio.get_event_loop().run_until_complete(ag.aclose()) + class CustomAsyncStreamWrapper: + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + wrapper = CustomAsyncStreamWrapper() + assert is_async_streaming_response({"stream": True}, wrapper) is True + assert is_async_streaming_response({}, wrapper) is False + assert is_async_streaming_response({"stream": False}, wrapper) is False + def test_is_sync_streaming_response_rejects_async_types(): """Async generators and coroutines should not be sync streaming.""" @@ -980,6 +1022,7 @@ async def _coro(): coro.close() + # --------------------------------------------------------------------------- # Streaming context: span_token kept alive during iteration # --------------------------------------------------------------------------- diff --git a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/dispatcher_wrapper.py b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/dispatcher_wrapper.py index 937000a541..1daadbc76d 100644 --- a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/dispatcher_wrapper.py +++ b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/dispatcher_wrapper.py @@ -70,6 +70,77 @@ # we use the regular OpenLLMetry instrumentations AVAILABLE_OPENLLMETRY_INSTRUMENTATIONS = ["OpenAI"] +# FR: marker role set on the LlamaIndex "*.workflow" span when it +# delegates full LLM attribute extraction to a child provider span. +# The FR backend LLM-usage extractor skips spans carrying this role so +# we don't double-count the wrapper + the child provider span. Safety +# findings are still emitted on the wrapper (that's where safety fires +# before the child provider span exists), which is why we ALSO need +# the model-attribution helpers below. +_FR_LLM_WRAPPER_ROLE_KEY = "fortifyroot.span.role" +_FR_LLM_WRAPPER_ROLE_VALUE = "llm_wrapper" + + +def _infer_provider_from_model(model) -> Optional[str]: + name = str(model or "").lower() + if "claude" in name or "anthropic" in name: + return "anthropic" + if "gpt" in name or name.startswith(("o1", "o3", "o4")): + return "openai" + return None + + +def _stamp_llm_model_for_safety(event: BaseEvent, span) -> None: + """Set gen_ai.request.model / gen_ai.system on a delegated wrapper + span so that safety findings emitted on it have model attribution. + + Leaves other attributes (prompts, params, usage) alone — those + belong to the child provider span. The FR backend LLM-usage + extractor skips wrapper spans by the ``fortifyroot.span.role`` + marker, so adding gen_ai.request.model here does NOT cause + double-counting. + """ + from opentelemetry.semconv._incubating.attributes import ( + gen_ai_attributes as GenAIAttributes, + ) + try: + model_dict = event.model_dict or {} + except Exception: # pragma: no cover — defensive; event shape may vary + return + if "llm" in model_dict: + model_dict = model_dict.get("llm", {}) + model = model_dict.get("model") if isinstance(model_dict, dict) else None + if model: + span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_MODEL, model) + provider = _infer_provider_from_model(model) + if provider: + span.set_attribute(GenAIAttributes.GEN_AI_SYSTEM, provider) + + +def _stamp_llm_response_model_for_safety(event: BaseEvent, span) -> None: + """Set gen_ai.response.model on a delegated wrapper span when the + LLM response carries a model different from the request (e.g. + Anthropic aliasing ``claude-4-sonnet`` → ``claude-sonnet-4``). + + Used for completion-location safety findings that are emitted on + the wrapper span after the response has arrived. + """ + from opentelemetry.semconv._incubating.attributes import ( + gen_ai_attributes as GenAIAttributes, + ) + response = getattr(event, "response", None) + raw = getattr(response, "raw", None) if response is not None else None + model = None + try: + if raw is not None: + model = getattr(raw, "model", None) + if not model and isinstance(raw, dict): + model = raw.get("model") + except Exception: # pragma: no cover + model = None + if model: + span.set_attribute(GenAIAttributes.GEN_AI_RESPONSE_MODEL, model) + CLASS_ANDMETHOD_NAME_FROM_ID_REGEX = re.compile(r"([a-zA-Z]+)\.([a-zA-Z_]+)-") STREAMING_END_EVENTS = ( LLMChatEndEvent, @@ -139,6 +210,17 @@ def _(self, event: LLMChatStartEvent): emit_chat_message_events(event) else: set_llm_chat_request(event, self.otel_span) + else: + # FR: even when the provider instrumentor owns the full LLM + # attribute set on its own child span, we still stamp + # ``gen_ai.request.model`` + ``gen_ai.system`` on this wrapper + # span so that safety violation events emitted here (via + # ``emit_deferred_findings``) carry model attribution. The + # wrapper span is marked ``fortifyroot.span.role="llm_wrapper"`` + # at creation time so the backend LLM-usage extractor skips + # it for event counting and avoids double-counting with the + # child provider span. + _stamp_llm_model_for_safety(event, self.otel_span) @update_span_for_event.register def _(self, event: LLMChatEndEvent): @@ -153,11 +235,20 @@ def _(self, event: LLMChatEndEvent): emit_chat_response_events(event) else: set_llm_chat_response(event, self.otel_span) # noqa: F821 + else: + # FR: completion-location safety findings (non-streaming) are + # emitted here for delegated providers too, so we need model + # attribution on the wrapper span for those events. + if not self.waiting_for_streaming: + apply_chat_end_safety(event, self.otel_span) + _stamp_llm_response_model_for_safety(event, self.otel_span) @update_span_for_event.register def _(self, event: LLMCompletionStartEvent): if not self.delegates_to_provider: apply_completion_start_span_attributes(event, self.otel_span) + else: + _stamp_llm_model_for_safety(event, self.otel_span) @update_span_for_event.register def _(self, event: LLMCompletionEndEvent): @@ -165,6 +256,13 @@ def _(self, event: LLMCompletionEndEvent): # FR: same as LLMChatEndEvent -- skip when streaming safety is active. if not self.waiting_for_streaming: apply_completion_end_safety(event, self.otel_span) + else: + # FR: delegated-provider completion safety still emits + # findings on the wrapper span; stamp response model on the + # wrapper so those findings carry model attribution. + if not self.waiting_for_streaming: + apply_completion_end_safety(event, self.otel_span) + _stamp_llm_response_model_for_safety(event, self.otel_span) @update_span_for_event.register def _(self, event: LLMPredictEndEvent): @@ -261,6 +359,14 @@ def new_span( span.set_attribute(SpanAttributes.TRACELOOP_SPAN_KIND, kind) span.set_attribute(SpanAttributes.TRACELOOP_ENTITY_NAME, span_name) + # FR: mark delegated LlamaIndex wrapper spans so the backend + # LLM-usage extractor skips them for event counting. Pair with + # ``_stamp_llm_model_for_safety()`` to keep safety findings + # emitted on this span (via emit_deferred_findings) correctly + # attributed to model / provider. See + # fr-system-tests st_phase_6.txt addendum 8. + if is_openllmetry_class: + span.set_attribute(_FR_LLM_WRAPPER_ROLE_KEY, _FR_LLM_WRAPPER_ROLE_VALUE) try: if should_send_prompts(): span.set_attribute( diff --git a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/safety.py b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/safety.py index 00d506af91..c0b2b6d6b2 100644 --- a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/safety.py +++ b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/safety.py @@ -49,6 +49,19 @@ logger = logging.getLogger(__name__) +def _infer_llm_provider(model) -> str | None: + model_name = str(model or "").lower() + # TODO(ST-6 follow-up): MVP support here is intentionally limited to the + # providers exercised in ST-6 (OpenAI + Anthropic). Expand this inference + # when we certify more LlamaIndex-backed providers so safety-emitted + # wrapper spans continue to carry provider/model attribution for them. + if "claude" in model_name or "anthropic" in model_name: + return "anthropic" + if "gpt" in model_name or model_name.startswith(("o1", "o3", "o4")): + return "openai" + return None + + def instrument_llm_safety_wrappers(): global _WRAPPERS_INSTALLED with _WRAPPERS_LOCK: @@ -259,7 +272,11 @@ def apply_completion_start_span_attributes(event, span): span.set_attribute( SpanAttributes.LLM_REQUEST_TYPE, LLMRequestTypeValues.COMPLETION.value ) - span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_MODEL, model_dict.get("model")) + model = model_dict.get("model") + span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_MODEL, model) + provider = _infer_llm_provider(model) + if provider: + span.set_attribute(GenAIAttributes.GEN_AI_SYSTEM, provider) span.set_attribute( GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE, model_dict.get("temperature"), @@ -513,7 +530,11 @@ def _set_completion_response_model_attributes(response, span): usage = get_object_value(raw, "usage") if usage is not None: completion_tokens = get_object_value(usage, "completion_tokens") + if completion_tokens is None: + completion_tokens = get_object_value(usage, "output_tokens") prompt_tokens = get_object_value(usage, "prompt_tokens") + if prompt_tokens is None: + prompt_tokens = get_object_value(usage, "input_tokens") total_tokens = get_object_value(usage, "total_tokens") if completion_tokens is not None: span.set_attribute( @@ -530,6 +551,11 @@ def _set_completion_response_model_attributes(response, span): SpanAttributes.LLM_USAGE_TOTAL_TOKENS, int(total_tokens), ) + elif prompt_tokens is not None and completion_tokens is not None: + span.set_attribute( + SpanAttributes.LLM_USAGE_TOTAL_TOKENS, + int(prompt_tokens) + int(completion_tokens), + ) def _mask_prompt_text( diff --git a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/span_utils.py b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/span_utils.py index 0bcba17b4f..329f86db05 100644 --- a/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/span_utils.py +++ b/packages/opentelemetry-instrumentation-llamaindex/opentelemetry/instrumentation/llamaindex/span_utils.py @@ -12,6 +12,19 @@ ) +def _infer_llm_provider(model) -> str | None: + model_name = str(model or "").lower() + # TODO(ST-6 follow-up): MVP support here is intentionally limited to the + # providers exercised in ST-6 (OpenAI + Anthropic). Expand this inference + # as additional LlamaIndex-backed providers are certified so delegated + # wrapper spans and backend LLMUsage attribution keep working for them too. + if "claude" in model_name or "anthropic" in model_name: + return "anthropic" + if "gpt" in model_name or model_name.startswith(("o1", "o3", "o4")): + return "openai" + return None + + @dont_throw def set_llm_chat_request(event, span) -> None: if not span.is_recording(): @@ -39,7 +52,11 @@ def set_llm_chat_request_model_attributes(event, span): if "llm" in model_dict: model_dict = model_dict.get("llm", {}) - span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_MODEL, model_dict.get("model")) + model = model_dict.get("model") + span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_MODEL, model) + provider = _infer_llm_provider(model) + if provider: + span.set_attribute(GenAIAttributes.GEN_AI_SYSTEM, provider) span.set_attribute( GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE, model_dict.get("temperature") ) @@ -93,6 +110,11 @@ def set_llm_chat_response_model_attributes(event, span): output_tokens = None total_tokens = None + # TODO(ST-6 follow-up): token extraction below currently covers the usage + # shapes we needed for the MVP providers/certified paths (OpenAI-style, + # Anthropic-style, and Cohere metadata fallback). Extend this branch as + # additional LlamaIndex-backed providers are added so token attribution + # stays correct without relying on backend heuristics. # Try OpenAI format first: raw.usage with completion_tokens, prompt_tokens usage = getattr(raw, "usage", None) or (raw.get("usage") if isinstance(raw, dict) else None) if usage: @@ -100,9 +122,13 @@ def set_llm_chat_response_model_attributes(event, span): output_tokens = usage.completion_tokens input_tokens = usage.prompt_tokens total_tokens = usage.total_tokens + elif hasattr(usage, "output_tokens"): + output_tokens = usage.output_tokens + input_tokens = usage.input_tokens + total_tokens = getattr(usage, "total_tokens", None) elif isinstance(usage, dict): - output_tokens = usage.get("completion_tokens") - input_tokens = usage.get("prompt_tokens") + output_tokens = usage.get("completion_tokens", usage.get("output_tokens")) + input_tokens = usage.get("prompt_tokens", usage.get("input_tokens")) total_tokens = usage.get("total_tokens") # Try Cohere format: raw.meta.tokens or raw.meta.billed_units diff --git a/packages/opentelemetry-instrumentation-llamaindex/tests/test_safety_hooks.py b/packages/opentelemetry-instrumentation-llamaindex/tests/test_safety_hooks.py index fd78994def..e76be06d80 100644 --- a/packages/opentelemetry-instrumentation-llamaindex/tests/test_safety_hooks.py +++ b/packages/opentelemetry-instrumentation-llamaindex/tests/test_safety_hooks.py @@ -392,6 +392,27 @@ def test_internal_helpers_cover_kwargs_and_non_recording_paths(): assert spans[0].attributes["gen_ai.usage.input_tokens"] == 1 assert spans[0].attributes["gen_ai.usage.output_tokens"] == 2 assert spans[0].attributes["llm.usage.total_tokens"] == 3 + + exporter, tracer = _test_span() + with patch("opentelemetry.instrumentation.llamaindex.safety.should_send_prompts", return_value=False): + with tracer.start_as_current_span("llamaindex.completion") as span: + apply_completion_end_safety( + LLMCompletionEndEvent( + prompt="prompt", + response=CompletionResponse( + text="secret", + raw={"model": "claude", "usage": {"input_tokens": 4, "output_tokens": 5}}, + ), + span_id="span-3", + ), + span, + ) + + spans = exporter.get_finished_spans() + assert spans[0].attributes["gen_ai.usage.input_tokens"] == 4 + assert spans[0].attributes["gen_ai.usage.output_tokens"] == 5 + assert spans[0].attributes["llm.usage.total_tokens"] == 9 + thinking_block = SimpleNamespace(block_type="thinking", content="idea") assert _block_text(thinking_block) == "idea" assert _set_block_text(thinking_block, "updated") is True diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py index d4c1178b96..f7c811030d 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py @@ -752,49 +752,59 @@ def _shared_attributes(self): @dont_throw def _process_complete_response(self): - _set_streaming_token_metrics( - self._request_kwargs, - self._complete_response, - self._span, - self._token_counter, - self._shared_attributes(), - ) + with self._cleanup_lock: + if self._cleanup_completed: + logger.debug("ChatStream complete response already processed, skipping") + return - # choice metrics - if self._choice_counter and self._complete_response.get("choices"): - _set_choice_counter_metrics( - self._choice_counter, - self._complete_response.get("choices"), + if not self._span or not self._span.is_recording(): + self._cleanup_completed = True + logger.debug("ChatStream span already ended before complete response") + return + + _set_streaming_token_metrics( + self._request_kwargs, + self._complete_response, + self._span, + self._token_counter, self._shared_attributes(), ) - # duration metrics - if self._start_time and isinstance(self._start_time, (float, int)): - duration = time.time() - self._start_time - else: - duration = None - if duration and isinstance(duration, (float, int)) and self._duration_histogram: - self._duration_histogram.record( - duration, attributes=self._shared_attributes() - ) - if self._streaming_time_to_generate and self._time_of_first_token: - self._streaming_time_to_generate.record( - time.time() - self._time_of_first_token, - attributes=self._shared_attributes(), - ) + # choice metrics + if self._choice_counter and self._complete_response.get("choices"): + _set_choice_counter_metrics( + self._choice_counter, + self._complete_response.get("choices"), + self._shared_attributes(), + ) - _set_response_attributes(self._span, self._complete_response) - if should_emit_events(): - for choice in self._complete_response.get("choices", []): - emit_event(_parse_choice_event(choice)) - else: - if should_send_prompts(): - _set_completions( - self._span, self._complete_response.get("choices")) + # duration metrics + if self._start_time and isinstance(self._start_time, (float, int)): + duration = time.time() - self._start_time + else: + duration = None + if duration and isinstance(duration, (float, int)) and self._duration_histogram: + self._duration_histogram.record( + duration, attributes=self._shared_attributes() + ) + if self._streaming_time_to_generate and self._time_of_first_token: + self._streaming_time_to_generate.record( + time.time() - self._time_of_first_token, + attributes=self._shared_attributes(), + ) + + _set_response_attributes(self._span, self._complete_response) + if should_emit_events(): + for choice in self._complete_response.get("choices", []): + emit_event(_parse_choice_event(choice)) + else: + if should_send_prompts(): + _set_completions( + self._span, self._complete_response.get("choices")) - self._span.set_status(Status(StatusCode.OK)) - self._span.end() - self._cleanup_completed = True + self._span.set_status(Status(StatusCode.OK)) + self._span.end() + self._cleanup_completed = True @dont_throw def _ensure_cleanup(self): diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/utils.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/utils.py index 1de4e63f3b..ee6eccb5a7 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/utils.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/utils.py @@ -4,12 +4,14 @@ import threading import traceback from contextlib import asynccontextmanager +from importlib import import_module from importlib.metadata import version from packaging import version as pkg_version from opentelemetry import context as context_api from opentelemetry._logs import Logger from opentelemetry.instrumentation.openai.shared.config import Config +from opentelemetry.instrumentation.utils import unwrap import openai @@ -160,6 +162,18 @@ def _handle_exception(e, func, logger): return async_wrapper if asyncio.iscoroutinefunction(func) else sync_wrapper +def unwrap_dotted_method(module, dotted_method): + """Unwrap methods registered as ``Class.method`` via wrap_function_wrapper.""" + try: + owner = import_module(module) + parts = dotted_method.split(".") + for part in parts[:-1]: + owner = getattr(owner, part) + unwrap(owner, parts[-1]) + except (AttributeError, ImportError, ModuleNotFoundError): + pass + + def run_async(method): try: loop = asyncio.get_running_loop() diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v0/__init__.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v0/__init__.py index 317fb7020b..ff7ce4aca1 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v0/__init__.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v0/__init__.py @@ -15,9 +15,11 @@ aembeddings_wrapper, embeddings_wrapper, ) -from opentelemetry.instrumentation.openai.utils import is_metrics_enabled +from opentelemetry.instrumentation.openai.utils import ( + is_metrics_enabled, + unwrap_dotted_method, +) from opentelemetry.instrumentation.openai.version import __version__ -from opentelemetry.instrumentation.utils import unwrap from opentelemetry.metrics import get_meter from opentelemetry.semconv._incubating.metrics import gen_ai_metrics as GenAIMetrics from opentelemetry.semconv_ai import Meters @@ -168,9 +170,9 @@ def _instrument(self, **kwargs): ) def _uninstrument(self, **kwargs): - unwrap("openai", "Completion.create") - unwrap("openai", "Completion.acreate") - unwrap("openai", "ChatCompletion.create") - unwrap("openai", "ChatCompletion.acreate") - unwrap("openai", "Embedding.create") - unwrap("openai", "Embedding.acreate") + unwrap_dotted_method("openai", "Completion.create") + unwrap_dotted_method("openai", "Completion.acreate") + unwrap_dotted_method("openai", "ChatCompletion.create") + unwrap_dotted_method("openai", "ChatCompletion.acreate") + unwrap_dotted_method("openai", "Embedding.create") + unwrap_dotted_method("openai", "Embedding.acreate") diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/__init__.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/__init__.py index 966127362d..22ddfb9fa9 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/__init__.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/__init__.py @@ -18,7 +18,10 @@ from opentelemetry.instrumentation.openai.shared.image_gen_wrappers import ( image_gen_metrics_wrapper, ) -from opentelemetry.instrumentation.openai.utils import is_metrics_enabled +from opentelemetry.instrumentation.openai.utils import ( + is_metrics_enabled, + unwrap_dotted_method, +) from opentelemetry.instrumentation.openai.v1.assistant_wrappers import ( assistants_create_wrapper, messages_list_wrapper, @@ -41,7 +44,6 @@ ) from opentelemetry.instrumentation.openai.version import __version__ -from opentelemetry.instrumentation.utils import unwrap from opentelemetry.metrics import get_meter from opentelemetry.semconv._incubating.metrics import gen_ai_metrics as GenAIMetrics from opentelemetry.semconv_ai import Meters @@ -349,31 +351,31 @@ def _instrument(self, **kwargs): ) def _uninstrument(self, **kwargs): - unwrap("openai.resources.chat.completions", "Completions.create") - unwrap("openai.resources.completions", "Completions.create") - unwrap("openai.resources.embeddings", "Embeddings.create") - unwrap("openai.resources.chat.completions", "AsyncCompletions.create") - unwrap("openai.resources.completions", "AsyncCompletions.create") - unwrap("openai.resources.embeddings", "AsyncEmbeddings.create") - unwrap("openai.resources.images", "Images.generate") + unwrap_dotted_method("openai.resources.chat.completions", "Completions.create") + unwrap_dotted_method("openai.resources.completions", "Completions.create") + unwrap_dotted_method("openai.resources.embeddings", "Embeddings.create") + unwrap_dotted_method("openai.resources.chat.completions", "AsyncCompletions.create") + unwrap_dotted_method("openai.resources.completions", "AsyncCompletions.create") + unwrap_dotted_method("openai.resources.embeddings", "AsyncEmbeddings.create") + unwrap_dotted_method("openai.resources.images", "Images.generate") # Beta APIs may not be available consistently in all versions try: - unwrap("openai.resources.beta.assistants", "Assistants.create") + unwrap_dotted_method("openai.resources.beta.assistants", "Assistants.create") uninstrument_additional_beta_safety_surfaces() - unwrap("openai.resources.beta.chat.completions", "Completions.parse") - unwrap("openai.resources.beta.chat.completions", "AsyncCompletions.parse") - unwrap("openai.resources.beta.threads.runs", "Runs.create") - unwrap("openai.resources.beta.threads.runs", "Runs.retrieve") - unwrap("openai.resources.beta.threads.runs", "Runs.create_and_stream") - unwrap("openai.resources.beta.threads.messages", "Messages.list") - unwrap("openai.resources.responses", "Responses.create") - unwrap("openai.resources.responses", "Responses.retrieve") - unwrap("openai.resources.responses", "Responses.cancel") - unwrap("openai.resources.responses", "AsyncResponses.create") - unwrap("openai.resources.responses", "AsyncResponses.retrieve") - unwrap("openai.resources.responses", "AsyncResponses.cancel") - unwrap("openai.resources.beta.realtime.realtime", "Realtime.connect") - unwrap("openai.resources.beta.realtime.realtime", "AsyncRealtime.connect") + unwrap_dotted_method("openai.resources.beta.chat.completions", "Completions.parse") + unwrap_dotted_method("openai.resources.beta.chat.completions", "AsyncCompletions.parse") + unwrap_dotted_method("openai.resources.beta.threads.runs", "Runs.create") + unwrap_dotted_method("openai.resources.beta.threads.runs", "Runs.retrieve") + unwrap_dotted_method("openai.resources.beta.threads.runs", "Runs.create_and_stream") + unwrap_dotted_method("openai.resources.beta.threads.messages", "Messages.list") + unwrap_dotted_method("openai.resources.responses", "Responses.create") + unwrap_dotted_method("openai.resources.responses", "Responses.retrieve") + unwrap_dotted_method("openai.resources.responses", "Responses.cancel") + unwrap_dotted_method("openai.resources.responses", "AsyncResponses.create") + unwrap_dotted_method("openai.resources.responses", "AsyncResponses.retrieve") + unwrap_dotted_method("openai.resources.responses", "AsyncResponses.cancel") + unwrap_dotted_method("openai.resources.beta.realtime.realtime", "Realtime.connect") + unwrap_dotted_method("openai.resources.beta.realtime.realtime", "AsyncRealtime.connect") except ImportError: pass diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/instrumentation_runtime.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/instrumentation_runtime.py index 8618e90114..9be9d0833f 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/instrumentation_runtime.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/instrumentation_runtime.py @@ -9,7 +9,7 @@ aruns_retrieve_wrapper, ) from opentelemetry.instrumentation.openai.v1.assistant_wrappers import messages_create_wrapper -from opentelemetry.instrumentation.utils import unwrap +from opentelemetry.instrumentation.openai.utils import unwrap_dotted_method def instrument_additional_beta_safety_surfaces(instrumentor, tracer): @@ -51,10 +51,10 @@ def instrument_additional_beta_safety_surfaces(instrumentor, tracer): def uninstrument_additional_beta_safety_surfaces(): - unwrap("openai.resources.beta.assistants", "AsyncAssistants.create") - unwrap("openai.resources.beta.threads.runs", "AsyncRuns.create") - unwrap("openai.resources.beta.threads.runs", "AsyncRuns.retrieve") - unwrap("openai.resources.beta.threads.runs", "AsyncRuns.create_and_stream") - unwrap("openai.resources.beta.threads.messages", "Messages.create") - unwrap("openai.resources.beta.threads.messages", "AsyncMessages.create") - unwrap("openai.resources.beta.threads.messages", "AsyncMessages.list") + unwrap_dotted_method("openai.resources.beta.assistants", "AsyncAssistants.create") + unwrap_dotted_method("openai.resources.beta.threads.runs", "AsyncRuns.create") + unwrap_dotted_method("openai.resources.beta.threads.runs", "AsyncRuns.retrieve") + unwrap_dotted_method("openai.resources.beta.threads.runs", "AsyncRuns.create_and_stream") + unwrap_dotted_method("openai.resources.beta.threads.messages", "Messages.create") + unwrap_dotted_method("openai.resources.beta.threads.messages", "AsyncMessages.create") + unwrap_dotted_method("openai.resources.beta.threads.messages", "AsyncMessages.list") diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/safety_registration.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/safety_registration.py index 0df0de6834..89dedffc92 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/safety_registration.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/v1/safety_registration.py @@ -27,7 +27,7 @@ responses_cancel_wrapper, responses_get_or_create_wrapper, ) -from opentelemetry.instrumentation.utils import unwrap +from opentelemetry.instrumentation.openai.utils import unwrap_dotted_method def instrument_beta_safety_wrappers( @@ -170,25 +170,25 @@ def instrument_beta_safety_wrappers( def uninstrument_beta_safety_wrappers(): - unwrap("openai.resources.beta.assistants", "Assistants.create") - unwrap("openai.resources.beta.assistants", "AsyncAssistants.create") - unwrap("openai.resources.beta.chat.completions", "Completions.parse") - unwrap("openai.resources.beta.chat.completions", "AsyncCompletions.parse") - unwrap("openai.resources.beta.threads.runs", "Runs.create") - unwrap("openai.resources.beta.threads.runs", "AsyncRuns.create") - unwrap("openai.resources.beta.threads.runs", "Runs.retrieve") - unwrap("openai.resources.beta.threads.runs", "AsyncRuns.retrieve") - unwrap("openai.resources.beta.threads.runs", "Runs.create_and_stream") - unwrap("openai.resources.beta.threads.runs", "AsyncRuns.create_and_stream") - unwrap("openai.resources.beta.threads.messages", "Messages.create") - unwrap("openai.resources.beta.threads.messages", "AsyncMessages.create") - unwrap("openai.resources.beta.threads.messages", "Messages.list") - unwrap("openai.resources.beta.threads.messages", "AsyncMessages.list") - unwrap("openai.resources.responses", "Responses.create") - unwrap("openai.resources.responses", "Responses.retrieve") - unwrap("openai.resources.responses", "Responses.cancel") - unwrap("openai.resources.responses", "AsyncResponses.create") - unwrap("openai.resources.responses", "AsyncResponses.retrieve") - unwrap("openai.resources.responses", "AsyncResponses.cancel") - unwrap("openai.resources.beta.realtime.realtime", "Realtime.connect") - unwrap("openai.resources.beta.realtime.realtime", "AsyncRealtime.connect") + unwrap_dotted_method("openai.resources.beta.assistants", "Assistants.create") + unwrap_dotted_method("openai.resources.beta.assistants", "AsyncAssistants.create") + unwrap_dotted_method("openai.resources.beta.chat.completions", "Completions.parse") + unwrap_dotted_method("openai.resources.beta.chat.completions", "AsyncCompletions.parse") + unwrap_dotted_method("openai.resources.beta.threads.runs", "Runs.create") + unwrap_dotted_method("openai.resources.beta.threads.runs", "AsyncRuns.create") + unwrap_dotted_method("openai.resources.beta.threads.runs", "Runs.retrieve") + unwrap_dotted_method("openai.resources.beta.threads.runs", "AsyncRuns.retrieve") + unwrap_dotted_method("openai.resources.beta.threads.runs", "Runs.create_and_stream") + unwrap_dotted_method("openai.resources.beta.threads.runs", "AsyncRuns.create_and_stream") + unwrap_dotted_method("openai.resources.beta.threads.messages", "Messages.create") + unwrap_dotted_method("openai.resources.beta.threads.messages", "AsyncMessages.create") + unwrap_dotted_method("openai.resources.beta.threads.messages", "Messages.list") + unwrap_dotted_method("openai.resources.beta.threads.messages", "AsyncMessages.list") + unwrap_dotted_method("openai.resources.responses", "Responses.create") + unwrap_dotted_method("openai.resources.responses", "Responses.retrieve") + unwrap_dotted_method("openai.resources.responses", "Responses.cancel") + unwrap_dotted_method("openai.resources.responses", "AsyncResponses.create") + unwrap_dotted_method("openai.resources.responses", "AsyncResponses.retrieve") + unwrap_dotted_method("openai.resources.responses", "AsyncResponses.cancel") + unwrap_dotted_method("openai.resources.beta.realtime.realtime", "Realtime.connect") + unwrap_dotted_method("openai.resources.beta.realtime.realtime", "AsyncRealtime.connect") diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py index a604dc9d40..8a94cd7f62 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py @@ -1,12 +1,16 @@ import asyncio +import logging from unittest.mock import patch import httpx import pytest +from openai.resources.chat.completions import Completions from openai.types.chat.chat_completion_message_tool_call import ( ChatCompletionMessageFunctionToolCall, ) +from wrapt import BoundFunctionWrapper from opentelemetry.sdk._logs import ReadableLogRecord +from opentelemetry.instrumentation.openai import OpenAIInstrumentor from opentelemetry.semconv._incubating.attributes import ( gen_ai_attributes as GenAIAttributes, ) @@ -17,6 +21,22 @@ from .utils import assert_request_contains_tracecontext, spy_decorator +def test_openai_v1_uninstrument_removes_chat_completion_wrapper(): + instrumentor = OpenAIInstrumentor() + if instrumentor.is_instrumented_by_opentelemetry: + instrumentor.uninstrument() + + assert not isinstance(Completions.create, BoundFunctionWrapper) + + instrumentor.instrument() + try: + assert isinstance(Completions.create, BoundFunctionWrapper) + finally: + instrumentor.uninstrument() + + assert not isinstance(Completions.create, BoundFunctionWrapper) + + @pytest.mark.vcr def test_chat(instrument_legacy, span_exporter, log_exporter, openai_client): openai_client.chat.completions.create( @@ -1756,6 +1776,43 @@ def test_chat_streaming_partial_consumption(instrument_legacy, span_exporter, lo ) +def test_chat_stream_complete_response_noops_after_cleanup( + tracer_provider, span_exporter, caplog +): + """Late stream finalization must not mutate an already-ended span.""" + from opentelemetry.instrumentation.openai.shared.chat_wrappers import ChatStream + + class _EmptyStream: + def __iter__(self): + return self + + def __next__(self): + raise StopIteration + + span = tracer_provider.get_tracer(__name__).start_span("openai.chat") + stream = ChatStream( + span, + _EmptyStream(), + start_time=0.0, + request_kwargs={"model": "gpt-3.5-turbo"}, + ) + + stream._ensure_cleanup() + assert len(span_exporter.get_finished_spans()) == 1 + + caplog.clear() + with caplog.at_level(logging.WARNING, logger="opentelemetry.sdk.trace"): + stream._process_complete_response() + + assert len(span_exporter.get_finished_spans()) == 1 + ended_span_messages = [ + record.getMessage() + for record in caplog.records + if "ended span" in record.getMessage() + ] + assert ended_span_messages == [] + + @pytest.mark.vcr def test_chat_streaming_exception_during_consumption(instrument_legacy, span_exporter, log_exporter, openai_client): """Test that streaming responses handle exceptions during consumption properly""" From 01df2a439949bb0c51202cb93cf38248367f572a Mon Sep 17 00:00:00 2001 From: Manasjyoti Sharma Date: Mon, 20 Apr 2026 20:52:01 +0530 Subject: [PATCH 2/2] fix: stabilize OpenAI instrumentation lifecycle and package tests --- .../instrumentation/openai/shared/__init__.py | 8 +- .../openai/shared/chat_wrappers.py | 22 +++++- .../tests/conftest.py | 2 - .../tests/traces/test_chat.py | 73 +++++++++++-------- .../tests/traces/test_chat_parse.py | 28 +------ .../tests/traces/test_completions.py | 32 ++------ .../tests/traces/test_embeddings.py | 24 ++---- .../tests/traces/utils.py | 31 ++++++++ 8 files changed, 115 insertions(+), 105 deletions(-) diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/__init__.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/__init__.py index a27b1bcf81..99c416cdb2 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/__init__.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/__init__.py @@ -27,6 +27,8 @@ logger = logging.getLogger(__name__) +LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA = "gen_ai.request.structured_output_schema" + def _set_span_attribute(span, name, value): if value is None or value == "": @@ -157,7 +159,7 @@ def _set_request_attributes(span, kwargs, instance=None): if schema: _set_span_attribute( span, - SpanAttributes.LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA, + LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA, json.dumps(schema), ) elif ( @@ -169,7 +171,7 @@ def _set_request_attributes(span, kwargs, instance=None): ): _set_span_attribute( span, - SpanAttributes.LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA, + LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA, json.dumps(response_format.model_json_schema()), ) else: @@ -185,7 +187,7 @@ def _set_request_attributes(span, kwargs, instance=None): if schema: _set_span_attribute( span, - SpanAttributes.LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA, + LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA, schema, ) diff --git a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py index f7c811030d..6f00a08a21 100644 --- a/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py +++ b/packages/opentelemetry-instrumentation-openai/opentelemetry/instrumentation/openai/shared/chat_wrappers.py @@ -104,7 +104,16 @@ def chat_wrapper( run_async(_handle_request(span, kwargs, instance)) try: start_time = time.time() - response = wrapped(*args, **kwargs) + token = context_api.attach( + context_api.set_value( + SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, + True, + ) + ) + try: + response = wrapped(*args, **kwargs) + finally: + context_api.detach(token) end_time = time.time() except Exception as e: # pylint: disable=broad-except end_time = time.time() @@ -205,7 +214,16 @@ async def achat_wrapper( try: start_time = time.time() - response = await wrapped(*args, **kwargs) + token = context_api.attach( + context_api.set_value( + SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY, + True, + ) + ) + try: + response = await wrapped(*args, **kwargs) + finally: + context_api.detach(token) end_time = time.time() except Exception as e: # pylint: disable=broad-except end_time = time.time() diff --git a/packages/opentelemetry-instrumentation-openai/tests/conftest.py b/packages/opentelemetry-instrumentation-openai/tests/conftest.py index a70ac5b47e..347a2691d0 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/conftest.py +++ b/packages/opentelemetry-instrumentation-openai/tests/conftest.py @@ -166,7 +166,6 @@ def instrument_with_content( Config.use_legacy_attributes = True Config.event_logger = None os.environ.pop(TRACELOOP_TRACE_CONTENT, None) - instrumentor.uninstrument() @pytest.fixture(scope="function") @@ -186,7 +185,6 @@ def instrument_with_no_content( Config.use_legacy_attributes = True Config.event_logger = None os.environ.pop(TRACELOOP_TRACE_CONTENT, None) - instrumentor.uninstrument() @pytest.fixture(autouse=True) diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py index 8a94cd7f62..f5441c2ef7 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat.py @@ -11,6 +11,7 @@ from wrapt import BoundFunctionWrapper from opentelemetry.sdk._logs import ReadableLogRecord from opentelemetry.instrumentation.openai import OpenAIInstrumentor +from opentelemetry.instrumentation.openai.shared.config import Config from opentelemetry.semconv._incubating.attributes import ( gen_ai_attributes as GenAIAttributes, ) @@ -18,23 +19,53 @@ from opentelemetry.trace import StatusCode from opentelemetry.instrumentation.openai.utils import is_reasoning_supported -from .utils import assert_request_contains_tracecontext, spy_decorator +from .utils import ( + assert_openai_exception_span, + assert_request_contains_tracecontext, + spy_decorator, +) + +def test_openai_v1_uninstrument_removes_chat_completion_wrapper( + tracer_provider, meter_provider +): + prior_config = { + "enrich_assistant": Config.enrich_assistant, + "exception_logger": Config.exception_logger, + "get_common_metrics_attributes": Config.get_common_metrics_attributes, + "upload_base64_image": Config.upload_base64_image, + "enable_trace_context_propagation": Config.enable_trace_context_propagation, + "use_legacy_attributes": Config.use_legacy_attributes, + } + instrumentor = OpenAIInstrumentor(**prior_config) + restore_session_instrumentation = instrumentor.is_instrumented_by_opentelemetry -def test_openai_v1_uninstrument_removes_chat_completion_wrapper(): - instrumentor = OpenAIInstrumentor() - if instrumentor.is_instrumented_by_opentelemetry: + if restore_session_instrumentation: instrumentor.uninstrument() assert not isinstance(Completions.create, BoundFunctionWrapper) - instrumentor.instrument() try: + instrumentor.instrument( + tracer_provider=tracer_provider, + meter_provider=meter_provider, + ) assert isinstance(Completions.create, BoundFunctionWrapper) finally: instrumentor.uninstrument() + if restore_session_instrumentation: + instrumentor.instrument( + tracer_provider=tracer_provider, + meter_provider=meter_provider, + ) + Config.enrich_assistant = prior_config["enrich_assistant"] + Config.exception_logger = prior_config["exception_logger"] + Config.get_common_metrics_attributes = prior_config["get_common_metrics_attributes"] + Config.upload_base64_image = prior_config["upload_base64_image"] + Config.enable_trace_context_propagation = prior_config["enable_trace_context_propagation"] + Config.use_legacy_attributes = prior_config["use_legacy_attributes"] - assert not isinstance(Completions.create, BoundFunctionWrapper) + assert isinstance(Completions.create, BoundFunctionWrapper) == restore_session_instrumentation @pytest.mark.vcr @@ -1563,18 +1594,7 @@ def test_chat_exception(instrument_legacy, span_exporter, openai_client): ) assert open_ai_span.attributes.get( SpanAttributes.LLM_IS_STREAMING) is False - assert open_ai_span.status.status_code == StatusCode.ERROR - assert open_ai_span.status.description.startswith("Error code: 401") - events = open_ai_span.events - assert len(events) == 1 - event = events[0] - assert event.name == "exception" - assert event.attributes["exception.type"] == "openai.AuthenticationError" - assert event.attributes["exception.message"].startswith("Error code: 401") - assert open_ai_span.attributes.get("error.type") == "AuthenticationError" - assert "Traceback (most recent call last):" in event.attributes["exception.stacktrace"] - assert "openai.AuthenticationError" in event.attributes["exception.stacktrace"] - assert "invalid_api_key" in event.attributes["exception.stacktrace"] + assert_openai_exception_span(open_ai_span) @pytest.mark.asyncio @@ -1603,18 +1623,7 @@ async def test_chat_async_exception(instrument_legacy, span_exporter, async_open ) assert open_ai_span.attributes.get( SpanAttributes.LLM_IS_STREAMING) is False - assert open_ai_span.status.status_code == StatusCode.ERROR - assert open_ai_span.status.description.startswith("Error code: 401") - events = open_ai_span.events - assert len(events) == 1 - event = events[0] - assert event.name == "exception" - assert event.attributes["exception.type"] == "openai.AuthenticationError" - assert event.attributes["exception.message"].startswith("Error code: 401") - assert "Traceback (most recent call last):" in event.attributes["exception.stacktrace"] - assert "openai.AuthenticationError" in event.attributes["exception.stacktrace"] - assert "invalid_api_key" in event.attributes["exception.stacktrace"] - assert open_ai_span.attributes.get("error.type") == "AuthenticationError" + assert_openai_exception_span(open_ai_span) @pytest.mark.vcr @@ -1648,7 +1657,7 @@ def test_chat_streaming_not_consumed(instrument_legacy, span_exporter, log_expor assert open_ai_span.end_time > open_ai_span.start_time assert open_ai_span.attributes.get( - SpanAttributes.LLM_REQUEST_MODEL) == "gpt-3.5-turbo" + GenAIAttributes.GEN_AI_REQUEST_MODEL) == "gpt-3.5-turbo" assert open_ai_span.attributes.get(SpanAttributes.LLM_IS_STREAMING) is True assert open_ai_span.attributes.get( SpanAttributes.LLM_REQUEST_TYPE) == "chat" @@ -1729,7 +1738,7 @@ def test_chat_streaming_partial_consumption(instrument_legacy, span_exporter, lo assert open_ai_span.end_time is not None assert open_ai_span.attributes.get( - SpanAttributes.LLM_REQUEST_MODEL) == "gpt-3.5-turbo" + GenAIAttributes.GEN_AI_REQUEST_MODEL) == "gpt-3.5-turbo" assert open_ai_span.attributes.get(SpanAttributes.LLM_IS_STREAMING) is True # Should have at least one event from the consumed chunk diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat_parse.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat_parse.py index 4a908d10ea..db936c35df 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat_parse.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_chat_parse.py @@ -9,6 +9,8 @@ from opentelemetry.sdk.trace import Span from opentelemetry.trace import StatusCode +from .utils import assert_openai_exception_span + class StructuredAnswer(BaseModel): rating: int @@ -548,18 +550,7 @@ def test_parsed_completion_exception( assert span.attributes.get(f"{GenAIAttributes.GEN_AI_PROMPT}.0.content") == "Tell me a joke about opentelemetry" assert span.attributes.get(f"{GenAIAttributes.GEN_AI_PROMPT}.0.role") == "user" - assert span.status.status_code == StatusCode.ERROR - assert span.status.description.startswith("Error code: 401") - events = span.events - assert len(events) == 1 - event = events[0] - assert event.name == "exception" - assert event.attributes["exception.type"] == "openai.AuthenticationError" - assert event.attributes["exception.message"].startswith("Error code: 401") - assert "Traceback (most recent call last):" in event.attributes["exception.stacktrace"] - assert "openai.AuthenticationError" in event.attributes["exception.stacktrace"] - assert "invalid_api_key" in event.attributes["exception.stacktrace"] - assert span.attributes.get("error.type") == "AuthenticationError" + assert_openai_exception_span(span) @pytest.mark.asyncio @@ -583,15 +574,4 @@ async def test_async_parsed_completion_exception( assert span.attributes.get(f"{GenAIAttributes.GEN_AI_PROMPT}.0.content") == "Tell me a joke about opentelemetry" assert span.attributes.get(f"{GenAIAttributes.GEN_AI_PROMPT}.0.role") == "user" - assert span.status.status_code == StatusCode.ERROR - assert span.status.description.startswith("Error code: 401") - events = span.events - assert len(events) == 1 - event = events[0] - assert event.name == "exception" - assert event.attributes["exception.type"] == "openai.AuthenticationError" - assert event.attributes["exception.message"].startswith("Error code: 401") - assert "Traceback (most recent call last):" in event.attributes["exception.stacktrace"] - assert "openai.AuthenticationError" in event.attributes["exception.stacktrace"] - assert "invalid_api_key" in event.attributes["exception.stacktrace"] - assert span.attributes.get("error.type") == "AuthenticationError" + assert_openai_exception_span(span) diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py index fd67cbf2bb..5a3fed82a0 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_completions.py @@ -9,7 +9,11 @@ from opentelemetry.trace import StatusCode from opentelemetry.semconv_ai import SpanAttributes -from .utils import assert_request_contains_tracecontext, spy_decorator +from .utils import ( + assert_openai_exception_span, + assert_request_contains_tracecontext, + spy_decorator, +) @pytest.mark.vcr @@ -899,18 +903,7 @@ def test_completion_exception(instrument_legacy, span_exporter, openai_client): open_ai_span.attributes[f"{GenAIAttributes.GEN_AI_PROMPT}.0.user"] == "Tell me a joke about opentelemetry" ) - assert open_ai_span.status.status_code == StatusCode.ERROR - assert open_ai_span.status.description.startswith("Error code: 401") - events = open_ai_span.events - assert len(events) == 1 - event = events[0] - assert event.name == "exception" - assert event.attributes["exception.type"] == "openai.AuthenticationError" - assert event.attributes["exception.message"].startswith("Error code: 401") - assert "Traceback (most recent call last):" in event.attributes["exception.stacktrace"] - assert "openai.AuthenticationError" in event.attributes["exception.stacktrace"] - assert "invalid_api_key" in event.attributes["exception.stacktrace"] - assert open_ai_span.attributes.get("error.type") == "AuthenticationError" + assert_openai_exception_span(open_ai_span) @pytest.mark.asyncio @@ -931,18 +924,7 @@ async def test_async_completion_exception(instrument_legacy, span_exporter, asyn open_ai_span.attributes[f"{GenAIAttributes.GEN_AI_PROMPT}.0.user"] == "Tell me a joke about opentelemetry" ) - assert open_ai_span.status.status_code == StatusCode.ERROR - assert open_ai_span.status.description.startswith("Error code: 401") - events = open_ai_span.events - assert len(events) == 1 - event = events[0] - assert event.name == "exception" - assert event.attributes["exception.type"] == "openai.AuthenticationError" - assert event.attributes["exception.message"].startswith("Error code: 401") - assert "Traceback (most recent call last):" in event.attributes["exception.stacktrace"] - assert "openai.AuthenticationError" in event.attributes["exception.stacktrace"] - assert "invalid_api_key" in event.attributes["exception.stacktrace"] - assert open_ai_span.attributes.get("error.type") == "AuthenticationError" + assert_openai_exception_span(open_ai_span) def assert_message_in_logs(log: ReadableLogRecord, event_name: str, expected_content: dict): diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/test_embeddings.py b/packages/opentelemetry-instrumentation-openai/tests/traces/test_embeddings.py index a56510ad48..2aab558d8e 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/test_embeddings.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/test_embeddings.py @@ -10,7 +10,11 @@ from opentelemetry.semconv_ai import SpanAttributes from opentelemetry.trace import StatusCode -from .utils import assert_request_contains_tracecontext, spy_decorator +from .utils import ( + assert_openai_exception_span, + assert_request_contains_tracecontext, + spy_decorator, +) @pytest.mark.vcr @@ -611,14 +615,7 @@ def test_embeddings_exception(instrument_legacy, span_exporter, openai_client): "openai.embeddings", ] open_ai_span = spans[0] - assert open_ai_span.status.status_code == StatusCode.ERROR - assert open_ai_span.status.description.startswith("Error code: 401") - events = open_ai_span.events - assert len(events) == 1 - event = events[0] - assert event.name == "exception" - assert event.attributes["exception.type"] == "openai.AuthenticationError" - assert event.attributes["exception.message"].startswith("Error code: 401") + assert_openai_exception_span(open_ai_span) @pytest.mark.asyncio @@ -635,14 +632,7 @@ async def test_async_embeddings_exception(instrument_legacy, span_exporter, asyn "openai.embeddings", ] open_ai_span = spans[0] - assert open_ai_span.status.status_code == StatusCode.ERROR - assert open_ai_span.status.description.startswith("Error code: 401") - events = open_ai_span.events - assert len(events) == 1 - event = events[0] - assert event.name == "exception" - assert event.attributes["exception.type"] == "openai.AuthenticationError" - assert event.attributes["exception.message"].startswith("Error code: 401") + assert_openai_exception_span(open_ai_span) def assert_message_in_logs(log: ReadableLogRecord, event_name: str, expected_content: dict): diff --git a/packages/opentelemetry-instrumentation-openai/tests/traces/utils.py b/packages/opentelemetry-instrumentation-openai/tests/traces/utils.py index 0db9c16c94..b370174d46 100644 --- a/packages/opentelemetry-instrumentation-openai/tests/traces/utils.py +++ b/packages/opentelemetry-instrumentation-openai/tests/traces/utils.py @@ -1,5 +1,6 @@ import httpx from opentelemetry.sdk.trace import Span +from opentelemetry.trace import StatusCode from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator from opentelemetry.trace.propagation import get_current_span from unittest.mock import MagicMock @@ -25,3 +26,33 @@ def assert_request_contains_tracecontext(request: httpx.Request, expected_span: assert request_span_context.trace_id == expected_span_context.trace_id assert request_span_context.span_id == expected_span_context.span_id + + +def assert_openai_exception_span(span: Span): + assert span.status.status_code == StatusCode.ERROR + assert span.status.description + + events = span.events + assert len(events) == 1 + + event = events[0] + assert event.name == "exception" + + exception_type = event.attributes["exception.type"] + assert exception_type in { + "openai.AuthenticationError", + "openai.APIConnectionError", + } + + assert event.attributes["exception.message"] == span.status.description + + error_type = span.attributes.get("error.type") + assert error_type in {"AuthenticationError", "APIConnectionError"} + assert error_type == exception_type.split(".")[-1] + + stacktrace = event.attributes["exception.stacktrace"] + assert "Traceback (most recent call last):" in stacktrace + assert exception_type in stacktrace + + if exception_type == "openai.AuthenticationError": + assert "invalid_api_key" in stacktrace