Initial implementation and tests

Author: drewby

instrumentation-genai/opentelemetry-instrumentation-openai-v2/src/opentelemetry/instrumentation/openai_v2/__init__.py

@@ -54,7 +54,12 @@
 from opentelemetry.trace import get_tracer
 
 from .instruments import Instruments
-from .patch import async_chat_completions_create, chat_completions_create
+from .patch import (
+    async_chat_completions_create,
+    async_embeddings_create,
+    chat_completions_create,
+    embeddings_create,
+)
 
 
 class OpenAIInstrumentor(BaseInstrumentor):
@@ -106,8 +111,27 @@ def _instrument(self, **kwargs):
             ),
         )
 
+        # Add instrumentation for the embeddings API
+        wrap_function_wrapper(
+            module="openai.resources.embeddings",
+            name="Embeddings.create",
+            wrapper=embeddings_create(
+                tracer, event_logger, instruments, is_content_enabled()
+            ),
+        )
+
+        wrap_function_wrapper(
+            module="openai.resources.embeddings",
+            name="AsyncEmbeddings.create",
+            wrapper=async_embeddings_create(
+                tracer, event_logger, instruments, is_content_enabled()
+            ),
+        )
+
     def _uninstrument(self, **kwargs):
         import openai  # pylint: disable=import-outside-toplevel
 
         unwrap(openai.resources.chat.completions.Completions, "create")
         unwrap(openai.resources.chat.completions.AsyncCompletions, "create")
+        unwrap(openai.resources.embeddings.Embeddings, "create")
+        unwrap(openai.resources.embeddings.AsyncEmbeddings, "create")

instrumentation-genai/opentelemetry-instrumentation-openai-v2/src/opentelemetry/instrumentation/openai_v2/patch.py

@@ -91,6 +91,7 @@ def traced_method(wrapped, instance, args, kwargs):
                     result,
                     span_attributes,
                     error_type,
+                    GenAIAttributes.GenAiOperationNameValues.CHAT.value,
                 )
 
     return traced_method
@@ -149,6 +150,137 @@ async def traced_method(wrapped, instance, args, kwargs):
                     result,
                     span_attributes,
                     error_type,
+                    GenAIAttributes.GenAiOperationNameValues.CHAT.value,
+                )
+
+    return traced_method
+
+
+def embeddings_create(
+    tracer: Tracer,
+    event_logger: EventLogger,
+    instruments: Instruments,
+    capture_content: bool,
+):
+    """Wrap the `create` method of the `Embeddings` class to trace it."""
+
+    def traced_method(wrapped, instance, args, kwargs):
+        span_attributes = {
+            **get_llm_request_attributes(
+                kwargs,
+                instance,
+                GenAIAttributes.GenAiOperationNameValues.EMBEDDINGS.value,
+            )
+        }
+
+        # Set embeddings dimensions if specified in the request
+        if "dimensions" in kwargs and kwargs["dimensions"] is not None:
+            span_attributes["gen_ai.embeddings.dimensions"] = kwargs[
+                "dimensions"
+            ]
+
+        span_name = f"{span_attributes[GenAIAttributes.GEN_AI_OPERATION_NAME]} {span_attributes[GenAIAttributes.GEN_AI_REQUEST_MODEL]}"
+        with tracer.start_as_current_span(
+            name=span_name,
+            kind=SpanKind.CLIENT,
+            attributes=span_attributes,
+            end_on_exit=False,
+        ) as span:
+            # Store the input for later use in the response attributes
+            input_text = kwargs.get("input", "")
+
+            start = default_timer()
+            result = None
+            error_type = None
+            try:
+                result = wrapped(*args, **kwargs)
+
+                if span.is_recording():
+                    _set_embeddings_response_attributes(
+                        span, result, event_logger, capture_content, input_text
+                    )
+
+                span.end()
+                return result
+
+            except Exception as error:
+                error_type = type(error).__qualname__
+                handle_span_exception(span, error)
+                raise
+            finally:
+                duration = max((default_timer() - start), 0)
+                _record_metrics(
+                    instruments,
+                    duration,
+                    result,
+                    span_attributes,
+                    error_type,
+                    GenAIAttributes.GenAiOperationNameValues.EMBEDDINGS.value,
+                )
+
+    return traced_method
+
+
+def async_embeddings_create(
+    tracer: Tracer,
+    event_logger: EventLogger,
+    instruments: Instruments,
+    capture_content: bool,
+):
+    """Wrap the `create` method of the `AsyncEmbeddings` class to trace it."""
+
+    async def traced_method(wrapped, instance, args, kwargs):
+        span_attributes = {
+            **get_llm_request_attributes(
+                kwargs,
+                instance,
+                GenAIAttributes.GenAiOperationNameValues.EMBEDDINGS.value,
+            )
+        }
+
+        # Set embeddings dimensions if specified in the request
+        if "dimensions" in kwargs and kwargs["dimensions"] is not None:
+            span_attributes["gen_ai.embeddings.dimensions"] = kwargs[
+                "dimensions"
+            ]
+
+        span_name = f"{span_attributes[GenAIAttributes.GEN_AI_OPERATION_NAME]} {span_attributes[GenAIAttributes.GEN_AI_REQUEST_MODEL]}"
+        with tracer.start_as_current_span(
+            name=span_name,
+            kind=SpanKind.CLIENT,
+            attributes=span_attributes,
+            end_on_exit=False,
+        ) as span:
+            # Store the input for later use in the response attributes
+            input_text = kwargs.get("input", "")
+
+            start = default_timer()
+            result = None
+            error_type = None
+            try:
+                result = await wrapped(*args, **kwargs)
+
+                if span.is_recording():
+                    _set_embeddings_response_attributes(
+                        span, result, event_logger, capture_content, input_text
+                    )
+
+                span.end()
+                return result
+
+            except Exception as error:
+                error_type = type(error).__qualname__
+                handle_span_exception(span, error)
+                raise
+            finally:
+                duration = max((default_timer() - start), 0)
+                _record_metrics(
+                    instruments,
+                    duration,
+                    result,
+                    span_attributes,
+                    error_type,
+                    GenAIAttributes.GenAiOperationNameValues.EMBEDDINGS.value,
                 )
 
     return traced_method
