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

opentelemetry/instrumentation/openai/shared/completion_wrappers.py

@@ -1,37 +1,44 @@
 import logging
 
 from opentelemetry import context as context_api
-
-from opentelemetry.semconv_ai import (
-    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
-    SpanAttributes,
-    LLMRequestTypeValues,
-)
-
-from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
-from opentelemetry.instrumentation.openai.utils import _with_tracer_wrapper, dont_throw
+from opentelemetry import trace
 from opentelemetry.instrumentation.openai.shared import (
     _set_client_attributes,
-    _set_request_attributes,
-    _set_span_attribute,
     _set_functions_attributes,
+    _set_request_attributes,
     _set_response_attributes,
+    _set_span_attribute,
+    _set_span_stream_usage,
     is_streaming_response,
-    should_send_prompts,
     model_as_dict,
-    should_record_stream_token_usage,
-    get_token_count_from_string,
-    _set_span_stream_usage,
     propagate_trace_context,
 )
-
-from opentelemetry.instrumentation.openai.utils import is_openai_v1
-
+from opentelemetry.instrumentation.openai.shared.config import Config
+from opentelemetry.semconv.attributes.error_attributes import ERROR_TYPE
+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,
+    is_openai_v1,
+    should_emit_events,
+    should_send_prompts,
+)
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+from opentelemetry.semconv_ai import (
+    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
+    LLMRequestTypeValues,
+    SpanAttributes,
+)
 from opentelemetry.trace import SpanKind
 from opentelemetry.trace.status import Status, StatusCode
 
-from opentelemetry.instrumentation.openai.shared.config import Config
-
 SPAN_NAME = "openai.completion"
 LLM_REQUEST_TYPE = LLMRequestTypeValues.COMPLETION
 
@@ -52,17 +59,27 @@ def completion_wrapper(tracer, wrapped, instance, args, kwargs):
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
     )
 
-    _handle_request(span, kwargs, instance)
-    response = wrapped(*args, **kwargs)
+    # Use the span as current context to ensure events get proper trace context
+    with trace.use_span(span, end_on_exit=False):
+        _handle_request(span, kwargs, instance)
 
-    if is_streaming_response(response):
-        # span will be closed after the generator is done
-        return _build_from_streaming_response(span, kwargs, response)
-    else:
-        _handle_response(response, span)
+        try:
+            response = wrapped(*args, **kwargs)
+        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
 
-    span.end()
-    return response
+        if is_streaming_response(response):
+            # span will be closed after the generator is done
+            return _build_from_streaming_response(span, kwargs, response)
+        else:
+            _handle_response(response, span, instance)
+
+        span.end()
+        return response
 
 
 @_with_tracer_wrapper
@@ -78,41 +95,66 @@ async def acompletion_wrapper(tracer, wrapped, instance, args, kwargs):
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
     )
 
-    _handle_request(span, kwargs, instance)
-    response = await wrapped(*args, **kwargs)
+    # Use the span as current context to ensure events get proper trace context
+    with trace.use_span(span, end_on_exit=False):
+        _handle_request(span, kwargs, instance)
 
-    if is_streaming_response(response):
-        # span will be closed after the generator is done
-        return _abuild_from_streaming_response(span, kwargs, response)
-    else:
-        _handle_response(response, span)
+        try:
+            response = await wrapped(*args, **kwargs)
+        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
 
-    span.end()
-    return response
+        if is_streaming_response(response):
+            # span will be closed after the generator is done
+            return _abuild_from_streaming_response(span, kwargs, response)
+        else:
+            _handle_response(response, span, instance)
+
+        span.end()
+        return response
 
 
 @dont_throw
 def _handle_request(span, kwargs, instance):
-    _set_request_attributes(span, kwargs)
-    if should_send_prompts():
-        _set_prompts(span, kwargs.get("prompt"))
-        _set_functions_attributes(span, kwargs.get("functions"))
+    _set_request_attributes(span, kwargs, instance)
+    if should_emit_events():
+        _emit_prompts_events(kwargs)
+    else:
+        if should_send_prompts():
+            _set_prompts(span, kwargs.get("prompt"))
+            _set_functions_attributes(span, kwargs.get("functions"))
     _set_client_attributes(span, instance)
     if Config.enable_trace_context_propagation:
         propagate_trace_context(span, kwargs)
 
 
+def _emit_prompts_events(kwargs):
+    prompt = kwargs.get("prompt")
+    if isinstance(prompt, list):
+        for p in prompt:
+            emit_event(MessageEvent(content=p))
+    elif isinstance(prompt, str):
+        emit_event(MessageEvent(content=prompt))
+
+
 @dont_throw
