opentelemetry-instrumentation-openai 0.34.1__py3-none-any.whl → 0.49.3__py3-none-any.whl

opentelemetry/instrumentation/openai/v1/assistant_wrappers.py

@@ -1,17 +1,30 @@
 import logging
 import time
+
 from opentelemetry import context as context_api
+from opentelemetry import trace
 from opentelemetry.instrumentation.openai.shared import (
     _set_span_attribute,
     model_as_dict,
 )
-from opentelemetry.trace import SpanKind
-from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
-
-from opentelemetry.semconv_ai import SpanAttributes, LLMRequestTypeValues
-
-from opentelemetry.instrumentation.openai.utils import _with_tracer_wrapper, dont_throw
 from opentelemetry.instrumentation.openai.shared.config import Config
+from opentelemetry.instrumentation.openai.shared.event_emitter import emit_event
+from opentelemetry.instrumentation.openai.shared.event_models import (
+    ChoiceEvent,
+    MessageEvent,
+)
+from opentelemetry.instrumentation.openai.utils import (
+    _with_tracer_wrapper,
+    dont_throw,
+    should_emit_events,
+)
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+from opentelemetry.semconv_ai import LLMRequestTypeValues, SpanAttributes
+from opentelemetry.trace import SpanKind, Status, StatusCode
 
 from openai._legacy_response import LegacyAPIResponse
 from openai.types.beta.threads.run import Run
@@ -45,17 +58,24 @@ def runs_create_wrapper(tracer, wrapped, instance, args, kwargs):
     thread_id = kwargs.get("thread_id")
     instructions = kwargs.get("instructions")
 
-    response = wrapped(*args, **kwargs)
-    response_dict = model_as_dict(response)
+    try:
+        response = wrapped(*args, **kwargs)
+        response_dict = model_as_dict(response)
 
-    runs[thread_id] = {
-        "start_time": time.time_ns(),
-        "assistant_id": kwargs.get("assistant_id"),
-        "instructions": instructions,
-        "run_id": response_dict.get("id"),
-    }
+        runs[thread_id] = {
+            "start_time": time.time_ns(),
+            "assistant_id": kwargs.get("assistant_id"),
+            "instructions": instructions,
+            "run_id": response_dict.get("id"),
+        }
 
-    return response
+        return response
+    except Exception as e:
+        runs[thread_id] = {
+            "exception": e,
+            "end_time": time.time_ns(),
+        }
+        raise
 
 
 @_with_tracer_wrapper
@@ -77,10 +97,16 @@ def runs_retrieve_wrapper(tracer, wrapped, instance, args, kwargs):
     if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
         return wrapped(*args, **kwargs)
 
-    response = wrapped(*args, **kwargs)
-    process_response(response)
-
-    return response
+    try:
+        response = wrapped(*args, **kwargs)
+        process_response(response)
+        return response
+    except Exception as e:
+        thread_id = kwargs.get("thread_id")
+        if thread_id in runs:
+            runs[thread_id]["exception"] = e
+            runs[thread_id]["end_time"] = time.time_ns()
+        raise
 
 
 @_with_tracer_wrapper
@@ -106,64 +132,114 @@ def messages_list_wrapper(tracer, wrapped, instance, args, kwargs):
         start_time=run.get("start_time"),
     )
 