@@ -160,15 +292,22 @@ def _record_metrics(
     result,
     span_attributes: dict,
     error_type: Optional[str],
+    operation_name: str,
 ):
+    """Generalized function to record metrics for both chat and embeddings operations."""
     common_attributes = {
-        GenAIAttributes.GEN_AI_OPERATION_NAME: GenAIAttributes.GenAiOperationNameValues.CHAT.value,
+        GenAIAttributes.GEN_AI_OPERATION_NAME: operation_name,
         GenAIAttributes.GEN_AI_SYSTEM: GenAIAttributes.GenAiSystemValues.OPENAI.value,
         GenAIAttributes.GEN_AI_REQUEST_MODEL: span_attributes[
             GenAIAttributes.GEN_AI_REQUEST_MODEL
         ],
     }
 
+    if "gen_ai.embeddings.dimensions" in span_attributes:
+        common_attributes["gen_ai.embeddings.dimensions"] = span_attributes[
+            "gen_ai.embeddings.dimensions"
+        ]
+
     if error_type:
         common_attributes["error.type"] = error_type
 
@@ -210,13 +349,18 @@ def _record_metrics(
             attributes=input_attributes,
         )
 
-        completion_attributes = {
+        token_count = (
+            result.usage.total_tokens
+            if operation_name
+            == GenAIAttributes.GenAiOperationNameValues.EMBEDDINGS.value
+            else result.usage.completion_tokens
+        )
+        attributes = {
             **common_attributes,
             GenAIAttributes.GEN_AI_TOKEN_TYPE: GenAIAttributes.GenAiTokenTypeValues.COMPLETION.value,
         }
         instruments.token_usage_histogram.record(
-            result.usage.completion_tokens,
-            attributes=completion_attributes,
+            token_count, attributes=attributes
         )
 
 
@@ -262,6 +406,54 @@ def _set_response_attributes(
         )
 
 
+def _set_embeddings_response_attributes(
+    span,
+    result,
+    event_logger: EventLogger,
+    capture_content: bool,
+    input_text: str,
+):
+    """Set attributes on the span based on the embeddings response."""
+    # Set the model name if available
+    set_span_attribute(
+        span, GenAIAttributes.GEN_AI_RESPONSE_MODEL, result.model
+    )
+
+    # Set embeddings dimensions if we can determine it from the response
+    if getattr(result, "data", None) and len(result.data) > 0:
+        first_embedding = result.data[0]
+        if getattr(first_embedding, "embedding", None):
+            set_span_attribute(
+                span,
+                "gen_ai.embeddings.dimensions",
+                len(first_embedding.embedding),
+            )
+
+    # Get the usage
+    if getattr(result, "usage", None):
+        set_span_attribute(
+            span,
+            GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS,
+            result.usage.prompt_tokens,
+        )
+        set_span_attribute(
+            span,
+            GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS,
+            result.usage.total_tokens,
+        )
+
+    # Emit event for embeddings if content capture is enabled
+    if capture_content:
+        input_event = Event(
+            name="gen_ai.embeddings",
+            attributes={
+                GenAIAttributes.GEN_AI_SYSTEM: GenAIAttributes.GenAiSystemValues.OPENAI.value
+            },
+            body={"content": input_text, "role": "user"},
+        )
+        event_logger.emit(input_event)
+
+
 class ToolCallBuffer:
     def __init__(self, index, tool_call_id, function_name):
         self.index = index

instrumentation-genai/opentelemetry-instrumentation-openai-v2/src/opentelemetry/instrumentation/openai_v2/utils.py

@@ -192,38 +192,65 @@ def get_llm_request_attributes(
         GenAIAttributes.GEN_AI_OPERATION_NAME: operation_name,
         GenAIAttributes.GEN_AI_SYSTEM: GenAIAttributes.GenAiSystemValues.OPENAI.value,
         GenAIAttributes.GEN_AI_REQUEST_MODEL: kwargs.get("model"),
-        GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE: kwargs.get("temperature"),
-        GenAIAttributes.GEN_AI_REQUEST_TOP_P: kwargs.get("p")
-        or kwargs.get("top_p"),
-        GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS: kwargs.get("max_tokens"),
-        GenAIAttributes.GEN_AI_REQUEST_PRESENCE_PENALTY: kwargs.get(
-            "presence_penalty"
-        ),
-        GenAIAttributes.GEN_AI_REQUEST_FREQUENCY_PENALTY: kwargs.get(
-            "frequency_penalty"
-        ),
-        GenAIAttributes.GEN_AI_OPENAI_REQUEST_SEED: kwargs.get("seed"),
     }
 
-    if (response_format := kwargs.get("response_format")) is not None:
-        # response_format may be string or object with a string in the `type` key
-        if isinstance(response_format, Mapping):
-            if (
-                response_format_type := response_format.get("type")
-            ) is not None:
+    # Add chat-specific attributes only for chat operations
+    if operation_name == GenAIAttributes.GenAiOperationNameValues.CHAT.value:
+        attributes.update(
+            {
+                GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE: kwargs.get(
+                    "temperature"
+                ),
+                GenAIAttributes.GEN_AI_REQUEST_TOP_P: kwargs.get("p")
+                or kwargs.get("top_p"),
+                GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS: kwargs.get(
+                    "max_tokens"
+                ),
+                GenAIAttributes.GEN_AI_REQUEST_PRESENCE_PENALTY: kwargs.get(
+                    "presence_penalty"
+                ),
+                GenAIAttributes.GEN_AI_REQUEST_FREQUENCY_PENALTY: kwargs.get(
+                    "frequency_penalty"
+                ),
+                GenAIAttributes.GEN_AI_OPENAI_REQUEST_SEED: kwargs.get("seed"),
+            }
+        )
+
+        if (response_format := kwargs.get("response_format")) is not None:
+            # response_format may be string or object with a string in the `type` key
+            if isinstance(response_format, Mapping):
+                if (
+                    response_format_type := response_format.get("type")
+                ) is not None:
+                    attributes[
+                        GenAIAttributes.GEN_AI_OPENAI_REQUEST_RESPONSE_FORMAT
+                    ] = response_format_type
+            else:
                 attributes[
                     GenAIAttributes.GEN_AI_OPENAI_REQUEST_RESPONSE_FORMAT
-                ] = response_format_type
-        else:
-            attributes[
-                GenAIAttributes.GEN_AI_OPENAI_REQUEST_RESPONSE_FORMAT
-            ] = response_format
+                ] = response_format
+
+        service_tier = kwargs.get("service_tier")
+        attributes[GenAIAttributes.GEN_AI_OPENAI_RESPONSE_SERVICE_TIER] = (
+            service_tier if service_tier != "auto" else None
+        )
+
+    # Add embeddings-specific attributes
+    elif (
+        operation_name
+        == GenAIAttributes.GenAiOperationNameValues.EMBEDDINGS.value
+    ):
+        # Add embedding dimensions if specified
+        if "dimensions" in kwargs and kwargs["dimensions"] is not None:
+            attributes["gen_ai.embeddings.dimensions"] = kwargs["dimensions"]
+
+        # Add encoding format if specified
+        if "encoding_format" in kwargs:
+            attributes["gen_ai.embeddings.encoding_format"] = kwargs[
+                "encoding_format"
+            ]
 
     set_server_address_and_port(client_instance, attributes)
-    service_tier = kwargs.get("service_tier")
-    attributes[GenAIAttributes.GEN_AI_OPENAI_RESPONSE_SERVICE_TIER] = (
-        service_tier if service_tier != "auto" else None
-    )
 
     # filter out None values
     return {k: v for k, v in attributes.items() if v is not None}