-def _handle_response(response, span):
+def _handle_response(response, span, instance=None):
     if is_openai_v1():
         response_dict = model_as_dict(response)
     else:
         response_dict = response
 
     _set_response_attributes(span, response_dict)
-
-    if should_send_prompts():
-        _set_completions(span, response_dict.get("choices"))
+    if should_emit_events():
+        for choice in response.choices:
+            emit_event(_parse_choice_event(choice))
+    else:
+        if should_send_prompts():
+            _set_completions(span, response_dict.get("choices"))
 
 
 def _set_prompts(span, prompt):
@@ -121,7 +163,7 @@ def _set_prompts(span, prompt):
 
     _set_span_attribute(
         span,
-        f"{SpanAttributes.LLM_PROMPTS}.0.user",
+        f"{GenAIAttributes.GEN_AI_PROMPT}.0.user",
         prompt[0] if isinstance(prompt, list) else prompt,
     )
 
@@ -133,7 +175,7 @@ def _set_completions(span, choices):
 
     for choice in choices:
         index = choice.get("index")
-        prefix = f"{SpanAttributes.LLM_COMPLETIONS}.{index}"
+        prefix = f"{GenAIAttributes.GEN_AI_COMPLETION}.{index}"
         _set_span_attribute(
             span, f"{prefix}.finish_reason", choice.get("finish_reason")
         )
@@ -142,7 +184,7 @@ def _set_completions(span, choices):
 
 @dont_throw
 def _build_from_streaming_response(span, request_kwargs, response):
-    complete_response = {"choices": [], "model": ""}
+    complete_response = {"choices": [], "model": "", "id": ""}
     for item in response:
         yield item
         _accumulate_streaming_response(complete_response, item)
@@ -151,8 +193,11 @@ def _build_from_streaming_response(span, request_kwargs, response):
 
     _set_token_usage(span, request_kwargs, complete_response)
 
-    if should_send_prompts():
-        _set_completions(span, complete_response.get("choices"))
+    if should_emit_events():
+        _emit_streaming_response_events(complete_response)
+    else:
+        if should_send_prompts():
+            _set_completions(span, complete_response.get("choices"))
 
     span.set_status(Status(StatusCode.OK))
     span.end()
@@ -160,7 +205,7 @@ def _build_from_streaming_response(span, request_kwargs, response):
 
 @dont_throw
 async def _abuild_from_streaming_response(span, request_kwargs, response):
-    complete_response = {"choices": [], "model": ""}
+    complete_response = {"choices": [], "model": "", "id": ""}
     async for item in response:
         yield item
         _accumulate_streaming_response(complete_response, item)
@@ -169,44 +214,42 @@ async def _abuild_from_streaming_response(span, request_kwargs, response):
 
     _set_token_usage(span, request_kwargs, complete_response)
 
-    if should_send_prompts():
-        _set_completions(span, complete_response.get("choices"))
+    if should_emit_events():
+        _emit_streaming_response_events(complete_response)
+    else:
+        if should_send_prompts():
+            _set_completions(span, complete_response.get("choices"))
 
     span.set_status(Status(StatusCode.OK))
     span.end()
 
 
-@dont_throw
-def _set_token_usage(span, request_kwargs, complete_response):
-    # use tiktoken calculate token usage
-    if should_record_stream_token_usage():
-        prompt_usage = -1
-        completion_usage = -1
-
-        # prompt_usage
-        if request_kwargs and request_kwargs.get("prompt"):
-            prompt_content = request_kwargs.get("prompt")
-            model_name = complete_response.get("model") or None
-
-            if model_name:
-                prompt_usage = get_token_count_from_string(prompt_content, model_name)
+def _emit_streaming_response_events(complete_response):
+    for i, choice in enumerate(complete_response["choices"]):
+        emit_event(
+            ChoiceEvent(
+                index=choice.get("index", i),
+                message={"content": choice.get("text"), "role": "assistant"},
+                finish_reason=choice.get("finish_reason", "unknown"),
+            )
+        )
 
-        # completion_usage
-        if complete_response.get("choices"):
-            completion_content = ""
-            model_name = complete_response.get("model") or None
 
-            for choice in complete_response.get("choices"):
-                if choice.get("text"):
-                    completion_content += choice.get("text")
+@dont_throw
+def _set_token_usage(span, request_kwargs, complete_response):
+    prompt_usage = -1
+    completion_usage = -1
 