-    i = 0
-    if assistants.get(run["assistant_id"]) is not None or Config.enrich_assistant:
-        if Config.enrich_assistant:
-            assistant = model_as_dict(
-                instance._client.beta.assistants.retrieve(run["assistant_id"])
+    # Use the span as current context to ensure events get proper trace context
+    with trace.use_span(span, end_on_exit=False):
+        if exception := run.get("exception"):
+            span.set_attribute(ERROR_TYPE, exception.__class__.__name__)
+            span.record_exception(exception)
+            span.set_status(Status(StatusCode.ERROR, str(exception)))
+            span.end()
+            return response
+
+        prompt_index = 0
+        if assistants.get(run["assistant_id"]) is not None or Config.enrich_assistant:
+            if Config.enrich_assistant:
+                assistant = model_as_dict(
+                    instance._client.beta.assistants.retrieve(run["assistant_id"])
+                )
+                assistants[run["assistant_id"]] = assistant
+            else:
+                assistant = assistants[run["assistant_id"]]
+
+            _set_span_attribute(
+                span,
+                GenAIAttributes.GEN_AI_SYSTEM,
+                "openai",
             )
-            assistants[run["assistant_id"]] = assistant
-        else:
-            assistant = assistants[run["assistant_id"]]
-
-        _set_span_attribute(
-            span,
-            SpanAttributes.LLM_SYSTEM,
-            "openai",
-        )
-        _set_span_attribute(
-            span,
-            SpanAttributes.LLM_REQUEST_MODEL,
-            assistant["model"],
-        )
-        _set_span_attribute(
-            span,
-            SpanAttributes.LLM_RESPONSE_MODEL,
-            assistant["model"],
-        )
-        _set_span_attribute(span, f"{SpanAttributes.LLM_PROMPTS}.{i}.role", "system")
-        _set_span_attribute(
-            span,
-            f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
-            assistant["instructions"],
-        )
-        i += 1
-    _set_span_attribute(span, f"{SpanAttributes.LLM_PROMPTS}.{i}.role", "system")
-    _set_span_attribute(
-        span, f"{SpanAttributes.LLM_PROMPTS}.{i}.content", run["instructions"]
-    )
-
-    for i, msg in enumerate(messages):
-        prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{i}"
-        content = msg.get("content")
-
-        _set_span_attribute(span, f"{prefix}.role", msg.get("role"))
-        _set_span_attribute(
-            span, f"{prefix}.content", content[0].get("text").get("value")
-        )
-
-    if run.get("usage"):
-        usage_dict = model_as_dict(run.get("usage"))
+            _set_span_attribute(
+                span,
+                GenAIAttributes.GEN_AI_REQUEST_MODEL,
+                assistant["model"],
+            )
+            _set_span_attribute(
+                span,
+                GenAIAttributes.GEN_AI_RESPONSE_MODEL,
+                assistant["model"],
+            )
+            if should_emit_events():
+                emit_event(MessageEvent(content=assistant["instructions"], role="system"))
+            else:
+                _set_span_attribute(
+                    span, f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.role", "system"
+                )
+                _set_span_attribute(
+                    span,
+                    f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.content",
+                    assistant["instructions"],
+                )
+            prompt_index += 1
         _set_span_attribute(
-            span,
-            SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
-            usage_dict.get("completion_tokens"),
+            span, f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.role", "system"
         )
         _set_span_attribute(
             span,
-            SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
-            usage_dict.get("prompt_tokens"),
+            f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.content",
+            run["instructions"],
         )
+        if should_emit_events():
+            emit_event(MessageEvent(content=run["instructions"], role="system"))
+        prompt_index += 1
+
+        completion_index = 0
+        for msg in messages:
+            prefix = f"{GenAIAttributes.GEN_AI_COMPLETION}.{completion_index}"
+            content = msg.get("content")
+
+            message_content = content[0].get("text").get("value")
+            message_role = msg.get("role")
+            if message_role in ["user", "system"]:
+                if should_emit_events():
+                    emit_event(MessageEvent(content=message_content, role=message_role))
+                else:
+                    _set_span_attribute(
+                        span,
+                        f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.role",
+                        message_role,
+                    )
+                    _set_span_attribute(
+                        span,
+                        f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.content",
+                        message_content,
+                    )
+                prompt_index += 1
+            else:
+                if should_emit_events():
+                    emit_event(
+                        ChoiceEvent(
+                            index=completion_index,
+                            message={"content": message_content, "role": message_role},
+                        )
+                    )
+                else:
+                    _set_span_attribute(span, f"{prefix}.role", msg.get("role"))
+                    _set_span_attribute(span, f"{prefix}.content", message_content)
+                    _set_span_attribute(
+                        span, f"gen_ai.response.{completion_index}.id", msg.get("id")
+                    )
+                completion_index += 1
+
+        if run.get("usage"):
+            usage_dict = model_as_dict(run.get("usage"))
+            _set_span_attribute(
+                span,
+                GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS,
+                usage_dict.get("completion_tokens"),
+            )
+            _set_span_attribute(
+                span,
+                GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS,
+                usage_dict.get("prompt_tokens"),
+            )
 
     span.end(run.get("end_time"))
 
@@ -184,47 +260,70 @@ def runs_create_and_stream_wrapper(tracer, wrapped, instance, args, kwargs):
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLMRequestTypeValues.CHAT.value},
     )
 