-            if model_name:
-                completion_usage = get_token_count_from_string(
-                    completion_content, model_name
-                )
+    # Use token usage from API response only
+    if complete_response.get("usage"):
+        usage = complete_response["usage"]
+        if usage.get("prompt_tokens"):
+            prompt_usage = usage["prompt_tokens"]
+        if usage.get("completion_tokens"):
+            completion_usage = usage["completion_tokens"]
 
-        # span record
-        _set_span_stream_usage(span, prompt_usage, completion_usage)
+    # span record
+    _set_span_stream_usage(span, prompt_usage, completion_usage)
 
 
 @dont_throw
@@ -215,6 +258,11 @@ def _accumulate_streaming_response(complete_response, item):
         item = model_as_dict(item)
 
     complete_response["model"] = item.get("model")
+    complete_response["id"] = item.get("id")
+
+    # capture usage information from the stream chunks
+    if item.get("usage"):
+        complete_response["usage"] = item.get("usage")
 
     for choice in item.get("choices"):
         index = choice.get("index")
@@ -228,3 +276,17 @@ def _accumulate_streaming_response(complete_response, item):
             complete_choice["text"] += choice.get("text")
 
     return complete_response
+
+
+def _parse_choice_event(choice) -> ChoiceEvent:
+    has_message = choice.text is not None
+    has_finish_reason = choice.finish_reason is not None
+
+    content = choice.text if has_message else None
+    finish_reason = choice.finish_reason if has_finish_reason else "unknown"
+
+    return ChoiceEvent(
+        index=choice.index,
+        message={"content": content, "role": "assistant"},
+        finish_reason=finish_reason,
+    )

opentelemetry/instrumentation/openai/shared/config.py

@@ -1,10 +1,15 @@
-from typing import Callable
+from typing import Callable, Optional
+
+from opentelemetry._logs import Logger
 
 
 class Config:
-    enrich_token_usage = False
     enrich_assistant = False
     exception_logger = None
     get_common_metrics_attributes: Callable[[], dict] = lambda: {}
-    upload_base64_image: Callable[[str, str, str], str] = lambda trace_id, span_id, base64_image_url: str
+    upload_base64_image: Callable[[str, str, str, str], str] = (
+        lambda trace_id, span_id, image_name, base64_string: str
+    )
     enable_trace_context_propagation: bool = True
+    use_legacy_attributes = True
+    event_logger: Optional[Logger] = None

opentelemetry/instrumentation/openai/shared/embeddings_wrappers.py

@@ -1,39 +1,49 @@
 import logging
 import time
+from collections.abc import Iterable
 
 from opentelemetry import context as context_api
-from opentelemetry.metrics import Counter, Histogram
-from opentelemetry.semconv_ai import (
-    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
-    SpanAttributes,
-    LLMRequestTypeValues,
-)
-
-from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
-from opentelemetry.instrumentation.openai.utils import (
-    dont_throw,
-    start_as_current_span_async,
-    _with_embeddings_telemetry_wrapper,
-)
 from opentelemetry.instrumentation.openai.shared import (
-    metric_shared_attributes,
+    OPENAI_LLM_USAGE_TOKEN_TYPES,
+    _get_openai_base_url,
     _set_client_attributes,
     _set_request_attributes,
-    _set_span_attribute,
     _set_response_attributes,
+    _set_span_attribute,
     _token_type,
-    should_send_prompts,
+    metric_shared_attributes,
     model_as_dict,
-    _get_openai_base_url,
-    OPENAI_LLM_USAGE_TOKEN_TYPES,
     propagate_trace_context,
 )
-
 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_embeddings_telemetry_wrapper,
+    dont_throw,
+    is_openai_v1,
+    should_emit_events,
+    should_send_prompts,
+    start_as_current_span_async,
+)
+from opentelemetry.instrumentation.utils import _SUPPRESS_INSTRUMENTATION_KEY
+from opentelemetry.metrics import Counter, Histogram
+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 (
+    SUPPRESS_LANGUAGE_MODEL_INSTRUMENTATION_KEY,
+    LLMRequestTypeValues,
+    SpanAttributes,
+)
+from opentelemetry.trace import SpanKind, Status, StatusCode
 
-from opentelemetry.instrumentation.openai.utils import is_openai_v1
-
-from opentelemetry.trace import SpanKind
+from openai._legacy_response import LegacyAPIResponse
+from openai.types.create_embedding_response import CreateEmbeddingResponse
 
 SPAN_NAME = "openai.embeddings"
 LLM_REQUEST_TYPE = LLMRequestTypeValues.EMBEDDING