-    i = 0
-    if assistants.get(assistant_id) is not None or Config.enrich_assistant:
-        if Config.enrich_assistant:
-            assistant = model_as_dict(
-                instance._client.beta.assistants.retrieve(assistant_id)
+    # Use the span as current context to ensure events get proper trace context
+    with trace.use_span(span, end_on_exit=False):
+        i = 0
+        if assistants.get(assistant_id) is not None or Config.enrich_assistant:
+            if Config.enrich_assistant:
+                assistant = model_as_dict(
+                    instance._client.beta.assistants.retrieve(assistant_id)
+                )
+                assistants[assistant_id] = assistant
+            else:
+                assistant = assistants[assistant_id]
+
+            _set_span_attribute(
+                span, GenAIAttributes.GEN_AI_REQUEST_MODEL, assistants[assistant_id]["model"]
+            )
+            _set_span_attribute(
+                span,
+                GenAIAttributes.GEN_AI_SYSTEM,
+                "openai",
             )
-            assistants[assistant_id] = assistant
+            _set_span_attribute(
+                span,
+                GenAIAttributes.GEN_AI_RESPONSE_MODEL,
+                assistants[assistant_id]["model"],
+            )
+            if should_emit_events():
+                emit_event(
+                    MessageEvent(
+                        content=assistants[assistant_id]["instructions"], role="system"
+                    )
+                )
+            else:
+                _set_span_attribute(
+                    span, f"{GenAIAttributes.GEN_AI_PROMPT}.{i}.role", "system"
+                )
+                _set_span_attribute(
+                    span,
+                    f"{GenAIAttributes.GEN_AI_PROMPT}.{i}.content",
+                    assistants[assistant_id]["instructions"],
+                )
+            i += 1
+        if should_emit_events():
+            emit_event(MessageEvent(content=instructions, role="system"))
         else:
-            assistant = assistants[assistant_id]
+            _set_span_attribute(span, f"{GenAIAttributes.GEN_AI_PROMPT}.{i}.role", "system")
+            _set_span_attribute(
+                span, f"{GenAIAttributes.GEN_AI_PROMPT}.{i}.content", instructions
+            )
 
-        _set_span_attribute(
-            span, SpanAttributes.LLM_REQUEST_MODEL, assistants[assistant_id]["model"]
+        from opentelemetry.instrumentation.openai.v1.event_handler_wrapper import (
+            EventHandleWrapper,
         )
-        _set_span_attribute(
-            span,
-            SpanAttributes.LLM_SYSTEM,
-            "openai",
-        )
-        _set_span_attribute(
-            span,
-            SpanAttributes.LLM_RESPONSE_MODEL,
-            assistants[assistant_id]["model"],
-        )
-        _set_span_attribute(span, f"{SpanAttributes.LLM_PROMPTS}.{i}.role", "system")
-        _set_span_attribute(
-            span,
-            f"{SpanAttributes.LLM_PROMPTS}.{i}.content",
-            assistants[assistant_id]["instructions"],
-        )
-        i += 1
-    _set_span_attribute(span, f"{SpanAttributes.LLM_PROMPTS}.{i}.role", "system")
-    _set_span_attribute(span, f"{SpanAttributes.LLM_PROMPTS}.{i}.content", instructions)
 
-    from opentelemetry.instrumentation.openai.v1.event_handler_wrapper import (
-        EventHandleWrapper,
-    )
-
-    kwargs["event_handler"] = EventHandleWrapper(
-        original_handler=kwargs["event_handler"], span=span
-    )
-
-    response = wrapped(*args, **kwargs)
+        kwargs["event_handler"] = EventHandleWrapper(
+            original_handler=kwargs["event_handler"],
+            span=span,
+        )
 
-    return response
+        try:
+            response = wrapped(*args, **kwargs)
+            return response
+        except Exception as e:
+            span.set_attribute(ERROR_TYPE, e.__class__.__name__)
+            span.record_exception(e)
+            span.set_status(Status(StatusCode.ERROR, str(e)))
+            span.end()
+            raise

opentelemetry/instrumentation/openai/v1/event_handler_wrapper.py

@@ -1,13 +1,18 @@
-from opentelemetry.instrumentation.openai.shared import (
-    _set_span_attribute,
+from opentelemetry.instrumentation.openai.shared import _set_span_attribute
+from opentelemetry.instrumentation.openai.shared.event_emitter import emit_event
+from opentelemetry.instrumentation.openai.shared.event_models import ChoiceEvent
+from opentelemetry.instrumentation.openai.utils import should_emit_events
+from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
 )
-from opentelemetry.semconv_ai import SpanAttributes
-from openai import AssistantEventHandler
+from opentelemetry.trace import Status, StatusCode
 from typing_extensions import override
 
+from openai import AssistantEventHandler
 
-class EventHandleWrapper(AssistantEventHandler):
 
+class EventHandleWrapper(AssistantEventHandler):
     _current_text_index = 0
     _prompt_tokens = 0
     _completion_tokens = 0
@@ -21,12 +26,12 @@ class EventHandleWrapper(AssistantEventHandler):
     def on_end(self):
         _set_span_attribute(
             self._span,
-            SpanAttributes.LLM_USAGE_PROMPT_TOKENS,
+            GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS,
             self._prompt_tokens,
         )
         _set_span_attribute(
             self._span,
-            SpanAttributes.LLM_USAGE_COMPLETION_TOKENS,
+            GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS,
             self._completion_tokens,
         )
         self._original_handler.on_end()
@@ -65,6 +70,9 @@ class EventHandleWrapper(AssistantEventHandler):
 
     @override
     def on_exception(self, exception: Exception):
+        self._span.set_attribute(ERROR_TYPE, exception.__class__.__name__)
+        self._span.record_exception(exception)
+        self._span.set_status(Status(StatusCode.ERROR, str(exception)))
         self._original_handler.on_exception(exception)
 
     @override
@@ -81,7 +89,22 @@ class EventHandleWrapper(AssistantEventHandler):
 
     @override
     def on_message_done(self, message):
+        _set_span_attribute(
+            self._span,
+            f"gen_ai.response.{self._current_text_index}.id",
+            message.id,
+        )
+        emit_event(
+            ChoiceEvent(
+                index=self._current_text_index,
+                message={
+                    "content": [item.model_dump() for item in message.content],
+                    "role": message.role,
+                },
+            )
+        )
         self._original_handler.on_message_done(message)
+        self._current_text_index += 1
 
     @override
     def on_text_created(self, text):
@@ -94,18 +117,17 @@ class EventHandleWrapper(AssistantEventHandler):
     @override
     def on_text_done(self, text):
         self._original_handler.on_text_done(text)
-        _set_span_attribute(
-            self._span,
-            f"{SpanAttributes.LLM_COMPLETIONS}.{self._current_text_index}.role",
-            "assistant",
-        )
-        _set_span_attribute(
-            self._span,
-            f"{SpanAttributes.LLM_COMPLETIONS}.{self._current_text_index}.content",
-            text.value,
-        )
-
-        self._current_text_index += 1
+        if not should_emit_events():
+            _set_span_attribute(
+                self._span,
+                f"{GenAIAttributes.GEN_AI_COMPLETION}.{self._current_text_index}.role",
+                "assistant",
+            )
+            _set_span_attribute(
+                self._span,
+                f"{GenAIAttributes.GEN_AI_COMPLETION}.{self._current_text_index}.content",
+                text.value,
+            )
 
     @override
     def on_image_file_done(self, image_file):