@@ -83,7 +93,12 @@ def embeddings_wrapper(
             if exception_counter:
                 exception_counter.add(1, attributes=attributes)
 
-            raise e
+            span.set_attribute(ERROR_TYPE, e.__class__.__name__)
+            span.record_exception(e)
+            span.set_status(Status(StatusCode.ERROR, str(e)))
+            span.end()
+
+            raise
 
         duration = end_time - start_time
 
@@ -124,6 +139,7 @@ async def aembeddings_wrapper(
         attributes={SpanAttributes.LLM_REQUEST_TYPE: LLM_REQUEST_TYPE.value},
     ) as span:
         _handle_request(span, kwargs, instance)
+
         try:
             # record time for duration
             start_time = time.time()
@@ -142,9 +158,15 @@ async def aembeddings_wrapper(
             if exception_counter:
                 exception_counter.add(1, attributes=attributes)
 
-            raise e
+            span.set_attribute(ERROR_TYPE, e.__class__.__name__)
+            span.record_exception(e)
+            span.set_status(Status(StatusCode.ERROR, str(e)))
+            span.end()
+
+            raise
 
         duration = end_time - start_time
+
         _handle_response(
             response,
             span,
@@ -160,10 +182,16 @@ async def aembeddings_wrapper(
 
 @dont_throw
 def _handle_request(span, kwargs, instance):
-    _set_request_attributes(span, kwargs)
-    if should_send_prompts():
-        _set_prompts(span, kwargs.get("input"))
+    _set_request_attributes(span, kwargs, instance)
+
+    if should_emit_events():
+        _emit_embeddings_message_event(kwargs.get("input"))
+    else:
+        if should_send_prompts():
+            _set_prompts(span, kwargs.get("input"))
+
     _set_client_attributes(span, instance)
+
     if Config.enable_trace_context_propagation:
         propagate_trace_context(span, kwargs)
 
@@ -194,6 +222,10 @@ def _handle_response(
     # span attributes
     _set_response_attributes(span, response_dict)
 
+    # emit events
+    if should_emit_events():
+        _emit_embeddings_choice_event(response)
+
 
 def _set_embeddings_metrics(
     instance,
@@ -219,7 +251,7 @@ def _set_embeddings_metrics(
                     continue
                 attributes_with_token_type = {
                     **shared_attributes,
-                    SpanAttributes.LLM_TOKEN_TYPE: _token_type(name),
+                    GenAIAttributes.GEN_AI_TOKEN_TYPE: _token_type(name),
                 }
                 token_counter.record(val, attributes=attributes_with_token_type)
 
@@ -241,10 +273,39 @@ def _set_prompts(span, prompt):
 
     if isinstance(prompt, list):
         for i, p in enumerate(prompt):
-            _set_span_attribute(span, f"{SpanAttributes.LLM_PROMPTS}.{i}.content", p)
+            _set_span_attribute(span, f"{GenAIAttributes.GEN_AI_PROMPT}.{i}.content", p)
     else:
         _set_span_attribute(
             span,
-            f"{SpanAttributes.LLM_PROMPTS}.0.content",
+            f"{GenAIAttributes.GEN_AI_PROMPT}.0.content",
             prompt,
         )
+
+
+def _emit_embeddings_message_event(embeddings) -> None:
+    if isinstance(embeddings, str):
+        emit_event(MessageEvent(content=embeddings))
+    elif isinstance(embeddings, Iterable):
+        for i in embeddings:
+            emit_event(MessageEvent(content=i))
+
+
+def _emit_embeddings_choice_event(response) -> None:
+    if isinstance(response, CreateEmbeddingResponse):
+        for embedding in response.data:
+            emit_event(
+                ChoiceEvent(
+                    index=embedding.index,
+                    message={"content": embedding.embedding, "role": "assistant"},
+                )
+            )
+
+    elif isinstance(response, LegacyAPIResponse):
+        parsed_response = response.parse()
+        for embedding in parsed_response.data:
+            emit_event(
+                ChoiceEvent(
+                    index=embedding.index,
+                    message={"content": embedding.embedding, "role": "assistant"},
+                )
+            )

opentelemetry/instrumentation/openai/shared/event_emitter.py

@@ -0,0 +1,108 @@
+from dataclasses import asdict
+from enum import Enum
+from typing import Union
+
+from opentelemetry._logs import LogRecord
+from opentelemetry.instrumentation.openai.shared.event_models import (
+    ChoiceEvent,
+    MessageEvent,
+)
+from opentelemetry.instrumentation.openai.utils import (
+    should_emit_events,
+    should_send_prompts,
+)
+from opentelemetry.semconv._incubating.attributes import (
+    gen_ai_attributes as GenAIAttributes,
+)
+
+from .config import Config
+
+
+class Roles(Enum):
+    USER = "user"
+    ASSISTANT = "assistant"
+    SYSTEM = "system"
+    TOOL = "tool"
+
+
+VALID_MESSAGE_ROLES = {role.value for role in Roles}
+"""The valid roles for naming the message event."""
+
+EVENT_ATTRIBUTES = {
+    GenAIAttributes.GEN_AI_SYSTEM: GenAIAttributes.GenAiSystemValues.OPENAI.value
+}
+"""The attributes to be used for the event."""
+
+
+def emit_event(event: Union[MessageEvent, ChoiceEvent]) -> None:
+    """
+    Emit an event to the OpenTelemetry SDK.
+
+    Args:
+        event: The event to emit.
+    """
+    if not should_emit_events():
+        return
+
+    if isinstance(event, MessageEvent):
+        _emit_message_event(event)
+    elif isinstance(event, ChoiceEvent):
+        _emit_choice_event(event)
+    else:
+        raise TypeError("Unsupported event type")
+
+
+def _emit_message_event(event: MessageEvent) -> None:
+    body = asdict(event)
+
+    if event.role in VALID_MESSAGE_ROLES:
+        name = "gen_ai.{}.message".format(event.role)
+        # According to the semantic conventions, the role is conditionally required if available
+        # and not equal to the "role" in the message name. So, remove the role from the body if
+        # it is the same as the in the event name.
+        body.pop("role", None)
+    else:
+        name = "gen_ai.user.message"
+
+    # According to the semantic conventions, only the assistant role has tool call
+    if event.role != Roles.ASSISTANT.value and event.tool_calls is not None:
+        del body["tool_calls"]
+    elif event.tool_calls is None:
+        del body["tool_calls"]
+
+    if not should_send_prompts():
+        del body["content"]
+        if body.get("tool_calls") is not None:
+            for tool_call in body["tool_calls"]:
+                tool_call["function"].pop("arguments", None)
+
+    log_record = LogRecord(
+        body=body,
+        attributes=EVENT_ATTRIBUTES,
+        event_name=name
+    )
+    Config.event_logger.emit(log_record)
+
+
+def _emit_choice_event(event: ChoiceEvent) -> None:
+    body = asdict(event)
+    if event.message["role"] == Roles.ASSISTANT.value:
+        # According to the semantic conventions, the role is conditionally required if available
+        # and not equal to "assistant", so remove the role from the body if it is "assistant".
+        body["message"].pop("role", None)
+
+    if event.tool_calls is None:
+        del body["tool_calls"]
+
+    if not should_send_prompts():
+        body["message"].pop("content", None)
+        if body.get("tool_calls") is not None:
+            for tool_call in body["tool_calls"]:
+                tool_call["function"].pop("arguments", None)
+
+    log_record = LogRecord(
+        body=body,
+        attributes=EVENT_ATTRIBUTES,
+        event_name="gen_ai.choice"
+    )
+    Config.event_logger.emit(log_record)

opentelemetry/instrumentation/openai/shared/event_models.py

@@ -0,0 +1,41 @@
+from dataclasses import dataclass
+from typing import Any, List, Literal, Optional, TypedDict
+
+
+class _FunctionToolCall(TypedDict):
+    function_name: str
+    arguments: Optional[dict[str, Any]]
+
+
+class ToolCall(TypedDict):
+    """Represents a tool call in the AI model."""
+
+    id: str
+    function: _FunctionToolCall
+    type: Literal["function"]
+
+
+class CompletionMessage(TypedDict):
+    """Represents a message in the AI model."""
+
+    content: Any
+    role: str = "assistant"
+
+
+@dataclass
+class MessageEvent:
+    """Represents an input event for the AI model."""
+
+    content: Any
+    role: str = "user"
+    tool_calls: Optional[List[ToolCall]] = None
+
+
+@dataclass
+class ChoiceEvent:
+    """Represents a completion event for the AI model."""
+
+    index: int
+    message: CompletionMessage
+    finish_reason: str = "unknown"
+    tool_calls: Optional[List[ToolCall]] = None